Update main.py

main.py

@@ -1,9 +1,25 @@
-from fastapi import FastAPI
+import os
+from fastapi import FastAPI, WebSocket
 
-from transformers import pipeline
+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.get("/api/")
-def index():
-    return {"output": "Hello world!"}
+@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)