DaniilAlpha commited on
Commit
aeb5107
·
1 Parent(s): 8de1402

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +21 -5
main.py CHANGED
@@ -2,7 +2,7 @@ from typing import Union
2
  from fastapi import FastAPI
3
  from fastapi.encoders import jsonable_encoder
4
  from fastapi.websockets import WebSocket, WebSocketDisconnect, WebSocketState
5
- from fastapi.responses import HTMLResponse, JSONResponse
6
 
7
  from answerer import Answerer
8
  from mapper import Mapper
@@ -17,6 +17,8 @@ answerer = Answerer(
17
  strategy="cpu bf16",
18
  ctx_limit=16*1024,
19
  )
 
 
20
 
21
  app = FastAPI()
22
 
@@ -63,6 +65,11 @@ HTML = """
63
  # def index():
64
  # return HTMLResponse(HTML)
65
 
 
 
 
 
 
66
  @app.post("/map")
67
  def map(query: Union[str, None], items: Union[list[str], None]):
68
  scores = mapper(query, items)
@@ -75,11 +82,20 @@ async def answer(ws: WebSocket):
75
  try: input = await ws.receive_text()
76
  except WebSocketDisconnect: return
77
 
78
- output = answerer(input, 32)
79
- async for el in output:
 
80
  if(ws.client_state == WebSocketState.CONNECTED):
81
  await ws.send_text(el)
82
  elif(ws.client_state == WebSocketState.DISCONNECTED):
83
  return
84
-
85
- await ws.close()
 
 
 
 
 
 
 
 
 
2
  from fastapi import FastAPI
3
  from fastapi.encoders import jsonable_encoder
4
  from fastapi.websockets import WebSocket, WebSocketDisconnect, WebSocketState
5
+ from fastapi.responses import JSONResponse
6
 
7
  from answerer import Answerer
8
  from mapper import Mapper
 
17
  strategy="cpu bf16",
18
  ctx_limit=16*1024,
19
  )
20
+
21
+ accelerator: Union[WebSocket, None] = None
22
 
23
  app = FastAPI()
24
 
 
65
  # def index():
66
  # return HTMLResponse(HTML)
67
 
68
+ @app.websocket("/accelerate")
69
+ async def answer(ws: WebSocket):
70
+ await ws.accept()
71
+ accelerator = ws
72
+
73
  @app.post("/map")
74
  def map(query: Union[str, None], items: Union[list[str], None]):
75
  scores = mapper(query, items)
 
82
  try: input = await ws.receive_text()
83
  except WebSocketDisconnect: return
84
 
85
+ if accelerator:
86
+ await accelerator.send_text(input)
87
+ output = await accelerator.receive_text()
88
  if(ws.client_state == WebSocketState.CONNECTED):
89
  await ws.send_text(el)
90
  elif(ws.client_state == WebSocketState.DISCONNECTED):
91
  return
92
+ else:
93
+ output = answerer(input, 32)
94
+ async for el in output:
95
+ if(ws.client_state == WebSocketState.CONNECTED):
96
+ await ws.send_text(el)
97
+ elif(ws.client_state == WebSocketState.DISCONNECTED):
98
+ return
99
+
100
+ await ws.close()
101
+