Spaces:
Paused
Paused
| from typing import Union | |
| from fastapi import FastAPI, WebSocket, WebSocketDisconnect | |
| from fastapi.responses import HTMLResponse, JSONResponse | |
| from answerer import Answerer | |
| from mapper import Mapper | |
| mapper = Mapper( | |
| repo="sentence-transformers", | |
| model="multi-qa-distilbert-cos-v1", | |
| ) | |
| answerer = Answerer( | |
| model="RWKV-5-World-3B-v2-20231118-ctx16k.pth", | |
| vocab="rwkv_vocab_v20230424", | |
| strategy="cpu bf16", | |
| ctx_limit=16*1024, | |
| ) | |
| 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("ws://127.0.0.1:2403/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) | |
| def map(query: Union[str, None], items: Union[list[str], None]): | |
| indices = mapper(query, items) | |
| return JSONResponse(indices) | |
| async def answer(ws: WebSocket): | |
| print("ws started!") | |
| await ws.accept() | |
| print("ws accepted!") | |
| try: input = await ws.receive_text() | |
| except WebSocketDisconnect: return | |
| print("input received!") | |
| await ws.send_text("OK!") | |
| output = answerer(input, 32) | |
| print("output created!") | |
| for el in output: | |
| print(f"sent: '{el}'") | |
| await ws.send_text(el) | |
| await ws.close() |