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()
input = await ws.receive_text()
output = answerer(input, 32)
async for el in output:
await ws.send_text(el)
await ws.close()
async def answerer_test():
for i in range(10000):
yield f"i"
@app.websocket("/answer_test")
async def answer(ws: WebSocket):
await ws.accept()
input = await ws.receive_text()
output = answerer_test()
for el in output:
print(el)
await ws.send_text(el)
await ws.close()