wehub-resource-sync
166 lines
4.2 KiB
Python
166 lines
4.2 KiB
Python
"""
|
|
Sparse ANN module tests
|
|
"""
|
|
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
|
|
from unittest.mock import patch
|
|
|
|
from scipy.sparse import random
|
|
from sklearn.preprocessing import normalize
|
|
|
|
from txtai.ann import SparseANNFactory
|
|
|
|
|
|
class TestSparse(unittest.TestCase):
|
|
"""
|
|
Sparse ANN tests.
|
|
"""
|
|
|
|
def testCustomBackend(self):
|
|
"""
|
|
Test resolving a custom backend
|
|
"""
|
|
|
|
self.assertIsNotNone(SparseANNFactory.create({"backend": "txtai.ann.IVFSparse"}))
|
|
|
|
def testCustomBackendNotFound(self):
|
|
"""
|
|
Test resolving an unresolvable backend
|
|
"""
|
|
|
|
with self.assertRaises(ImportError):
|
|
SparseANNFactory.create({"backend": "notfound.ann"})
|
|
|
|
def testIVFSparse(self):
|
|
"""
|
|
Test IVFSparse backend
|
|
"""
|
|
|
|
# Generate test record
|
|
insert = self.generate(500, 30522)
|
|
append = self.generate(500, 30522)
|
|
|
|
# Count of records
|
|
count = insert.shape[0] + append.shape[0]
|
|
|
|
# Create ANN
|
|
path = os.path.join(tempfile.gettempdir(), "ivfsparse")
|
|
ann = SparseANNFactory.create({"backend": "ivfsparse", "ivfsparse": {"nlist": 2, "nprobe": 2, "sample": 1.0}})
|
|
|
|
# Test indexing
|
|
ann.index(insert)
|
|
ann.append(append)
|
|
|
|
# Validate search results
|
|
results = [x[0] for x in ann.search(insert[5], 10)[0]]
|
|
self.assertIn(5, results)
|
|
|
|
# Validate save/load/delete
|
|
ann.save(path)
|
|
ann.load(path)
|
|
|
|
# Validate count
|
|
self.assertEqual(ann.count(), count)
|
|
|
|
# Test delete
|
|
ann.delete([0])
|
|
self.assertEqual(ann.count(), count - 1)
|
|
|
|
# Re-validate search results
|
|
results = [x[0] for x in ann.search(append[0], 10)[0]]
|
|
self.assertIn(insert.shape[0], results)
|
|
|
|
# Close ANN
|
|
ann.close()
|
|
|
|
# Test cluster pruning
|
|
ann = SparseANNFactory.create({"backend": "ivfsparse", "ivfsparse": {"nlist": 15, "nprobe": 1, "sample": 1.0}})
|
|
ann.index(insert)
|
|
self.assertLess(len(ann.blocks), 15)
|
|
ann.close()
|
|
|
|
def testIVFSparseTopnOverLimit(self):
|
|
"""
|
|
Test IVFSparse topn when limit exceeds the number of indexed documents
|
|
"""
|
|
|
|
# Generate a small dataset (5 documents)
|
|
data = self.generate(5, 30522)
|
|
|
|
ann = SparseANNFactory.create({"backend": "ivfsparse"})
|
|
ann.index(data)
|
|
|
|
# Search with limit (10) greater than document count (5)
|
|
results = ann.search(data[0], 10)
|
|
self.assertGreater(len(results[0]), 0)
|
|
|
|
# Batch search with multiple queries exceeding document count
|
|
results = ann.search(data, 10)
|
|
self.assertEqual(len(results), data.shape[0])
|
|
for result in results:
|
|
self.assertGreater(len(result), 0)
|
|
|
|
ann.close()
|
|
|
|
@patch("sqlalchemy.orm.Query.limit")
|
|
def testPGSparse(self, query):
|
|
"""
|
|
Test Sparse Postgres backend
|
|
"""
|
|
|
|
# Generate test record
|
|
data = self.generate(1, 30522)
|
|
|
|
# Mock database query
|
|
query.return_value = [(x, -1.0) for x in range(data.shape[0])]
|
|
|
|
# Create ANN
|
|
path = os.path.join(tempfile.gettempdir(), "pgsparse.sqlite")
|
|
ann = SparseANNFactory.create({"backend": "pgsparse", "dimensions": 30522, "pgsparse": {"url": f"sqlite:///{path}", "schema": "txtai"}})
|
|
|
|
# Test indexing
|
|
ann.index(data)
|
|
ann.append(data)
|
|
|
|
# Validate search results
|
|
self.assertEqual(ann.search(data, 1), [[(0, 1.0)]])
|
|
|
|
# Validate save/load/delete
|
|
ann.save(None)
|
|
ann.load(None)
|
|
|
|
# Validate count
|
|
self.assertEqual(ann.count(), 2)
|
|
|
|
# Test delete
|
|
ann.delete([0])
|
|
self.assertEqual(ann.count(), 1)
|
|
|
|
# Test > 1000 dimensions
|
|
data = random(1, 30522, format="csr", density=0.1)
|
|
ann.index(data)
|
|
self.assertEqual(ann.count(), 1)
|
|
|
|
# Close ANN
|
|
ann.close()
|
|
|
|
def generate(self, m, n):
|
|
"""
|
|
Generates random normalized sparse data.
|
|
|
|
Args:
|
|
m, n: shape of the matrix
|
|
|
|
Returns:
|
|
csr matrix
|
|
"""
|
|
|
|
# Generate random csr matrix
|
|
data = random(m, n, format="csr")
|
|
|
|
# Normalize and return
|
|
return normalize(data)
|