from transformers import Pipeline
from sentence_transformers import SentenceTransformer
import torch

class SentimentModelPipe(Pipeline):

    def __init__(self, **kwargs):
        Pipeline.__init__(self, **kwargs)
        self.smodel = SentenceTransformer(kwargs.get("embedding_model", "sentence-transformers/all-MiniLM-L6-v2"))

    def _sanitize_parameters(self, **kw):
        return {}, {}, {}

    def preprocess(self, inputs):
        return self.smodel.encode(inputs, convert_to_tensor=True)
    
    def postprocess(self, outputs):
        return outputs.argmax(1).item()
    
    def _forward(self, tensor):
        with torch.no_grad():
            out = self.model(tensor)
        return out