import os
from fastapi import FastAPI, WebSocket

from answerer import Answerer

print(f"HELLO_WORLD: {os.environ['HELLO_WORLD']}")

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()

@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)