Update main.py

main.py

@@ -2,7 +2,7 @@ from typing import Union
 from fastapi import FastAPI
 from fastapi.encoders import jsonable_encoder
 from fastapi.websockets import WebSocket, WebSocketDisconnect, WebSocketState
-from fastapi.responses import HTMLResponse, JSONResponse
+from fastapi.responses import JSONResponse
 
 from answerer import Answerer
 from mapper import Mapper
@@ -17,6 +17,8 @@ answerer = Answerer(
   strategy="cpu bf16",
   ctx_limit=16*1024,
 )
+
+accelerator: Union[WebSocket, None] = None
   
 app = FastAPI()
 
@@ -63,6 +65,11 @@ HTML = """
 # def index():
 #   return HTMLResponse(HTML)
 
+@app.websocket("/accelerate")
+async def answer(ws: WebSocket):
+  await ws.accept()
+  accelerator = ws
+
 @app.post("/map")
 def map(query: Union[str, None], items: Union[list[str], None]):
   scores = mapper(query, items)
@@ -75,11 +82,20 @@ async def answer(ws: WebSocket):
   try: input = await ws.receive_text()
   except WebSocketDisconnect: return
 
-  output = answerer(input, 32)
-  async for el in output:
+  if accelerator:
+    await accelerator.send_text(input)
+    output = await accelerator.receive_text()
     if(ws.client_state == WebSocketState.CONNECTED):
       await ws.send_text(el)
     elif(ws.client_state == WebSocketState.DISCONNECTED):
       return
-    
-  await ws.close()
+  else:
+    output = answerer(input, 32)
+    async for el in output:
+      if(ws.client_state == WebSocketState.CONNECTED):
+        await ws.send_text(el)
+      elif(ws.client_state == WebSocketState.DISCONNECTED):
+        return
+      
+  await ws.close()
+