1  """
 2  Summary module tests
 3  """
 4  
 5  import unittest
 6  
 7  from txtai.pipeline import Summary
 8  
 9  
10  class TestSummary(unittest.TestCase):
11      """
12      Summary tests.
13      """
14  
15      @classmethod
16      def setUpClass(cls):
17          """
18          Create single summary instance.
19          """
20  
21          cls.text = (
22              "Search is the base of many applications. Once data starts to pile up, users want to be able to find it. It's the foundation "
23              "of the internet and an ever-growing challenge that is never solved or done. The field of Natural Language Processing (NLP) is "
24              "rapidly evolving with a number of new developments. Large-scale general language models are an exciting new capability "
25              "allowing us to add amazing functionality quickly with limited compute and people. Innovation continues with new models "
26              "and advancements coming in at what seems a weekly basis. This article introduces txtai, an AI-powered search engine "
27              "that enables Natural Language Understanding (NLU) based search in any application."
28          )
29  
30          cls.summary = Summary("t5-small")
31  
32      def testSummary(self):
33          """
34          Test summarization of text
35          """
36  
37          self.assertEqual(self.summary(self.text, minlength=15, maxlength=15), "the field of natural language processing (NLP) is rapidly evolving")
38  
39      def testSummaryBatch(self):
40          """
41          Test batch summarization of text
42          """
43  
44          summaries = self.summary([self.text, self.text], maxlength=15)
45          self.assertEqual(len(summaries), 2)
46  
47      def testSummaryNoLength(self):
48          """
49          Test summary with no max length set
50          """
51  
52          self.assertEqual(
53              self.summary(self.text + self.text),
54              "search is the base of many applications. Once data starts to pile up, users want to be able to find it. "
55              + "Large-scale general language models are an exciting new capability allowing us to add amazing functionality quickly "
56              + "with limited compute and people.",
57          )
58  
59      def testSummaryShort(self):
60          """
61          Test that summarization is skipped
62          """
63  
64          self.assertEqual(self.summary("Text", maxlength=15), "Text")