Create api.py

api.py

@@ -0,0 +1,44 @@
+from typing import Union
+
+from fastapi import FastAPI
+from epub2txt import epub2txt
+import config
+
+model = llm.model(config.CKP_PATH)
+prompt = os.getenv("prompt")
+
+from fastllm_pytools import llm
+
+app = FastAPI()
+
+
+@app.websockets("/{f_name}/ws")
+async def read_root(*,
+                    websocket: WebSocket,
+                   f_name: str):
+    await websocket.accept()
+    ch_list = epub2txt(f_name, outputlist=True)
+    chapter_titles = epub2txt.content_titles
+    title = epub2txt.title
+
+    idx = 0
+    sm_list = []
+    for text in ch_list[2:]:
+        idx += 1
+        docs = []
+        for i in range(0, len(text)//2000+1, 2000):
+            t = text[i:i+2048]
+            if len(t) > 0:
+                docs.append(model.response(prompt+t))
+                await websocket.send_text(f"chsum: {docs[-1]}")
+        hist = docs[0]
+        for doc in docs[1:]:
+            hist = model.response(prompt+"\n"+hist+"\n"+doc)
+            await websocket.send_text(f"draft_sum: {hist}")
+        sm_list.append(hist)
+        mdobj_str = f"# {title}\n\n"
+        for ct, sm in zip(chapter_titles[2:], sm_list):
+            mdobj_str += f"## {ct}\n\n{sm}\n\n\n"
+    await websocket.send_text(f"output: {mdobj_str}")
+
+# uvicorn api:app --reload