start and call ws api

app.py

@@ -2,18 +2,13 @@ import os
 import pip
 import gradio as gr
 from epub2txt import epub2txt
+from websocket import create_connection
 
-if not os.path.exists("./chatglm2-6b-int4.flm"):
+if not os.path.exists(os.getenv("checkpoint_path")):
     os.system("git clone --recurse-submodules https://github.com/ztxz16/fastllm.git")
-    os.system("cd fastllm; mkdir build; cd build; cmake ..; make -j")
-    # os.system("cd fastllm; mkdir build; cd build; cmake ..; make -j; cd tools; python setup.py install --user --prefix=")
+    os.system("cd fastllm; mkdir build; cd build; cmake ..; make -j; cd tools; python setup.py install --user --prefix=")
     # os.system("wget https://huggingface.co/huangyuyang/chatglm2-6b-int4.flm/resolve/main/chatglm2-6b-int4.flm")
-    pip.main(['install', '-e', 'fastllm/build/tools/', '--user'])
-    # exit()
-    
-from fastllm_pytools import llm
-model = llm.model("./chatglm2-6b-int4.flm")
-prompt = os.getenv("prompt")
+    os.system("uvicorn api:app --reload")
 
 class GUI:
     def __init__(self, *args, **kwargs):
@@ -22,21 +17,35 @@ class GUI:
                 gr.Markdown(scale=2).attach_load_event(self.hello, None)
                 gr.LoginButton()
                 gr.LogoutButton()
-            self.out = gr.Markdown()
+            out = gr.Markdown()
             inp = gr.File(file_types=['.epub'])
-            inp.change(self.process, inp, self.out)
+            inp.change(self.process, inp, out)
+            self.ws = None
+            self.out = []
         demo.launch()
 
     def process(self, file, profile: gr.OAuthProfile | None):
         if profile is None:
             return gr.update(value='Login to access the tool.')
-        ch_list = epub2txt(file.name, outputlist=True)
+
+        h = ''
         chapter_titles = epub2txt.content_titles
         title = epub2txt.title
+        if self.ws is None:
+            self.ws = create_connection(f"ws://localhost/{file.name}/ws")
+        res = self.ws.recv()
+        if 'chsum: ' in res:
+            self.out.append(res.remove("chsum: "))
+        elif 'draft_sum: ' in res:
+            h = res[11:]
+        elif 'output: ' in res:
+            self.ws.close()
+            self.ws = None
+            self.out = []
+            return gr.update(value=res)
         return gr.update(
-            value=f"# {title}\n\n" + model.response(prompt+ch_list[2]))
-            # value=f"# {title}\n\n" + "\n\n".join(
-            #     [f"## {ct}\n\n{c}" for ct, c in zip(chapter_titles, ch_list)]))
+            value=f"# {title}\n\n" +h+ "\n\n".join(
+                [f"## {ct}\n\n{c}" for ct, c in zip(chapter_titles, self.out)]))
 
     def hello(self, profile: gr.OAuthProfile | None):
         if profile is None: