# from gliner import GLiNER

# class EndpointHandler:
#     def __init__(self, path=""):
#         # Use the provided path for loading the model
#         self.model = GLiNER.from_pretrained(path)

#     def __call__(self, data):
#         try:
#             text = data.get("text", "")
#             labels = data.get("labels", [])
#             if not text or not labels:
#                 return {"error": "Please provide 'text' and 'labels'"}
#             entities = self.model.predict_entities(text, labels)
#             return {"entities": entities}
#         except Exception as e:
#             return {"error": str(e)}

from gliner import GLiNER
import torch

class EndpointHandler:
    def __init__(self, path=""):
        # Load without device_map, then move to GPU
        self.model = GLiNER.from_pretrained(path)  # Remove device_map="cuda"
        self.model = self.model.to("cuda")
        self.model.eval()  # Lock for inference

    def __call__(self, data):
    # If data is wrapped in 'inputs' (as Hugging Face does), unwrap it
        if isinstance(data, dict) and "inputs" in data:
            data = data["inputs"]

        text = data.get("text", "")
        labels = data.get("labels", [])

        if not text or not labels:
            return {"error": "Please provide 'text' and 'labels'"}

        entities = self.model.predict_entities(text, labels)
        return {"entities": entities}