txtai/test/python/testembeddings.py

"""
Embeddings module tests
"""

import json
import os
import tempfile
import unittest

from unittest.mock import patch

import numpy as np

from txtai.embeddings import Embeddings, Reducer
from txtai.serialize import SerializeFactory


class TestEmbeddings(unittest.TestCase):
    """
    Embeddings tests.
    """

    @classmethod
    def setUpClass(cls):
        """
        Initialize test data.
        """

        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",
        ]

        # Create embeddings model, backed by sentence-transformers & transformers
        cls.embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2"})

    @classmethod
    def tearDownClass(cls):
        """
        Cleanup data.
        """

        if cls.embeddings:
            cls.embeddings.close()

    def testArchive(self):
        """
        Test saving and loading a compressed index archive
        """

        # Index data with sparse keyword vectors and content
        embeddings = Embeddings({"keyword": True, "content": True})
        embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])

        # Generate temp file paths
        archive = os.path.join(tempfile.gettempdir(), "embeddings.archive.tar.gz")
        index = os.path.join(tempfile.gettempdir(), "embeddings.archive")

        # Save to an archive, load it and save back to the same archive
        embeddings.save(archive)
        embeddings.close()

        embeddings = Embeddings().load(archive)
        self.assertTrue(embeddings.exists(archive))
        embeddings.save(archive)

        # Save to a directory, update data and save to the archive again
        embeddings.save(index)
        embeddings.upsert([(0, "Feel good story: baby panda born", None)])
        embeddings.save(archive)
        embeddings.close()

        # Directory has the original data
        embeddings = Embeddings().load(index)
        self.assertEqual(embeddings.count(), len(self.data))
        self.assertEqual(embeddings.search("lottery ticket", 1)[0]["id"], "4")
        embeddings.close()

        # Archive has the updated data
        embeddings = Embeddings().load(archive)
        self.assertEqual(embeddings.count(), len(self.data))
        self.assertEqual(embeddings.search("feel good story", 1)[0]["id"], "0")
        embeddings.close()

    def testAutoId(self):
        """
        Test auto id generation
        """

        # Default sequence id
        embeddings = Embeddings()
        embeddings.index(self.data)

        uid = embeddings.search(self.data[4], 1)[0][0]
        self.assertEqual(uid, 4)

        # UUID
        embeddings = Embeddings(autoid="uuid4")
        embeddings.index(self.data)

        uid = embeddings.search(self.data[4], 1)[0][0]
        self.assertEqual(len(uid), 36)

    def testColumns(self):
        """
        Test custom text/object columns
        """

        embeddings = Embeddings({"keyword": True, "columns": {"text": "value"}})
        data = [{"value": x} for x in self.data]
        embeddings.index([(uid, text, None) for uid, text in enumerate(data)])

        # Run search
        uid = embeddings.search("lottery", 1)[0][0]
        self.assertEqual(uid, 4)

    def testContext(self):
        """
        Test embeddings context manager
        """

        # Generate temp file path
        index = os.path.join(tempfile.gettempdir(), "embeddings.context")

        with Embeddings() as embeddings:
            embeddings.index(self.data)
            embeddings.save(index)

        with Embeddings().load(index) as embeddings:
            uid = embeddings.search(self.data[4], 1)[0][0]
            self.assertEqual(uid, 4)

    def testDefaults(self):
        """
        Test default configuration
        """

        # Run index with no config which will fall back to default configuration
        embeddings = Embeddings()
        embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])

        self.assertEqual(embeddings.count(), 6)

    def testDelete(self):
        """
        Test delete
        """

        # Create an index for the list of text
        self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])

        # Delete best match
        self.embeddings.delete([4])

        # Search for best match
        uid = self.embeddings.search("feel good story", 1)[0][0]

        self.assertEqual(self.embeddings.count(), 5)
        self.assertEqual(uid, 5)

    def testDense(self):
        """
        Test dense alias
        """

        # Dense flag is an alias for path
        embeddings = Embeddings(dense="sentence-transformers/nli-mpnet-base-v2")
        embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])

        self.assertEqual(embeddings.count(), 6)

    def testEmpty(self):
        """
        Test empty index
        """

        # Test search against empty index
        embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2"})
        self.assertEqual(embeddings.search("test"), [])

        # Test index with no data
        embeddings.index([])
        self.assertIsNone(embeddings.ann)

        # Test upsert with no data
        embeddings.index([(0, "this is a test", None)])
        embeddings.upsert([])
        self.assertIsNotNone(embeddings.ann)

    def testEmptySave(self):
        """
        Test saving an empty index
        """

        for content in [False, True]:
            # Keyword index, no data ever indexed
            embeddings = Embeddings({"keyword": True, "content": content})
            embeddings.index([])

            # Generate temp file path
            index = os.path.join(tempfile.gettempdir(), f"embeddings.emptysave.{content}")

            # Test save/load
            embeddings.save(index)
            embeddings.load(index)

            # Validate index is still empty
            self.assertEqual(embeddings.count(), 0)
            self.assertEqual(embeddings.search("test"), [])

    def testEmptyString(self):
        """
        Test empty string indexing
        """

        # Test empty string
        self.embeddings.index([(0, "", None)])
        self.assertTrue(self.embeddings.search("test"))

        # Test empty string with dict
        self.embeddings.index([(0, {"text": ""}, None)])
        self.assertTrue(self.embeddings.search("test"))

    def testNoneText(self):
        """
        Test dict documents where the text field is present but set to None
        """

        # None text mixed with real text
        self.embeddings.index([(0, {"text": "alpha content"}, None), (1, {"text": None}, None), (2, {"text": "gamma content"}, None)])

        # Only the two documents with real text should be indexed
        self.assertEqual(self.embeddings.count(), 2)
        self.assertEqual(self.embeddings.search("alpha content", 1)[0][0], 0)
        self.assertEqual(self.embeddings.search("gamma content", 1)[0][0], 2)

    def testExternal(self):
        """
        Test embeddings backed by external vectors
        """

        def transform(data):
            embeddings = []
            for text in data:
                # Create dummy embedding using sum and mean of character ordinals
                ordinals = [ord(c) for c in text]
                embeddings.append(np.array([sum(ordinals), np.mean(ordinals)]))

            return embeddings

        # Index data using simple embeddings transform method
        embeddings = Embeddings({"method": "external", "transform": transform})
        embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])

        # Run search
        uid = embeddings.search(self.data[4], 1)[0][0]
        self.assertEqual(uid, 4)

    def testExternalPrecomputed(self):
        """
        Test embeddings backed by external pre-computed vectors
        """

        # Test with no transform function
        data = np.random.rand(5, 10).astype(np.float32)

        embeddings = Embeddings({"method": "external"})
        embeddings.index([(uid, row, None) for uid, row in enumerate(data)])

        # Run search
        uid = embeddings.search(data[4], 1)[0][0]
        self.assertEqual(uid, 4)

    def testHybrid(self):
        """
        Test hybrid search
        """

        # Build data array
        data = [(uid, text, None) for uid, text in enumerate(self.data)]

        # Index data with sparse + dense vectors
        embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "hybrid": True})
        embeddings.index(data)

        # Run search
        uid = embeddings.search("feel good story", 1)[0][0]
        self.assertEqual(uid, 4)

        # Generate temp file path
        index = os.path.join(tempfile.gettempdir(), "embeddings.hybrid")

        # Test load/save
        embeddings.save(index)
        embeddings.load(index)

        # Run search
        uid = embeddings.search("feel good story", 1)[0][0]
        self.assertEqual(uid, 4)

        # Index data with sparse + dense vectors and unnormalized scores
        embeddings.config["scoring"]["normalize"] = False
        embeddings.index(data)

        # Run search
        uid = embeddings.search("feel good story", 1)[0][0]
        self.assertEqual(uid, 4)

        # Index data with sparse + dense vectors and bb25 normalization
        embeddings.config["scoring"]["normalize"] = "bb25"
        embeddings.index(data)

        # Run search
        uid = embeddings.search("canada intact iceberg a", 1)[0][0]
        self.assertEqual(uid, 1)

        # Test upsert
        data[0] = (0, "Feel good story: baby panda born", None)
        embeddings.upsert([data[0]])

        uid = embeddings.search("feel good story", 1)[0][0]
        self.assertEqual(uid, 0)

    def testIds(self):
        """
        Test legacy config ids loading
        """

        # Create an index for the list of text
        self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])

        # Generate temp file path
        index = os.path.join(tempfile.gettempdir(), "embeddings.ids")

        # Save index
        self.embeddings.save(index)

        # Set ids on config to simulate legacy ids format
        with open(f"{index}/config.json", "r", encoding="utf-8") as handle:
            config = json.load(handle)
            config["ids"] = list(range(len(self.data)))

        with open(f"{index}/config.json", "w", encoding="utf-8") as handle:
            json.dump(config, handle, default=str, indent=2)

        # Reload index
        self.embeddings.load(index)

        # Run search
        uid = self.embeddings.search("feel good story", 1)[0][0]
        self.assertEqual(uid, 4)

        # Check that ids is not in config
        self.assertTrue("ids" not in self.embeddings.config)

    @patch.dict(os.environ, {"ALLOW_PICKLE": "True"})
    def testIdsPickle(self):
        """
        Test legacy pickle ids
        """

        # Create an index for the list of text
        self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])

        # Generate temp file path
        index = os.path.join(tempfile.gettempdir(), "embeddings.idspickle")

        # Save index
        self.embeddings.save(index)

        # Create ids as pickle
        path = os.path.join(tempfile.gettempdir(), "embeddings.idspickle", "ids")
        serializer = SerializeFactory.create("pickle", allowpickle=True)
        serializer.save(self.embeddings.ids.ids, path)

        with self.assertWarns(RuntimeWarning):
            self.embeddings.load(index)

        # Run search
        uid = self.embeddings.search("feel good story", 1)[0][0]
        self.assertEqual(uid, 4)

    def testIndex(self):
        """
        Test index
        """

        # Create an index for the list of text
        self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])

        # Search for best match
        uid = self.embeddings.search("feel good story", 1)[0][0]

        self.assertEqual(uid, 4)

    def testKeyword(self):
        """
        Test keyword only (sparse) search
        """

        # Build data array
        data = [(uid, text, None) for uid, text in enumerate(self.data)]

        # Index data with sparse keyword vectors
        embeddings = Embeddings({"keyword": True})
        embeddings.index(data)

        # Run search
        uid = embeddings.search("lottery ticket", 1)[0][0]
        self.assertEqual(uid, 4)

        # Test count method
        self.assertEqual(embeddings.count(), len(data))

        # Generate temp file path
        index = os.path.join(tempfile.gettempdir(), "embeddings.keyword")

        # Test load/save
        embeddings.save(index)
        embeddings.load(index)

        # Run search
        uid = embeddings.search("lottery ticket", 1)[0][0]
        self.assertEqual(uid, 4)

        # Update data
        data[0] = (0, "Feel good story: baby panda born", None)
        embeddings.upsert([data[0]])

        # Search for best match
        uid = embeddings.search("feel good story", 1)[0][0]
        self.assertEqual(uid, 0)

    def testLimitBindParameter(self):
        """
        Test a LIMIT bind parameter in a content SQL query
        """

        embeddings = Embeddings({"keyword": True, "content": True})
        embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])

        # A ":n" LIMIT bind parameter must not crash on the candidate-count parse (str vs int)
        results = embeddings.search("select id from txtai order by id limit :n", parameters={"n": 2})
        self.assertEqual(len(results), 2)

        # A plain integer limit still works
        self.assertEqual(len(embeddings.search("select id from txtai order by id limit 3")), 3)

    def testOffset(self):
        """
        Test an OFFSET clause with a similar clause
        """

        # Longer documents score lower, which sorts the results by id
        embeddings = Embeddings({"keyword": True, "content": True})
        embeddings.index([(uid, "apple " + "pie " * uid, None) for uid in range(10)])

        # First page
        results = embeddings.search("select id from txtai where similar('apple') limit 3")
        self.assertEqual([x["id"] for x in results], ["0", "1", "2"])

        # Second page must not be empty
        results = embeddings.search("select id from txtai where similar('apple') limit 3 offset 3")
        self.assertEqual([x["id"] for x in results], ["3", "4", "5"])

        # An explicit ORDER BY score DESC pages the same way
        results = embeddings.search("select id from txtai where similar('apple') ORDER BY score DESC limit 3 offset 3")
        self.assertEqual([x["id"] for x in results], ["3", "4", "5"])

        # Selecting the score column pages the same way
        results = embeddings.search("select id, score from txtai where similar('apple') limit 3 offset 3")
        self.assertEqual([x["id"] for x in results], ["3", "4", "5"])

        # A tie-breaker after score DESC pages the same way
        results = embeddings.search("select id from txtai where similar('apple') order by score desc, id limit 3 offset 3")
        self.assertEqual([x["id"] for x in results], ["3", "4", "5"])

        # A bind parameter offset is not a number, so it does not add candidates
        results = embeddings.search("select id from txtai where similar('apple') limit 3 offset :n", parameters={"n": 0})
        self.assertEqual([x["id"] for x in results], ["0", "1", "2"])

        # Each query in a batch gets its own page
        queries = [f"select id from txtai where similar('apple') limit 3{offset}" for offset in [" offset 3", " offset 6", ""]]
        results = embeddings.batchsearch(queries)
        self.assertEqual([[x["id"] for x in result] for result in results], [["3", "4", "5"], ["6", "7", "8"], ["0", "1", "2"]])

    def testOffsetOrderBy(self):
        """
        Test an OFFSET clause with a similar clause and an ORDER BY other than score DESC
        """

        embeddings = Embeddings({"keyword": True, "content": True})
        embeddings.index([(uid, {"text": "apple " + "pie " * uid, "n": uid}, None) for uid in range(40)])

        # Every page sorts the same 30 candidates, the default of 10x the query limit
        for orderby in ["n desc", "n desc, score desc", "score asc"]:
            query = f"select n from txtai where similar('apple') and n >= 0 order by {orderby} limit 3"
            self.assertEqual([x["n"] for x in embeddings.search(query)], [29, 28, 27])
            self.assertEqual([x["n"] for x in embeddings.search(f"{query} offset 3")], [26, 25, 24])

        # A select alias named score replaces the similarity score in the order
        for select, orderby in [("n as score", " order by score desc"), ("n as score", ""), ("n score", ""), ('n as "score"', "")]:
            query = f"select {select} from txtai where similar('apple') and n >= 0{orderby} limit 3"
            self.assertEqual([x["score"] for x in embeddings.search(query)], [29, 28, 27])
            self.assertEqual([x["score"] for x in embeddings.search(f"{query} offset 3")], [26, 25, 24])

    def testQuantize(self):
        """
        Test scalar quantization
        """

        for ann in ["faiss", "numpy", "torch"]:
            # Index data with 1-bit scalar quantization
            embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "quantize": 1, "backend": ann})
            embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])

            # Search for best match
            uid = embeddings.search("feel good story", 1)[0][0]
            self.assertEqual(uid, 4)

    def testReducer(self):
        """
        Test reducer model
        """

        # Test model with single PCA component
        data = np.random.rand(5, 5).astype(np.float32)
        reducer = Reducer(data, 1)

        # Generate query and keep original data to ensure it changes
        query = np.random.rand(1, 5).astype(np.float32)
        original = query.copy()

        # Run test
        reducer(query)
        self.assertFalse(np.array_equal(query, original))

        # Test model with multiple PCA components
        reducer = Reducer(data, 3)

        # Generate query and keep original data to ensure it changes
        query = np.random.rand(5).astype(np.float32)
        original = query.copy()

        # Run test
        reducer(query)
        self.assertFalse(np.array_equal(query, original))

    @patch.dict(os.environ, {"ALLOW_PICKLE": "True"})
    def testReducerLegacy(self):
        """
        Test reducer model with legacy model format
        """

        # Test model with single PCA component
        data = np.random.rand(5, 5).astype(np.float32)
        reducer = Reducer(data, 1)

        # Save legacy format
        path = os.path.join(tempfile.gettempdir(), "reducer")
        serializer = SerializeFactory.create("pickle", allowpickle=True)
        serializer.save(reducer.model, path)

        # Load legacy format
        reducer = Reducer()
        reducer.load(path)

        # Generate query and keep original data to ensure it changes
        query = np.random.rand(1, 5).astype(np.float32)
        original = query.copy()

        # Run test
        reducer(query)
        self.assertFalse(np.array_equal(query, original))

    def testSave(self):
        """
        Test save
        """

        # Create an index for the list of text
        self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])

        # Generate temp file path
        index = os.path.join(tempfile.gettempdir(), "embeddings.base")

        self.embeddings.save(index)
        self.embeddings.load(index)

        # Search for best match
        uid = self.embeddings.search("feel good story", 1)[0][0]

        self.assertEqual(uid, 4)

        # Test offsets still work after save/load
        self.embeddings.upsert([(0, "Looking out into the dreadful abyss", None)])
        self.assertEqual(self.embeddings.count(), len(self.data))

    def testShortcuts(self):
        """
        Test embeddings creation shortcuts
        """

        tests = [
            ({"keyword": True}, ["scoring"]),
            ({"keyword": "sif"}, ["scoring"]),
            ({"sparse": True}, ["scoring"]),
            ({"dense": True}, ["ann"]),
            ({"hybrid": True}, ["ann", "scoring"]),
            ({"hybrid": "tfidf"}, ["ann", "scoring"]),
            ({"hybrid": "sparse"}, ["ann", "scoring"]),
            ({"graph": True}, ["graph"]),
        ]

        for config, checks in tests:
            embeddings = Embeddings(config)
            embeddings.index(["test"])

            for attr in checks:
                self.assertIsNotNone(getattr(embeddings, attr))

    def testSimilarity(self):
        """
        Test similarity
        """

        # Get best matching id
        uid = self.embeddings.similarity("feel good story", self.data)[0][0]

        self.assertEqual(uid, 4)

    def testSparse(self):
        """
        Test sparse vector search
        """

        # Build data array
        data = [(uid, text, None) for uid, text in enumerate(self.data)]

        # Index data with sparse vectors
        embeddings = Embeddings({"sparse": "sparse-encoder-testing/splade-bert-tiny-nq"})
        embeddings.index(data)

        # Run search
        uid = embeddings.search("lottery ticket", 1)[0][0]
        self.assertEqual(uid, 4)

        # Test count method
        self.assertEqual(embeddings.count(), len(data))

        # Generate temp file path
        index = os.path.join(tempfile.gettempdir(), "embeddings.sparse")

        # Test load/save
        embeddings.save(index)
        embeddings.load(index)

        # Run search
        uid = embeddings.search("lottery ticket", 1)[0][0]
        self.assertEqual(uid, 4)

        # Test similarity
        uid = embeddings.similarity("lottery ticket", self.data)[0][0]
        self.assertEqual(uid, 4)

        # Update data
        data[0] = (0, "Feel good story: baby panda born", None)
        embeddings.upsert([data[0]])

        # Search for best match
        uid = embeddings.search("feel good story", 1)[0][0]
        self.assertEqual(uid, 0)

    def testSubindex(self):
        """
        Test subindex
        """

        # Build data array
        data = [(uid, text, None) for uid, text in enumerate(self.data)]

        # Disable top-level indexing and create subindex
        embeddings = Embeddings({"defaults": False, "indexes": {"index1": {"path": "sentence-transformers/nli-mpnet-base-v2"}}})
        embeddings.index(data)

        # Test transform
        self.assertEqual(embeddings.transform("feel good story").shape, (768,))
        self.assertEqual(embeddings.transform("feel good story", index="index1").shape, (768,))
        with self.assertRaises(KeyError):
            embeddings.transform("feel good story", index="index2")

        # Run search
        uid = embeddings.search("feel good story", 1)[0][0]
        self.assertEqual(uid, 4)

        # Generate temp file path
        index = os.path.join(tempfile.gettempdir(), "embeddings.subindex")

        # Test load/save
        embeddings.save(index)
        embeddings.load(index)

        # Run search
        uid = embeddings.search("feel good story", 1)[0][0]
        self.assertEqual(uid, 4)

        # Update data
        data[0] = (0, "Feel good story: baby panda born", None)
        embeddings.upsert([data[0]])

        # Search for best match
        uid = embeddings.search("feel good story", 10)[0][0]
        self.assertEqual(uid, 0)

        # Check missing text is set to id when top-level indexing is disabled
        embeddings.upsert([(embeddings.count(), {"content": "empty text"}, None)])
        uid = embeddings.search(f"{embeddings.count() - 1}", 1)[0][0]
        self.assertEqual(uid, embeddings.count() - 1)

        # Close embeddings
        embeddings.close()

    def testSubindexVectors(self):
        """
        Test subindexes with the same model path and different vectors settings
        """

        # Build data array
        data = [(uid, text, None) for uid, text in enumerate(self.data)]

        # Subindexes share a models cache. The same model path with different vectors settings must load separate models.
        path = "neuml/colbert-bert-tiny"
        embeddings = Embeddings(
            {
                "defaults": False,
                "indexes": {
                    "index1": {"path": path, "vectors": {"muvera": {"repetitions": 1}}},
                    "index2": {"path": path, "vectors": {"muvera": {"repetitions": 2}}},
                },
            }
        )
        embeddings.index(data)

        # MUVERA output dimensions = repetitions * 2^5 * 16
        self.assertEqual(embeddings.transform("feel good story", index="index1").shape, (512,))
        self.assertEqual(embeddings.transform("feel good story", index="index2").shape, (1024,))
        self.assertEqual(embeddings.indexes["index1"].config["dimensions"], 512)
        self.assertEqual(embeddings.indexes["index2"].config["dimensions"], 1024)
        self.assertIsNot(embeddings.indexes["index1"].model.model, embeddings.indexes["index2"].model.model)

        # Close embeddings
        embeddings.close()

        # Per-encode settings share the same loaded model
        embeddings = Embeddings(
            {
                "defaults": False,
                "indexes": {
                    "index1": {"path": path, "maxlength": 32, "vectors": {"muvera": {"repetitions": 1}}},
                    "index2": {"path": path, "maxlength": 64, "vectors": {"muvera": {"repetitions": 1}}},
                },
            }
        )
        embeddings.index(data)
        self.assertIs(embeddings.indexes["index1"].model.model, embeddings.indexes["index2"].model.model)

        # Close embeddings
        embeddings.close()

    def testTerms(self):
        """
        Test extracting keyword terms from queries
        """

        embeddings = Embeddings({"keyword": True, "content": True})
        embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])

        # Plain keyword query is returned unchanged
        self.assertEqual(embeddings.terms("feel good story"), "feel good story")

        # SQL query with a similar() clause extracts its terms
        self.assertEqual(embeddings.terms("select id, text from txtai where similar('feel good story')"), "feel good story")

        # SQL query without a similar() clause has no keyword terms (previously raised KeyError)
        self.assertEqual(embeddings.terms("select id, text from txtai where id = 1"), None)

    def testTruncate(self):
        """
        Test dimensionality truncation
        """

        # Truncate vectors to a specified number of dimensions
        embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "dimensionality": 750, "vectors": {"revision": "main"}})
        embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])

        # Search for best match
        uid = embeddings.search("feel good story", 1)[0][0]
        self.assertEqual(uid, 4)

    def testUpsert(self):
        """
        Test upsert
        """

        # Build data array
        data = [(uid, text, None) for uid, text in enumerate(self.data)]

        # Reset embeddings for test
        self.embeddings.ann = None
        self.embeddings.ids = None

        # Create an index for the list of text
        self.embeddings.upsert(data)

        # Update data
        data[0] = (0, "Feel good story: baby panda born", None)
        self.embeddings.upsert([data[0]])

        # Search for best match
        uid = self.embeddings.search("feel good story", 1)[0][0]

        self.assertEqual(uid, 0)

    @patch("os.cpu_count")
    def testWords(self, cpucount):
        """
        Test embeddings backed by word vectors
        """

        # Mock CPU count
        cpucount.return_value = 1

        # Create dataset
        data = [(x, row.split(), None) for x, row in enumerate(self.data)]

        # Create embeddings model, backed by word vectors
        embeddings = Embeddings({"path": "neuml/glove-6B-quantized", "scoring": "bm25", "pca": 3, "quantize": True})

        # Call scoring and index methods
        embeddings.score(data)
        embeddings.index(data)

        # Test search
        self.assertIsNotNone(embeddings.search("win", 1))

        # Generate temp file path
        index = os.path.join(tempfile.gettempdir(), "embeddings.wordvectors")

        # Test save/load
        embeddings.save(index)
        embeddings.load(index)

        # Test search
        self.assertIsNotNone(embeddings.search("win", 1))

    @patch("os.cpu_count")
    def testWordsUpsert(self, cpucount):
        """
        Test embeddings backed by word vectors with upserts
        """

        # Mock CPU count
        cpucount.return_value = 1

        # Create dataset
        data = [(x, row.split(), None) for x, row in enumerate(self.data)]

        # Create embeddings model, backed by word vectors
        embeddings = Embeddings({"path": "neuml/glove-6B/model.sqlite", "scoring": "bm25", "pca": 3})

        # Call scoring and index methods
        embeddings.score(data)
        embeddings.index(data)

        # Now upsert and override record
        data = [(0, "win win", None)]

        # Update scoring and run upsert
        embeddings.score(data)
        embeddings.upsert(data)

        # Test search after upsert
        uid = embeddings.search("win", 1)[0][0]
        self.assertEqual(uid, 0)