|
|
from asyncio import sleep |
|
|
from typing import Optional |
|
|
from fastapi import FastAPI |
|
|
from fastapi.encoders import jsonable_encoder |
|
|
from fastapi.websockets import WebSocket, WebSocketDisconnect |
|
|
from fastapi.responses import HTMLResponse, JSONResponse |
|
|
from websockets import ConnectionClosed |
|
|
|
|
|
from accelerator import Accelerator |
|
|
from answerer import Answerer |
|
|
from mapper import Mapper |
|
|
|
|
|
try: mapper = Mapper("sentence-transformers/multi-qa-distilbert-cos-v1") |
|
|
except Exception as e: print(f"ERROR! cannot load Mapper model!\n{e}") |
|
|
|
|
|
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)"> |
|
|
<textarea id="prompt"></textarea> |
|
|
<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> |
|
|
""" |
|
|
|
|
|
@app.get("/") |
|
|
def index(): |
|
|
return HTMLResponse(HTML) |
|
|
|
|
|
@app.websocket("/accelerate") |
|
|
async def answer(ws: WebSocket): |
|
|
await accelerator.connect(ws) |
|
|
while accelerator.connected(): |
|
|
await sleep(10) |
|
|
|
|
|
@app.post("/map") |
|
|
def map(query: Optional[str], items: Optional[list[str]]): |
|
|
scores = mapper(query, items) |
|
|
return JSONResponse(jsonable_encoder(scores)) |
|
|
|
|
|
async def handle_answerer_local(ws: WebSocket, input: str): |
|
|
output = answerer(input, 128) |
|
|
el: str |
|
|
async for el in output: pass |
|
|
await ws.send_text(el) |
|
|
|
|
|
async def handle_answerer_accelerated(ws: WebSocket, input: str): |
|
|
output = await accelerator.accelerate(input) |
|
|
if output: await ws.send_text(output) |
|
|
else: await handle_answerer_local(ws, input) |
|
|
|
|
|
@app.websocket("/answer") |
|
|
async def answer(ws: WebSocket): |
|
|
await ws.accept() |
|
|
|
|
|
try: |
|
|
input = await ws.receive_text() |
|
|
if accelerator.connected(): await handle_answerer_accelerated(ws, input) |
|
|
else: await handle_answerer_local(ws, input) |
|
|
except ConnectionClosed: return |
|
|
except WebSocketDisconnect: return |
|
|
|
|
|
await ws.close() |
|
|
|
