Spaces:
Paused
Paused
Commit
·
aeb5107
1
Parent(s):
8de1402
Update main.py
Browse files
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
|
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 |
-
|
79 |
-
|
|
|
80 |
if(ws.client_state == WebSocketState.CONNECTED):
|
81 |
await ws.send_text(el)
|
82 |
elif(ws.client_state == WebSocketState.DISCONNECTED):
|
83 |
return
|
84 |
-
|
85 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
+
|