from transformers import AutoModelForSequenceClassification, AutoTokenizer

class PreTrainedPipeline():
    def __init__(self, path):
        """
        Initialize model
        """
        self.model = AutoModelForSequenceClassification.from_pretrained("garrettbaber/twitter-roberta-base-fear-intensity")
        self.tokenizer = AutoTokenizer.from_pretrained("garrettbaber/twitter-roberta-base-fear-intensity")

    def __call__(self, inputs):
        """
        Args:
            inputs (:obj:`np.array`):
                The raw waveform of audio received. By default at 16KHz.
        Return:
            A :obj:`dict`:. The object return should be liked {"text": "XXX"} containing
            the detected text from the input audio.
        """
        tokens = self.tokenizer(inputs, return_tensors="pt")
        outputs = self.model(**tokens)
        logits = outputs.get("logits")
        rawScore = logits.tolist().pop().pop()
        
        return {
            "score": f"{rawScore:.3f}"
        }