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