from typing import Dict, List, Any
from gliner import GLiNER

class EndpointHandler:
    def __init__(self, path=""):
        # Initialize the GLiNER model
        self.model = GLiNER.from_pretrained("urchade/gliner_multi-v2.1")

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        Args:
            data (Dict[str, Any]): The input data including:
                - "inputs": The text input from which to extract information.
                - "labels": The labels to predict entities for.

        Returns:
            List[Dict[str, Any]]: The extracted entities from the text, formatted as required.
        """
        # Get inputs and labels
        inputs = data.get("inputs", "")
        labels = ["party", "document title"]
        print('labels',labels)
        # Predict entities using GLiNER
        entities = self.model.predict_entities(inputs, labels)

        # Format the results to match the expected output structure
        formatted_results = []
        for entity in entities:
            formatted_entity = {
               entity["label"]: entity["text"],
            }
            print(formatted_entity)
            formatted_results.append(formatted_entity)

        return formatted_results