from typing import Dict, List, Any
from setfit import SetFitModel


class EndpointHandler:
    def __init__(self, path=""):
        # load model
        self.model = SetFitModel.from_pretrained(path)
        # ag_news id to label mapping
        self.id2label = {0: "Absent", 1: "Present"}

    # def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
    #     """
    #      data args:
    #           inputs (:obj: `str`)
    #     Return:
    #           A :obj:`list` | `dict`: will be serialized and returned
    #     """
    #     # get inputs
    #     inputs = data.pop("inputs", data)
    #     if isinstance(inputs, str):
    #         inputs = [inputs]

    #     # run normal prediction
    #     scores = self.model.predict_proba(inputs)[0]

    #     return [{"label": self.id2label[i], "score": score.item()} for i, score in enumerate(scores)]

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        data args:
            inputs (:obj: `List[str]`) - List of strings
        Return:
            A :obj:`list` of dicts: each dict contains 'label' and 'score' for each input string
        """
        # get inputs
        inputs = data.pop("inputs", data)
        if not isinstance(inputs, list):
            raise ValueError("Input must be a list of strings")

        # run normal prediction
        all_scores = self.model.predict_proba(inputs)  # This returns a list of score arrays

        # Format the results for each input string
        results = []
        for scores in all_scores:
            results.append([
                {"label": self.id2label[i], "score": score.item()} for i, score in enumerate(scores)
            ])

        return results