Add scores in output

pipeline.py

@@ -25,6 +25,8 @@ class ExtSummPipeline(Pipeline):
             strategy_args : any
                 Parameters of the strategy.
 
+            out_scores : bool
+                If True, the score for each sentence is returned.
         Outputs
         -------
             selected_sents : list[str]
@@ -32,6 +34,8 @@ class ExtSummPipeline(Pipeline):
 
             selected_idxs : list[int]
                 List of the indexes of the selected sentences in the original input
+
+            sents_scores : Tensor (optional)
     """
 
     
@@ -44,6 +48,8 @@ class ExtSummPipeline(Pipeline):
             postprocess_kwargs["strategy"] = kwargs["strategy"]
         if "strategy_args" in kwargs:
             postprocess_kwargs["strategy_args"] = kwargs["strategy_args"]
+        if "out_scores" in kwargs:
+            postprocess_kwargs["out_scores"] = kwargs["out_scores"]
 
         return {}, {}, postprocess_kwargs
 
@@ -95,7 +101,11 @@ class ExtSummPipeline(Pipeline):
         return { "predictions": out_predictions, "sentences": sentences }
 
 
-    def postprocess(self, args, strategy: str="count", strategy_args=3):
+    def postprocess(self, args, strategy: str="count", strategy_args=3, out_scores=False):
         predictions = args["predictions"]
         sentences = args["sentences"]
-        return select(sentences, predictions, strategy, strategy_args)
+        out = select(sentences, predictions, strategy, strategy_args)
+
+        if out_scores: out += (predictions,)
+
+        return out