DaniilAlpha commited on
Commit
26790f6
·
1 Parent(s): 185b88a

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +23 -3
main.py CHANGED
@@ -2,7 +2,12 @@ from fastapi import FastAPI, WebSocket
2
  from fastapi.responses import HTMLResponse
3
 
4
  from answerer import Answerer
 
5
 
 
 
 
 
6
  answerer = Answerer(
7
  model="RWKV-5-World-3B-v2-20231118-ctx16k.pth",
8
  vocab="rwkv_vocab_v20230424",
@@ -52,16 +57,31 @@ HTML = """
52
  """
53
 
54
  @app.get("/")
55
- async def index():
56
- return HTMLResponse(HTML)
 
 
 
 
57
 
58
  @app.websocket("/answer")
59
  async def answer(ws: WebSocket):
60
  await ws.accept()
 
 
61
 
62
  input = await ws.receive_text()
 
 
 
 
 
63
  output = answerer(input, 32)
 
 
 
64
  for el in output:
65
- await ws.send_text(el)
 
66
 
67
  await ws.close()
 
2
  from fastapi.responses import HTMLResponse
3
 
4
  from answerer import Answerer
5
+ from mapper import Mapper
6
 
7
+ mapper = Mapper(
8
+ repo="sentence-transformers",
9
+ model="multi-qa-distilbert-cos-v1",
10
+ )
11
  answerer = Answerer(
12
  model="RWKV-5-World-3B-v2-20231118-ctx16k.pth",
13
  vocab="rwkv_vocab_v20230424",
 
57
  """
58
 
59
  @app.get("/")
60
+ def index():
61
+ return HTMLResponse(HTML)
62
+
63
+ @app.get("/map")
64
+ def map():
65
+ return HTMLResponse(HTML)
66
 
67
  @app.websocket("/answer")
68
  async def answer(ws: WebSocket):
69
  await ws.accept()
70
+
71
+ print("ws accepted!")
72
 
73
  input = await ws.receive_text()
74
+
75
+ print("input received!")
76
+
77
+ await ws.send_text("OK!")
78
+
79
  output = answerer(input, 32)
80
+
81
+ print("output created!")
82
+
83
  for el in output:
84
+ print(f"sent: '{el}'")
85
+ await ws.send_text(el)
86
 
87
  await ws.close()