chore: import upstream snapshot with attribution

test/python/testmodels/testpooling.py

@@ -0,0 +1,123 @@
+"""
+Pooling module tests
+"""
+
+import unittest
+
+from txtai.models import Models, ClsPooling, LastPooling, MeanPooling, PoolingFactory
+
+
+class TestPooling(unittest.TestCase):
+    """
+    Pooling tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Initialize device
+        """
+
+        # Device id
+        cls.device = Models.deviceid(True)
+
+    def testCLS(self):
+        """
+        Test CLS pooling
+        """
+
+        # Test CLS pooling
+        pooling = PoolingFactory.create({"path": "flax-sentence-embeddings/multi-qa_v1-MiniLM-L6-cls_dot", "device": self.device})
+        self.assertEqual(type(pooling), ClsPooling)
+
+        pooling = PoolingFactory.create({"method": "clspooling", "path": "sentence-transformers/nli-mpnet-base-v2", "device": self.device})
+        self.assertEqual(type(pooling), ClsPooling)
+
+        # Test CLS pooling encoding
+        self.assertEqual(pooling.encode(["test"])[0].shape, (768,))
+
+    def testLast(self):
+        """
+        Test last pooling
+        """
+
+        # Test last pooling
+        pooling = PoolingFactory.create({"path": "neuml/bert-tiny-sts-last-pooling", "device": self.device})
+        self.assertEqual(type(pooling), LastPooling)
+
+        pooling = PoolingFactory.create({"method": "lastpooling", "path": "sentence-transformers/nli-mpnet-base-v2", "device": self.device})
+        self.assertEqual(type(pooling), LastPooling)
+
+        # Test last pooling encoding
+        self.assertEqual(pooling.encode(["test"])[0].shape, (768,))
+
+    def testLength(self):
+        """
+        Test pooling with max_seq_length
+        """
+
+        # Test reading max_seq_length parmaeter
+        pooling = PoolingFactory.create({"path": "sentence-transformers/nli-mpnet-base-v2", "device": self.device, "maxlength": True})
+        self.assertEqual(pooling.maxlength, 75)
+
+        # Test specified maxlength
+        pooling = PoolingFactory.create({"path": "sentence-transformers/nli-mpnet-base-v2", "device": self.device, "maxlength": 256})
+        self.assertEqual(pooling.maxlength, 256)
+
+        # Test max_seq_length is ignored when parameter is omitted
+        pooling = PoolingFactory.create({"path": "sentence-transformers/nli-mpnet-base-v2", "device": self.device})
+        self.assertEqual(pooling.maxlength, 512)
+
+        # Test maxlength when max_seq_length not present
+        pooling = PoolingFactory.create({"path": "hf-internal-testing/tiny-random-gpt2", "device": self.device, "maxlength": True})
+        self.assertEqual(pooling.maxlength, 1024)
+
+    def testMean(self):
+        """
+        Test mean pooling
+        """
+
+        # Test mean pooling
+        pooling = PoolingFactory.create({"path": "sentence-transformers/nli-mpnet-base-v2", "device": self.device})
+        self.assertEqual(type(pooling), MeanPooling)
+
+        pooling = PoolingFactory.create(
+            {"method": "meanpooling", "path": "flax-sentence-embeddings/multi-qa_v1-MiniLM-L6-cls_dot", "device": self.device}
+        )
+        self.assertEqual(type(pooling), MeanPooling)
+
+    def testMuvera(self):
+        """
+        Test late pooling with MUVERA fixed dimensional encoding
+        """
+
+        # Test MUVERA encoding
+        for model in ["neuml/colbert-bert-tiny", "neuml/pylate-bert-tiny"]:
+            # Test defaults
+            pooling = PoolingFactory.create({"path": model, "device": self.device})
+            self.assertEqual(pooling.encode(["test"], category="query").shape, (1, 10240))
+
+            # Test custom settings
+            pooling = PoolingFactory.create(
+                {"path": model, "device": self.device, "modelargs": {"muvera": {"repetitions": 5, "hashes": 2, "projection": 8}}}
+            )
+            self.assertEqual(pooling.encode(["test"], category="data").shape, (1, 160))
+
+    def testPrompts(self):
+        """
+        Test instruction prompts
+        """
+
+        # Load model with prompts
+        pooling = PoolingFactory.create({"path": "neuml/bert-tiny-prompts", "device": self.device, "loadprompts": True})
+
+        # Test prompts are prepended
+        self.assertEqual(pooling.preencode(["abc"], "query")[0], "query: abc")
+        self.assertEqual(pooling.preencode(["text"], "data")[0], "document: text")
+
+        # Load model with prompts disabled (default)
+        pooling = PoolingFactory.create({"path": "neuml/bert-tiny-prompts", "device": self.device})
+
+        # Test that prompts are not prepended
+        self.assertEqual(pooling.preencode(["abc"], "query")[0], "abc")
+        self.assertEqual(pooling.preencode(["text"], "data")[0], "text")