chore: import upstream snapshot with attribution

test/python/testvectors/testdense/testcustom.py

@@ -0,0 +1,51 @@
+"""
+Custom module tests
+"""
+
+import os
+import unittest
+
+import numpy as np
+
+from txtai.vectors import VectorsFactory
+
+
+class TestCustom(unittest.TestCase):
+    """
+    Custom vectors tests
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create custom vectors instance.
+        """
+
+        cls.model = VectorsFactory.create({"method": "txtai.vectors.HFVectors", "path": "sentence-transformers/nli-mpnet-base-v2"}, None)
+
+    def testIndex(self):
+        """
+        Test transformers indexing
+        """
+
+        # Generate enough volume to test batching
+        documents = [(x, "This is a test", None) for x in range(1000)]
+
+        ids, dimension, batches, stream = self.model.index(documents)
+
+        self.assertEqual(len(ids), 1000)
+        self.assertEqual(dimension, 768)
+        self.assertEqual(batches, 2)
+        self.assertIsNotNone(os.path.exists(stream))
+
+        # Test shape of serialized embeddings
+        with open(stream, "rb") as queue:
+            self.assertEqual(np.load(queue).shape, (500, 768))
+
+    def testNotFound(self):
+        """
+        Test unresolvable vector backend
+        """
+
+        with self.assertRaises(ImportError):
+            VectorsFactory.create({"method": "notfound.vectors"})

test/python/testvectors/testdense/testexternal.py

@@ -0,0 +1,54 @@
+"""
+External module tests
+"""
+
+import os
+import unittest
+
+import numpy as np
+
+from txtai.vectors import External, VectorsFactory
+
+
+class TestExternal(unittest.TestCase):
+    """
+    External vectors tests
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create External vectors instance.
+        """
+
+        cls.model = VectorsFactory.create({"method": "external"}, None)
+
+    def testIndex(self):
+        """
+        Test indexing with external vectors
+        """
+
+        # Generate dummy data
+        data = np.random.rand(1000, 768).astype(np.float32)
+
+        # Generate enough volume to test batching
+        documents = [(x, data[x], None) for x in range(1000)]
+
+        ids, dimension, batches, stream = self.model.index(documents)
+
+        self.assertEqual(len(ids), 1000)
+        self.assertEqual(dimension, 768)
+        self.assertEqual(batches, 2)
+        self.assertIsNotNone(os.path.exists(stream))
+
+        # Test shape of serialized embeddings
+        with open(stream, "rb") as queue:
+            self.assertEqual(np.load(queue).shape, (500, 768))
+
+    def testMethod(self):
+        """
+        Test method is derived when transform function passed
+        """
+
+        model = VectorsFactory.create({"transform": lambda x: x}, None)
+        self.assertTrue(isinstance(model, External))

test/python/testvectors/testdense/testhuggingface.py

@@ -0,0 +1,99 @@
+"""
+Huggingface module tests
+"""
+
+import os
+import unittest
+
+import numpy as np
+
+from txtai.vectors import VectorsFactory
+
+
+class TestHFVectors(unittest.TestCase):
+    """
+    HFVectors tests
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create HFVectors instance.
+        """
+
+        cls.model = VectorsFactory.create({"path": "sentence-transformers/nli-mpnet-base-v2"}, None)
+
+    def testIndex(self):
+        """
+        Test transformers indexing
+        """
+
+        # Generate enough volume to test batching
+        documents = [(x, "This is a test", None) for x in range(1000)]
+
+        ids, dimension, batches, stream = self.model.index(documents)
+
+        self.assertEqual(len(ids), 1000)
+        self.assertEqual(dimension, 768)
+        self.assertEqual(batches, 2)
+        self.assertIsNotNone(os.path.exists(stream))
+
+        # Test shape of serialized embeddings
+        with open(stream, "rb") as queue:
+            self.assertEqual(np.load(queue).shape, (500, 768))
+
+    def testText(self):
+        """
+        Test transformers text conversion
+        """
+
+        self.model.tokenize = True
+        self.assertEqual(self.model.prepare("Y 123 This is a test!"), "test")
+        self.assertEqual(self.model.prepare(["This", "is", "a", "test"]), "This is a test")
+
+        self.model.tokenize = False
+        self.assertEqual(self.model.prepare("Y 123 This is a test!"), "Y 123 This is a test!")
+        self.assertEqual(self.model.prepare(["This", "is", "a", "test"]), "This is a test")
+
+    def testTransform(self):
+        """
+        Test transformers transform
+        """
+
+        # Sample documents: one where tokenizer changes text and one with no changes to text
+        documents = [(0, "This is a test and has no tokenization", None), (1, "test tokenization", None)]
+
+        # Run with tokenization enabled
+        self.model.tokenize = True
+        embeddings1 = [self.model.transform(d) for d in documents]
+
+        # Run with tokenization disabled
+        self.model.tokenize = False
+        embeddings2 = [self.model.transform(d) for d in documents]
+
+        self.assertFalse(np.array_equal(embeddings1[0], embeddings2[0]))
+        self.assertTrue(np.array_equal(embeddings1[1], embeddings2[1]))
+
+    def testTransformArray(self):
+        """
+        Test transformers skips transforming NumPy arrays
+        """
+
+        # Generate data and run through vector model
+        data1 = np.random.rand(5, 5).astype(np.float32)
+        data2 = self.model.transform((0, data1, None))
+
+        # Test transform method returns original data
+        self.assertTrue(np.array_equal(data1, data2))
+
+    def testTransformLong(self):
+        """
+        Test transformers transform on long text
+        """
+
+        # Sample documents: short text and longer text
+        documents = [(0, "This is long text " * 512, None), (1, "This is short text", None)]
+
+        # Run transform and ensure it completes without errors
+        embeddings = [self.model.transform(d) for d in documents]
+        self.assertIsNotNone(embeddings)

test/python/testvectors/testdense/testlitellm.py

@@ -0,0 +1,83 @@
+"""
+LiteLLM module tests
+"""
+
+import json
+import os
+import unittest
+
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from threading import Thread
+
+import numpy as np
+
+from txtai.vectors import VectorsFactory
+
+
+class RequestHandler(BaseHTTPRequestHandler):
+    """
+    Test HTTP handler.
+    """
+
+    def do_POST(self):
+        """
+        POST request handler.
+        """
+
+        # Generate mock response
+        response = [[0.0] * 768]
+        response = json.dumps(response).encode("utf-8")
+
+        self.send_response(200)
+        self.send_header("content-type", "application/json")
+        self.send_header("content-length", len(response))
+        self.end_headers()
+
+        self.wfile.write(response)
+        self.wfile.flush()
+
+
+class TestLiteLLM(unittest.TestCase):
+    """
+    LiteLLM vectors tests
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create mock http server.
+        """
+
+        cls.httpd = HTTPServer(("127.0.0.1", 8004), RequestHandler)
+
+        server = Thread(target=cls.httpd.serve_forever, daemon=True)
+        server.start()
+
+    @classmethod
+    def tearDownClass(cls):
+        """
+        Shutdown mock http server.
+        """
+
+        cls.httpd.shutdown()
+
+    def testIndex(self):
+        """
+        Test indexing with LiteLLM vectors
+        """
+
+        # LiteLLM vectors instance
+        model = VectorsFactory.create(
+            {"path": "huggingface/sentence-transformers/all-MiniLM-L6-v2", "vectors": {"api_base": "http://127.0.0.1:8004"}}, None
+        )
+
+        ids, dimension, batches, stream = model.index([(0, "test", None)])
+
+        self.assertEqual(len(ids), 1)
+        self.assertEqual(dimension, 768)
+        self.assertEqual(batches, 1)
+        self.assertIsNotNone(os.path.exists(stream))
+
+        # Test shape of serialized embeddings
+        with open(stream, "rb") as queue:
+            self.assertEqual(np.load(queue).shape, (1, 768))

test/python/testvectors/testdense/testlitert.py

@@ -0,0 +1,42 @@
+"""
+LiteRT module tests
+"""
+
+import os
+import unittest
+
+import numpy as np
+
+from txtai.vectors import VectorsFactory
+
+
+class TestLiteRT(unittest.TestCase):
+    """
+    LiteRT vectors tests
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create LiteRT instance.
+        """
+
+        cls.model = VectorsFactory.create(
+            {"path": "neuml/bert-hash-nano-embeddings-litert/bert-hash-nano-embeddings-int4.tflite", "gpu": False}, None
+        )
+
+    def testIndex(self):
+        """
+        Test indexing with LiteRT vectors
+        """
+
+        ids, dimension, batches, stream = self.model.index([(0, "test", None)])
+
+        self.assertEqual(len(ids), 1)
+        self.assertEqual(dimension, 128)
+        self.assertEqual(batches, 1)
+        self.assertIsNotNone(os.path.exists(stream))
+
+        # Test shape of serialized embeddings
+        with open(stream, "rb") as queue:
+            self.assertEqual(np.load(queue).shape, (1, 128))

test/python/testvectors/testdense/testllama.py

@@ -0,0 +1,40 @@
+"""
+Llama module tests
+"""
+
+import os
+import unittest
+
+import numpy as np
+
+from txtai.vectors import VectorsFactory
+
+
+class TestLlamaCpp(unittest.TestCase):
+    """
+    llama.cpp vectors tests
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create LlamaCpp instance.
+        """
+
+        cls.model = VectorsFactory.create({"path": "nomic-ai/nomic-embed-text-v1.5-GGUF/nomic-embed-text-v1.5.Q2_K.gguf"}, None)
+
+    def testIndex(self):
+        """
+        Test indexing with LlamaCpp vectors
+        """
+
+        ids, dimension, batches, stream = self.model.index([(0, "test", None)])
+
+        self.assertEqual(len(ids), 1)
+        self.assertEqual(dimension, 768)
+        self.assertEqual(batches, 1)
+        self.assertIsNotNone(os.path.exists(stream))
+
+        # Test shape of serialized embeddings
+        with open(stream, "rb") as queue:
+            self.assertEqual(np.load(queue).shape, (1, 768))

test/python/testvectors/testdense/testm2v.py

@@ -0,0 +1,40 @@
+"""
+Model2Vec module tests
+"""
+
+import os
+import unittest
+
+import numpy as np
+
+from txtai.vectors import VectorsFactory
+
+
+class TestModel2Vec(unittest.TestCase):
+    """
+    Model2vec vectors tests
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create Model2Vec instance.
+        """
+
+        cls.model = VectorsFactory.create({"path": "minishlab/potion-base-8M"}, None)
+
+    def testIndex(self):
+        """
+        Test indexing with Model2Vec vectors
+        """
+
+        ids, dimension, batches, stream = self.model.index([(0, "test", None)])
+
+        self.assertEqual(len(ids), 1)
+        self.assertEqual(dimension, 256)
+        self.assertEqual(batches, 1)
+        self.assertIsNotNone(os.path.exists(stream))
+
+        # Test shape of serialized embeddings
+        with open(stream, "rb") as queue:
+            self.assertEqual(np.load(queue).shape, (1, 256))

test/python/testvectors/testdense/testsbert.py

@@ -0,0 +1,61 @@
+"""
+Sentence Transformers module tests
+"""
+
+import os
+import platform
+import unittest
+
+from unittest.mock import patch
+
+import numpy as np
+
+from txtai.vectors import VectorsFactory
+
+
+class TestSTVectors(unittest.TestCase):
+    """
+    STVectors tests
+    """
+
+    def testIndex(self):
+        """
+        Test indexing with sentence-transformers vectors
+        """
+
+        model = VectorsFactory.create({"method": "sentence-transformers", "path": "paraphrase-MiniLM-L3-v2"}, None)
+        ids, dimension, batches, stream = model.index([(0, "test", None)])
+
+        self.assertEqual(len(ids), 1)
+        self.assertEqual(dimension, 384)
+        self.assertEqual(batches, 1)
+        self.assertIsNotNone(os.path.exists(stream))
+
+        # Test shape of serialized embeddings
+        with open(stream, "rb") as queue:
+            self.assertEqual(np.load(queue).shape, (1, 384))
+
+    @unittest.skipIf(platform.system() == "Darwin", "Torch memory sharing not supported on macOS")
+    @patch("torch.cuda.device_count")
+    def testMultiGPU(self, count):
+        """
+        Test multiple gpu encoding
+        """
+
+        # Mock accelerator count
+        count.return_value = 2
+
+        model = VectorsFactory.create({"method": "sentence-transformers", "path": "paraphrase-MiniLM-L3-v2", "gpu": "all"}, None)
+        ids, dimension, batches, stream = model.index([(0, "test", None)])
+
+        self.assertEqual(len(ids), 1)
+        self.assertEqual(dimension, 384)
+        self.assertEqual(batches, 1)
+        self.assertIsNotNone(os.path.exists(stream))
+
+        # Test shape of serialized embeddings
+        with open(stream, "rb") as queue:
+            self.assertEqual(np.load(queue).shape, (1, 384))
+
+        # Close the multiprocessing pool
+        model.close()

test/python/testvectors/testdense/testvectors.py

@@ -0,0 +1,68 @@
+"""
+Vectors module tests
+"""
+
+import os
+import tempfile
+import unittest
+
+import numpy as np
+
+from txtai.vectors import Vectors, Recovery
+
+
+class TestVectors(unittest.TestCase):
+    """
+    Vectors tests.
+    """
+
+    def testNotImplemented(self):
+        """
+        Test exceptions for non-implemented methods
+        """
+
+        vectors = Vectors(None, None, None)
+
+        self.assertRaises(NotImplementedError, vectors.load, None)
+        self.assertRaises(NotImplementedError, vectors.encode, None)
+
+    def testNormalize(self):
+        """
+        Test batch normalize and single input normalize are equal
+        """
+
+        vectors = Vectors(None, None, None)
+
+        # Generate data
+        data1 = np.random.rand(5, 5).astype(np.float32)
+        data2 = data1.copy()
+
+        # Keep original data to ensure it changed
+        original = data1.copy()
+
+        # Normalize data
+        vectors.normalize(data1)
+        for x in data2:
+            vectors.normalize(x)
+
+        # Test both data arrays are the same and changed from original
+        self.assertTrue(np.allclose(data1, data2))
+        self.assertFalse(np.allclose(data1, original))
+
+    def testRecovery(self):
+        """
+        Test vectors recovery failure
+        """
+
+        # Checkpoint directory
+        checkpoint = os.path.join(tempfile.gettempdir(), "recovery")
+        os.makedirs(checkpoint, exist_ok=True)
+
+        # Create empty file
+        # pylint: disable=R1732
+        f = open(os.path.join(checkpoint, "id"), "w", encoding="utf-8")
+        f.close()
+
+        # Create the recovery instance with an empty checkpoint file
+        recovery = Recovery(checkpoint, "id", np.load)
+        self.assertIsNone(recovery())

test/python/testvectors/testdense/testwordvectors.py

@@ -0,0 +1,163 @@
+"""
+WordVectors module tests
+"""
+
+import os
+import tempfile
+import unittest
+
+from unittest.mock import patch
+
+import numpy as np
+
+from huggingface_hub.errors import HFValidationError
+from txtai.vectors import VectorsFactory
+from txtai.vectors.dense.words import create, transform, WordVectors
+
+
+class TestWordVectors(unittest.TestCase):
+    """
+    Vectors tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Sets the pretrained model to use
+        """
+
+        # Test with pretrained glove quantized vectors
+        cls.path = "neuml/glove-6B-quantized"
+
+    @patch("os.cpu_count")
+    def testIndex(self, cpucount):
+        """
+        Test word vectors indexing
+        """
+
+        # Mock CPU count
+        cpucount.return_value = 1
+
+        # Generate data
+        documents = [(x, "This is a test", None) for x in range(1000)]
+
+        model = VectorsFactory.create({"path": self.path, "parallel": True}, None)
+
+        ids, dimension, batches, stream = model.index(documents, 1)
+
+        self.assertEqual(len(ids), 1000)
+        self.assertEqual(dimension, 300)
+        self.assertEqual(batches, 1000)
+        self.assertIsNotNone(os.path.exists(stream))
+
+        # Test shape of serialized embeddings
+        with open(stream, "rb") as queue:
+            self.assertEqual(np.load(queue).shape, (1, 300))
+
+    @patch("os.cpu_count")
+    def testIndexBatch(self, cpucount):
+        """
+        Test word vectors indexing with batch size set
+        """
+
+        # Mock CPU count
+        cpucount.return_value = 1
+
+        # Generate data
+        documents = [(x, "This is a test", None) for x in range(1000)]
+
+        model = VectorsFactory.create({"path": self.path, "parallel": True}, None)
+
+        ids, dimension, batches, stream = model.index(documents, 512)
+
+        self.assertEqual(len(ids), 1000)
+        self.assertEqual(dimension, 300)
+        self.assertEqual(batches, 2)
+        self.assertIsNotNone(os.path.exists(stream))
+
+        # Test shape of serialized embeddings
+        with open(stream, "rb") as queue:
+            self.assertEqual(np.load(queue).shape, (512, 300))
+            self.assertEqual(np.load(queue).shape, (488, 300))
+
+    def testIndexSerial(self):
+        """
+        Test word vector indexing in single process mode
+        """
+
+        # Generate data
+        documents = [(x, "This is a test", None) for x in range(1000)]
+
+        model = VectorsFactory.create({"path": self.path, "parallel": False}, None)
+
+        ids, dimension, batches, stream = model.index(documents, 1)
+
+        self.assertEqual(len(ids), 1000)
+        self.assertEqual(dimension, 300)
+        self.assertEqual(batches, 1000)
+        self.assertIsNotNone(os.path.exists(stream))
+
+        # Test shape of serialized embeddings
+        with open(stream, "rb") as queue:
+            self.assertEqual(np.load(queue).shape, (1, 300))
+
+    def testIndexSerialBatch(self):
+        """
+        Test word vector indexing in single process mode with batch size set
+        """
+
+        # Generate data
+        documents = [(x, "This is a test", None) for x in range(1000)]
+
+        model = VectorsFactory.create({"path": self.path, "parallel": False}, None)
+
+        ids, dimension, batches, stream = model.index(documents, 512)
+
+        self.assertEqual(len(ids), 1000)
+        self.assertEqual(dimension, 300)
+        self.assertEqual(batches, 2)
+        self.assertIsNotNone(os.path.exists(stream))
+
+        # Test shape of serialized embeddings
+        with open(stream, "rb") as queue:
+            self.assertEqual(np.load(queue).shape, (512, 300))
+            self.assertEqual(np.load(queue).shape, (488, 300))
+
+    def testLookup(self):
+        """
+        Test word vector lookup
+        """
+
+        model = VectorsFactory.create({"path": self.path}, None)
+        self.assertEqual(model.lookup(["txtai", "embeddings", "sentence"]).shape, (3, 300))
+
+    def testMultiprocess(self):
+        """
+        Test multiprocess helper methods
+        """
+
+        create({"path": self.path}, None)
+
+        uid, vector = transform((0, "test", None))
+        self.assertEqual(uid, 0)
+        self.assertEqual(vector.shape, (300,))
+
+    def testNoExist(self):
+        """
+        Test loading model that doesn't exist
+        """
+
+        # Test non-existent local path raises an exception
+        with self.assertRaises((IOError, HFValidationError)):
+            VectorsFactory.create({"method": "words", "path": os.path.join(tempfile.gettempdir(), "noexist")}, None)
+
+        # Test non-existent hub path is handled properly
+        self.assertFalse(WordVectors.ismodel("invalid/user/noexist"))
+
+    def testTransform(self):
+        """
+        Test word vector transform
+        """
+
+        model = VectorsFactory.create({"path": self.path}, None)
+        self.assertEqual(len(model.transform((None, ["txtai"], None))), 300)