from typing import Dict, List, Any from setfit import SetFitModel import json class EndpointHandler: def __init__(self, path=""): # load model self.model = SetFitModel.from_pretrained(path) with open('/repository/label_config.json', 'r') as file: raw = json.load(file) self.id2label = {int(k): v for k, v in raw.items()} 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] results = [ {"label": self.id2label[i], "score": score.item()} for i, score in enumerate(scores) ] max_element = max(results, key=lambda x: x['score']) return max_element