from transformers import Pipeline
import torch

class DeclarativePipeline(Pipeline):

    def _sanitize_parameters(self, **kwargs):
        preprocess_kwargs = {}
        return preprocess_kwargs, {}, {}


    def preprocess(self, inputs):
        question = inputs.split("#")[0]
        answer = inputs.split("#")[1]
        return self.tokenizer(
            question,
            answer,
            add_special_tokens=True,
            max_length=512,
            padding='max_length',
            truncation='only_first',
            return_attention_mask=True,
            return_tensors=self.framework
        )

    def _forward(self, model_inputs):
        self.model.eval()
        with torch.no_grad():
            generate_ids = self.model.generate(
                    input_ids=model_inputs["input_ids"],
                    attention_mask=model_inputs["attention_mask"],
                    max_length=128,
                    num_beams=4,
                    num_return_sequences=1,
                    no_repeat_ngram_size=2,
                    early_stopping=True,
                )
        return generate_ids

    def postprocess(self, model_outputs):
        preds = [
            self.tokenizer.decode(gen_id,
                                  skip_special_tokens=True,
                                  clean_up_tokenization_spaces=True)
            for gen_id in model_outputs
        ]
        return "".join(preds)