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