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-small-long-nli",
   "id2label": {
     "0": "irrelevant",

pipeline.py

@@ -3,15 +3,8 @@ 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}
@@ -82,10 +75,29 @@ class DocumentSentenceRelevancePipeline(Pipeline):
         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]
+            sentences_dict = {
+                "sentence": sentences,
+                "label": output["sentences"]["label"],
+                "score": output["sentences"]["score"]
+            }
+            # Create the final output structure
+            final_output = {
+                "document": output["document"],
+                "sentences": [
+                    {
+                        "sentence": sent,
+                        "label": label,
+                        "score": score
+                    }
+                    for sent, label, score in zip(
+                        sentences_dict["sentence"],
+                        sentences_dict["label"],
+                        sentences_dict["score"]
+                    )
+                ]
+            }
+            pipeline_outputs.append(final_output)
+        return pipeline_outputs
 
     def postprocess(self, model_outputs, threshold = 0.5):
         doc_logits = model_outputs.doc_logits
@@ -96,14 +108,23 @@ class DocumentSentenceRelevancePipeline(Pipeline):
         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]
+        sentence_probabilities = sentence_probabilities.squeeze()
+        
+        if len(sentence_best_class.shape) == 0:
+            sentence_best_class = sentence_best_class.unsqueeze(0)
+            sentence_probabilities = sentence_probabilities.unsqueeze(0)
+        
+        batch_indices = torch.arange(len(sentence_best_class))
+        sentence_scores = sentence_probabilities[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}