chore: import upstream snapshot with attribution

test/python/testpipeline/testaudio/testaudiomixer.py

@@ -0,0 +1,29 @@
+"""
+AudioMixer module tests
+"""
+
+import unittest
+
+import numpy as np
+
+from txtai.pipeline import AudioMixer
+
+
+class TestAudioStream(unittest.TestCase):
+    """
+    AudioStream tests.
+    """
+
+    def testAudioStream(self):
+        """
+        Test mixing audio streams
+        """
+
+        audio1 = np.random.rand(2, 5000), 100
+        audio2 = np.random.rand(2, 5000), 100
+
+        mixer = AudioMixer()
+        audio, rate = mixer((audio1, audio2))
+
+        self.assertEqual(audio.shape, (2, 5000))
+        self.assertEqual(rate, 100)

test/python/testpipeline/testaudio/testaudiostream.py

@@ -0,0 +1,37 @@
+"""
+AudioStream module tests
+"""
+
+import unittest
+
+from unittest.mock import patch
+
+import soundfile as sf
+
+from txtai.pipeline import AudioStream
+
+# pylint: disable=C0411
+from utils import Utils
+
+
+class TestAudioStream(unittest.TestCase):
+    """
+    AudioStream tests.
+    """
+
+    @patch("sounddevice.play")
+    def testAudioStream(self, play):
+        """
+        Test playing audio
+        """
+
+        play.return_value = True
+
+        # Read audio data
+        audio, rate = sf.read(Utils.PATH + "/Make_huge_profits.wav")
+
+        stream = AudioStream()
+        self.assertIsNotNone(stream([(audio, rate), AudioStream.COMPLETE]))
+
+        # Wait for completion
+        stream.wait()

test/python/testpipeline/testaudio/testmicrophone.py

@@ -0,0 +1,81 @@
+"""
+Microphone module tests
+"""
+
+import unittest
+
+from unittest.mock import patch
+
+import numpy as np
+import soundfile as sf
+
+from txtai.pipeline import Microphone
+
+# pylint: disable=C0411
+from utils import Utils
+
+
+class TestMicrophone(unittest.TestCase):
+    """
+    Microphone tests.
+    """
+
+    # pylint: disable=C0115,C0116
+    @patch("sounddevice.RawInputStream")
+    def testMicrophone(self, inputstream):
+        """
+        Test listening to microphone
+        """
+
+        class RawInputStream:
+            def __init__(self, **kwargs):
+                self.args = kwargs
+
+                # Read audio data
+                self.index, self.passes = 0, 0
+                audio, self.samplerate = sf.read(Utils.PATH + "/Make_huge_profits.wav")
+
+                # Convert data to PCM
+                self.audio = self.int16(audio)
+
+                # Start with random data to test that speech is not detected
+                self.data = np.concatenate((self.audio * 50, np.zeros(shape=self.audio.shape, dtype=np.int16)))
+
+            def start(self):
+                pass
+
+            def stop(self):
+                pass
+
+            def read(self, size):
+                # Get chunk
+                chunk = self.data[self.index : self.index + size]
+                self.index += size
+
+                # Initial pass is random data, 2nd pass is speech data
+                if self.index > len(self.data):
+                    if not self.passes:
+                        self.index, self.passes = 0, self.passes + 1
+                        self.data = self.audio
+                    elif self.index >= len(self.audio) * 10:
+                        # Break out of loop if speech continues to not be detected
+                        raise IOError("Data exhausted")
+
+                return chunk, False
+
+            def int16(self, data):
+                i = np.iinfo(np.int16)
+                absmax = 2 ** (i.bits - 1)
+                offset = i.min + absmax
+                return (data * absmax + offset).clip(i.min, i.max).astype(np.int16)
+
+        # Mock input stream
+        inputstream.side_effect = RawInputStream
+
+        # Create microphone pipeline and read data
+        pipeline = Microphone()
+        data, rate = pipeline()
+
+        # Validate sample rate and length of data
+        self.assertEqual(len(data), 91220)
+        self.assertEqual(rate, 16000)

test/python/testpipeline/testaudio/testtexttoaudio.py

@@ -0,0 +1,26 @@
+"""
+TextToAudio module tests
+"""
+
+import unittest
+
+from txtai.pipeline import TextToAudio
+
+
+class TestTextToAudio(unittest.TestCase):
+    """
+    TextToAudio tests.
+    """
+
+    def testTextToAudio(self):
+        """
+        Test generating audio for text
+        """
+
+        tta = TextToAudio("hf-internal-testing/tiny-random-MusicgenForConditionalGeneration")
+
+        # Check that data is generated
+        audio, rate = tta("This is a test")
+
+        self.assertGreater(len(audio), 0)
+        self.assertEqual(rate, 24000)

test/python/testpipeline/testaudio/testtexttospeech.py

@@ -0,0 +1,82 @@
+"""
+TextToSpeech module tests
+"""
+
+import unittest
+
+from unittest.mock import patch
+
+from txtai.pipeline import TextToSpeech
+
+
+class TestTextToSpeech(unittest.TestCase):
+    """
+    TextToSpeech tests.
+    """
+
+    def testESPnet(self):
+        """
+        Test generating speech for text with an ESPnet model
+        """
+
+        tts = TextToSpeech()
+
+        # Check that data is generated
+        speech, rate = tts("This is a test")
+
+        self.assertGreater(len(speech), 0)
+        self.assertEqual(rate, 22050)
+
+    def testKokoro(self):
+        """
+        Test generating speech for text with a Kokoro model
+        """
+
+        tts = TextToSpeech("neuml/kokoro-int8-onnx", maxtokens=2)
+
+        # Check that data is generated
+        speech, rate = tts("This is a test")
+
+        self.assertGreater(len(speech), 0)
+        self.assertEqual(rate, 22050)
+
+    @patch("onnxruntime.get_available_providers")
+    @patch("torch.cuda.is_available")
+    def testProviders(self, cuda, providers):
+        """
+        Test that GPU provider is detected
+        """
+
+        # Test CUDA and onnxruntime-gpu installed
+        cuda.return_value = True
+        providers.return_value = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+
+        tts = TextToSpeech()
+        self.assertEqual(tts.providers()[0][0], "CUDAExecutionProvider")
+
+    def testSpeechT5(self):
+        """
+        Test generating speech for text with a SpeechT5 model
+        """
+
+        tts = TextToSpeech("neuml/txtai-speecht5-onnx")
+
+        # Check that data is generated
+        speech, rate = tts("This is a test")
+
+        self.assertGreater(len(speech), 0)
+        self.assertEqual(rate, 22050)
+
+    def testStreaming(self):
+        """
+        Test streaming speech generation
+        """
+
+        tts = TextToSpeech()
+
+        # Check that data is generated
+        speech, rate = list(tts("This is a test. And another".split(), stream=True))[0]
+
+        # Check that data is generated
+        self.assertGreater(len(speech), 0)
+        self.assertEqual(rate, 22050)

test/python/testpipeline/testaudio/testtranscription.py

@@ -0,0 +1,104 @@
+"""
+Transcription module tests
+"""
+
+import unittest
+
+import numpy as np
+import soundfile as sf
+
+from scipy import signal
+
+from txtai.pipeline import Transcription
+
+# pylint: disable=C0411
+from utils import Utils
+
+
+class TestTranscription(unittest.TestCase):
+    """
+    Transcription tests.
+    """
+
+    def testArray(self):
+        """
+        Test audio data to text transcription
+        """
+
+        transcribe = Transcription()
+
+        # Read audio data
+        raw, samplerate = sf.read(Utils.PATH + "/Make_huge_profits.wav")
+
+        self.assertEqual(transcribe((raw, samplerate)), "Make huge profits without working make up to one hundred thousand dollars a day")
+        self.assertEqual(transcribe(raw, samplerate), "Make huge profits without working make up to one hundred thousand dollars a day")
+
+    def testChunks(self):
+        """
+        Test splitting transcription into chunks
+        """
+
+        transcribe = Transcription()
+
+        result = transcribe(Utils.PATH + "/Make_huge_profits.wav", join=False)[0]
+
+        self.assertIsInstance(result["raw"], np.ndarray)
+        self.assertIsNotNone(result["rate"])
+        self.assertEqual(result["text"], "Make huge profits without working make up to one hundred thousand dollars a day")
+
+    def testFile(self):
+        """
+        Test audio file 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"
+        )
+
+    def testGenerateArguments(self):
+        """
+        Test transcription with generation keyword arguments
+        """
+
+        transcribe = Transcription()
+
+        # Read audio data
+        raw, samplerate = sf.read(Utils.PATH + "/Make_huge_profits.wav")
+
+        self.assertEqual(
+            transcribe(raw, samplerate, language="English", task="transcribe"),
+            "Make huge profits without working make up to one hundred thousand dollars a day",
+        )
+
+    def testResample(self):
+        """
+        Test resampled audio file to text transcription
+        """
+
+        transcribe = Transcription()
+
+        # Read audio data
+        raw, samplerate = sf.read(Utils.PATH + "/Make_huge_profits.wav")
+
+        # Resample for testing
+        samples = round(len(raw) * float(22050) / samplerate)
+        raw, samplerate = signal.resample(raw, samples), 22050
+
+        self.assertEqual(transcribe(raw, samplerate), "Make huge profits without working make up to one hundred thousand dollars a day")
+
+    def testStereo(self):
+        """
+        Test audio file in stereo to text transcription
+        """
+
+        transcribe = Transcription()
+
+        # Read audio data
+        raw, samplerate = sf.read(Utils.PATH + "/Make_huge_profits.wav")
+
+        # Convert mono to stereo
+        raw = np.column_stack((raw, raw))
+
+        self.assertEqual(transcribe(raw, samplerate), "Make huge profits without working make up to one hundred thousand dollars a day")

test/python/testpipeline/testdata/testfiletohtml.py

@@ -0,0 +1,24 @@
+"""
+FileToHTML module tests
+"""
+
+import os
+import unittest
+
+from unittest.mock import patch
+
+from txtai.pipeline.data.filetohtml import Tika
+
+
+class TestFileToHTML(unittest.TestCase):
+    """
+    FileToHTML tests.
+    """
+
+    @patch.dict(os.environ, {"TIKA_JAVA": "1112444abc"})
+    def testTika(self):
+        """
+        Test the Tika.available returns False when Java is not available
+        """
+
+        self.assertFalse(Tika.available())

test/python/testpipeline/testdata/testtabular.py

@@ -0,0 +1,123 @@
+"""
+Tabular module tests
+"""
+
+import unittest
+
+from txtai.pipeline import Tabular
+
+# pylint: disable=C0411
+from utils import Utils
+
+
+class TestTabular(unittest.TestCase):
+    """
+    Tabular tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create single tabular instance
+        """
+
+        cls.tabular = Tabular("id", ["text"])
+
+    def testContent(self):
+        """
+        Test parsing additional content
+        """
+
+        tabular = Tabular("id", ["text"], True)
+
+        row = {"id": 0, "text": "This is a test", "flag": 1}
+
+        # When content is enabled, both (uid, text, tags) and (uid, data, tags) rows are generated
+        # given that data doesn't necessarily include the text to index
+        rows = tabular([row])
+        uid, data, _ = rows[1]
+
+        # Data should contain the entire input row
+        self.assertEqual(uid, 0)
+        self.assertEqual(data, row)
+
+        # Only select flag field
+        tabular.content = ["flag"]
+        rows = tabular([row])
+        uid, data, _ = rows[1]
+
+        # Data should only contain a single field, flag
+        self.assertEqual(uid, 0)
+        self.assertTrue(list(data.keys()) == ["flag"])
+        self.assertEqual(data["flag"], 1)
+
+    def testCSV(self):
+        """
+        Test parsing a CSV file
+        """
+
+        rows = self.tabular([Utils.PATH + "/tabular.csv"])
+        uid, text, _ = rows[0][0]
+
+        self.assertEqual(uid, 0)
+        self.assertEqual(text, "The first sentence")
+
+    def testDict(self):
+        """
+        Test parsing a dict
+        """
+
+        rows = self.tabular([{"id": 0, "text": "This is a test"}])
+        uid, text, _ = rows[0]
+
+        self.assertEqual(uid, 0)
+        self.assertEqual(text, "This is a test")
+
+    def testInvalid(self):
+        """
+        Test invalid file paths
+        """
+
+        with self.assertRaises(ValueError):
+            self.tabular([Utils.PATH + "/article.pdf"])
+
+        with self.assertRaises(ValueError):
+            self.tabular(["https://invalid.path"])
+
+    def testList(self):
+        """
+        Test parsing a list
+        """
+
+        rows = self.tabular([[{"id": 0, "text": "This is a test"}]])
+        uid, text, _ = rows[0][0]
+
+        self.assertEqual(uid, 0)
+        self.assertEqual(text, "This is a test")
+
+    def testMissingColumns(self):
+        """
+        Test rows with uneven or missing columns
+        """
+
+        tabular = Tabular("id", ["text"], True)
+
+        rows = tabular([{"id": 0, "text": "This is a test", "metadata": "meta"}, {"id": 1, "text": "This is a test"}])
+
+        # When content is enabled both (id, text, tag) and (id, data, tag) tuples are generated given that
+        # data doesn't necessarily include the text to index
+        _, data, _ = rows[3]
+
+        self.assertIsNone(data["metadata"])
+
+    def testNoColumns(self):
+        """
+        Test creating text without specifying columns
+        """
+
+        tabular = Tabular("id")
+        rows = tabular([{"id": 0, "text": "This is a test", "summary": "Describes text in more detail"}])
+        uid, text, _ = rows[0]
+
+        self.assertEqual(uid, 0)
+        self.assertEqual(text, "This is a test. Describes text in more detail")

test/python/testpipeline/testdata/testtextractor.py

@@ -0,0 +1,253 @@
+"""
+Textractor module tests
+"""
+
+import platform
+import unittest
+
+from txtai.pipeline import Textractor
+
+# pylint: disable=C0411
+from utils import Utils
+
+
+class TestTextractor(unittest.TestCase):
+    """
+    Textractor tests.
+    """
+
+    def testClean(self):
+        """
+        Test text cleaning method
+        """
+
+        # Default text cleaning
+        textractor = Textractor()
+        self.assertEqual(textractor(" a  b  c "), "a b c")
+
+        # Require text to be minlength
+        textractor = Textractor(minlength=10)
+        self.assertEqual(textractor(" a  b  c "), None)
+
+        # Disable text cleaning
+        textractor = Textractor(cleantext=False, minlength=10)
+        self.assertEqual(textractor(" a  b  c "), " a  b  c ")
+
+    def testChonkie(self):
+        """
+        Test a chonkie chunker
+        """
+
+        # Test chonkie chunking
+        textractor = Textractor(chunker="sentence", chunk_size=5, chunk_overlap=0)
+        self.assertEqual(textractor("This is a test. And another test."), ["This is a test.", "And another test."])
+
+        # Test bad chunker throws an exception
+        with self.assertRaises(AttributeError):
+            textractor = Textractor(chunker="badchunker")
+
+    def testDefault(self):
+        """
+        Test default text extraction
+        """
+
+        # Text input
+        textractor = Textractor(backend=None)
+        text = textractor(Utils.PATH + "/tabular.csv")
+        self.assertEqual(len(text), 125)
+
+        # Markdown input
+        textractor = Textractor(sections=True)
+        sections = textractor("# Heading 1\nText1\n\n# Heading 2\nText2\n")
+
+        # Check number of sections is as expected
+        self.assertEqual(len(sections), 2)
+
+    @unittest.skipIf(platform.system() == "Darwin", "Docling skipped on macOS to avoid MPS issues")
+    def testDocling(self):
+        """
+        Test docling backend
+        """
+
+        textractor = Textractor(backend="docling")
+
+        # Extract text and check for Markdown formatting
+        text = textractor(Utils.PATH + "/article.pdf")
+        self.assertTrue("## Introducing txtai" in text)
+
+    def testLines(self):
+        """
+        Test extraction to lines
+        """
+
+        textractor = Textractor(lines=True)
+
+        # Extract text as lines
+        lines = textractor(Utils.PATH + "/article.pdf")
+
+        # Check number of lines is as expected
+        self.assertEqual(len(lines), 35)
+
+    def testLiteParse(self):
+        """
+        Test liteparse backend
+        """
+
+        textractor = Textractor(backend="liteparse")
+
+        # Extract text and check for Markdown formatting
+        text = textractor(Utils.PATH + "/article.pdf")
+        self.assertTrue("# Introducing txtai" in text)
+
+    def testHTML(self):
+        """
+        Test HTML to Markdown
+        """
+
+        # Headings
+        self.assertMarkdown("<h1>This is a test</h1>", "# This is a test")
+        self.assertMarkdown("<h6>This is a test</h6>", "###### This is a test")
+
+        # Blockquotes
+        self.assertMarkdown("<blockquote>This is a test</blockquote>", "> This is a test")
+
+        # Lists
+        self.assertMarkdown("<ul><li>Test1</li><li>Test2</li></ul>", "- Test1\n- Test2")
+        self.assertMarkdown("<ol><li>Test1</li><li>Test2</li></ol>", "1. Test1\n2. Test2")
+
+        # Code
+        self.assertMarkdown("<code>This is a test</code>", "```\nThis is a test\n```")
+        self.assertMarkdown("<pre>This is a test</pre>", "```\nThis is a test\n```")
+
+        # Tables
+        self.assertMarkdown(
+            "<table><tr><th>Header1</th><th>Header2</th></tr><tr><td>Test1</td><td>Test2</td></tr></table>",
+            "|Header1|Header2|\n|---|---|\n|Test1|Test2|",
+        )
+
+        # Ignore list
+        self.assertMarkdown("<aside>This is a test</aside>", "")
+
+        # Text formatting
+        self.assertMarkdown("<p>This is a test</p>", "This is a test")
+        self.assertMarkdown("<p>This is a <b>test</b</p>", "This is a **test**")
+        self.assertMarkdown("<p>This is a <strong>test</strong></p>", "This is a **test**")
+        self.assertMarkdown("<p>This is a <i>test</i></p>", "This is a *test*")
+        self.assertMarkdown("<p>This is a <em>test</em></p>", "This is a *test*")
+        self.assertMarkdown("<p>This is a <a href='link'>test</a>", "This is a [test](link)")
+
+        # Collapse to outer tag
+        self.assertMarkdown("<p>This is a <strong><em>test</em></strong></p>", "This is a **test**")
+        self.assertMarkdown("<p>This is a <em><strong>test</strong></em></p>", "This is a *test*")
+
+    def testParagraphs(self):
+        """
+        Test extraction to paragraphs
+        """
+
+        textractor = Textractor(paragraphs=True)
+
+        # Extract text as paragraphs
+        paragraphs = textractor(Utils.PATH + "/article.pdf")
+
+        # Check number of paragraphs is as expected
+        self.assertEqual(len(paragraphs), 11)
+
+    def testSections(self):
+        """
+        Test extraction to sections
+        """
+
+        textractor = Textractor(sections=True)
+
+        # Extract as sections
+        sections = textractor(Utils.PATH + "/document.pdf")
+
+        # Check number of sections is as expected
+        self.assertEqual(len(sections), 3)
+
+    def testSentences(self):
+        """
+        Test 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 testSingle(self):
+        """
+        Test 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), 2471)
+
+    def testTable(self):
+        """
+        Test table extraction
+        """
+
+        textractor = Textractor()
+
+        # Extract text as a single block
+        for name in ["document.docx", "spreadsheet.xlsx"]:
+            text = textractor(f"{Utils.PATH}/{name}")
+
+            # Check for table header
+            self.assertTrue("|---|" in text)
+
+    def testTikaFlag(self):
+        """
+        Test legacy tika flag
+        """
+
+        textractor = Textractor(tika=True)
+        self.assertIsNotNone(textractor.html)
+
+        textractor = Textractor(tika=False)
+        self.assertIsNone(textractor.html)
+
+    def testTuples(self):
+        """
+        Test output tuples
+        """
+
+        # Default text cleaning
+        textractor = Textractor(tuples=True)
+
+        path, text = textractor(Utils.PATH + "/article.pdf")
+        self.assertEqual(path, Utils.PATH + "/article.pdf")
+        self.assertEqual(len(text), 2471)
+
+    def testURL(self):
+        """
+        Test parsing a remote URL
+        """
+
+        # Test parsing URLs for each backend
+        for backend in ["docling", "liteparse", "tika"]:
+            textractor = Textractor(backend=backend)
+            text = textractor("https://github.com/neuml/txtai")
+            self.assertTrue("txtai is an all-in-one AI framework" in text)
+
+    def assertMarkdown(self, html, expected):
+        """
+        Helper method to assert generated markdown is as expected.
+
+        Args:
+            html: input html snippet
+            expected: expected markdown text
+        """
+
+        textractor = Textractor()
+        self.assertEqual(textractor(f"<html><body>{html}</body></html>"), expected)

test/python/testpipeline/testdata/testtokenizer.py

@@ -0,0 +1,68 @@
+"""
+Tokenizer module tests
+"""
+
+import unittest
+
+from txtai.pipeline import Tokenizer
+
+
+class TestTokenizer(unittest.TestCase):
+    """
+    Tokenizer tests.
+    """
+
+    def testAlphanumTokenize(self):
+        """
+        Test alphanumeric tokenization
+        """
+
+        # Alphanumeric tokenization through backwards compatible static method
+        self.assertEqual(Tokenizer.tokenize("Y this is a test!"), ["test"])
+        self.assertEqual(Tokenizer.tokenize("abc123 ABC 123"), ["abc123", "abc"])
+
+    def testEmptyTokenize(self):
+        """
+        Test handling empty and None inputs
+        """
+
+        # Test that parser can handle empty or None strings
+        self.assertEqual(Tokenizer.tokenize(""), [])
+        self.assertEqual(Tokenizer.tokenize(None), None)
+
+    def testStandardTokenize(self):
+        """
+        Test standard tokenization
+        """
+
+        # Default standard tokenizer parameters
+        tokenizer = Tokenizer()
+
+        # Define token tests
+        tests = [
+            ("Y this is a test!", ["y", "this", "is", "a", "test"]),
+            ("abc123 ABC 123", ["abc123", "abc", "123"]),
+            ("Testing hy-phenated words", ["testing", "hy", "phenated", "words"]),
+            ("111-111-1111", ["111", "111", "1111"]),
+            ("Test.1234", ["test", "1234"]),
+        ]
+
+        # Run through tests
+        for test, result in tests:
+            # Unicode Text Segmentation per Unicode Annex #29
+            self.assertEqual(tokenizer(test), result)
+
+    def testNgramTokenize(self):
+        """
+        Test ngram tokenization
+        """
+
+        # Standard ngram tokenization
+        tokenizer = Tokenizer(lowercase=True, ngrams=3)
+        result = tokenizer("NGRAM TEST")
+        self.assertIn("ngr", result)
+
+        # Case sensitive ngram tokenization
+        tokenizer = Tokenizer(lowercase=False, ngrams=3)
+        result = tokenizer("NGRAM TEST")
+        self.assertIn("NGR", result)

test/python/testpipeline/testdata/testurlretrieve.py

@@ -0,0 +1,115 @@
+"""
+URLRetrieve module tests
+"""
+
+import contextlib
+import unittest
+
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from threading import Thread
+from urllib.request import build_opener
+
+from txtai.pipeline import URLRetrieve
+from txtai.pipeline.data.urlretrieve import SafeRedirectHandler
+
+
+class RequestHandler(BaseHTTPRequestHandler):
+    """
+    Test HTTP handler.
+    """
+
+    def do_GET(self):
+        """
+        GET request handler.
+        """
+
+        if self.path == "/valid":
+            redirect = "https://github.com/neuml/txtai"
+        elif self.path == "/invalid":
+            redirect = "http://127.0.0.1"
+        else:
+            redirect = None
+
+        if redirect:
+            self.send_response(301)
+            self.send_header("Location", redirect)
+            self.end_headers()
+        else:
+            response = "test".encode("utf-8")
+
+            self.send_response(200)
+            self.send_header("content-type", "text/plain")
+            self.send_header("content-length", len(response))
+            self.end_headers()
+
+            self.wfile.write(response)
+            self.wfile.flush()
+
+
+class TestURLRetrieve(unittest.TestCase):
+    """
+    URLRetrieve tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create mock http server
+        """
+
+        cls.httpd = HTTPServer(("127.0.0.1", 8006), RequestHandler)
+
+        server = Thread(target=cls.httpd.serve_forever, daemon=True)
+        server.start()
+
+    @classmethod
+    def tearDownClass(cls):
+        """
+        Shutdown mock http server.
+        """
+
+        cls.httpd.shutdown()
+
+    def testRedirect(self):
+        """
+        Test redirects
+        """
+
+        urlretrieve = URLRetrieve(safeopen=True)
+
+        # Test redirect handler
+        opener = build_opener(SafeRedirectHandler(urlretrieve))
+
+        # Test valid direct
+        with contextlib.closing(opener.open("http://127.0.0.1:8006/valid")) as connection:
+            self.assertTrue("txtai is an all-in-one AI framework" in str(connection.read()))
+
+        # Test invalid redirect
+        with self.assertRaises(IOError):
+            contextlib.closing(opener.open("http://127.0.0.1:8006/invalid"))
+
+    def testRetrieve(self):
+        """
+        Test retrieval
+        """
+
+        urlretrieve = URLRetrieve()
+        data = urlretrieve("http://127.0.0.1:8006/data")
+        self.assertEqual(data, b"test")
+
+    def testSafeopen(self):
+        """
+        Test safeopen checks
+        """
+
+        urlretrieve = URLRetrieve(safeopen=True)
+
+        # Verify that local ip addresses fail
+        with self.assertRaises(IOError):
+            urlretrieve("http://127.0.0.1")
+
+        with self.assertRaises(IOError):
+            urlretrieve("https://127.0.0.1")
+
+        with self.assertRaises(IOError):
+            urlretrieve("https://[::1]")

test/python/testpipeline/testimage/testcaption.py

@@ -0,0 +1,36 @@
+"""
+Caption module tests
+"""
+
+import unittest
+
+from PIL import Image
+
+from transformers import AutoModelForImageTextToText, AutoImageProcessor, AutoTokenizer
+from txtai.pipeline import Caption
+
+# pylint: disable=C0411
+from utils import Utils
+
+
+class TestCaption(unittest.TestCase):
+    """
+    Caption tests.
+    """
+
+    def testCaption(self):
+        """
+        Test captions
+        """
+
+        caption = Caption()
+        self.assertEqual(caption(Image.open(Utils.PATH + "/books.jpg")), "a book shelf filled with books and a stack of books")
+
+        # Load passing models directly
+        path = "ydshieh/vit-gpt2-coco-en"
+        model = AutoModelForImageTextToText.from_pretrained(path)
+        tokenizer = AutoTokenizer.from_pretrained(path)
+        processor = AutoImageProcessor.from_pretrained(path)
+
+        caption = Caption((model, tokenizer, processor))
+        self.assertEqual(caption(Image.open(Utils.PATH + "/books.jpg")), "a book shelf filled with books and a stack of books")

test/python/testpipeline/testimage/testimagehash.py

@@ -0,0 +1,74 @@
+"""
+ImageHash module tests
+"""
+
+import unittest
+
+from PIL import Image
+
+from txtai.pipeline import ImageHash
+
+# pylint: disable=C0411
+from utils import Utils
+
+
+class TestImageHash(unittest.TestCase):
+    """
+    ImageHash tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Caches an image to hash
+        """
+
+        cls.image = Image.open(Utils.PATH + "/books.jpg")
+
+    def testArray(self):
+        """
+        Test numpy return type
+        """
+
+        ihash = ImageHash(strings=False)
+        self.assertEqual(ihash(self.image).shape, (64,))
+
+    def testAverage(self):
+        """
+        Test average hash
+        """
+
+        ihash = ImageHash("average")
+        self.assertIn(ihash(self.image), ["0859dd04bfbfbf00", "0859dd04ffbfbf00"])
+
+    def testColor(self):
+        """
+        Test color hash
+        """
+
+        ihash = ImageHash("color")
+        self.assertIn(ihash(self.image), ["1ffffe02000e000c0e0000070000", "1ff8fe03000e00070e0000070000"])
+
+    def testDifference(self):
+        """
+        Test difference hash
+        """
+
+        ihash = ImageHash("difference")
+        self.assertEqual(ihash(self.image), "d291996d6969686a")
+
+    def testPerceptual(self):
+        """
+        Test perceptual hash
+        """
+
+        ihash = ImageHash("perceptual")
+        self.assertEqual(ihash(self.image), "8be8418577b331b9")
+
+    def testWavelet(self):
+        """
+        Test wavelet hash
+        """
+
+        ihash = ImageHash("wavelet")
+        self.assertEqual(ihash(Utils.PATH + "/books.jpg"), "68015d85bfbf3f00")

test/python/testpipeline/testimage/testobjects.py

@@ -0,0 +1,40 @@
+"""
+Objects module tests
+"""
+
+import unittest
+
+from txtai.pipeline import Objects
+
+# pylint: disable=C0411
+from utils import Utils
+
+
+class TestObjects(unittest.TestCase):
+    """
+    Object detection tests.
+    """
+
+    def testClassification(self):
+        """
+        Test object detection using an image classification model
+        """
+
+        objects = Objects(classification=True, threshold=0.3)
+        self.assertEqual(objects(Utils.PATH + "/books.jpg")[0][0], "library")
+
+    def testDetection(self):
+        """
+        Test object detection using an object detection model
+        """
+
+        objects = Objects()
+        self.assertEqual(objects(Utils.PATH + "/books.jpg")[0][0], "book")
+
+    def testFlatten(self):
+        """
+        Test object detection using an object detection model, flatten to return only objects
+        """
+
+        objects = Objects()
+        self.assertEqual(objects(Utils.PATH + "/books.jpg", flatten=True)[0], "book")

test/python/testpipeline/testllm/testgenerator.py

@@ -0,0 +1,24 @@
+"""
+Generator module tests
+"""
+
+import unittest
+
+from txtai.pipeline import Generator
+
+
+class TestGenerator(unittest.TestCase):
+    """
+    Sequences tests.
+    """
+
+    def testGeneration(self):
+        """
+        Test text pipeline generation
+        """
+
+        model = Generator("hf-internal-testing/tiny-random-gpt2")
+        start = "Hello, how are"
+
+        # Test that text is generated
+        self.assertIsNotNone(model(start))

test/python/testpipeline/testllm/testlitellm.py

@@ -0,0 +1,115 @@
+"""
+LiteLLM module tests
+"""
+
+import json
+import os
+import time
+import unittest
+import uuid
+
+from unittest.mock import patch
+
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from threading import Thread
+
+from txtai.pipeline import LLM
+
+
+class RequestHandler(BaseHTTPRequestHandler):
+    """
+    Test HTTP handler.
+    """
+
+    def do_POST(self):
+        """
+        POST request handler.
+        """
+
+        # Parse input headers
+        length = int(self.headers["content-length"])
+        data = json.loads(self.rfile.read(length))
+
+        if data.get("stream"):
+            # Mock streaming response
+            content = "application/octet-stream"
+            response = (
+                "data: "
+                + json.dumps(
+                    {
+                        "id": str(uuid.uuid4()),
+                        "object": "chat.completion.chunk",
+                        "created": int(time.time() * 1000),
+                        "model": "test",
+                        "choices": [{"id": 0, "delta": {"content": "blue"}}],
+                    }
+                )
+                + "\n\ndata: [DONE]\n\n"
+            )
+        else:
+            # Mock standard response
+            content = "application/json"
+            response = json.dumps(
+                {
+                    "id": str(uuid.uuid4()),
+                    "object": "chat.completion",
+                    "created": int(time.time() * 1000),
+                    "model": "test",
+                    "choices": [{"id": 0, "message": {"role": "assistant", "content": "blue"}, "finish_reason": "stop"}],
+                }
+            )
+
+        # Encode response as bytes
+        response = response.encode("utf-8")
+
+        self.send_response(200)
+        self.send_header("content-type", content)
+        self.send_header("content-length", len(response))
+        self.end_headers()
+
+        self.wfile.write(response)
+        self.wfile.flush()
+
+
+class TestLiteLLM(unittest.TestCase):
+    """
+    LiteLLM tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create mock http server.
+        """
+
+        cls.httpd = HTTPServer(("127.0.0.1", 8000), RequestHandler)
+
+        server = Thread(target=cls.httpd.serve_forever, daemon=True)
+        server.start()
+
+    @classmethod
+    def tearDownClass(cls):
+        """
+        Shutdown mock http server.
+        """
+
+        cls.httpd.shutdown()
+
+    @patch.dict(os.environ, {"OPENAI_API_KEY": "test"})
+    def testGeneration(self):
+        """
+        Test generation with LiteLLM
+        """
+
+        # Test model generation with LiteLLM
+        model = LLM("openai/gpt-4o", api_base="http://127.0.0.1:8000")
+        self.assertEqual(model("The sky is"), "blue")
+
+        # Test default role
+        self.assertEqual(model("The sky is", defaultrole="user"), "blue")
+
+        # Test streaming
+        self.assertEqual(" ".join(x for x in model("The sky is", stream=True)), "blue")
+
+        # Test vision
+        self.assertEqual(model.isvision(), False)

test/python/testpipeline/testllm/testlitert.py

@@ -0,0 +1,31 @@
+"""
+LiteRT module tests
+"""
+
+import unittest
+
+from txtai.pipeline import LLM
+
+
+class TestLiteRT(unittest.TestCase):
+    """
+    LiteRT tests.
+    """
+
+    def testGeneration(self):
+        """
+        Test generation with LiteRT
+        """
+
+        # Test model generation with LiteRT
+        model = LLM("neuml/gemma-4-tiny-random-litert-lm/gemma-4-tiny-random.litertlm", mtp=False, maxlength=25)
+
+        # Test standard
+        self.assertIsNotNone(model("Hello"))
+
+        # Test streaming
+        self.assertIsNotNone(list(model("Hello", stream=True)))
+
+        # Test CPU fallback
+        model = LLM("neuml/gemma-4-tiny-random-litert-lm/gemma-4-tiny-random.litertlm", mtp=True, maxlength=25)
+        self.assertIsNotNone(model("Hello"))

test/python/testpipeline/testllm/testllama.py

@@ -0,0 +1,76 @@
+"""
+Llama module tests
+"""
+
+import unittest
+
+from unittest.mock import patch
+
+from txtai.pipeline import LLM
+
+
+class TestLlama(unittest.TestCase):
+    """
+    llama.cpp tests.
+    """
+
+    @patch("llama_cpp.Llama")
+    def testContext(self, llama):
+        """
+        Test n_ctx with llama.cpp
+        """
+
+        class Llama:
+            """
+            Mock llama.cpp instance to test invalid context
+            """
+
+            def __init__(self, **kwargs):
+                if kwargs.get("n_ctx") == 0 or kwargs.get("n_ctx", 0) >= 10000:
+                    raise ValueError("Failed to create context")
+
+                # Save parameters
+                self.params = kwargs
+
+        # Mock llama.cpp instance
+        llama.side_effect = Llama
+
+        # Model to test
+        path = "TheBloke/TinyLlama-1.1B-Chat-v0.3-GGUF/tinyllama-1.1b-chat-v0.3.Q2_K.gguf"
+
+        # Test omitting n_ctx falls back to default settings
+        llm = LLM(path)
+        self.assertNotIn("n_ctx", llm.generator.llm.params)
+
+        # Test n_ctx=0 falls back to default settings
+        llm = LLM(path, n_ctx=0)
+        self.assertNotIn("n_ctx", llm.generator.llm.params)
+
+        # Test n_ctx manually set
+        llm = LLM(path, n_ctx=1024)
+        self.assertEqual(llm.generator.llm.params["n_ctx"], 1024)
+
+        # Mock a value for n_ctx that's too big
+        with self.assertRaises(ValueError):
+            llm = LLM(path, n_ctx=10000)
+
+    def testGeneration(self):
+        """
+        Test generation with llama.cpp
+        """
+
+        # Test model generation with llama.cpp
+        model = LLM("TheBloke/TinyLlama-1.1B-Chat-v0.3-GGUF/tinyllama-1.1b-chat-v0.3.Q2_K.gguf", chat_format="chatml")
+
+        # Test with prompt
+        self.assertEqual(model("2 + 2 = ", maxlength=10, seed=0, stop=["."], defaultrole="prompt")[0], "4")
+
+        # Test with list of messages
+        messages = [{"role": "system", "content": "You are a helpful assistant. You answer math problems."}, {"role": "user", "content": "2+2?"}]
+        self.assertIsNotNone(model(messages, maxlength=10, seed=0, stop=["."]))
+
+        # Test default role
+        self.assertIsNotNone(model("2 + 2 = ", maxlength=10, seed=0, stop=["."], defaultrole="user"))
+
+        # Test streaming
+        self.assertEqual(" ".join(x for x in model("2 + 2 = ", maxlength=10, stream=True, seed=0, stop=["."], defaultrole="prompt"))[0], "4")

test/python/testpipeline/testllm/testllm.py

@@ -0,0 +1,185 @@
+"""
+LLM module tests
+"""
+
+import unittest
+
+import torch
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from txtai.pipeline import LLM, Generation
+
+# pylint: disable=C0411
+from utils import Utils
+
+
+class TestLLM(unittest.TestCase):
+    """
+    LLM tests.
+    """
+
+    def testArguments(self):
+        """
+        Test pipeline keyword arguments
+        """
+
+        start = "Hello, how are"
+
+        # Test that text is generated with custom parameters
+        model = LLM("hf-internal-testing/tiny-random-gpt2", task="language-generation", dtype="torch.float32")
+        self.assertIsNotNone(model(start))
+
+        model = LLM("hf-internal-testing/tiny-random-gpt2", task="language-generation", dtype=torch.float32)
+        self.assertIsNotNone(model(start))
+
+    def testBatchSize(self):
+        """
+        Test batch size
+        """
+
+        model = LLM("sshleifer/tiny-gpt2")
+        self.assertIsNotNone(model(["Hello, how are"] * 2, batch_size=2))
+
+    def testCustom(self):
+        """
+        Test custom LLM framework
+        """
+
+        model = LLM("hf-internal-testing/tiny-random-gpt2", task="language-generation", method="txtai.pipeline.HFGeneration")
+        self.assertIsNotNone(model("Hello, how are"))
+
+    def testCustomNotFound(self):
+        """
+        Test resolving an unresolvable LLM framework
+        """
+
+        with self.assertRaises(ImportError):
+            LLM("hf-internal-testing/tiny-random-gpt2", method="notfound.generation")
+
+    def testDefaultRole(self):
+        """
+        Test default role
+        """
+
+        model = LLM("hf-internal-testing/tiny-random-LlamaForCausalLM")
+        generator = model.generator
+
+        # Validate that the LLM supports chat messages
+        self.assertEqual(model.ischat(), True)
+
+        messages = [
+            ("Hello", list),
+            ("\n<|im_start|>Hello<|im_end|>", str),
+            ("<|start|>Hello<|end|>", str),
+            ("<|start_of_role|>system<|end_of_role|>", str),
+            ("[INST]Hello[/INST]", str),
+        ]
+
+        for message, expected in messages:
+            # Test auto detection of formats
+            self.assertEqual(type(generator.format([message], "auto")[0]), expected)
+
+            # Test always setting user chat messages
+            self.assertEqual(type(generator.format([message], "user")[0]), list)
+
+            # Test always keeping as prompt text
+            self.assertEqual(type(generator.format([message], "prompt")[0]), str)
+
+    def testExternal(self):
+        """
+        Test externally loaded model
+        """
+
+        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2")
+        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
+
+        model = LLM((model, tokenizer), template="{text}")
+        start = "Hello, how are"
+
+        # Test that text is generated
+        self.assertIsNotNone(model(start))
+
+    def testMaxLength(self):
+        """
+        Test max length
+        """
+
+        model = LLM("sshleifer/tiny-gpt2")
+        self.assertIsInstance(model("Hello, how are", maxlength=10), str)
+
+    def testNotImplemented(self):
+        """
+        Test exceptions for non-implemented methods
+        """
+
+        generation = Generation()
+        self.assertRaises(NotImplementedError, generation.stream, None, None, None, None)
+
+    def testStop(self):
+        """
+        Test stop strings
+        """
+
+        model = LLM("sshleifer/tiny-gpt2")
+        self.assertIsNotNone(model("Hello, how are", stop=["you"]))
+
+    def testStream(self):
+        """
+        Test streaming generation
+        """
+
+        model = LLM("sshleifer/tiny-gpt2")
+        self.assertIsInstance(" ".join(x for x in model("Hello, how are", stream=True)), str)
+
+    def testStripThink(self):
+        """
+        Test stripthink parameter
+        """
+
+        # pylint: disable=W0613
+        def execute1(*args, **kwargs):
+            return ["<think>test</think>you"]
+
+        def execute2(*args, **kwargs):
+            return ["<|channel|>final<|message|> you"]
+
+        model = LLM("hf-internal-testing/tiny-random-LlamaForCausalLM")
+
+        for method in [execute1, execute2]:
+            # Override execute method
+            model.generator.execute = method
+            self.assertEqual(model("Hello, how are", stripthink=True), "you")
+            self.assertEqual(model("Hello, how are", stripthink=False), method()[0])
+
+    def testStripThinkStream(self):
+        """
+        Test stripthink parameter with streaming output
+        """
+
+        # pylint: disable=W0613
+        def execute1(*args, **kwargs):
+            yield from "<think>test</think>you"
+
+        def execute2(*args, **kwargs):
+            yield from "<|channel|>final<|message|>you"
+
+        model = LLM("hf-internal-testing/tiny-random-LlamaForCausalLM")
+
+        for method in [execute1, execute2]:
+            # Override execute method
+            model.generator.execute = method
+            self.assertEqual("".join(model("Hello, how are", stripthink=True, stream=True)), "you")
+            self.assertEqual("".join(model("Hello, how are", stripthink=False, stream=True)), "".join(list(method())))
+
+    def testVision(self):
+        """
+        Test vision LLM
+        """
+
+        model = LLM("neuml/tiny-random-qwen2vl")
+        result = model(
+            [{"role": "user", "content": [{"type": "text", "text": "What is in this image?"}, {"type": "image", "image": Utils.PATH + "/books.jpg"}]}]
+        )
+
+        self.assertIsNotNone(result)

test/python/testpipeline/testllm/testopencode.py

@@ -0,0 +1,71 @@
+"""
+OpenCode module tests
+"""
+
+import json
+import unittest
+
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from threading import Thread
+
+from txtai.pipeline import LLM
+
+
+class RequestHandler(BaseHTTPRequestHandler):
+    """
+    Test HTTP handler.
+    """
+
+    def do_POST(self):
+        """
+        POST request handler.
+        """
+
+        # Mock response
+        content = "application/json"
+        response = json.dumps({"id": "0", "parts": [{"type": "text", "text": "blue"}]})
+
+        # Encode response as bytes
+        response = response.encode("utf-8")
+
+        self.send_response(200)
+        self.send_header("content-type", content)
+        self.send_header("content-length", len(response))
+        self.end_headers()
+
+        self.wfile.write(response)
+        self.wfile.flush()
+
+
+class TestOpenCode(unittest.TestCase):
+    """
+    OpenCode tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create mock http server.
+        """
+
+        cls.httpd = HTTPServer(("127.0.0.1", 8005), RequestHandler)
+
+        server = Thread(target=cls.httpd.serve_forever, daemon=True)
+        server.start()
+
+    @classmethod
+    def tearDownClass(cls):
+        """
+        Shutdown mock http server.
+        """
+
+        cls.httpd.shutdown()
+
+    def testGeneration(self):
+        """
+        Test generation with OpenCode
+        """
+
+        # Test model generation with LiteLLM
+        model = LLM("opencode/big-pickle", url="http://127.0.0.1:8005")
+        self.assertEqual(model("The sky is"), "blue")

test/python/testpipeline/testllm/testrag.py

@@ -0,0 +1,225 @@
+"""
+RAG module tests
+"""
+
+import platform
+import unittest
+
+from txtai.embeddings import Embeddings
+from txtai.pipeline import Questions, RAG, Similarity
+
+
+class TestRAG(unittest.TestCase):
+    """
+    RAG tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create single rag instance.
+        """
+
+        cls.data = [
+            "Giants hit 3 HRs to down Dodgers",
+            "Giants 5 Dodgers 4 final",
+            "Dodgers drop Game 2 against the Giants, 5-4",
+            "Blue Jays beat Red Sox final score 2-1",
+            "Red Sox lost to the Blue Jays, 2-1",
+            "Blue Jays at Red Sox is over. Score: 2-1",
+            "Phillies win over the Braves, 5-0",
+            "Phillies 5 Braves 0 final",
+            "Final: Braves lose to the Phillies in the series opener, 5-0",
+            "Lightning goaltender pulled, lose to Flyers 4-1",
+            "Flyers 4 Lightning 1 final",
+            "Flyers win 4-1",
+        ]
+
+        # Create embeddings model, backed by sentence-transformers & transformers
+        cls.embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2"})
+
+        # Create rag instance
+        cls.rag = RAG(cls.embeddings, "distilbert-base-cased-distilled-squad")
+
+    @classmethod
+    def tearDownClass(cls):
+        """
+        Cleanup data.
+        """
+
+        if cls.embeddings:
+            cls.embeddings.close()
+
+    def testAnswer(self):
+        """
+        Test qa extraction with an answer
+        """
+
+        questions = ["What team won the game?", "What was score?"]
+
+        # pylint: disable=C3001
+        execute = lambda query: self.rag([(question, query, question, False) for question in questions], self.data)
+
+        answers = execute("Red Sox - Blue Jays")
+        self.assertEqual("Blue Jays", answers[0][1])
+        self.assertEqual("2-1", answers[1][1])
+
+        # Ad-hoc questions
+        question = "What hockey team won?"
+
+        answers = self.rag([(question, question, question, False)], self.data)
+        self.assertEqual("Flyers", answers[0][1])
+
+    def testEmptyQuery(self):
+        """
+        Test an empty queries list
+        """
+
+        self.assertEqual(self.rag.query(None, None), [])
+
+    def testNoAnswer(self):
+        """
+        Test qa extraction with no answer
+        """
+
+        question = ""
+
+        answers = self.rag([(question, question, question, False)], self.data)
+        self.assertIsNone(answers[0][1])
+
+        question = "abcdef"
+        answers = self.rag([(question, question, question, False)], self.data)
+        self.assertIsNone(answers[0][1])
+
+    @unittest.skipIf(platform.system() == "Darwin", "Quantized models not supported on macOS")
+    def testQuantize(self):
+        """
+        Test qa extraction backed by a quantized model
+        """
+
+        rag = RAG(self.embeddings, "distilbert-base-cased-distilled-squad", True)
+
+        question = "How many home runs?"
+
+        answers = rag([(question, question, question, True)], self.data)
+        self.assertIsNotNone(answers[0][1])
+
+    def testOutputs(self):
+        """
+        Test output formatting rules
+        """
+
+        question = "How many home runs?"
+
+        # Test flatten to list of answers
+        rag = RAG(self.embeddings, "distilbert-base-cased-distilled-squad", output="flatten")
+        answers = rag([(question, question, question, True)], self.data)
+        self.assertTrue(answers[0].startswith("Giants hit 3 HRs"))
+
+        # Test reference field
+        rag = RAG(self.embeddings, "distilbert-base-cased-distilled-squad", output="reference")
+        answers = rag([(question, question, question, True)], self.data)
+        self.assertTrue(self.data[answers[0][2]].startswith("Giants hit 3 HRs"))
+
+    def testPrompt(self):
+        """
+        Test a user prompt with templating
+        """
+
+        embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": True})
+        embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
+
+        rag = RAG(
+            embeddings,
+            "google/flan-t5-small",
+            template="""
+              Answer the following question and return a number.
+              Question: {question}
+              Context:{context}""",
+            output="flatten",
+        )
+
+        self.assertEqual(rag("How many HRs"), "3")
+
+    def testPromptTemplates(self):
+        """
+        Test system and user prompt templates
+        """
+
+        rag = RAG(
+            self.embeddings,
+            "sshleifer/tiny-gpt2",
+            system="You are a friendly assistant",
+            template="""
+              Answer the following question and return a number.
+              Question: {question}
+              Context:{context}""",
+        )
+
+        prompts = rag.prompts(["How many HRs?"], [self.data])[0]
+        self.assertEqual([x["role"] for x in prompts], ["system", "user"])
+
+    def testSearch(self):
+        """
+        Test qa extraction with an embeddings search for context
+        """
+
+        embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": True})
+        embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
+
+        rag = RAG(embeddings, "distilbert-base-cased-distilled-squad")
+
+        question = "How many home runs?"
+
+        answers = rag([(question, question, question, True)])
+        self.assertTrue(answers[0][1].startswith("Giants hit 3 HRs"))
+
+    def testSimilarity(self):
+        """
+        Test qa extraction using a Similarity pipeline to build context
+        """
+
+        # Create rag instance
+        rag = RAG(Similarity("prajjwal1/bert-medium-mnli"), Questions("distilbert-base-cased-distilled-squad"))
+
+        question = "How many home runs?"
+
+        answers = rag([(question, "HRs", question, True)], self.data)
+        self.assertTrue(answers[0][1].startswith("Giants hit 3 HRs"))
+
+    def testSnippet(self):
+        """
+        Test qa extraction with a full answer snippet
+        """
+
+        question = "How many home runs?"
+
+        answers = self.rag([(question, question, question, True)], self.data)
+        self.assertTrue(answers[0][1].startswith("Giants hit 3 HRs"))
+
+    def testSnippetEmpty(self):
+        """
+        Test snippet method can handle empty parameters
+        """
+
+        self.assertEqual(self.rag.snippets(["name"], [None], [None], [None]), [("name", None)])
+
+    def testStringInput(self):
+        """
+        Test with single string input
+        """
+
+        result = self.rag("How many home runs?", self.data)
+        self.assertEqual(result["answer"], "3")
+
+    def testTasks(self):
+        """
+        Test loading models with task parameter
+        """
+
+        for task, model in [
+            ("language-generation", "hf-internal-testing/tiny-random-gpt2"),
+            ("sequence-sequence", "hf-internal-testing/tiny-random-t5"),
+        ]:
+            rag = RAG(self.embeddings, model, task=task)
+            self.assertIsNotNone(rag)

test/python/testpipeline/testllm/testsequences.py

@@ -0,0 +1,21 @@
+"""
+Sequences module tests
+"""
+
+import unittest
+
+from txtai.pipeline import Sequences
+
+
+class TestSequences(unittest.TestCase):
+    """
+    Sequences tests.
+    """
+
+    def testGeneration(self):
+        """
+        Test text2text pipeline generation
+        """
+
+        model = Sequences("t5-small")
+        self.assertEqual(model("Testing the model", prefix="translate English to German: "), "Das Modell zu testen")

test/python/testpipeline/testtext/testentity.py

@@ -0,0 +1,63 @@
+"""
+Entity module tests
+"""
+
+import unittest
+
+from txtai.pipeline import Entity
+
+
+class TestEntity(unittest.TestCase):
+    """
+    Entity tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create entity instance.
+        """
+
+        cls.entity = Entity("dslim/bert-base-NER")
+
+    def testEntity(self):
+        """
+        Test entity
+        """
+
+        # Run entity extraction
+        entities = self.entity("Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg")
+        self.assertEqual([e[0] for e in entities], ["Canada", "Manhattan"])
+
+    def testEntityFlatten(self):
+        """
+        Test entity with flattened output
+        """
+
+        # Test flatten
+        entities = self.entity("Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg", flatten=True)
+        self.assertEqual(entities, ["Canada", "Manhattan"])
+
+        # Test flatten with join
+        entities = self.entity(
+            "Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg", flatten=True, join=True
+        )
+        self.assertEqual(entities, "Canada Manhattan")
+
+    def testEntityTypes(self):
+        """
+        Test entity type filtering
+        """
+
+        # Run entity extraction
+        entities = self.entity("Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg", labels=["PER"])
+        self.assertFalse(entities)
+
+    def testGliner(self):
+        """
+        Test entity pipeline with a GLiNER model
+        """
+
+        entity = Entity("neuml/gliner-bert-tiny")
+        entities = entity("My name is John Smith.", flatten=True)
+        self.assertEqual(entities, ["John Smith"])

test/python/testpipeline/testtext/testlabels.py

@@ -0,0 +1,85 @@
+"""
+Labels module tests
+"""
+
+import unittest
+
+from txtai.pipeline import Labels
+
+
+class TestLabels(unittest.TestCase):
+    """
+    Labels tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create single labels instance.
+        """
+
+        cls.data = [
+            "US tops 5 million confirmed virus cases",
+            "Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
+            "Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
+            "The National Park Service warns against sacrificing slower friends in a bear attack",
+            "Maine man wins $1M from $25 lottery ticket",
+            "Make huge profits without work, earn up to $100,000 a day",
+        ]
+
+        cls.labels = Labels("prajjwal1/bert-medium-mnli")
+
+    def testLabel(self):
+        """
+        Test labels with single text input
+        """
+
+        self.assertEqual(self.labels("This is the best sentence ever", ["positive", "negative"])[0][0], 0)
+
+    def testLabelFlatten(self):
+        """
+        Test labels with single text input, flattened to top text labels
+        """
+
+        self.assertEqual(self.labels("This is the best sentence ever", ["positive", "negative"], flatten=True)[0], "positive")
+
+    def testLabelBatch(self):
+        """
+        Test labels with multiple text inputs
+        """
+
+        results = [l[0][0] for l in self.labels(["This is the best sentence ever", "This is terrible"], ["positive", "negative"])]
+        self.assertEqual(results, [0, 1])
+
+    def testLabelBatchFlatten(self):
+        """
+        Test labels with multiple text inputs, flattened to top text labels
+        """
+
+        results = [l[0] for l in self.labels(["This is the best sentence ever", "This is terrible"], ["positive", "negative"], flatten=True)]
+        self.assertEqual(results, ["positive", "negative"])
+
+    def testLabelFixed(self):
+        """
+        Test labels with a fixed label text classification model
+        """
+
+        labels = Labels(dynamic=False)
+
+        # Get index of "POSITIVE" label
+        index = labels.labels().index("POSITIVE")
+
+        # Verify results
+        self.assertEqual(labels("This is the best sentence ever")[0][0], index)
+        self.assertEqual(labels("This is the best sentence ever", multilabel=True)[0][0], index)
+
+    def testLabelFixedFlatten(self):
+        """
+        Test labels with a fixed label text classification model, flattened to top text labels
+        """
+
+        labels = Labels(dynamic=False)
+
+        # Verify results
+        self.assertEqual(labels("This is the best sentence ever", flatten=True)[0], "POSITIVE")
+        self.assertEqual(labels("This is the best sentence ever", multilabel=True, flatten=True)[0], "POSITIVE")

test/python/testpipeline/testtext/testreranker.py

@@ -0,0 +1,42 @@
+"""
+Reranker module tests
+"""
+
+import unittest
+
+from txtai import Embeddings
+from txtai.pipeline import Reranker, Similarity
+
+
+class TestReranker(unittest.TestCase):
+    """
+    Reranker tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create single labels instance.
+        """
+
+        cls.data = [
+            "US tops 5 million confirmed virus cases",
+            "Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
+            "Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
+            "The National Park Service warns against sacrificing slower friends in a bear attack",
+            "Maine man wins $1M from $25 lottery ticket",
+            "Make huge profits without work, earn up to $100,000 a day",
+        ]
+
+    def testRanker(self):
+        """
+        Test re-ranking pipeline
+        """
+
+        embeddings = Embeddings(content=True)
+        embeddings.index(self.data)
+
+        similarity = Similarity("neuml/colbert-bert-tiny", lateencode=True)
+
+        ranker = Reranker(embeddings, similarity)
+        self.assertEqual(ranker("lottery winner")[0]["id"], "4")

test/python/testpipeline/testtext/testsimilarity.py

@@ -0,0 +1,105 @@
+"""
+Similarity module tests
+"""
+
+import unittest
+
+from txtai.pipeline import Similarity
+
+
+class TestSimilarity(unittest.TestCase):
+    """
+    Similarity tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create single labels instance.
+        """
+
+        cls.data = [
+            "US tops 5 million confirmed virus cases",
+            "Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
+            "Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
+            "The National Park Service warns against sacrificing slower friends in a bear attack",
+            "Maine man wins $1M from $25 lottery ticket",
+            "Make huge profits without work, earn up to $100,000 a day",
+        ]
+
+        cls.similarity = Similarity("prajjwal1/bert-medium-mnli")
+
+    def testCrossEncoder(self):
+        """
+        Test cross-encoder similarity model
+        """
+
+        similarity = Similarity("cross-encoder/ms-marco-MiniLM-L-2-v2", crossencode=True)
+        uid = similarity("Who won the lottery?", self.data)[0][0]
+        self.assertEqual(self.data[uid], self.data[4])
+
+    def testCrossEncoderBatch(self):
+        """
+        Test cross-encoder similarity model with multiple inputs
+        """
+
+        similarity = Similarity("cross-encoder/ms-marco-MiniLM-L-2-v2", crossencode=True)
+        results = [r[0][0] for r in similarity(["Who won the lottery?", "Where did an iceberg collapse?"], self.data)]
+        self.assertEqual(results, [4, 1])
+
+    def testLateEncoder(self):
+        """
+        Test late-encoder similarity model
+        """
+
+        similarity = Similarity("neuml/pylate-bert-tiny", lateencode=True)
+        uid = similarity("Who won the lottery?", self.data)[0][0]
+        self.assertEqual(self.data[uid], self.data[4])
+
+        # Test encode method
+        # pylint: disable=E1101
+        self.assertEqual(similarity.encode(["Who won the lottery?"], "data").shape, (1, 8, 128))
+
+    def testLateEncoderBatch(self):
+        """
+        Test late-encoder similarity model with multiple inputs
+        """
+
+        similarity = Similarity("neuml/colbert-bert-tiny", lateencode=True)
+        results = [r[0][0] for r in similarity(["Who won the lottery?", "Where did an iceberg collapse?"], self.data)]
+        self.assertEqual(results, [4, 1])
+
+    def testSimilarity(self):
+        """
+        Test similarity with single query
+        """
+
+        uid = self.similarity("feel good story", self.data)[0][0]
+        self.assertEqual(self.data[uid], self.data[4])
+
+    def testSimilarityBatch(self):
+        """
+        Test similarity with multiple queries
+        """
+
+        results = [r[0][0] for r in self.similarity(["feel good story", "climate change"], self.data)]
+        self.assertEqual(results, [4, 1])
+
+    def testSimilarityFixed(self):
+        """
+        Test similarity with a fixed label text classification model
+        """
+
+        similarity = Similarity(dynamic=False)
+
+        # Test with query as label text and label id
+        self.assertLessEqual(similarity("negative", ["This is the best sentence ever"])[0][1], 0.1)
+        self.assertLessEqual(similarity("0", ["This is the best sentence ever"])[0][1], 0.1)
+
+    def testSimilarityLong(self):
+        """
+        Test similarity with long text
+        """
+
+        uid = self.similarity("other", ["Very long text " * 1000, "other text"])[0][0]
+        self.assertEqual(uid, 1)

test/python/testpipeline/testtext/testsummary.py

@@ -0,0 +1,64 @@
+"""
+Summary module tests
+"""
+
+import unittest
+
+from txtai.pipeline import Summary
+
+
+class TestSummary(unittest.TestCase):
+    """
+    Summary 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("t5-small")
+
+    def testSummary(self):
+        """
+        Test summarization of text
+        """
+
+        self.assertEqual(self.summary(self.text, minlength=15, maxlength=15), "the field of natural language processing (NLP) is rapidly evolving")
+
+    def testSummaryBatch(self):
+        """
+        Test batch summarization of text
+        """
+
+        summaries = self.summary([self.text, self.text], maxlength=15)
+        self.assertEqual(len(summaries), 2)
+
+    def testSummaryNoLength(self):
+        """
+        Test summary with no max length set
+        """
+
+        self.assertEqual(
+            self.summary(self.text + self.text),
+            "search is the base of many applications. Once data starts to pile up, users want to be able to find it. "
+            + "Large-scale general language models are an exciting new capability allowing us to add amazing functionality quickly "
+            + "with limited compute and people.",
+        )
+
+    def testSummaryShort(self):
+        """
+        Test that summarization is skipped
+        """
+
+        self.assertEqual(self.summary("Text", maxlength=15), "Text")

test/python/testpipeline/testtext/testtranslation.py

@@ -0,0 +1,175 @@
+"""
+Translation module tests
+"""
+
+import unittest
+import time
+
+import requests
+
+from txtai.pipeline import Translation
+
+
+class TestTranslation(unittest.TestCase):
+    """
+    Translation tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create single translation instance.
+        """
+
+        cls.translate = Translation()
+
+        # Preload list of models. Handle HF Hub errors.
+        complete, wait = False, 1
+        while not complete:
+            try:
+                cls.translate.lookup("en", "es")
+                complete = True
+            except requests.exceptions.HTTPError:
+                # Exponential backoff
+                time.sleep(wait)
+
+                # Wait up to 16 seconds
+                wait = min(wait * 2, 16)
+
+    def testDetect(self):
+        """
+        Test language detection
+        """
+
+        test = ["This is a test language detection."]
+        language = self.translate.detect(test)
+
+        self.assertListEqual(language, ["en"])
+
+    def testDetectWithCustomFunc(self):
+        """
+        Test language detection with custom function
+        """
+
+        def dummy_func(text):
+            return ["en" for x in text]
+
+        translate = Translation(langdetect=dummy_func)
+
+        test = ["This is a test language detection."]
+        language = translate.detect(test)
+
+        self.assertListEqual(language, ["en"])
+
+    def testLongTranslation(self):
+        """
+        Test a translation longer than max tokenization length
+        """
+
+        text = "This is a test translation to Spanish. " * 100
+        translation = self.translate(text, "es")
+
+        # Validate translation text
+        self.assertIsNotNone(translation)
+
+    def testM2M100Translation(self):
+        """
+        Test a translation using M2M100 models
+        """
+
+        text = self.translate("This is a test translation to Croatian", "hr")
+
+        # Validate translation text
+        self.assertEqual(text, "Ovo je testni prijevod na hrvatski")
+
+    def testMarianTranslation(self):
+        """
+        Test a translation using Marian models
+        """
+
+        text = "This is a test translation into Spanish"
+        translation = self.translate(text, "es")
+
+        # Validate translation text
+        self.assertEqual(translation, "Esta es una traducción de prueba al español")
+
+        # Validate translation back
+        translation = self.translate(translation, "en")
+        self.assertEqual(translation, text)
+
+    def testNoLang(self):
+        """
+        Test no matching language id
+        """
+
+        self.assertIsNone(self.translate.langid([], "zz"))
+
+    def testNoModel(self):
+        """
+        Test no known available model found
+        """
+
+        self.assertEqual(self.translate.modelpath("zz", "en"), "Helsinki-NLP/opus-mt-mul-en")
+
+    def testNoTranslation(self):
+        """
+        Test translation skipped when text already in destination language
+        """
+
+        text = "This is a test translation to English"
+        translation = self.translate(text, "en")
+
+        # Validate no translation
+        self.assertEqual(text, translation)
+
+    def testShowmodelsChunked(self):
+        """
+        Test a long translation with showmodels flag. When text is chunked
+        by the tokenizer, results should still be properly concatenated as
+        a 3-tuple (translation, language, model) rather than a malformed tuple.
+        """
+
+        text = "This is a test translation to Spanish. " * 100
+        result = self.translate(text, "es", showmodels=True)
+
+        # Result should be a tuple of exactly 3 elements
+        self.assertIsInstance(result, tuple)
+        self.assertEqual(len(result), 3)
+
+        translation, language, modelpath = result
+
+        # Translation should be a single string, not a nested tuple
+        self.assertIsInstance(translation, str)
+        self.assertIsNotNone(translation)
+        self.assertGreater(len(translation), 0)
+
+        # Language and model should be valid strings
+        self.assertEqual(language, "en")
+        self.assertIsInstance(modelpath, str)
+
+    def testTranslationWithShowmodels(self):
+        """
+        Test a translation using Marian models and showmodels flag to return
+        model and language.
+        """
+
+        text = "This is a test translation into Spanish"
+        result = self.translate(text, "es", showmodels=True)
+
+        translation, language, modelpath = result
+        # Validate translation text
+        self.assertEqual(translation, "Esta es una traducción de prueba al español")
+        # Validate detected language
+        self.assertEqual(language, "en")
+        # Validate model
+        self.assertEqual(modelpath, "Helsinki-NLP/opus-mt-en-es")
+
+        # Validate translation back
+        result = self.translate(translation, "en", showmodels=True)
+
+        translation, language, modelpath = result
+        self.assertEqual(translation, text)
+        # Validate detected language
+        self.assertEqual(language, "es")
+        # Validate model
+        self.assertEqual(modelpath, "Helsinki-NLP/opus-mt-es-en")

test/python/testpipeline/testtrain/testonnx.py

@@ -0,0 +1,161 @@
+"""
+ONNX module tests
+"""
+
+import os
+import tempfile
+import unittest
+
+from unittest.mock import patch
+
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.linear_model import LogisticRegression
+from sklearn.pipeline import Pipeline
+
+from txtai.embeddings import Embeddings
+from txtai.models import OnnxModel
+from txtai.pipeline import HFOnnx, HFTrainer, Labels, MLOnnx, Questions
+
+
+class TestOnnx(unittest.TestCase):
+    """
+    ONNX tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create default datasets.
+        """
+
+        cls.data = [{"text": "Dogs", "label": 0}, {"text": "dog", "label": 0}, {"text": "Cats", "label": 1}, {"text": "cat", "label": 1}] * 100
+
+    def testDefault(self):
+        """
+        Test exporting an ONNX model with default parameters
+        """
+
+        # Export model to ONNX, use default parameters
+        onnx = HFOnnx()
+        model = onnx("google/bert_uncased_L-2_H-128_A-2")
+
+        # Validate model has data
+        self.assertGreater(len(model), 0)
+
+        # Validate model device properly works
+        self.assertEqual(OnnxModel(model).device, -1)
+
+    def testClassification(self):
+        """
+        Test exporting a classification model to ONNX and running inference
+        """
+
+        path = "google/bert_uncased_L-2_H-128_A-2"
+
+        trainer = HFTrainer()
+        model, tokenizer = trainer(path, self.data)
+
+        # Output file path
+        output = os.path.join(tempfile.gettempdir(), "onnx")
+
+        # Export model to ONNX
+        onnx = HFOnnx()
+        model = onnx((model, tokenizer), "text-classification", output, True)
+
+        # Test classification
+        labels = Labels((model, path), dynamic=False)
+        self.assertEqual(labels("cat")[0][0], 1)
+
+    @patch("onnxruntime.get_available_providers")
+    @patch("torch.cuda.is_available")
+    def testPooling(self, cuda, providers):
+        """
+        Test exporting a pooling model to ONNX and running inference
+        """
+
+        path = "sentence-transformers/paraphrase-MiniLM-L3-v2"
+
+        # Export model to ONNX
+        onnx = HFOnnx()
+        model = onnx(path, "pooling", quantize=True)
+
+        # Test no CUDA and onnxruntime installed
+        cuda.return_value = False
+        providers.return_value = ["CPUExecutionProvider"]
+
+        embeddings = Embeddings({"path": model, "tokenizer": path})
+        self.assertEqual(embeddings.similarity("animal", ["dog", "book", "rug"])[0][0], 0)
+
+        # Test no CUDA and onnxruntime-gpu installed
+        cuda.return_value = False
+        providers.return_value = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+
+        embeddings = Embeddings({"path": model, "tokenizer": path})
+        self.assertIsNotNone(embeddings)
+
+        # Test CUDA and only onnxruntime installed
+        cuda.return_value = True
+        providers.return_value = ["CPUExecutionProvider"]
+
+        embeddings = Embeddings({"path": model, "tokenizer": path})
+        self.assertIsNotNone(embeddings)
+
+        # Test CUDA and onnxruntime-gpu installed
+        cuda.return_value = True
+        providers.return_value = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+
+        embeddings = Embeddings({"path": model, "tokenizer": path})
+        self.assertIsNotNone(embeddings)
+
+    def testQA(self):
+        """
+        Test exporting a QA model to ONNX and running inference
+        """
+
+        path = "distilbert-base-cased-distilled-squad"
+
+        # Export model to ONNX
+        onnx = HFOnnx()
+        model = onnx(path, "question-answering")
+
+        questions = Questions((model, path))
+        self.assertEqual(questions(["What is the price?"], ["The price is $30"])[0], "$30")
+
+    def testScikit(self):
+        """
+        Test exporting a scikit-learn model to ONNX and running inference
+        """
+
+        # pylint: disable=W0613
+        def tokenizer(inputs, **kwargs):
+            if isinstance(inputs, str):
+                inputs = [inputs]
+
+            return {"input_ids": [[x] for x in inputs]}
+
+        # Train a scikit-learn model
+        model = Pipeline([("tfidf", TfidfVectorizer()), ("lr", LogisticRegression())])
+        model.fit([x["text"] for x in self.data], [x["label"] for x in self.data])
+
+        # Export model to ONNX
+        onnx = MLOnnx()
+        model = onnx(model)
+
+        # Test classification
+        labels = Labels((model, tokenizer), dynamic=False)
+        self.assertEqual(labels("cat")[0][0], 1)
+
+    def testZeroShot(self):
+        """
+        Test exporting a zero shot classification model to ONNX and running inference
+        """
+
+        path = "prajjwal1/bert-medium-mnli"
+
+        # Export model to ONNX
+        onnx = HFOnnx()
+        model = onnx(path, "zero-shot-classification", quantize=True)
+
+        # Test zero shot classification
+        labels = Labels((model, path))
+        self.assertEqual(labels("That is great news", ["negative", "positive"])[0][0], 1)

test/python/testpipeline/testtrain/testquantization.py

@@ -0,0 +1,35 @@
+"""
+Quantization module tests
+"""
+
+import platform
+import unittest
+
+from transformers import AutoModel
+
+from txtai.pipeline import HFModel, HFPipeline
+
+
+class TestQuantization(unittest.TestCase):
+    """
+    Quantization tests.
+    """
+
+    @unittest.skipIf(platform.system() == "Darwin", "Quantized models not supported on macOS")
+    def testModel(self):
+        """
+        Test quantizing a model through HFModel.
+        """
+
+        model = HFModel(quantize=True, gpu=False)
+        model = model.prepare(AutoModel.from_pretrained("google/bert_uncased_L-2_H-128_A-2"))
+        self.assertIsNotNone(model)
+
+    @unittest.skipIf(platform.system() == "Darwin", "Quantized models not supported on macOS")
+    def testPipeline(self):
+        """
+        Test quantizing a model through HFPipeline.
+        """
+
+        pipeline = HFPipeline("text-classification", "google/bert_uncased_L-2_H-128_A-2", True, False)
+        self.assertIsNotNone(pipeline)

test/python/testpipeline/testtrain/testtrainer.py

@@ -0,0 +1,335 @@
+"""
+Trainer module tests
+"""
+
+import os
+import unittest
+import tempfile
+
+import torch
+
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+from txtai.data import Data
+from txtai.pipeline import HFTrainer, Labels, Questions, Sequences
+
+
+class TestTrainer(unittest.TestCase):
+    """
+    Trainer tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create default datasets.
+        """
+
+        cls.data = [{"text": "Dogs", "label": 0}, {"text": "dog", "label": 0}, {"text": "Cats", "label": 1}, {"text": "cat", "label": 1}] * 100
+
+    def testBasic(self):
+        """
+        Test training a model with basic parameters
+        """
+
+        trainer = HFTrainer()
+        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", self.data)
+
+        labels = Labels((model, tokenizer), dynamic=False)
+        self.assertEqual(labels("cat")[0][0], 1)
+
+    def testCLM(self):
+        """
+        Test training a model with causal language modeling
+        """
+
+        trainer = HFTrainer()
+
+        # Test default parameters
+        model, _ = trainer("hf-internal-testing/tiny-random-gpt2", self.data, maxlength=16, task="language-generation")
+        self.assertIsNotNone(model)
+
+        # Test pack merging
+        model, _ = trainer("hf-internal-testing/tiny-random-gpt2", self.data, maxlength=16, task="language-generation", merge="pack")
+        self.assertIsNotNone(model)
+
+        # Test no merging
+        model, _ = trainer("hf-internal-testing/tiny-random-gpt2", self.data, maxlength=16, task="language-generation", merge=None)
+        self.assertIsNotNone(model)
+
+    def testCustom(self):
+        """
+        Test training a model with custom parameters
+        """
+
+        # pylint: disable=E1120
+        model = AutoModelForSequenceClassification.from_pretrained("google/bert_uncased_L-2_H-128_A-2")
+        tokenizer = AutoTokenizer.from_pretrained("google/bert_uncased_L-2_H-128_A-2")
+
+        trainer = HFTrainer()
+        model, tokenizer = trainer(
+            (model, tokenizer),
+            self.data,
+            self.data,
+            columns=("text", "label"),
+            do_eval=True,
+            output_dir=os.path.join(tempfile.gettempdir(), "trainer"),
+        )
+
+        labels = Labels((model, tokenizer), dynamic=False)
+        self.assertEqual(labels("cat")[0][0], 1)
+
+    def testDataFrame(self):
+        """
+        Test training a model with a mock pandas DataFrame
+        """
+
+        class TestDataFrame:
+            """
+            Test DataFrame
+            """
+
+            def __init__(self, data):
+                # Get list of columns
+                self.columns = list(data[0].keys())
+
+                # Build columnar data view
+                self.data = {}
+                for column in self.columns:
+                    self.data[column] = Values([row[column] for row in data])
+
+            def __getitem__(self, column):
+                return self.data[column]
+
+        class Values:
+            """
+            Test values list
+            """
+
+            def __init__(self, values):
+                self.values = list(values)
+
+            def __getitem__(self, index):
+                return self.values[index]
+
+            def unique(self):
+                """
+                Returns a list of unique values.
+
+                Returns:
+                    unique list of values
+                """
+
+                return set(self.values)
+
+        # Mock DataFrame
+        df = TestDataFrame(self.data)
+
+        trainer = HFTrainer()
+        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", df)
+
+        labels = Labels((model, tokenizer), dynamic=False)
+        self.assertEqual(labels("cat")[0][0], 1)
+
+    def testDataset(self):
+        """
+        Test training a model with a mock Hugging Face Dataset
+        """
+
+        class TestDataset(torch.utils.data.Dataset):
+            """
+            Test Dataset
+            """
+
+            def __init__(self, data):
+                self.data = data
+                self.unique = lambda _: [0, 1]
+
+            def __len__(self):
+                return len(self.data)
+
+            def __getitem__(self, index):
+                return self.data[index]
+
+            def column_names(self):
+                """
+                Returns column names for this dataset
+
+                Returns:
+                    list of columns
+                """
+
+                return ["text", "label"]
+
+            # pylint: disable=W0613
+            def map(self, fn, batched, num_proc, remove_columns):
+                """
+                Map each dataset row using fn.
+
+                Args:
+                    fn: function
+                    batched: batch records
+
+                Returns:
+                    updated Dataset
+                """
+
+                self.data = [fn(x) for x in self.data]
+                return self
+
+        ds = TestDataset(self.data)
+
+        trainer = HFTrainer()
+        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", ds)
+
+        labels = Labels((model, tokenizer), dynamic=False)
+        self.assertEqual(labels("cat")[0][0], 1)
+
+    def testEmpty(self):
+        """
+        Test an empty training data object
+        """
+
+        self.assertIsNone(Data(None, None, None).process(None))
+
+    def testKD(self):
+        """
+        Test knowledge distillation
+        """
+
+        # Base model
+        trainer = HFTrainer()
+        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", self.data)
+
+        # Train with knowledge distillation
+        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", self.data, teacher=(model, tokenizer))
+
+        labels = Labels((model, tokenizer), dynamic=False)
+        self.assertEqual(labels("cat")[0][0], 1)
+
+    def testMLM(self):
+        """
+        Test training a model with masked language modeling.
+        """
+
+        trainer = HFTrainer()
+        model, _ = trainer("hf-internal-testing/tiny-random-bert", self.data, task="language-modeling")
+
+        # Test model completed successfully
+        self.assertIsNotNone(model)
+
+    def testMultiLabel(self):
+        """
+        Test training model with labels provided as a list
+        """
+
+        data = []
+        for x in self.data:
+            data.append({"text": x["text"], "label": [0.0, 1.0] if x["label"] else [1.0, 0.0]})
+
+        trainer = HFTrainer()
+        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", data)
+
+        labels = Labels((model, tokenizer), dynamic=False)
+        self.assertEqual(labels("cat")[0][0], 1)
+
+    def testPEFT(self):
+        """
+        Test training a model with causal language modeling and PEFT
+        """
+
+        trainer = HFTrainer()
+        model, _ = trainer(
+            "hf-internal-testing/tiny-random-gpt2",
+            self.data,
+            maxlength=16,
+            task="language-generation",
+            quantize=True,
+            lora=True,
+        )
+
+        # Test model completed successfully
+        self.assertIsNotNone(model)
+
+    def testQA(self):
+        """
+        Test training a QA model
+        """
+
+        # Training data
+        data = [
+            {"question": "What ingredient?", "context": "1 can whole tomatoes", "answers": "tomatoes"},
+            {"question": "What ingredient?", "context": "Crush 1 tomato", "answers": "tomato"},
+            {"question": "What ingredient?", "context": "1 yellow onion", "answers": "onion"},
+            {"question": "What ingredient?", "context": "Unwrap 2 red onions", "answers": "onions"},
+            {"question": "What ingredient?", "context": "1 red pepper", "answers": "pepper"},
+            {"question": "What ingredient?", "context": "Clean 3 red peppers", "answers": "peppers"},
+            {"question": "What ingredient?", "context": "1 clove garlic", "answers": "garlic"},
+            {"question": "What ingredient?", "context": "Unwrap 3 cloves of garlic", "answers": "garlic"},
+            {"question": "What ingredient?", "context": "3 pieces of ginger", "answers": "ginger"},
+            {"question": "What ingredient?", "context": "Peel 1 orange", "answers": "orange"},
+            {"question": "What ingredient?", "context": "1/2 lb beef", "answers": "beef"},
+            {"question": "What ingredient?", "context": "Roast 3 lbs of beef", "answers": "beef"},
+            {"question": "What ingredient?", "context": "1 pack of chicken", "answers": "chicken"},
+            {"question": "What ingredient?", "context": "Forest through the trees", "answers": None},
+        ]
+
+        trainer = HFTrainer()
+        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", data, data, task="question-answering", num_train_epochs=40)
+
+        questions = Questions((model, tokenizer), gpu=True)
+        self.assertTrue("onion" in questions(["What ingredient?"], ["Peel 1 onion"])[0])
+
+    def testRegression(self):
+        """
+        Test training a model with a regression (continuous) output
+        """
+
+        data = []
+        for x in self.data:
+            data.append({"text": x["text"], "label": x["label"] + 0.1})
+
+        trainer = HFTrainer()
+        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", data)
+
+        labels = Labels((model, tokenizer), dynamic=False)
+
+        # Regression tasks return a single entry with the regression output
+        self.assertGreater(labels("cat")[0][1], 0.5)
+
+    def testRTD(self):
+        """
+        Test training a language model with replaced token detection
+        """
+
+        # Save directory
+        output = os.path.join(tempfile.gettempdir(), "trainer.rtd")
+
+        trainer = HFTrainer()
+        model, _ = trainer("hf-internal-testing/tiny-random-electra", self.data, task="token-detection", output_dir=output)
+
+        # Test model completed successfully
+        self.assertIsNotNone(model)
+
+        # Test output directories exist
+        self.assertTrue(os.path.exists(os.path.join(output, "generator")))
+        self.assertTrue(os.path.exists(os.path.join(output, "discriminator")))
+
+    def testSeqSeq(self):
+        """
+        Test training a sequence-sequence model
+        """
+
+        data = [
+            {"source": "Running again", "target": "Sleeping again"},
+            {"source": "Run", "target": "Sleep"},
+            {"source": "running", "target": "sleeping"},
+        ]
+
+        trainer = HFTrainer()
+        model, tokenizer = trainer("t5-small", data, task="sequence-sequence", prefix="translate Run to Sleep: ", learning_rate=1e-3)
+
+        # Run run-sleep translation
+        sequences = Sequences((model, tokenizer))
+        result = sequences("translate Run to Sleep: run")
+        self.assertEqual(result.lower(), "sleep")