Update handler.py

374b849 verified about 2 years ago

1.54 kB

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"]
        # Predict entities using GLiNER
        entities = self.model.predict_entities(inputs, labels)

        # Initialize a dictionary to store organized entities
            organized_entities = {label: {"labels": [], "scores": []} for label in labels}

            for entity in entities:
                label = entity['label']
                text = entity['text']
                score = entity['score']

                # Append text and score to the corresponding label
                organized_entities[label]["labels"].append(text)
                organized_entities[label]["scores"].append(score)

            # Store organized entities in document metadata
            doc.meta["entities"] = organized_entities

        return {"documents": documents}

	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"]
	# Predict entities using GLiNER
	entities = self.model.predict_entities(inputs, labels)

	# Initialize a dictionary to store organized entities
	organized_entities = {label: {"labels": [], "scores": []} for label in labels}

	for entity in entities:
	label = entity['label']
	text = entity['text']
	score = entity['score']

	# Append text and score to the corresponding label
	organized_entities[label]["labels"].append(text)
	organized_entities[label]["scores"].append(score)

	# Store organized entities in document metadata
	doc.meta["entities"] = organized_entities

	return {"documents": documents}