import json from typing import Any, Dict, List import sklearn import os import joblib import numpy as np class PreTrainedPipeline(): def __init__(self, path: str): # load the model self.model = joblib.load((os.path.join(path, "pipeline.pkl")) def __call__(self, inputs: str) -> List[Dict[str, float]]: """ Args: inputs (:obj:`str`): a string containing some text Return: A :obj:`list`:. The object returned should be a list of one list like [[{"label": 0.9939950108528137}]] containing: - "label": A string representing what the label/class is. There can be multiple labels. - "score": A score between 0 and 1 describing how confident the model is for this label/class. """ predictions = self.model.predict_proba([inputs]) labels = [] for cls in predictions[0]: labels.append({ "label": f"LABEL_{cls}", "score": predictions[0][cls], }) return labels