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