File size: 2,305 Bytes
eb6c6fa
ef18a5d
 
 
c7f3620
9a6620f
71f5570
1e28d04
26790f6
9a6620f
9c604c5
1e28d04
2bc4710
1e28d04
 
 
 
aeb5107
71f5570
1e28d04
9a6620f
 
a3ccca5
 
 
 
 
 
6942628
05dc8f9
 
 
a3ccca5
 
a05f54e
a3ccca5
 
 
 
7cab91c
9e7b790
a3ccca5
 
9e7b790
 
 
 
 
a3ccca5
 
 
9e7b790
 
 
 
a3ccca5
 
 
 
 
 
bc2d927
 
 
26790f6
aeb5107
 
71f5570
aeb5107
d7562ab
eb6c6fa
ef18a5d
 
a3ccca5
1e28d04
 
 
3518fdf
64e57d0
 
26790f6
71f5570
 
ef18a5d
71f5570
ef18a5d
 
aeb5107
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
from typing import Union
from fastapi import FastAPI
from fastapi.encoders import jsonable_encoder
from fastapi.websockets import WebSocket, WebSocketDisconnect, WebSocketState
from fastapi.responses import HTMLResponse, JSONResponse

from accelerator import Accelerator
from answerer import Answerer
from mapper import Mapper

# mapper = Mapper("multi-qa-distilbert-cos-v1")
answerer = Answerer(
  model="RWKV-5-World-3B-v2-20231118-ctx16k",
  vocab="rwkv_vocab_v20230424",
  strategy="cpu bf16",
  ctx_limit=16*1024,
)

accelerator = Accelerator()
  
app = FastAPI()

HTML = """
<!DOCTYPE HTML>

<html>

<body>
  <form action="" onsubmit="ask(event)">
    <input id="prompt" type="text" autocomplete="off" />
    <br>
    <input type="submit" value="SEND" />
  </form>

  <p id="output"></p>
  <script>
    const prompt = document.getElementById("prompt");
    const output = document.getElementById("output");

    const ws = new WebSocket("wss://daniilalpha-answerer-api.hf.space/answer"); 
    ws.onmessage = (e) => answer(e.data);

    function ask(event) {
      if(ws.readyState != 1) {
        answer("websocket is not connected!");
        return;
      }

      ws.send(prompt.value);
      event.preventDefault();
    }

    function answer(value) {
      output.innerHTML = value;
    }
  </script>
</body>

</html>
"""

@app.get("/")
def index():
  return HTMLResponse(HTML)

@app.websocket("/accelerate")
async def answer(ws: WebSocket):
  await accelerator.connect(ws)

@app.post("/map")
def map(query: Union[str, None], items: Union[list[str], None]):
  scores = mapper(query, items)
  return JSONResponse(jsonable_encoder(scores))

@app.websocket("/answer")
async def answer(ws: WebSocket):
  await ws.accept()
      
  try: input = await ws.receive_text()
  except WebSocketDisconnect: return

  if accelerator.connected():
    output = await accelerator.accelerate(input)
    if(ws.client_state == WebSocketState.CONNECTED):
      await ws.send_text(output)
    elif(ws.client_state == WebSocketState.DISCONNECTED):
      return
  else:
    output = answerer(input, 32)
    async for el in output:
      if(ws.client_state == WebSocketState.CONNECTED):
        await ws.send_text(el)
      elif(ws.client_state == WebSocketState.DISCONNECTED):
        return
      
  await ws.close()