from peft import PeftModel, PeftConfig
from transformers import AutoModelForSequenceClassification, Pipeline, AutoTokenizer

config = PeftConfig.from_pretrained("ctu-aic/lora-xlm-roberta-large-squad2-csfever_v2-f1")
model = AutoModelForSequenceClassification.from_pretrained(config.base_model_name_or_path)
model = PeftModel.from_pretrained(model, config)
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)

#pipeline for NLI
class NliPipeline(Pipeline):
    def _sanitize_parameters(self, **kwargs):
        preprocess_kwargs = {}
        if "evidence" in kwargs:
            preprocess_kwargs["evidence"] = kwargs["evidence"]
        return preprocess_kwargs, {}, {}

    def preprocess(self, claim, evidence=""):
        model_input = self.tokenizer(claim, evidence, return_tensors=self.framework, truncation=True)
        return model_input
    def _forward(self, model_inputs):
        outputs = self.model(**model_inputs)
        return outputs

    def postprocess(self, model_outputs):
        logits = model_outputs.logits

        predictions = torch.argmax(logits, dim=-1)
        return {"logits": logits, "label": int(predictions[0])}

nli_pipeline = NliPipeline(model=model, tokenizer=tokenizer)

nli_pipeline("claim", "evidence")