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


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

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        # get inputs
        inputs = data.pop("inputs", data)
        labels = data.pop("labels", None)

        prediction = self.model.predict_entities(inputs, labels)
        return prediction