from threading import Thread
from fastapi import FastAPI, WebSocket
from fastapi.responses import HTMLResponse
from answerer import Answerer
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 = """
"""
@app.get("/")
async def index():
return HTMLResponse(HTML)
@app.websocket("/answer")
async def answer(ws: WebSocket):
await ws.accept()
print("ws connected!")
async def callback(output: str):
print(f"sent: '{output}'")
await ws.send_text(output)
input = await ws.receive_text()
answerer_thread = Thread(
target=answerer,
args=(input, 32),
kwargs={"callback": callback}
)
answerer_thread.start()
answerer_thread.join()
print("output ready!")
await ws.close()
print("ws closed!")