from fastapi import FastAPI, WebSocket
from fastapi.responses import HTMLResponse
from answerer import Answerer
answerer = Answerer(
repo="BlinkDL/rwkv-5-world",
filename="RWKV-5-World-3B-v2-20231118-ctx16k",
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()
input = await ws.receive_text()
stream = answerer(input, 256)
for el in stream:
await ws.send_text(el)