from typing import Any, Dict, Tuple
from transformers import Pipeline
from transformers.pipelines.base import GenericTensor
from transformers.utils import ModelOutput
from typing import Union,List

class SimilarPipeline(Pipeline):
    def __init__(self, max_length=512,*args, **kwargs):
        super().__init__(*args, **kwargs)
        self.max_length = max_length

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

    def preprocess(self, input: Union[Tuple[str],List[Tuple[str]]], **preprocess_parameters: Dict) -> Dict[str, GenericTensor]:
        if isinstance(input, list):
            a = list(map(lambda x: x[0], input))
            b = list(map(lambda x: x[1], input))
        else:
            a = input[0]
            b = input[1]    
        tensors = self.tokenizer(
            a,
            b,
            max_length=self.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        )
        return tensors
    
    def _forward(self, input_tensors: Dict[str, GenericTensor], **forward_parameters: Dict) -> ModelOutput:
        _,logits = self.model(**input_tensors)
        return logits.tolist()
    
    def postprocess(
            self, 
            model_outputs: ModelOutput, 
            **postprocess_parameters: Dict
        ) -> Any:
        return model_outputs