Add tests for v3 functionality, closes #71

test/python/testsummary.py

@@ -0,0 +1,45 @@
+"""
+Summary module tests
+"""
+
+import unittest
+
+from txtai.pipeline import Summary
+
+
+class TestSummary(unittest.TestCase):
+    """
+    Labels tests
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create single summary instance
+        """
+
+        cls.text = (
+            "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 "
+            "of the internet and an ever-growing challenge that is never solved or done. The field of Natural Language Processing (NLP) is "
+            "rapidly evolving with a number of new developments. Large-scale general language models are an exciting new capability "
+            "allowing us to add amazing functionality quickly with limited compute and people. Innovation continues with new models "
+            "and advancements coming in at what seems a weekly basis. This article introduces txtai, an AI-powered search engine "
+            "that enables Natural Language Understanding (NLU) based search in any application."
+        )
+
+        cls.summary = Summary()
+
+    def testSummary(self):
+        """
+        Test summarization of text
+        """
+
+        self.assertEqual(self.summary(self.text, minlength=10, maxlength=10), "Search is the foundation of the internet")
+
+    def testSummaryBatch(self):
+        """
+        Test batch summarization of text
+        """
+
+        summaries = self.summary([self.text, self.text], maxlength=10)
+        self.assertEqual(len(summaries), 2)

test/python/testtextractor.py

@@ -0,0 +1,55 @@
+"""
+Summary module tests
+"""
+
+import unittest
+
+from txtai.pipeline import Textractor
+
+# pylint: disable = C0411
+from utils import Utils
+
+
+class TestSummary(unittest.TestCase):
+    """
+    Labels tests
+    """
+
+    def testSingle(self):
+        """
+        Tests a single extraction with no tokenization of the results
+        """
+
+        textractor = Textractor()
+
+        # Extract text as a single block
+        text = textractor(Utils.PATH + "/article.pdf")
+
+        # Check length of text is as expected
+        self.assertEqual(len(text), 2301)
+
+    def testSentences(self):
+        """
+        Tests extraction to sentences
+        """
+
+        textractor = Textractor(sentences=True)
+
+        # Extract text as sentences
+        sentences = textractor(Utils.PATH + "/article.pdf")
+
+        # Check number of sentences is as expected
+        self.assertEqual(len(sentences), 17)
+
+    def testParagraphs(self):
+        """
+        Tests extraction to paragraphs
+        """
+
+        textractor = Textractor(paragraphs=True)
+
+        # Extract text as sentences
+        paragraphs = textractor(Utils.PATH + "/article.pdf")
+
+        # Check number of paragraphs is as expected
+        self.assertEqual(len(paragraphs), 13)

test/python/testtranscription.py

@@ -0,0 +1,27 @@
+"""
+Transcription module tests
+"""
+
+import unittest
+
+from txtai.pipeline import Transcription
+
+# pylint: disable = C0411
+from utils import Utils
+
+
+class TestTranscription(unittest.TestCase):
+    """
+    Transcription tests
+    """
+
+    def testTranscription(self):
+        """
+        Test audio to text transcription
+        """
+
+        transcribe = Transcription()
+
+        self.assertEqual(
+            transcribe(Utils.PATH + "/Make_huge_profits.wav"), "Make huge profits without working make up to one hundred thousand dollars a day"
+        )

test/python/testtranslation.py

@@ -0,0 +1,68 @@
+"""
+Translation module tests
+"""
+
+import unittest
+
+from txtai.pipeline import Translation
+
+
+class TestTranslation(unittest.TestCase):
+    """
+    Translation tests
+    """
+
+    def testMarianTranslation(self):
+        """
+        Tests a translation using Marian models
+        """
+
+        translate = Translation()
+
+        text = "This is a test translation into Spanish"
+        translation = translate(text, "es")
+
+        # Validate translation text
+        self.assertEqual(translation, "Esta es una traducción de prueba al español")
+
+        # Validate translation back
+        translation = translate(translation, "en")
+        self.assertEqual(translation, text)
+
+    def testM2M100Translation(self):
+        """
+        Tests a translation using M2M100 models
+        """
+
+        translate = Translation()
+
+        text = translate("This is a test translation to Croatian", "hr")
+
+        # Validate translation text
+        self.assertEqual(text, "Ovo je testni prijevod na hrvatski")
+
+    def testNoTranslation(self):
+        """
+        Tests translation skipped when text already in destination language
+        """
+
+        translate = Translation()
+
+        text = "This is a test translation to English"
+        translation = translate(text, "en")
+
+        # Validate no translation
+        self.assertEqual(text, translation)
+
+    def testLongTranslation(self):
+        """
+        Tests a translation longer than max tokenization length.
+        """
+
+        translate = Translation()
+
+        text = "This is a test translation to Spanish. " * 100
+        translation = translate(text, "es")
+
+        # Validate translation text
+        self.assertIsNotNone(translation)

test/python/testworkflow.py

@@ -0,0 +1,95 @@
+"""
+Workflow module tests
+"""
+
+import unittest
+
+from txtai.embeddings import Documents, Embeddings
+from txtai.pipeline import Summary, Transcription, Translation, Textractor
+from txtai.workflow import Workflow, Task, FileTask, ImageTask, WorkflowTask
+
+# pylint: disable = C0411
+from utils import Utils
+
+
+class TestWorkFlow(unittest.TestCase):
+    """
+    Workflow tests
+    """
+
+    def testBaseWorkflow(self):
+        """
+        Tests a basic workflow
+        """
+
+        translate = Translation()
+
+        # Workflow that translate text to Spanish
+        workflow = Workflow([Task(lambda x: translate(x, "es"))])
+
+        results = list(workflow(["The sky is blue", "Forest through the trees"]))
+
+        self.assertEqual(len(results), 2)
+
+    def testComplexWorkflow(self):
+        """
+        Tests a complex workflow
+        """
+
+        textractor = Textractor(paragraphs=True, minlength=150, join=True)
+        summary = Summary()
+        transcribe = Transcription()
+        translate = Translation()
+
+        embeddings = Embeddings({"method": "transformers", "path": "sentence-transformers/paraphrase-xlm-r-multilingual-v1"})
+        documents = Documents()
+
+        def index(x):
+            documents.add(x)
+            return x
+
+        # Extract text and summarize articles
+        articles = Workflow([FileTask(textractor), Task(lambda x: summary(x, maxlength=15))])
+
+        # Complex workflow that handles text extraction and audio transcription
+        # Results are translated to Spanish and loaded into an embeddings index
+        tasks = [WorkflowTask(articles, r".\.pdf$"), FileTask(transcribe, r"\.wav$"), Task(lambda x: translate(x, "es")), Task(index, unpack=False)]
+
+        data = [
+            "article.pdf",
+            "Beijing_mobilises.wav",
+            "Canadas_last_fully.wav",
+            "Maine_man_wins_1_mil.wav",
+            "Make_huge_profits.wav",
+            "The_National_Park.wav",
+            "US_tops_5_million.wav",
+        ]
+
+        # Convert file paths to data tuples
+        data = [(x, "file:///%s/%s" % (Utils.PATH, element), None) for x, element in enumerate(data)]
+
+        # Execute workflow, discard results as they are streamed
+        workflow = Workflow(tasks)
+        for _ in workflow(data):
+            pass
+
+        # Build the embeddings index
+        embeddings.index(documents)
+
+        # Cleanup temporary storage
+        documents.close()
+
+        # Run search and validate result
+        index, _ = embeddings.search("buscar texto", 1)[0]
+        self.assertEqual(index, 0)
+
+    def testImageWorkflow(self):
+        """
+        Tests an image task
+        """
+
+        workflow = Workflow([ImageTask()])
+
+        results = list(workflow(["file://" + Utils.PATH + "/books.jpg"]))
+
+        self.assertEqual(results[0].size, (1024, 682))

test/python/utils.py

@@ -0,0 +1,11 @@
+"""
+Utils module
+"""
+
+
+class Utils:
+    """
+    Utility constants and methods
+    """
+
+    PATH = "/tmp/txtai"