Upload custom_pipeline.py

custom_pipeline.py

@@ -0,0 +1,38 @@
+from transformers import AutoTokenizer, Pipeline
+import torch
+
+class PairTextClassificationPipeline(Pipeline):
+    def __init__(self, model, tokenizer=None, **kwargs):
+        # Initialize tokenizer first
+        if tokenizer is None:
+            tokenizer = AutoTokenizer.from_pretrained(model.config._name_or_path)
+        # Make sure we store the tokenizer before calling super().__init__
+        self.tokenizer = tokenizer
+        super().__init__(model=model, tokenizer=tokenizer, **kwargs)
+        self.prompt = "<pad> Determine if the hypothesis is true given the premise?\n\nPremise: {text1}\n\nHypothesis: {text2}"
+        
+    def _sanitize_parameters(self, **kwargs):
+        preprocess_kwargs = {}
+        return preprocess_kwargs, {}, {}
+
+    def preprocess(self, inputs):
+        # Expect inputs to be list of (Premise, Hypothesis) tuples
+        pair_dict = {'text1': inputs[0], 'text2': inputs[1]}
+        formatted_prompt = self.prompt.format(**pair_dict)
+        model_inputs = self.tokenizer(
+            formatted_prompt,
+            return_tensors='pt', 
+            padding=True
+        )
+        return model_inputs
+
+    def _forward(self, model_inputs):
+        model_outputs = self.model(**model_inputs)
+        return model_outputs
+
+    def postprocess(self, model_outputs):
+        logits = model_outputs.logits
+        logits = logits[:, 0, :] # tok_cls
+        transformed_probs = torch.softmax(logits, dim=-1)
+        raw_scores = transformed_probs[:, 1] # probability of class 1
+        return raw_scores.item()