Upload DocumentSentenceRelevancePipeline

config.json

@@ -7,6 +7,15 @@
     "AutoModel": "modeling.MultiHeadModel"
   },
   "classifier_dropout": 0.1,
+  "custom_pipelines": {
+    "context-relevance": {
+      "impl": "pipeline.DocumentSentenceRelevancePipeline",
+      "pt": [
+        "AutoModel"
+      ],
+      "tf": []
+    }
+  },
   "encoder_name": "tasksource/deberta-base-long-nli",
   "id2label": {
     "0": "irrelevant",

pipeline.py

@@ -0,0 +1,109 @@
+from transformers import Pipeline
+import torch
+from typing import Union
+
+
+
+def convert_to_list(data):
+    first_list = next(iter(data.values()))    
+    return [
+        {key: values[i] for key, values in data.items()}
+        for i in range(len(first_list))
+    ]
+
+class DocumentSentenceRelevancePipeline(Pipeline):
+    def _sanitize_parameters(self, **kwargs):
+        threshold = kwargs.get("threshold", 0.5)
+        return {}, {}, {"threshold": threshold}
+
+    def preprocess(self, inputs):
+        question = inputs.get("question", "")
+        context = inputs.get("context", [""])
+        response = inputs.get("response", "")
+
+        q_enc = self.tokenizer(question, add_special_tokens=True, truncation=False, padding=False)
+        r_enc = self.tokenizer(response, add_special_tokens=True, truncation=False, padding=False)
+
+        question_ids = q_enc["input_ids"]
+        response_ids = r_enc["input_ids"]
+
+        document_sentences_ids = []
+        for s in context:
+            s_enc = self.tokenizer(s, add_special_tokens=True, truncation=False, padding=False)
+            document_sentences_ids.append(s_enc["input_ids"])
+
+        ids = question_ids + response_ids
+        pair_ids = []
+        for s_ids in document_sentences_ids:
+            pair_ids.extend(s_ids)
+
+        total_length = len(ids) + len(pair_ids)
+        if total_length > self.tokenizer.model_max_length:
+            num_tokens_to_remove = total_length - self.tokenizer.model_max_length
+            ids, pair_ids, _ = self.tokenizer.truncate_sequences(
+                ids=ids,
+                pair_ids=pair_ids,
+                num_tokens_to_remove=num_tokens_to_remove,
+                truncation_strategy="only_second",
+                stride=0,
+            )
+        combined_ids = ids + pair_ids
+        token_types = [0]*len(ids) + [1]*len(pair_ids)
+        attention_mask = [1]*len(combined_ids)
+
+        sentence_positions = []
+        current_pos = len(ids)
+        found_sentences = 0
+
+        for i, tok_id in enumerate(pair_ids):
+            if tok_id == self.tokenizer.cls_token_id:
+                sentence_positions.append(current_pos + i)
+                found_sentences += 1
+
+        input_ids = torch.tensor([combined_ids], dtype=torch.long)
+        attention_mask = torch.tensor([attention_mask], dtype=torch.long)
+        token_type_ids = torch.tensor([token_types], dtype=torch.long)
+        sentence_positions = torch.tensor([sentence_positions], dtype=torch.long)
+
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "token_type_ids": token_type_ids,
+            "sentence_positions": sentence_positions
+        }
+
+    def _forward(self, model_inputs):
+        return self.model(**model_inputs)
+
+    def __call__(self, inputs: Union[dict[str, str], list[dict[str, str]]], **kwargs):
+        if isinstance(inputs, dict):
+            inputs = [inputs]
+        model_outputs = super().__call__(inputs, **kwargs)
+        pipeline_outputs = []
+        for i, output in enumerate(model_outputs):
+            sentences = inputs[i]["context"]
+            output["sentences"]["sentence"] = sentences
+            output['sentences'] = convert_to_list(output['sentences'])
+            pipeline_outputs.append(output)
+        return pipeline_outputs if len(pipeline_outputs) > 1 else pipeline_outputs[0]
+
+    def postprocess(self, model_outputs, threshold = 0.5):
+        doc_logits = model_outputs.doc_logits
+        sent_logits = model_outputs.sent_logits
+        document_probabilities = torch.softmax(doc_logits, dim=-1)
+        sentence_probabilities = torch.softmax(sent_logits, dim=-1)
+        
+        document_best_class = (document_probabilities[:, 1] > threshold).long()
+        sentence_best_class = (sentence_probabilities[:, :, 1] > threshold).long()
+        document_score = document_probabilities[:, document_best_class]
+        sentence_best_class = sentence_best_class.squeeze()
+        batch_indices = torch.arange(sentence_probabilities.size(1))
+        sentence_scores = sentence_probabilities.squeeze()[batch_indices, sentence_best_class]
+        best_document_label = document_best_class.numpy().item()
+        best_document_label = self.model.config.id2label[best_document_label]
+
+        best_sentence_labels = sentence_best_class.numpy().tolist()
+        best_sentence_labels = [self.model.config.id2label[label] for label in best_sentence_labels]
+        document_output = {"label": best_document_label, "score": document_score.numpy().item()}
+        sentence_output = {"label": best_sentence_labels, "score": sentence_scores.numpy().tolist()}
+        return {"document": document_output, "sentences": sentence_output}