chore: import upstream snapshot with attribution

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)