from transformers import Pipeline
import torch

class TBCP(Pipeline):
    def _sanitize_parameters(self, **kwargs):
        postprocess_kwargs = {}
        if "text_pair" in kwargs:
            postprocess_kwargs["top_k"] = kwargs["top_k"]
        return {}, {}, postprocess_kwargs

    def preprocess(self, text):
        return self.tokenizer(text, return_tensors="pt")

    def _forward(self, model_inputs):
        return self.model(**model_inputs)

    def postprocess(self, model_outputs,top_k = None):
        logits = model_outputs.logits
        probabilities = torch.nn.functional.softmax(logits, dim=-1)
        
        best_class = probabilities.argmax().item()
        label = self.model.config.id2label[best_class]
        score = probabilities.squeeze()[best_class].item()
        logits = logits.squeeze().tolist()
        return {"label": label, "score": score, "logits": logits}