File size: 785 Bytes
93b0f77 b073019 93b0f77 |
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 |
import json
import asyncio
import websockets
from transformers import pipeline
extractor5 = pipeline("text2text-generation", model="occ_extract")
async def occ_extractor(websocket, path):
try:
while True:
data = await websocket.recv()
payload = json.loads(data)
intent = payload["prompt"]
label=extractor5(intent)[0]["generated_text"]
if label=="":
label="No occupation detected"
await websocket.send(json.dumps(label))
except websockets.ConnectionClosed:
print("Connection closed")
async def start_server():
server2 = await websockets.serve(occ_extractor, "0.0.0.0", 8766)
print("Server task started")
await asyncio.Future()
asyncio.run(start_server())
|