Spaces:
Paused
Paused
| 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 accelerator import Accelerator | |
| from answerer import Answerer | |
| from mapper import Mapper | |
| mapper = Mapper("multi-qa-distilbert-cos-v1") | |
| answerer = Answerer( | |
| model="RWKV-5-World-3B-v2-20231118-ctx16k", | |
| vocab="rwkv_vocab_v20230424", | |
| strategy="cpu bf16", | |
| ctx_limit=16*1024, | |
| ) | |
| accelerator = Accelerator() | |
| app = FastAPI() | |
| HTML = """ | |
| <!DOCTYPE HTML> | |
| <html> | |
| <body> | |
| <form action="" onsubmit="ask(event)"> | |
| <input id="prompt" type="text" autocomplete="off" /> | |
| <br> | |
| <input type="submit" value="SEND" /> | |
| </form> | |
| <p id="output"></p> | |
| <script> | |
| const prompt = document.getElementById("prompt"); | |
| const output = document.getElementById("output"); | |
| const ws = new WebSocket("wss://daniilalpha-answerer-api.hf.space/answer"); | |
| ws.onmessage = (e) => answer(e.data); | |
| function ask(event) { | |
| if(ws.readyState != 1) { | |
| answer("websocket is not connected!"); | |
| return; | |
| } | |
| ws.send(prompt.value); | |
| event.preventDefault(); | |
| } | |
| function answer(value) { | |
| output.innerHTML = value; | |
| } | |
| </script> | |
| </body> | |
| </html> | |
| """ | |
| def index(): | |
| return HTMLResponse(HTML) | |
| async def answer(ws: WebSocket): | |
| await accelerator.connect(ws) | |
| def map(query: Union[str, None], items: Union[list[str], None]): | |
| scores = mapper(query, items) | |
| return JSONResponse(jsonable_encoder(scores)) | |
| async def answer(ws: WebSocket): | |
| await ws.accept() | |
| try: input = await ws.receive_text() | |
| except WebSocketDisconnect: return | |
| if accelerator.connected(): | |
| output = await accelerator.accelerate(input) | |
| if(ws.client_state == WebSocketState.CONNECTED): | |
| await ws.send_text(output) | |
| elif(ws.client_state == WebSocketState.DISCONNECTED): | |
| return | |
| 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() | |