chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,679 @@
|
||||
"""
|
||||
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
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
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 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 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 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 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 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)
|
||||
Reference in New Issue
Block a user