Merge pull request #1295 from Bdysj/pad-summary-batches

Pad batched inputs in the Summary pipeline

src/python/txtai/pipeline/text/summary.py

@@ -91,14 +91,20 @@ class Summary(HFModel):
         if prefix:
             inputs = [prefix + x for x in inputs]
 
-        # Tokenize inputs
-        tokens = self.tokenizer(inputs, truncation=kwargs.pop("truncation", False), return_tensors="pt").to(self.device)
+        truncation = kwargs.pop("truncation", False)
 
-        # Generate outputs
-        outputs = self.model.generate(**tokens, **kwargs)
+        results = []
+        for batch in self.batch(inputs, self.batchsize):
+            # Tokenize inputs, padding texts of different lengths within a batch
+            tokens = self.tokenizer(batch, padding=True, truncation=truncation, return_tensors="pt").to(self.device)
 
-        # Decode and return
-        return [self.tokenizer.decode(x, skip_special_tokens=True) for x in outputs]
+            # Generate outputs
+            outputs = self.model.generate(**tokens, **kwargs)
+
+            # Decode outputs
+            results.extend(self.tokenizer.decode(x, skip_special_tokens=True) for x in outputs)
+
+        return results
 
     def clean(self, text):
         """

test/python/testpipeline/testtext/testsummary.py

@@ -44,6 +44,24 @@ class TestSummary(unittest.TestCase):
         summaries = self.summary([self.text, self.text], maxlength=15)
         self.assertEqual(len(summaries), 2)
 
+    def testSummaryBatchLengths(self):
+        """
+        Test batch summarization of texts with different lengths
+        """
+
+        texts = [self.text, self.text + " " + self.text, self.text]
+        summaries = self.summary(texts, maxlength=15)
+        self.assertEqual(len(summaries), 3)
+        self.assertTrue(all(summary and summary not in texts for summary in summaries))
+
+        # Inputs that span multiple batches keep their order
+        batchsize = self.summary.batchsize
+        try:
+            self.summary.batchsize = 2
+            self.assertEqual(self.summary(texts, maxlength=15), summaries)
+        finally:
+            self.summary.batchsize = batchsize
+
     def testSummaryNoLength(self):
         """
         Test summary with no max length set