NORMA-DEV/Gliner_haystack

1from typing import Dict, List, Any2from gliner import GLiNER3 4class EndpointHandler:5    def __init__(self, path=""):6        # Initialize the GLiNER model7        self.model = GLiNER.from_pretrained("urchade/gliner_multi-v2.1")8 9    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:10        """11        Args:12            data (Dict[str, Any]): The input data including:13                - "inputs": The text input from which to extract information.14                - "labels": The labels to predict entities for.15 16        Returns:17            List[Dict[str, Any]]: The extracted entities from the text, formatted as required.18        """19        # Get inputs and labels20        inputs = data.get("inputs", "")21        labels = ["party", "document title"]22        # Predict entities using GLiNER23        entities = self.model.predict_entities(inputs, labels)24 25        # Initialize a dictionary to store organized entities26            organized_entities = {label: {"labels": [], "scores": []} for label in labels}27 28            for entity in entities:29                label = entity['label']30                text = entity['text']31                score = entity['score']32 33                # Append text and score to the corresponding label34                organized_entities[label]["labels"].append(text)35                organized_entities[label]["scores"].append(score)36 37            # Store organized entities in document metadata38            doc.meta["entities"] = organized_entities39 40        return {"documents": documents}41