From e765be3f6c772ac75aff1ddf4bc8e304b45f9d2c Mon Sep 17 00:00:00 2001 From: davidmezzetti <561939+davidmezzetti@users.noreply.github.com> Date: Mon, 14 Jul 2025 07:48:23 -0400 Subject: [PATCH] Refactor of ann and vectors packages for sparse vectors. Closes #930. Closes #934. Closes #935. Closes #936. Closes #937. --- docs/embeddings/configuration/scoring.md | 23 +- setup.py | 7 +- src/python/txtai/ann/__init__.py | 9 +- src/python/txtai/ann/dense/__init__.py | 11 + src/python/txtai/ann/{ => dense}/annoy.py | 2 +- src/python/txtai/ann/{ => dense}/factory.py | 2 +- src/python/txtai/ann/{ => dense}/faiss.py | 2 +- src/python/txtai/ann/{ => dense}/hnsw.py | 2 +- src/python/txtai/ann/{ => dense}/numpy.py | 4 +- src/python/txtai/ann/{ => dense}/pgvector.py | 26 +- src/python/txtai/ann/{ => dense}/sqlite.py | 2 +- src/python/txtai/ann/{ => dense}/torch.py | 0 src/python/txtai/ann/sparse/__init__.py | 7 + src/python/txtai/ann/sparse/factory.py | 61 +++++ .../ivf.py => ann/sparse/ivfsparse.py} | 212 +++++----------- src/python/txtai/ann/sparse/pgsparse.py | 44 ++++ src/python/txtai/embeddings/base.py | 82 +++++-- src/python/txtai/embeddings/index/indexes.py | 19 +- .../txtai/embeddings/index/transform.py | 32 +-- .../txtai/pipeline/audio/texttospeech.py | 2 +- src/python/txtai/scoring/__init__.py | 1 - src/python/txtai/scoring/base.py | 32 ++- src/python/txtai/scoring/factory.py | 20 +- src/python/txtai/scoring/pgtext.py | 4 +- src/python/txtai/scoring/sparse.py | 232 ++++++++---------- src/python/txtai/scoring/tfidf.py | 7 +- src/python/txtai/util/__init__.py | 1 + src/python/txtai/util/sparsearray.py | 62 +++++ src/python/txtai/vectors/__init__.py | 12 +- src/python/txtai/vectors/base.py | 103 +++++++- src/python/txtai/vectors/dense/__init__.py | 12 + .../txtai/vectors/{ => dense}/external.py | 6 +- .../txtai/vectors/{ => dense}/factory.py | 4 +- .../txtai/vectors/{ => dense}/huggingface.py | 6 +- .../txtai/vectors/{ => dense}/litellm.py | 4 +- src/python/txtai/vectors/{ => dense}/llama.py | 4 +- src/python/txtai/vectors/{ => dense}/m2v.py | 4 +- src/python/txtai/vectors/{ => dense}/sbert.py | 37 ++- src/python/txtai/vectors/{ => dense}/words.py | 10 +- src/python/txtai/vectors/recovery.py | 9 +- src/python/txtai/vectors/sparse/__init__.py | 7 + src/python/txtai/vectors/sparse/base.py | 76 ++++++ src/python/txtai/vectors/sparse/factory.py | 55 +++++ src/python/txtai/vectors/sparse/sbert.py | 30 +++ test/python/testann/__init__.py | 0 .../{testann.py => testann/testdense.py} | 10 +- test/python/testann/testsparse.py | 131 ++++++++++ test/python/testembeddings.py | 4 + test/python/testoptional.py | 32 ++- test/python/testscoring/__init__.py | 0 .../testkeyword.py} | 99 +------- test/python/testscoring/testsparse.py | 136 ++++++++++ test/python/testvectors/testdense/__init__.py | 0 .../testvectors/{ => testdense}/testcustom.py | 0 .../{ => testdense}/testexternal.py | 0 .../{ => testdense}/testhuggingface.py | 0 .../{ => testdense}/testlitellm.py | 0 .../testvectors/{ => testdense}/testllama.py | 0 .../testvectors/{ => testdense}/testm2v.py | 0 .../testvectors/{ => testdense}/testsbert.py | 0 .../{ => testdense}/testvectors.py | 2 +- .../{ => testdense}/testwordvectors.py | 2 +- .../python/testvectors/testsparse/__init__.py | 0 .../testvectors/testsparse/testsbert.py | 32 +++ .../testvectors/testsparse/testvectors.py | 40 +++ 65 files changed, 1249 insertions(+), 526 deletions(-) create mode 100644 src/python/txtai/ann/dense/__init__.py rename src/python/txtai/ann/{ => dense}/annoy.py (98%) rename src/python/txtai/ann/{ => dense}/factory.py (98%) rename src/python/txtai/ann/{ => dense}/faiss.py (99%) rename src/python/txtai/ann/{ => dense}/hnsw.py (99%) rename src/python/txtai/ann/{ => dense}/numpy.py (98%) rename src/python/txtai/ann/{ => dense}/pgvector.py (93%) rename src/python/txtai/ann/{ => dense}/sqlite.py (99%) rename src/python/txtai/ann/{ => dense}/torch.py (100%) create mode 100644 src/python/txtai/ann/sparse/__init__.py create mode 100644 src/python/txtai/ann/sparse/factory.py rename src/python/txtai/{scoring/ivf.py => ann/sparse/ivfsparse.py} (58%) create mode 100644 src/python/txtai/ann/sparse/pgsparse.py create mode 100644 src/python/txtai/util/sparsearray.py create mode 100644 src/python/txtai/vectors/dense/__init__.py rename src/python/txtai/vectors/{ => dense}/external.py (93%) rename src/python/txtai/vectors/{ => dense}/factory.py (97%) rename src/python/txtai/vectors/{ => dense}/huggingface.py (90%) rename src/python/txtai/vectors/{ => dense}/litellm.py (95%) rename src/python/txtai/vectors/{ => dense}/llama.py (96%) rename src/python/txtai/vectors/{ => dense}/m2v.py (95%) rename src/python/txtai/vectors/{ => dense}/sbert.py (62%) rename src/python/txtai/vectors/{ => dense}/words.py (97%) create mode 100644 src/python/txtai/vectors/sparse/__init__.py create mode 100644 src/python/txtai/vectors/sparse/base.py create mode 100644 src/python/txtai/vectors/sparse/factory.py create mode 100644 src/python/txtai/vectors/sparse/sbert.py create mode 100644 test/python/testann/__init__.py rename test/python/{testann.py => testann/testdense.py} (98%) create mode 100644 test/python/testann/testsparse.py create mode 100644 test/python/testscoring/__init__.py rename test/python/{testscoring.py => testscoring/testkeyword.py} (78%) create mode 100644 test/python/testscoring/testsparse.py create mode 100644 test/python/testvectors/testdense/__init__.py rename test/python/testvectors/{ => testdense}/testcustom.py (100%) rename test/python/testvectors/{ => testdense}/testexternal.py (100%) rename test/python/testvectors/{ => testdense}/testhuggingface.py (100%) rename test/python/testvectors/{ => testdense}/testlitellm.py (100%) rename test/python/testvectors/{ => testdense}/testllama.py (100%) rename test/python/testvectors/{ => testdense}/testm2v.py (100%) rename test/python/testvectors/{ => testdense}/testsbert.py (100%) rename test/python/testvectors/{ => testdense}/testvectors.py (96%) rename test/python/testvectors/{ => testdense}/testwordvectors.py (98%) create mode 100644 test/python/testvectors/testsparse/__init__.py create mode 100644 test/python/testvectors/testsparse/testsbert.py create mode 100644 test/python/testvectors/testsparse/testvectors.py diff --git a/docs/embeddings/configuration/scoring.md b/docs/embeddings/configuration/scoring.md index 7e9c7618..2c579d12 100644 --- a/docs/embeddings/configuration/scoring.md +++ b/docs/embeddings/configuration/scoring.md @@ -28,17 +28,30 @@ Additional settings for Postgres full-text keyword indexes. ### sparse ```yaml path: sparse vector model path +vectormethod: vector embeddings method gpu: boolean|int|string|device -encodebatch: int -modelargs: additional model init args +batch: Sets the transform batch size +encodebatch: Sets the encode batch size +vectors: additional model init args encodeargs: additional encode() args -ivf: #IVFFlat index configuration +backend: ivfsparse|pgsparse +``` + +Sparse vector scoring options. The sparse scoring instance combines a sparse vector model with a sparse approximate nearest neighbor index (ANN). + +#### ivfsparse +```yaml +ivfsparse: sample: percent of data to use for model training (0.0 - 1.0) - nclusters: desired number of clusters (int) + nlist: desired number of clusters (int) nprobe: search probe setting (int) ``` -Sparse vector scoring options. +Inverted file (IVF) index with flat vector file storage and sparse array support. + +#### pgsparse + +Sparse ANN backed by Postgres. Supports same options as the [pgvector](../ann/#pgvector) ANN. ## terms ```yaml diff --git a/setup.py b/setup.py index c6c960d1..ddf07f00 100644 --- a/setup.py +++ b/setup.py @@ -31,8 +31,10 @@ extras["agent"] = ["mcpadapt>=0.1.0", "smolagents>=1.17"] extras["ann"] = [ "annoy>=1.16.3", + "scikit-learn>=0.23.1", "hnswlib>=0.5.0", - "pgvector>=0.3.0", + "pgvector>=0.4.1", + "scipy>=1.4.1", "sqlalchemy>=2.0.20", "sqlite-vec>=0.1.1", ] @@ -94,13 +96,14 @@ extras["pipeline"] = ( + extras["pipeline-train"] ) -extras["scoring"] = ["sentence-transformers>=5.0.0", "sqlalchemy>=2.0.20"] +extras["scoring"] = ["sqlalchemy>=2.0.20"] extras["vectors"] = [ "litellm>=1.37.16", "llama-cpp-python>=0.2.75", "model2vec>=0.3.0", "scikit-learn>=0.23.1", + "scipy>=1.4.1", "sentence-transformers>=5.0.0", "skops>=0.9.0", "staticvectors>=0.2.0", diff --git a/src/python/txtai/ann/__init__.py b/src/python/txtai/ann/__init__.py index ad25b6f3..8a54db19 100644 --- a/src/python/txtai/ann/__init__.py +++ b/src/python/txtai/ann/__init__.py @@ -2,11 +2,6 @@ ANN imports """ -from .annoy import Annoy from .base import ANN -from .factory import ANNFactory -from .faiss import Faiss -from .hnsw import HNSW -from .numpy import NumPy -from .pgvector import PGVector -from .torch import Torch +from .dense import * +from .sparse import * diff --git a/src/python/txtai/ann/dense/__init__.py b/src/python/txtai/ann/dense/__init__.py new file mode 100644 index 00000000..804b2481 --- /dev/null +++ b/src/python/txtai/ann/dense/__init__.py @@ -0,0 +1,11 @@ +""" +Dense ANN imports +""" + +from .annoy import Annoy +from .factory import ANNFactory +from .faiss import Faiss +from .hnsw import HNSW +from .numpy import NumPy +from .pgvector import PGVector +from .torch import Torch diff --git a/src/python/txtai/ann/annoy.py b/src/python/txtai/ann/dense/annoy.py similarity index 98% rename from src/python/txtai/ann/annoy.py rename to src/python/txtai/ann/dense/annoy.py index 14551317..eca8e865 100644 --- a/src/python/txtai/ann/annoy.py +++ b/src/python/txtai/ann/dense/annoy.py @@ -10,7 +10,7 @@ try: except ImportError: ANNOY = False -from .base import ANN +from ..base import ANN # pylint: disable=W0223 diff --git a/src/python/txtai/ann/factory.py b/src/python/txtai/ann/dense/factory.py similarity index 98% rename from src/python/txtai/ann/factory.py rename to src/python/txtai/ann/dense/factory.py index cb0a4892..4beeccd3 100644 --- a/src/python/txtai/ann/factory.py +++ b/src/python/txtai/ann/dense/factory.py @@ -2,7 +2,7 @@ Factory module """ -from ..util import Resolver +from ...util import Resolver from .annoy import Annoy from .faiss import Faiss diff --git a/src/python/txtai/ann/faiss.py b/src/python/txtai/ann/dense/faiss.py similarity index 99% rename from src/python/txtai/ann/faiss.py rename to src/python/txtai/ann/dense/faiss.py index 6144794e..cbef53bf 100644 --- a/src/python/txtai/ann/faiss.py +++ b/src/python/txtai/ann/dense/faiss.py @@ -11,7 +11,7 @@ from faiss import omp_set_num_threads from faiss import index_factory, IO_FLAG_MMAP, METRIC_INNER_PRODUCT, read_index, write_index from faiss import index_binary_factory, read_index_binary, write_index_binary, IndexBinaryIDMap -from .base import ANN +from ..base import ANN if platform.system() == "Darwin": # Workaround for a Faiss issue causing segmentation faults on macOS. See txtai FAQ for more. diff --git a/src/python/txtai/ann/hnsw.py b/src/python/txtai/ann/dense/hnsw.py similarity index 99% rename from src/python/txtai/ann/hnsw.py rename to src/python/txtai/ann/dense/hnsw.py index 4460238c..f101ade1 100644 --- a/src/python/txtai/ann/hnsw.py +++ b/src/python/txtai/ann/dense/hnsw.py @@ -13,7 +13,7 @@ try: except ImportError: HNSWLIB = False -from .base import ANN +from ..base import ANN class HNSW(ANN): diff --git a/src/python/txtai/ann/numpy.py b/src/python/txtai/ann/dense/numpy.py similarity index 98% rename from src/python/txtai/ann/numpy.py rename to src/python/txtai/ann/dense/numpy.py index 0ab0a4ae..e8fb433e 100644 --- a/src/python/txtai/ann/numpy.py +++ b/src/python/txtai/ann/dense/numpy.py @@ -4,9 +4,9 @@ NumPy module import numpy as np -from ..serialize import SerializeFactory +from ...serialize import SerializeFactory -from .base import ANN +from ..base import ANN class NumPy(ANN): diff --git a/src/python/txtai/ann/pgvector.py b/src/python/txtai/ann/dense/pgvector.py similarity index 93% rename from src/python/txtai/ann/pgvector.py rename to src/python/txtai/ann/dense/pgvector.py index b5068e94..5b78cc1b 100644 --- a/src/python/txtai/ann/pgvector.py +++ b/src/python/txtai/ann/dense/pgvector.py @@ -18,7 +18,7 @@ try: except ImportError: PGVECTOR = False -from .base import ANN +from ..base import ANN class PGVector(ANN): @@ -111,7 +111,7 @@ class PGVector(ANN): self.schema() # Table name - table = self.setting("table", "vectors") + table = self.setting("table", self.defaulttable()) # Get embedding column and index settings column, index = self.column() @@ -147,7 +147,7 @@ class PGVector(ANN): self.close() # Create engine - self.engine = create_engine(self.setting("url", os.environ.get("ANN_URL")), poolclass=StaticPool, echo=False) + self.engine = create_engine(self.url(), poolclass=StaticPool, echo=False) self.connection = self.engine.connect() # Start database session @@ -191,6 +191,26 @@ class PGVector(ANN): args = (sql, parameters) if self.engine.dialect.name == "postgresql" else (text("SELECT 1"),) self.database.execute(*args) + def defaulttable(self): + """ + Returns the default table name. + + Returns: + default table name + """ + + return "vectors" + + def url(self): + """ + Reads the database url parameter. + + Returns: + database url + """ + + return self.setting("url", os.environ.get("ANN_URL")) + def column(self): """ Gets embedding column and index definitions for the current settings. diff --git a/src/python/txtai/ann/sqlite.py b/src/python/txtai/ann/dense/sqlite.py similarity index 99% rename from src/python/txtai/ann/sqlite.py rename to src/python/txtai/ann/dense/sqlite.py index 7c00c85c..0dc185e9 100644 --- a/src/python/txtai/ann/sqlite.py +++ b/src/python/txtai/ann/dense/sqlite.py @@ -13,7 +13,7 @@ try: except ImportError: SQLITEVEC = False -from .base import ANN +from ..base import ANN class SQLite(ANN): diff --git a/src/python/txtai/ann/torch.py b/src/python/txtai/ann/dense/torch.py similarity index 100% rename from src/python/txtai/ann/torch.py rename to src/python/txtai/ann/dense/torch.py diff --git a/src/python/txtai/ann/sparse/__init__.py b/src/python/txtai/ann/sparse/__init__.py new file mode 100644 index 00000000..616cdd4e --- /dev/null +++ b/src/python/txtai/ann/sparse/__init__.py @@ -0,0 +1,7 @@ +""" +Sparse ANN imports +""" + +from .factory import SparseANNFactory +from .ivfsparse import IVFSparse +from .pgsparse import PGSparse diff --git a/src/python/txtai/ann/sparse/factory.py b/src/python/txtai/ann/sparse/factory.py new file mode 100644 index 00000000..d7b42b21 --- /dev/null +++ b/src/python/txtai/ann/sparse/factory.py @@ -0,0 +1,61 @@ +""" +Factory module +""" + +from ...util import Resolver + +from .ivfsparse import IVFSparse +from .pgsparse import PGSparse + + +class SparseANNFactory: + """ + Methods to create Sparse ANN indexes. + """ + + @staticmethod + def create(config): + """ + Create an Sparse ANN. + + Args: + config: index configuration parameters + + Returns: + Sparse ANN + """ + + # ANN instance + ann = None + backend = config.get("backend", "ivfsparse") + + # Create ANN instance + if backend == "ivfsparse": + ann = IVFSparse(config) + elif backend == "pgsparse": + ann = PGSparse(config) + else: + ann = SparseANNFactory.resolve(backend, config) + + # Store config back + config["backend"] = backend + + return ann + + @staticmethod + def resolve(backend, config): + """ + Attempt to resolve a custom backend. + + Args: + backend: backend class + config: index configuration parameters + + Returns: + ANN + """ + + try: + return Resolver()(backend)(config) + except Exception as e: + raise ImportError(f"Unable to resolve sparse ann backend: '{backend}'") from e diff --git a/src/python/txtai/scoring/ivf.py b/src/python/txtai/ann/sparse/ivfsparse.py similarity index 58% rename from src/python/txtai/scoring/ivf.py rename to src/python/txtai/ann/sparse/ivfsparse.py index 41c9f051..c14e98fd 100644 --- a/src/python/txtai/scoring/ivf.py +++ b/src/python/txtai/ann/sparse/ivfsparse.py @@ -1,5 +1,5 @@ """ -IVF module +IVFSparse module """ import math @@ -8,23 +8,25 @@ import numpy as np # Conditional import try: - from scipy.sparse import csr_matrix, vstack + from scipy.sparse import vstack from sklearn.cluster import MiniBatchKMeans from sklearn.metrics import pairwise_distances_argmin_min from sklearn.utils.extmath import safe_sparse_dot - SKLEARN = True + IVFSPARSE = True except ImportError: - SKLEARN = False + IVFSPARSE = False -from ..serialize import SerializeFactory +from ...serialize import SerializeFactory +from ...util import SparseArray +from ..base import ANN -class IVFFlat: +class IVFSparse(ANN): """ - Inverted file (IVF) index with flat vector storage and sparse array support. + Inverted file (IVF) index with flat vector file storage and sparse array support. - IVFFlat builds an IVF index and enables approximate nearest neighbor (ANN) search. + IVFSparse builds an IVF index and enables approximate nearest neighbor (ANN) search. This index is modeled after Faiss and supports many of the same parameters. @@ -32,18 +34,10 @@ class IVFFlat: """ def __init__(self, config): - """ - Create a new IVFFlat instance. + super().__init__(config) - Args: - config: index configuration - """ - - if not SKLEARN: - raise ImportError('IVFFlat is not available - install "scoring" extra to enable') - - # Index configuration - self.config = config if isinstance(config, dict) else {} + if not IVFSPARSE: + raise ImportError('IVFSparse is not available - install "ann" extra to enable') # Cluster centroids, if computed self.centroids = None @@ -57,17 +51,9 @@ class IVFFlat: # Deleted ids self.deletes = None - def index(self, data): - """ - Builds a new IVFFlat index. - - Args: - data: input data - config: index configuration - """ - + def index(self, embeddings): # Compute model training size - train, sample = data, self.config.get("sample") + train, sample = embeddings, self.setting("sample") if sample: # Get sample for training rng = np.random.default_rng(0) @@ -75,7 +61,7 @@ class IVFFlat: train = train[indices] # Get number of clusters. Note that final number of clusters could be lower due to filtering duplicate centroids. - clusters = self.nclusters(data.shape[0], train.shape[0]) + clusters = self.nlist(embeddings.shape[0], train.shape[0]) # A single cluster is an exact search if clusters > 1: @@ -85,103 +71,75 @@ class IVFFlat: # Find closest points to each cluster center and use those as centroids # Filter out duplicate centroids indices, _ = pairwise_distances_argmin_min(kmeans.cluster_centers_, train, metric="cosine") - self.centroids = data[np.unique(indices)] + self.centroids = embeddings[np.unique(indices)] # Sort into clusters - ids = self.aggregate(data) + ids = self.aggregate(embeddings) # Sort clusters by id self.ids = dict(sorted(ids.items(), key=lambda x: x[0])) # Create cluster data blocks - self.blocks = {k: data[v] for k, v in self.ids.items()} + self.blocks = {k: embeddings[v] for k, v in self.ids.items()} # Initialize deletes self.deletes = [] - def append(self, data): - """ - Appends elements to an existing index. - - Args: - data: input data - """ + # Add id offset and index build metadata + self.config["offset"] = embeddings.shape[0] + self.metadata({"clusters": clusters}) + def append(self, embeddings): # Get offset offset = self.size() # Sort into clusters and merge - for cluster, ids in self.aggregate(data).items(): + for cluster, ids in self.aggregate(embeddings).items(): # Add new ids self.ids[cluster].extend([x + offset for x in ids]) # Add new data - self.blocks[cluster] = vstack([self.blocks[cluster], data[ids]]) + self.blocks[cluster] = vstack([self.blocks[cluster], embeddings[ids]]) + + # Update id offset and index metadata + self.config["offset"] += embeddings.shape[0] + self.metadata() def delete(self, ids): - """ - Mark ids as deleted. This prevents deleted results from showing up in search results. - The data is not removed from the underlying data structures. - - Args: - ids: ids to delete - """ - # Set index ids as deleted self.deletes.extend(ids) - def search(self, query, limit): - """ - Searches IVFFlat index for query. Returns topn results. + def search(self, queries, limit): + results = [] + for query in queries: + if self.centroids is not None: + # Approximate search + indices, _ = self.topn(query, self.centroids, self.nprobe()) - Args: - query: query array - limit: maximum results + # Stack into single ids list + ids = np.concatenate([self.ids[x] for x in indices if x in self.ids]) - Returns: - query results - """ + # Stack data rows + data = vstack([self.blocks[x] for x in indices if x in self.blocks]) + else: + # Exact search + ids, data = np.array(self.ids[0]), self.blocks[0] - if self.centroids is not None: - # Approximate search - indices, _ = self.topn(query, self.centroids, self.nprobe()) + # Get deletes + deletes = np.argwhere(np.isin(ids, self.deletes)).ravel() - # Stack into single ids list - ids = np.concatenate([self.ids[x] for x in indices if x in self.ids]) + # Calculate similarity + indices, scores = self.topn(query, data, limit, deletes) - # Stack data rows - data = vstack([self.blocks[x] for x in indices if x in self.blocks]) - else: - # Exact search - ids, data = np.array(self.ids[0]), self.blocks[0] + # Map data ids and return + results.append(list(zip(ids[indices].tolist(), scores.tolist()))) - # Get deletes - deletes = np.argwhere(np.isin(ids, self.deletes)).ravel() - - # Calculate similarity - indices, scores = self.topn(query, data, limit, deletes) - - # Map data ids and return - return list(zip(ids[indices].tolist(), scores.tolist())) + return results def count(self): - """ - Number of elements in the IVFFlat index. - - Returns: - count - """ - return self.size() - len(self.deletes) def load(self, path): - """ - Loads an IVFFlat index from path. - - Args: - path: file path - """ - # Create streaming serializer and limit read size to a byte at a time to ensure # only msgpack data is consumed serializer = SerializeFactory.create("msgpack", streaming=True, read_size=1) @@ -192,7 +150,7 @@ class IVFFlat: header = next(unpacker) # Read cluster centroids, if available - self.centroids = SparseArray().read(f) if header["centroids"] else None + self.centroids = SparseArray().load(f) if header["centroids"] else None # Read cluster ids self.ids = dict(next(unpacker)) @@ -200,23 +158,18 @@ class IVFFlat: # Read cluster data blocks self.blocks = {} for key in self.ids: - self.blocks[key] = SparseArray().read(f) + self.blocks[key] = SparseArray().load(f) # Read deletes self.deletes = next(unpacker) def save(self, path): - """ - Saves an IVFFlat index to path. This format uses a combination of msgpack and NumPy serialization. - Sparse arrays are saved as groups of NumPy arrays. - - IVFFlat storage format: - - header msgpack - - centroids sparse array (optional based on header parameters) - - cluster ids msgpack - - cluster data blocks list of sparse arrays - - deletes msgpack - """ + # IVFSparse storage format: + # - header msgpack + # - centroids sparse array (optional based on header parameters) + # - cluster ids msgpack + # - cluster data blocks list of sparse arrays + # - deletes msgpack # Create message pack serializer serializer = SerializeFactory.create("msgpack") @@ -227,14 +180,14 @@ class IVFFlat: # Write cluster centroids, if available if self.centroids is not None: - SparseArray().write(self.centroids, f) + SparseArray().save(f, self.centroids) # Write cluster id mapping serializer.savestream(list(self.ids.items()), f) # Write cluster data blocks for block in self.blocks.values(): - SparseArray().write(block, f) + SparseArray().save(f, block) # Write deletes serializer.savestream(self.deletes, f) @@ -296,10 +249,10 @@ class IVFFlat: return indices[0], scores[0] - def nclusters(self, count, train): + def nlist(self, count, train): """ - Calculates the number of clusters for this IVFFlat index. Note that the final number of clusters could be - lower as duplicate cluster centroids are filtered out. + Calculates the number of clusters for this IVFSparse index. Note that the final number of clusters + could be lower as duplicate cluster centroids are filtered out. Args: count: initial dataset size @@ -313,7 +266,7 @@ class IVFFlat: default = 1 if count <= 5000 else self.cells(train) # Number of clusters to create - return self.config.get("nclusters", default) + return self.setting("nlist", default) def nprobe(self): """ @@ -327,11 +280,11 @@ class IVFFlat: size = self.size() default = 6 if size <= 5000 else self.cells(size) // 64 - return self.config.get("nprobe", default) + return self.setting("nprobe", default) def cells(self, count): """ - Calculates the number of IVF cells for an IVFFlat index. + Calculates the number of IVF cells for an IVFSparse index. Args: count: number of rows @@ -353,38 +306,3 @@ class IVFFlat: """ return sum(len(x) for x in self.ids.values()) - - -class SparseArray: - """ - Methods to read and write sparse arrays to file. - """ - - def read(self, f): - """ - Reads a sparse array from input file. - - Args: - f: input file handle - - Returns: - sparse array - """ - - data, indices, indptr, shape = np.load(f), np.load(f), np.load(f), np.load(f) - - # Read sparse array - return csr_matrix((data, indices, indptr), shape=shape) - - def write(self, array, f): - """ - Writes a sparse array to file. - - Args: - array: sparse array - f: output file handle - """ - - # Write sparse array to file - for x in [array.data, array.indices, array.indptr, array.shape]: - np.save(f, x, allow_pickle=False) diff --git a/src/python/txtai/ann/sparse/pgsparse.py b/src/python/txtai/ann/sparse/pgsparse.py new file mode 100644 index 00000000..1a971f21 --- /dev/null +++ b/src/python/txtai/ann/sparse/pgsparse.py @@ -0,0 +1,44 @@ +""" +PGSparse module +""" + +import os + +# Conditional import +try: + from pgvector import SparseVector + from pgvector.sqlalchemy import SPARSEVEC + + PGSPARSE = True +except ImportError: + PGSPARSE = False + +from ..dense import PGVector + + +class PGSparse(PGVector): + """ + Builds a Sparse ANN index backed by a Postgres database. + """ + + def __init__(self, config): + if not PGSPARSE: + raise ImportError('PGSparse is not available - install "ann" extra to enable') + + super().__init__(config) + + # Quantization not supported + self.qbits = None + + def defaulttable(self): + return "svectors" + + def url(self): + return self.setting("url", os.environ.get("SCORING_URL", os.environ.get("ANN_URL"))) + + def column(self): + return SPARSEVEC(self.config["dimensions"]), "sparsevec_ip_ops" + + def prepare(self, data): + # Wrap as sparse vector + return SparseVector(data) diff --git a/src/python/txtai/embeddings/base.py b/src/python/txtai/embeddings/base.py index 8868503f..3fdbfac7 100644 --- a/src/python/txtai/embeddings/base.py +++ b/src/python/txtai/embeddings/base.py @@ -6,8 +6,6 @@ import json import os import tempfile -import numpy as np - from ..ann import ANNFactory from ..archive import ArchiveFactory from ..cloud import CloudFactory @@ -323,7 +321,7 @@ class Embeddings: self.defaults() # Get vector model - model = self.indexes.model(index) if index and self.indexes else self.model if self.model else self.indexes.model() + model = self.findmodel(index) # Convert documents into embeddings embeddings = model.batchtransform(Stream(self)(documents), category) @@ -437,8 +435,11 @@ class Embeddings: queries = self.batchtransform(((None, query, None) for query in queries), "query") data = self.batchtransform(((None, row, None) for row in data), "data") + # Get vector model + model = self.findmodel() + # Dot product on normalized vectors is equal to cosine similarity - scores = np.dot(queries, data.T).tolist() + scores = model.dot(queries, data) # Add index and sort desc based on score return [sorted(enumerate(score), key=lambda x: x[1], reverse=True) for score in scores] @@ -711,13 +712,13 @@ class Embeddings: def issparse(self): """ - Checks if this instance has an associated scoring instance with term indexing enabled. + Checks if this instance has an associated sparse keyword or sparse vectors scoring index. Returns: - True if term index is enabled, False otherwise + True if scoring has an associated sparse keyword/vector index, False otherwise """ - return self.scoring and self.scoring.hasterms() + return self.scoring and self.scoring.issparse() def isdense(self): """ @@ -737,7 +738,25 @@ class Embeddings: True if term weighting is enabled, False otherwise """ - return self.scoring and not self.scoring.hasterms() + return self.scoring and self.scoring.isweighted() + + def findmodel(self, index=None): + """ + Finds the primary vector model used by this instance. + + Returns: + Vectors + """ + + return ( + self.indexes.findmodel(index) + if index and self.indexes + else ( + self.model + if self.model + else self.scoring.findmodel() if self.scoring and self.scoring.findmodel() else self.indexes.findmodel() if self.indexes else None + ) + ) def configure(self, config): """ @@ -755,7 +774,7 @@ class Embeddings: # Create scoring instance for word vectors term weighting scoring = self.config.get("scoring") if self.config else None - self.scoring = self.createscoring() if scoring and (not isinstance(scoring, dict) or not scoring.get("terms")) else None + self.scoring = self.createscoring() if scoring and not self.hassparse() else None # Dense vectors - transforms data to embeddings vectors self.model = self.loadvectors() if self.config else None @@ -791,9 +810,8 @@ class Embeddings: # Initialize ANN, will be created after index transformations complete self.ann = None - # Create scoring only if term indexing is enabled - scoring = self.config.get("scoring") - if scoring and isinstance(scoring, dict) and self.config["scoring"].get("terms"): + # Create scoring only if the scoring config is for a sparse index + if self.hassparse(): self.scoring = self.createscoring() # Create subindexes, if necessary @@ -832,24 +850,29 @@ class Embeddings: Logic to derive default sparse index configuration. """ - # Get sparse index configuration - index, method = None, None - for x in ["keyword", "sparse", "hybrid"]: + # Check for keyword and hybrid parameters + method = None + for x in ["keyword", "hybrid"]: value = self.config.get(x) if value: - index, method = (x, value if isinstance(value, str) else "bm25") + method = value if isinstance(value, str) else "bm25" - # Default scoring configuration - self.config["scoring"] = {"method": method, "terms": True, "normalize": True} + # Enable dense index when hybrid enabled + if x == "hybrid": + self.config["dense"] = True - # Sparse vector models - if index == "sparse" or "/" in method: - # Derive sparse model path - path = self.config["sparse"] if index == "sparse" else method - path = path if isinstance(path, str) else "prithivida/Splade_PP_en_v2" + sparse = self.config.get("sparse", {}) + if sparse or method == "sparse": + # Sparse vector configuration + sparse = {"path": self.config.get("sparse")} if isinstance(sparse, str) else {} if isinstance(sparse, bool) else sparse + sparse["path"] = sparse.get("path", "prithivida/Splade_PP_en_v2") # Merge in sparse parameters - self.config["scoring"] = {**self.config["scoring"], **{"method": "sparse", "path": path}} + self.config["scoring"] = {**{"method": "sparse"}, **sparse} + + elif method: + # Sparse keyword configuration + self.config["scoring"] = {"method": method, "terms": True, "normalize": True} def defaultallowed(self): """ @@ -1051,6 +1074,17 @@ class Embeddings: return None + def hassparse(self): + """ + Checks is this embeddings database has an associated sparse index. + + Returns: + True if this embeddings has an associated scoring index + """ + + # Create scoring only if scoring is a sparse keyword/vector index + return ScoringFactory.issparse(self.config.get("scoring")) + def columns(self, config): """ Adds custom text/object column information if it's provided. diff --git a/src/python/txtai/embeddings/index/indexes.py b/src/python/txtai/embeddings/index/indexes.py index f9e5f12d..721e945f 100644 --- a/src/python/txtai/embeddings/index/indexes.py +++ b/src/python/txtai/embeddings/index/indexes.py @@ -25,6 +25,7 @@ class Indexes: self.indexes = indexes self.documents = None + self.checkpoint = None # Transform columns columns = embeddings.config.get("columns", {}) @@ -83,29 +84,34 @@ class Indexes: return list(self.indexes.keys())[0] - def model(self, index=None): + def findmodel(self, index=None): """ - Lookups a vector model. If index is empty, the first vector model is returned. + Finds a vector model. If index is empty, the first vector model is returned. + + Args: + index: index name to match Returns: Vectors """ # Find vector model - matches = [self.indexes[index]] if index else [index.model for index in self.indexes.values() if index.model] + matches = [self.indexes[index].findmodel()] if index else [index.findmodel() for index in self.indexes.values() if index.findmodel()] return matches[0] if matches else None - def insert(self, documents, index=None): + def insert(self, documents, index=None, checkpoint=None): """ Inserts a batch of documents into each subindex. Args: documents: list of (id, data, tags) index: indexid offset + checkpoint: optional checkpoint directory, enables indexing restart """ if not self.documents: self.documents = Documents() + self.checkpoint = checkpoint # Create batch containing documents added to parent index batch = [] @@ -139,12 +145,13 @@ class Indexes: Builds each subindex. """ - for index in self.indexes.values(): - index.index(self.documents) + for name, index in self.indexes.items(): + index.index(self.documents, checkpoint=f"{self.checkpoint}/{name}" if self.checkpoint else None) # Reset document stream self.documents.close() self.documents = None + self.checkpoint = None def upsert(self): """ diff --git a/src/python/txtai/embeddings/index/transform.py b/src/python/txtai/embeddings/index/transform.py index ce7e27d0..3d7a4d0e 100644 --- a/src/python/txtai/embeddings/index/transform.py +++ b/src/python/txtai/embeddings/index/transform.py @@ -2,8 +2,6 @@ Transform module """ -import os - import numpy as np from .action import Action @@ -92,29 +90,11 @@ class Transform: (document ids, dimensions, embeddings) """ - # Consume stream and transform documents to vectors - ids, dimensions, batches, stream = self.model.index(self.stream(documents), self.batch, self.checkpoint) + # Determine dtype + dtype = np.uint8 if self.qbits else np.float32 - # Check that embeddings are available and load as a memmap - embeddings = None - if ids: - # Determine dtype - dtype = np.uint8 if self.qbits else np.float32 - - # Write batches - embeddings = np.memmap(buffer, dtype=dtype, shape=(len(ids), dimensions), mode="w+") - with open(stream, "rb") as queue: - x = 0 - for _ in range(batches): - batch = np.load(queue) - embeddings[x : x + batch.shape[0]] = batch - x += batch.shape[0] - - # Remove temporary file (if checkpointing is disabled) - if not self.checkpoint: - os.remove(stream) - - return (ids, dimensions, embeddings) + # Transform documents into vectors + return self.model.vectors(self.stream(documents), self.batch, self.checkpoint, buffer, dtype) def ids(self, documents): """ @@ -211,11 +191,11 @@ class Transform: # Load batch into scoring if self.scoring: - self.scoring.insert(batch, self.offset) + self.scoring.insert(batch, self.offset, self.checkpoint) # Load batch into subindex documents stream if self.indexes: - self.indexes.insert(batch, self.offset) + self.indexes.insert(batch, self.offset, self.checkpoint) # Load batch into graph if self.graph: diff --git a/src/python/txtai/pipeline/audio/texttospeech.py b/src/python/txtai/pipeline/audio/texttospeech.py index af23a388..82b1ee7b 100644 --- a/src/python/txtai/pipeline/audio/texttospeech.py +++ b/src/python/txtai/pipeline/audio/texttospeech.py @@ -407,7 +407,7 @@ class SpeechT5(SpeechPipeline): self.vocoder = ort.InferenceSession(cached_file(path_or_repo_id=path, filename="decoder_postnet_and_vocoder.onnx"), providers=providers) self.processor = SpeechT5Processor.from_pretrained(path) - self.defaultspeaker = np.load(cached_file(path_or_repo_id=path, filename="speaker.npy")) + self.defaultspeaker = np.load(cached_file(path_or_repo_id=path, filename="speaker.npy"), allow_pickle=False) # Max number of input tokens model can handle self.maxtokens = maxtokens diff --git a/src/python/txtai/scoring/__init__.py b/src/python/txtai/scoring/__init__.py index e1816992..cb7cdc3d 100644 --- a/src/python/txtai/scoring/__init__.py +++ b/src/python/txtai/scoring/__init__.py @@ -5,7 +5,6 @@ Scoring imports from .base import Scoring from .bm25 import BM25 from .factory import ScoringFactory -from .ivf import IVFFlat from .pgtext import PGText from .sif import SIF from .sparse import Sparse diff --git a/src/python/txtai/scoring/base.py b/src/python/txtai/scoring/base.py index 63503ad9..ddd4e51b 100644 --- a/src/python/txtai/scoring/base.py +++ b/src/python/txtai/scoring/base.py @@ -24,13 +24,17 @@ class Scoring: self.text = columns.get("text", "text") self.object = columns.get("object", "object") - def insert(self, documents, index=None): + # Vector model, if available + self.model = None + + def insert(self, documents, index=None, checkpoint=None): """ Inserts documents into the scoring index. Args: documents: list of (id, dict|text|tokens, tags) index: indexid offset + checkpoint: optional checkpoint directory, enables indexing restart """ raise NotImplementedError @@ -143,16 +147,36 @@ class Scoring: raise NotImplementedError - def hasterms(self): + def findmodel(self): """ - Check if this scoring instance has an associated terms index. + Returns the associated vector model used by this scoring instance, if any. Returns: - True if this scoring instance has an associated terms index. + associated vector model + """ + + return self.model + + def issparse(self): + """ + Check if this scoring instance has an associated sparse keyword or sparse vector index. + + Returns: + True if this index has an associated sparse index """ raise NotImplementedError + def isweighted(self): + """ + Check if this scoring instance is for term weighting (i.e.) it has no associated sparse index. + + Returns: + True if this index is for term weighting + """ + + return not self.issparse() + def isnormalized(self): """ Check if this scoring instance returns normalized scores. diff --git a/src/python/txtai/scoring/factory.py b/src/python/txtai/scoring/factory.py index 868db626..70cccf7d 100644 --- a/src/python/txtai/scoring/factory.py +++ b/src/python/txtai/scoring/factory.py @@ -22,7 +22,7 @@ class ScoringFactory: Factory method to construct a Scoring instance. Args: - config: scoring configuration parameters - supports bm25, sif, tfidf + config: scoring configuration parameters models: models cache Returns: @@ -58,6 +58,24 @@ class ScoringFactory: return scoring + @staticmethod + def issparse(config): + """ + Checks if this scoring configuration builds a sparse index. + + Args: + config: scoring configuration + + Returns: + True if this config is for a sparse index + """ + + # Types that are always a sparse index + indexes = ["pgtext", "sparse"] + + # True if this config is for a sparse index + return config and isinstance(config, dict) and (config.get("method") in indexes or config.get("terms")) + @staticmethod def resolve(backend, config): """ diff --git a/src/python/txtai/scoring/pgtext.py b/src/python/txtai/scoring/pgtext.py index 4569151a..a32cc630 100644 --- a/src/python/txtai/scoring/pgtext.py +++ b/src/python/txtai/scoring/pgtext.py @@ -36,7 +36,7 @@ class PGText(Scoring): # Language self.language = self.config.get("language", "english") - def insert(self, documents, index=None): + def insert(self, documents, index=None, checkpoint=None): # Initialize tables self.initialize(recreate=True) @@ -108,7 +108,7 @@ class PGText(Scoring): self.database.close() self.engine.dispose() - def hasterms(self): + def issparse(self): return True def isnormalized(self): diff --git a/src/python/txtai/scoring/sparse.py b/src/python/txtai/scoring/sparse.py index 864186f0..717a0909 100644 --- a/src/python/txtai/scoring/sparse.py +++ b/src/python/txtai/scoring/sparse.py @@ -2,20 +2,13 @@ Sparse module """ -# Conditional import -try: - from scipy.sparse import csr_matrix, vstack - from sentence_transformers import SparseEncoder - from sklearn.preprocessing import normalize +from queue import Queue +from threading import Thread - SPARSE = True -except ImportError: - SPARSE = False - -from ..models import Models +from ..ann import SparseANNFactory +from ..vectors import SparseVectorsFactory from .base import Scoring -from .ivf import IVFFlat class Sparse(Scoring): @@ -23,185 +16,166 @@ class Sparse(Scoring): Sparse vector scoring. """ + # End of stream message + COMPLETE = 1 + def __init__(self, config=None, models=None): super().__init__(config) - if not SPARSE: - raise ImportError('Sparse encoder is not available - install "scoring" extra to enable') + # Vector configuration + config = {k: v for k, v in config.items() if k != "method"} + if "vectormethod" in config: + config["method"] = config["vectormethod"] - # Models cache - self.models = models + # Load the SparseVectors model + self.model = SparseVectorsFactory.create(config, models) - # Pool parameter - self.pool = None + # Sparse ANN + self.ann = None - # Sparse encoder - self.model = self.loadmodel() + # Encoding processing parameters + self.batch = self.config.get("batch", 1024) + self.thread, self.queue, self.data = None, None, None - # Encode batch size - controls underlying model batch size when encoding vectors - self.encodebatch = config.get("encodebatch", 32) + def insert(self, documents, index=None, checkpoint=None): + # Start processing thread, if necessary + self.start(checkpoint) - # Index backend - self.backend = None - - # Input queue - data to be indexed - self.queue = None - - def insert(self, documents, index=None): data = [] - for _, document, _ in documents: + for uid, document, tags in documents: # Extract text, if necessary if isinstance(document, dict): document = document.get(self.text, document.get(self.object)) if document is not None: # Add data - data.append(" ".join(document) if isinstance(document, list) else document) + data.append((uid, " ".join(document) if isinstance(document, list) else document, tags)) - # Encode and normalize data - data = self.encode(data) - - # Create or add to existing sparse vectors - self.queue = vstack((self.queue, data)) if self.queue is not None else data + # Add batch of data + self.queue.put(data) def delete(self, ids): - self.backend.delete(ids) + self.ann.delete(ids) def index(self, documents=None): # Insert documents, if provided if documents: self.insert(documents) - # Create index - if self.queue is not None: - # Create a new index instance - self.backend = IVFFlat(self.config.get("ivf")) - self.backend.index(self.queue) - - # Clear queue - self.queue = None + # Create ANN, if there is pending data + embeddings = self.stop() + if embeddings is not None: + self.ann = SparseANNFactory.create(self.config) + self.ann.index(embeddings) def upsert(self, documents=None): # Insert documents, if provided if documents: self.insert(documents) - # Upsert index - if self.backend and self.queue is not None: - self.backend.append(self.queue) + # Check for existing index and pending data + if self.ann: + embeddings = self.stop() + if embeddings is not None: + self.ann.append(embeddings) else: self.index() - # Clear queue - self.queue = None - def weights(self, tokens): # Not supported return None def search(self, query, limit=3): - return self.backend.search(self.encode([query]), limit) + return self.batchsearch([query], limit)[0] def batchsearch(self, queries, limit=3, threads=True): - return [self.search(query, limit) for query in queries] + # Convert queries to embedding vectors + embeddings = self.model.batchtransform((None, query, None) for query in queries) + + # Run ANN search + return self.ann.search(embeddings, limit) def count(self): - return self.backend.count() + return self.ann.count() def load(self, path): - # Read IVFFlat index - self.backend = IVFFlat(self.config.get("ivf")) - self.backend.load(path) + self.ann = SparseANNFactory.create(self.config) + self.ann.load(path) def save(self, path): - # Write IVFFlat index - if self.backend is not None: - self.backend.save(path) + # Save Sparse ANN + if self.ann: + self.ann.save(path) def close(self): - # Close pool before model is closed - if self.pool: - self.model.stop_multi_process_pool(self.pool) - self.pool = None + # Close Sparse ANN + if self.ann: + self.ann.close() - self.model, self.backend, self.queue = None, None, None + # Clear parameters + self.model, self.ann, self.thread, self.queue = None, None, None, None - def hasterms(self): + def issparse(self): return True def isnormalized(self): return True - def loadmodel(self): + def start(self, checkpoint): """ - Loads the sparse encoder model. - - Returns: - SparseEncoder - """ - - # Model path - path = self.config.get("path") - - # Check if model is cached - if self.models and path in self.models: - return self.models[path] - - # Get target device - gpu, pool = self.config.get("gpu", True), False - - # Default mode uses a single GPU. Setting to all spawns a process per GPU. - if isinstance(gpu, str) and gpu == "all": - # Get number of accelerator devices available - devices = Models.acceleratorcount() - - # Enable multiprocessing pooling only when multiple devices are available - gpu, pool = devices <= 1, devices > 1 - - # Tensor device id - deviceid = Models.deviceid(gpu) - - # Additional model arguments - modelargs = self.config.get("modelargs", {}) - - # Build embeddings with sentence-transformers - model = SparseEncoder(path, device=Models.device(deviceid), **modelargs) - - # Start process pool for multiple GPUs - if pool: - self.pool = model.start_multi_process_pool() - - # Store model in cache - if self.models is not None and path: - self.models[path] = model - - # Return model - return model - - def encode(self, data): - """ - Encodes a batch of data using the Sparse Encoder model. + Starts an encoding processing thread. Args: - data: input data - - Returns: - encoded data as a SciPy CSR Matrix + checkpoint: checkpoint directory """ - # Additional encoding arguments - encodeargs = self.config.get("encodeargs", {}) + if not self.thread: + self.queue = Queue(5) + self.thread = Thread(target=self.encode, args=(checkpoint,)) + self.thread.start() - # Encode data - data = self.model.encode(data, pool=self.pool, batch_size=self.encodebatch, **encodeargs) + def stop(self): + """ + Stops an encoding processing thread. Return processed results. - # Get data attributes - data = data.cpu().coalesce() - indices = data.indices().numpy() - values = data.values().numpy() + Returns: + results + """ - # Convert to CSR Matrix - matrix = csr_matrix((values, indices), shape=data.size()) + results = None + if self.thread: + # Send EOS message + self.queue.put(Sparse.COMPLETE) - # Normalize and return - return normalize(matrix) + self.thread.join() + self.thread, self.queue = None, None + + # Get return value + results = self.data + self.data = None + + return results + + def encode(self, checkpoint): + """ + Encodes streaming data. + + Args: + checkpoint: checkpoint directory + """ + + # Streaming encoding of data + _, dimensions, self.data = self.model.vectors(self.stream(), self.batch, checkpoint) + + # Save number of dimensions + self.config["dimensions"] = dimensions + + def stream(self): + """ + Streams data from an input queue until end of stream message received. + """ + + batch = self.queue.get() + while batch != Sparse.COMPLETE: + yield from batch + batch = self.queue.get() diff --git a/src/python/txtai/scoring/tfidf.py b/src/python/txtai/scoring/tfidf.py index 83ab2163..e1c5f29c 100644 --- a/src/python/txtai/scoring/tfidf.py +++ b/src/python/txtai/scoring/tfidf.py @@ -55,7 +55,7 @@ class TFIDF(Scoring): self.normalize = self.config.get("normalize") self.avgscore = None - def insert(self, documents, index=None): + def insert(self, documents, index=None, checkpoint=None): # Insert documents, calculate word frequency, total tokens and total documents for uid, document, tags in documents: # Extract text, if necessary @@ -232,7 +232,7 @@ class TFIDF(Scoring): if self.terms: self.terms.close() - def hasterms(self): + def issparse(self): return self.terms is not None def isnormalized(self): @@ -349,6 +349,9 @@ class TFIDF(Scoring): resolved results """ + # Convert to Python values + scores = [(x, float(score)) for x, score in scores] + if self.documents: return [{"id": x, "text": self.documents[x], "score": score} for x, score in scores] diff --git a/src/python/txtai/util/__init__.py b/src/python/txtai/util/__init__.py index ef78ad70..4cbc088d 100644 --- a/src/python/txtai/util/__init__.py +++ b/src/python/txtai/util/__init__.py @@ -3,4 +3,5 @@ Utility imports """ from .resolver import Resolver +from .sparsearray import SparseArray from .template import TemplateFormatter diff --git a/src/python/txtai/util/sparsearray.py b/src/python/txtai/util/sparsearray.py new file mode 100644 index 00000000..dbd66a05 --- /dev/null +++ b/src/python/txtai/util/sparsearray.py @@ -0,0 +1,62 @@ +""" +SparseArray module +""" + +import numpy as np + +# Conditional import +try: + from scipy.sparse import csr_matrix + + SCIPY = True +except ImportError: + SCIPY = False + + +class SparseArray: + """ + Methods to load and save sparse arrays to file. + """ + + def __init__(self): + """ + Creates a SparseArray instance. + """ + + if not SCIPY: + raise ImportError("SciPy is not available - install scipy to enable") + + def load(self, f): + """ + Loads a sparse array from file. + + Args: + f: input file handle + + Returns: + sparse array + """ + + # Load raw data + data, indices, indptr, shape = ( + np.load(f, allow_pickle=False), + np.load(f, allow_pickle=False), + np.load(f, allow_pickle=False), + np.load(f, allow_pickle=False), + ) + + # Load data into sparse array + return csr_matrix((data, indices, indptr), shape=shape) + + def save(self, f, array): + """ + Saves a sparse array to file. + + Args: + f: output file handle + array: sparse array + """ + + # Save sparse array to file + for x in [array.data, array.indices, array.indptr, array.shape]: + np.save(f, x, allow_pickle=False) diff --git a/src/python/txtai/vectors/__init__.py b/src/python/txtai/vectors/__init__.py index fba44cb3..13c64781 100644 --- a/src/python/txtai/vectors/__init__.py +++ b/src/python/txtai/vectors/__init__.py @@ -1,14 +1,8 @@ """ -Vectors import +Vectors imports """ from .base import Vectors -from .external import External -from .factory import VectorsFactory -from .huggingface import HFVectors -from .litellm import LiteLLM -from .llama import LlamaCpp -from .m2v import Model2Vec +from .dense import * from .recovery import Recovery -from .sbert import STVectors -from .words import WordVectors +from .sparse import * diff --git a/src/python/txtai/vectors/base.py b/src/python/txtai/vectors/base.py index 989ea0d6..733d6055 100644 --- a/src/python/txtai/vectors/base.py +++ b/src/python/txtai/vectors/base.py @@ -70,12 +70,13 @@ class Vectors: raise NotImplementedError - def encode(self, data): + def encode(self, data, category=None): """ Encodes a batch of data using vector model. Args: data: batch of data + category: optional category for instruction-based embeddings Return: transformed data @@ -116,14 +117,14 @@ class Vectors: checkpoint: optional checkpoint directory, enables indexing restart Returns: - (ids, dimensions, stream) + (ids, dimensions, batches, stream) """ ids, dimensions, batches, stream = [], None, 0, None # Generate recovery config if checkpoint is set vectorsid = self.vectorsid() if checkpoint else None - recovery = Recovery(checkpoint, vectorsid) if checkpoint else None + recovery = Recovery(checkpoint, vectorsid, self.loadembeddings) if checkpoint else None # Convert all documents to embedding arrays, stream embeddings to disk to control memory usage with self.spool(checkpoint, vectorsid) as output: @@ -148,6 +149,42 @@ class Vectors: return (ids, dimensions, batches, stream) + def vectors(self, documents, batchsize=500, checkpoint=None, buffer=None, dtype=None): + """ + Bulk encodes documents into vectors using index(). Return the data as a mmap-ed array. + + Args: + documents: list of (id, data, tags) + batchsize: index batch size + checkpoint: optional checkpoint directory, enables indexing restart + buffer: file path used for memmap buffer + dtype: dtype for buffer + + Returns: + (ids, dimensions, embeddings) + """ + + # Consume stream and transform documents to vectors + ids, dimensions, batches, stream = self.index(documents, batchsize, checkpoint) + + # Check that embeddings are available and load as a memmap + embeddings = None + if ids: + # Write batches + embeddings = np.memmap(buffer, dtype=dtype, shape=(len(ids), dimensions), mode="w+") + with open(stream, "rb") as queue: + x = 0 + for _ in range(batches): + batch = self.loadembeddings(queue) + embeddings[x : x + batch.shape[0]] = batch + x += batch.shape[0] + + # Remove temporary file (if checkpointing is disabled) + if not checkpoint: + os.remove(stream) + + return (ids, dimensions, embeddings) + def close(self): """ Closes this vectors instance. @@ -188,7 +225,22 @@ class Vectors: if documents and isinstance(documents[0], np.ndarray): return np.array(documents, dtype=np.float32) - return self.vectorize(documents) + return self.vectorize(documents, category) + + def dot(self, queries, data): + """ + Calculates the dot product similarity between queries and documents. This method + assumes each of the inputs are normalized. + + Args: + queries: queries + data: search data + + Returns: + dot product scores + """ + + return np.dot(queries, data.T).tolist() def vectorsid(self): """ @@ -246,10 +298,10 @@ class Vectors: # Attempt to read embeddings from a recovery file embeddings = recovery() if recovery else None - embeddings = self.vectorize(documents) if embeddings is None else embeddings + embeddings = self.vectorize(documents, "data") if embeddings is None else embeddings if embeddings is not None: dimensions = embeddings.shape[1] - np.save(output, embeddings) + self.saveembeddings(output, embeddings) return (ids, dimensions) @@ -299,7 +351,7 @@ class Vectors: return data - def vectorize(self, data): + def vectorize(self, data, category=None): """ Runs data vectorization, which consists of the following steps. @@ -310,13 +362,17 @@ class Vectors: Args: data: input data + category: category for instruction-based embeddings Returns: embeddings vectors """ + # Default instruction category + category = category if category else "query" + # Transform data into vectors - embeddings = self.encode(data) + embeddings = self.encode(data, category) if embeddings is not None: # Truncate embeddings, if necessary @@ -324,7 +380,7 @@ class Vectors: embeddings = self.truncate(embeddings) # Normalize data - self.normalize(embeddings) + embeddings = self.normalize(embeddings) # Apply quantization, if necessary if self.qbits: @@ -332,6 +388,30 @@ class Vectors: return embeddings + def loadembeddings(self, f): + """ + Loads embeddings from file. + + Args: + f: file to load from + + Returns: + embeddings + """ + + return np.load(f, allow_pickle=False) + + def saveembeddings(self, f, embeddings): + """ + Saves embeddings to output. + + Args: + f: output file + embeddings: embeddings to save + """ + + np.save(f, embeddings, allow_pickle=False) + def truncate(self, embeddings): """ Truncates embeddings to the configured dimensionality. @@ -354,6 +434,9 @@ class Vectors: Args: embeddings: input embeddings + + Returns: + embeddings """ # Calculation is different for matrices vs vectors @@ -362,6 +445,8 @@ class Vectors: else: embeddings /= np.linalg.norm(embeddings) + return embeddings + def quantize(self, embeddings): """ Quantizes embeddings using scalar quantization. diff --git a/src/python/txtai/vectors/dense/__init__.py b/src/python/txtai/vectors/dense/__init__.py new file mode 100644 index 00000000..c7831436 --- /dev/null +++ b/src/python/txtai/vectors/dense/__init__.py @@ -0,0 +1,12 @@ +""" +Dense vectors imports +""" + +from .external import External +from .factory import VectorsFactory +from .huggingface import HFVectors +from .litellm import LiteLLM +from .llama import LlamaCpp +from .m2v import Model2Vec +from .sbert import STVectors +from .words import WordVectors diff --git a/src/python/txtai/vectors/external.py b/src/python/txtai/vectors/dense/external.py similarity index 93% rename from src/python/txtai/vectors/external.py rename to src/python/txtai/vectors/dense/external.py index bb2d7571..f8b10b74 100644 --- a/src/python/txtai/vectors/external.py +++ b/src/python/txtai/vectors/dense/external.py @@ -6,9 +6,9 @@ import types import numpy as np -from ..util import Resolver +from ...util import Resolver -from .base import Vectors +from ..base import Vectors class External(Vectors): @@ -25,7 +25,7 @@ class External(Vectors): def loadmodel(self, path): return None - def encode(self, data): + def encode(self, data, category=None): # Call external transform function, if available and data not already an array # Batching is handed by the external transform function if self.transform and data and not isinstance(data[0], np.ndarray): diff --git a/src/python/txtai/vectors/factory.py b/src/python/txtai/vectors/dense/factory.py similarity index 97% rename from src/python/txtai/vectors/factory.py rename to src/python/txtai/vectors/dense/factory.py index 5ef28fe3..ec6d88b4 100644 --- a/src/python/txtai/vectors/factory.py +++ b/src/python/txtai/vectors/dense/factory.py @@ -2,7 +2,7 @@ Factory module """ -from ..util import Resolver +from ...util import Resolver from .external import External from .huggingface import HFVectors @@ -15,7 +15,7 @@ from .words import WordVectors class VectorsFactory: """ - Methods to create Vectors models. + Methods to create dense vector models. """ @staticmethod diff --git a/src/python/txtai/vectors/huggingface.py b/src/python/txtai/vectors/dense/huggingface.py similarity index 90% rename from src/python/txtai/vectors/huggingface.py rename to src/python/txtai/vectors/dense/huggingface.py index 3db9520c..bbb563d9 100644 --- a/src/python/txtai/vectors/huggingface.py +++ b/src/python/txtai/vectors/dense/huggingface.py @@ -2,9 +2,9 @@ Hugging Face module """ -from ..models import Models, PoolingFactory +from ...models import Models, PoolingFactory -from .base import Vectors +from ..base import Vectors class HFVectors(Vectors): @@ -39,6 +39,6 @@ class HFVectors(Vectors): } ) - def encode(self, data): + def encode(self, data, category=None): # Encode data using vectors model return self.model.encode(data, batch=self.encodebatch) diff --git a/src/python/txtai/vectors/litellm.py b/src/python/txtai/vectors/dense/litellm.py similarity index 95% rename from src/python/txtai/vectors/litellm.py rename to src/python/txtai/vectors/dense/litellm.py index 20d812c9..ee6292ea 100644 --- a/src/python/txtai/vectors/litellm.py +++ b/src/python/txtai/vectors/dense/litellm.py @@ -12,7 +12,7 @@ try: except ImportError: LITELLM = False -from .base import Vectors +from ..base import Vectors class LiteLLM(Vectors): @@ -57,7 +57,7 @@ class LiteLLM(Vectors): def loadmodel(self, path): return None - def encode(self, data): + def encode(self, data, category=None): # Call external embeddings API using LiteLLM # Batching is handled server-side response = api.embedding(model=self.config.get("path"), input=data, **self.config.get("vectors", {})) diff --git a/src/python/txtai/vectors/llama.py b/src/python/txtai/vectors/dense/llama.py similarity index 96% rename from src/python/txtai/vectors/llama.py rename to src/python/txtai/vectors/dense/llama.py index ac6d3fd5..9e5c7470 100644 --- a/src/python/txtai/vectors/llama.py +++ b/src/python/txtai/vectors/dense/llama.py @@ -16,7 +16,7 @@ try: except ImportError: LLAMA_CPP = False -from .base import Vectors +from ..base import Vectors class LlamaCpp(Vectors): @@ -58,7 +58,7 @@ class LlamaCpp(Vectors): # Create llama.cpp instance return Llama(path, n_ctx=0, verbose=modelargs.pop("verbose", False), embedding=True, **modelargs) - def encode(self, data): + def encode(self, data, category=None): # Generate embeddings and return as a NumPy array # llama-cpp-python has it's own batching built-in using n_batch parameter return np.array(self.model.embed(data), dtype=np.float32) diff --git a/src/python/txtai/vectors/m2v.py b/src/python/txtai/vectors/dense/m2v.py similarity index 95% rename from src/python/txtai/vectors/m2v.py rename to src/python/txtai/vectors/dense/m2v.py index 46c2f18d..fa44e4a7 100644 --- a/src/python/txtai/vectors/m2v.py +++ b/src/python/txtai/vectors/dense/m2v.py @@ -15,7 +15,7 @@ try: except ImportError: MODEL2VEC = False -from .base import Vectors +from ..base import Vectors class Model2Vec(Vectors): @@ -59,7 +59,7 @@ class Model2Vec(Vectors): def loadmodel(self, path): return StaticModel.from_pretrained(path) - def encode(self, data): + def encode(self, data, category=None): # Additional model arguments modelargs = self.config.get("vectors", {}) diff --git a/src/python/txtai/vectors/sbert.py b/src/python/txtai/vectors/dense/sbert.py similarity index 62% rename from src/python/txtai/vectors/sbert.py rename to src/python/txtai/vectors/dense/sbert.py index e98bc909..93906c8f 100644 --- a/src/python/txtai/vectors/sbert.py +++ b/src/python/txtai/vectors/dense/sbert.py @@ -1,5 +1,5 @@ """ -SentenceTransformers module +Sentence Transformers module """ # Conditional import @@ -10,9 +10,9 @@ try: except ImportError: SENTENCE_TRANSFORMERS = False -from ..models import Models +from ...models import Models -from .base import Vectors +from ..base import Vectors class STVectors(Vectors): @@ -48,8 +48,8 @@ class STVectors(Vectors): # Additional model arguments modelargs = self.config.get("vectors", {}) - # Build embeddings with sentence-transformers - model = SentenceTransformer(path, device=Models.device(deviceid), **modelargs) + # Load sentence-transformers encoder + model = self.loadencoder(path, device=Models.device(deviceid), **modelargs) # Start process pool for multiple GPUs if pool: @@ -58,9 +58,15 @@ class STVectors(Vectors): # Return model return model - def encode(self, data): - # Encode with sentence transformers - return self.model.encode(data, pool=self.pool, batch_size=self.encodebatch) + def encode(self, data, category=None): + # Get encode method based on input category + encode = self.model.encode_query if category == "query" else self.model.encode_document if category == "data" else self.model.encode + + # Additional encoding arguments + encodeargs = self.config.get("encodeargs", {}) + + # Encode with sentence transformers encoder + return encode(data, pool=self.pool, batch_size=self.encodebatch, **encodeargs) def close(self): # Close pool before model is closed in parent method @@ -69,3 +75,18 @@ class STVectors(Vectors): self.pool = None super().close() + + def loadencoder(self, path, device, **kwargs): + """ + Loads the embeddings encoder model from path. + + Args: + path: model path + device: tensor device + kwargs: additional keyword args + + Returns: + embeddings encoder + """ + + return SentenceTransformer(path, device=device, **kwargs) diff --git a/src/python/txtai/vectors/words.py b/src/python/txtai/vectors/dense/words.py similarity index 97% rename from src/python/txtai/vectors/words.py rename to src/python/txtai/vectors/dense/words.py index 81e6700a..b748e990 100644 --- a/src/python/txtai/vectors/words.py +++ b/src/python/txtai/vectors/dense/words.py @@ -22,9 +22,9 @@ try: except ImportError: STATICVECTORS = False -from ..pipeline import Tokenizer +from ...pipeline import Tokenizer -from .base import Vectors +from ..base import Vectors # Logging configuration logger = logging.getLogger(__name__) @@ -128,7 +128,7 @@ class WordVectors(Vectors): def loadmodel(self, path): return StaticVectors(path) - def encode(self, data): + def encode(self, data, category=None): # Iterate over each data element, tokenize (if necessary) and build an aggregated embeddings vector embeddings = [] for tokens in data: @@ -181,14 +181,14 @@ class WordVectors(Vectors): embeddings.append(embedding) if len(embeddings) == batchsize: - np.save(output, np.array(embeddings, dtype=np.float32)) + np.save(output, np.array(embeddings, dtype=np.float32), allow_pickle=False) batches += 1 embeddings = [] # Final embeddings batch if embeddings: - np.save(output, np.array(embeddings, dtype=np.float32)) + np.save(output, np.array(embeddings, dtype=np.float32), allow_pickle=False) batches += 1 return (ids, dimensions, batches, stream) diff --git a/src/python/txtai/vectors/recovery.py b/src/python/txtai/vectors/recovery.py index 74930404..f0dfb7c8 100644 --- a/src/python/txtai/vectors/recovery.py +++ b/src/python/txtai/vectors/recovery.py @@ -5,24 +5,23 @@ Recovery module import os import shutil -import numpy as np - class Recovery: """ Vector embeddings recovery. This class handles streaming embeddings from a vector checkpoint file. """ - def __init__(self, checkpoint, vectorsid): + def __init__(self, checkpoint, vectorsid, load): """ Creates a Recovery instance. Args: checkpoint: checkpoint directory vectorsid: vectors uid for current configuration + load: load embeddings method """ - self.spool, self.path = None, None + self.spool, self.path, self.load = None, None, load # Get unique file id path = f"{checkpoint}/{vectorsid}" @@ -46,7 +45,7 @@ class Recovery: """ try: - return np.load(self.spool) if self.spool else None + return self.load(self.spool) if self.spool else None except EOFError: # End of spool file, cleanup self.spool.close() diff --git a/src/python/txtai/vectors/sparse/__init__.py b/src/python/txtai/vectors/sparse/__init__.py new file mode 100644 index 00000000..dfce5a8a --- /dev/null +++ b/src/python/txtai/vectors/sparse/__init__.py @@ -0,0 +1,7 @@ +""" +Sparse vectors imports +""" + +from .base import SparseVectors +from .factory import SparseVectorsFactory +from .sbert import SparseSTVectors diff --git a/src/python/txtai/vectors/sparse/base.py b/src/python/txtai/vectors/sparse/base.py new file mode 100644 index 00000000..11736919 --- /dev/null +++ b/src/python/txtai/vectors/sparse/base.py @@ -0,0 +1,76 @@ +""" +SparseVectors module +""" + +# Conditional import +try: + from scipy.sparse import csr_matrix, vstack + from sklearn.preprocessing import normalize + from sklearn.utils.extmath import safe_sparse_dot + + SPARSE = True +except ImportError: + SPARSE = False + +from ...util import SparseArray +from ..base import Vectors + + +# pylint: disable=W0223 +class SparseVectors(Vectors): + """ + Base class for sparse vector models. Vector models transform input content into sparse arrays. + """ + + def __init__(self, config, scoring, models): + # Check before parent constructor since it calls loadmodel + if not SPARSE: + raise ImportError('SparseVectors is not available - install "vectors" extra to enable') + + super().__init__(config, scoring, models) + + def encode(self, data, category=None): + # Encode data to embeddings + embeddings = super().encode(data, category) + + # Get sparse torch vector attributes + embeddings = embeddings.cpu().coalesce() + indices = embeddings.indices().numpy() + values = embeddings.values().numpy() + + # Return as SciPy CSR Matrix + return csr_matrix((values, indices), shape=embeddings.size()) + + def vectors(self, documents, batchsize=500, checkpoint=None, buffer=None, dtype=None): + # Run indexing + ids, dimensions, batches, stream = self.index(documents, batchsize, checkpoint) + + # Rebuild sparse array + embeddings = None + with open(stream, "rb") as queue: + for _ in range(batches): + # Read in array batch + data = self.loadembeddings(queue) + embeddings = vstack((embeddings, data)) if embeddings is not None else data + + # Return sparse array + return (ids, dimensions, embeddings) + + def dot(self, queries, data): + return safe_sparse_dot(queries, data.T, dense_output=True).tolist() + + def loadembeddings(self, f): + return SparseArray().load(f) + + def saveembeddings(self, f, embeddings): + SparseArray().save(f, embeddings) + + def truncate(self, embeddings): + raise ValueError("Truncate is not supported for sparse vectors") + + def normalize(self, embeddings): + # Normalize in place using method that supports sparse vectors + return normalize(embeddings, copy=False) + + def quantize(self, embeddings): + raise ValueError("Quantize is not supported for sparse vectors") diff --git a/src/python/txtai/vectors/sparse/factory.py b/src/python/txtai/vectors/sparse/factory.py new file mode 100644 index 00000000..e9240716 --- /dev/null +++ b/src/python/txtai/vectors/sparse/factory.py @@ -0,0 +1,55 @@ +""" +Factory module +""" + +from ...util import Resolver + +from .sbert import SparseSTVectors + + +class SparseVectorsFactory: + """ + Methods to create sparse vector models. + """ + + @staticmethod + def create(config, models=None): + """ + Create a Vectors model instance. + + Args: + config: vector configuration + models: models cache + + Returns: + Vectors + """ + + # Get vector method + method = config.get("method", "sentence-transformers") + + # Sentence Transformers vectors + if method == "sentence-transformers": + return SparseSTVectors(config, None, models) if config and config.get("path") else None + + # Resolve custom method + return SparseVectorsFactory.resolve(method, config, models) if method else None + + @staticmethod + def resolve(backend, config, models): + """ + Attempt to resolve a custom backend. + + Args: + backend: backend class + config: vector configuration + models: models cache + + Returns: + Vectors + """ + + try: + return Resolver()(backend)(config, None, models) + except Exception as e: + raise ImportError(f"Unable to resolve sparse vectors backend: '{backend}'") from e diff --git a/src/python/txtai/vectors/sparse/sbert.py b/src/python/txtai/vectors/sparse/sbert.py new file mode 100644 index 00000000..4c76e0a6 --- /dev/null +++ b/src/python/txtai/vectors/sparse/sbert.py @@ -0,0 +1,30 @@ +""" +Sparse Sentence Transformers module +""" + +# Conditional import +try: + from sentence_transformers import SparseEncoder + + SENTENCE_TRANSFORMERS = True +except ImportError: + SENTENCE_TRANSFORMERS = False + +from ..dense.sbert import STVectors +from .base import SparseVectors + + +class SparseSTVectors(SparseVectors, STVectors): + """ + Builds sparse vectors using sentence-transformers (aka SBERT). + """ + + def __init__(self, config, scoring, models): + # Check before parent constructor since it calls loadmodel + if not SENTENCE_TRANSFORMERS: + raise ImportError('sentence-transformers is not available - install "vectors" extra to enable') + + super().__init__(config, scoring, models) + + def loadencoder(self, path, device, **kwargs): + return SparseEncoder(path, device=device, **kwargs) diff --git a/test/python/testann/__init__.py b/test/python/testann/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/python/testann.py b/test/python/testann/testdense.py similarity index 98% rename from test/python/testann.py rename to test/python/testann/testdense.py index de31ebe9..5f689a11 100644 --- a/test/python/testann.py +++ b/test/python/testann/testdense.py @@ -1,5 +1,5 @@ """ -ANN module tests +Dense ANN module tests """ import os @@ -17,9 +17,9 @@ from txtai.serialize import SerializeFactory # pylint: disable=R0904 -class TestANN(unittest.TestCase): +class TestDense(unittest.TestCase): """ - ANN tests. + Dense ANN tests. """ def testAnnoy(self): @@ -79,10 +79,10 @@ class TestANN(unittest.TestCase): # pylint: disable=C0415, W0611 # Force reload of class - name = "txtai.ann.faiss" + name = "txtai.ann.dense.faiss" module = sys.modules[name] del sys.modules[name] - import txtai.ann.faiss + import txtai.ann.dense.faiss # Run tests self.runTests("faiss") diff --git a/test/python/testann/testsparse.py b/test/python/testann/testsparse.py new file mode 100644 index 00000000..403f6f39 --- /dev/null +++ b/test/python/testann/testsparse.py @@ -0,0 +1,131 @@ +""" +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() + + @patch("sqlalchemy.orm.Query.limit") + def testPGSparse(self, query): + """ + Test Sparse Postgres backend + """ + + # Generate test record + data = self.generate(1, 240) + + # 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": 240, "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) + + # 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) diff --git a/test/python/testembeddings.py b/test/python/testembeddings.py index 3e157d92..27816d91 100644 --- a/test/python/testembeddings.py +++ b/test/python/testembeddings.py @@ -486,6 +486,10 @@ class TestEmbeddings(unittest.TestCase): 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]]) diff --git a/test/python/testoptional.py b/test/python/testoptional.py index 54f42511..f1a88b2e 100644 --- a/test/python/testoptional.py +++ b/test/python/testoptional.py @@ -107,8 +107,9 @@ class TestOptional(unittest.TestCase): Test missing ANN dependencies """ - from txtai.ann import ANNFactory + from txtai.ann import ANNFactory, SparseANNFactory + # Test dense methods with self.assertRaises(ImportError): ANNFactory.create({"backend": "annoy"}) @@ -121,6 +122,13 @@ class TestOptional(unittest.TestCase): with self.assertRaises(ImportError): ANNFactory.create({"backend": "sqlite"}) + # Test sparse methods + with self.assertRaises(ImportError): + SparseANNFactory.create({"backend": "ivfsparse"}) + + with self.assertRaises(ImportError): + SparseANNFactory.create({"backend": "pgsparse"}) + def testApi(self): """ Test missing api dependencies @@ -294,24 +302,20 @@ class TestOptional(unittest.TestCase): Test missing scoring dependencies """ - from txtai.scoring import IVFFlat, ScoringFactory + from txtai.scoring import ScoringFactory with self.assertRaises(ImportError): ScoringFactory.create({"method": "pgtext"}) - with self.assertRaises(ImportError): - ScoringFactory.create({"method": "sparse"}) - - with self.assertRaises(ImportError): - IVFFlat({}) - def testVectors(self): """ Test missing vector dependencies """ - from txtai.vectors import VectorsFactory + from txtai.vectors import SparseVectors, VectorsFactory, SparseVectorsFactory + from txtai.util import SparseArray + # Test dense vectors with self.assertRaises(ImportError): VectorsFactory.create({"method": "litellm", "path": "huggingface/sentence-transformers/all-MiniLM-L6-v2"}, None) @@ -331,6 +335,16 @@ class TestOptional(unittest.TestCase): model = VectorsFactory.create({"path": "sentence-transformers/all-MiniLM-L6-v2"}, None) self.assertIsNotNone(model) + # Test sparse vectors + with self.assertRaises(ImportError): + SparseVectors(None, None, None) + + with self.assertRaises(ImportError): + SparseVectorsFactory.create({"method": "sentence-transformers", "path": "sparse-encoder-testing/splade-bert-tiny-nq"}, None) + + with self.assertRaises(ImportError): + SparseArray() + def testWorkflow(self): """ Test missing workflow dependencies diff --git a/test/python/testscoring/__init__.py b/test/python/testscoring/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/python/testscoring.py b/test/python/testscoring/testkeyword.py similarity index 78% rename from test/python/testscoring.py rename to test/python/testscoring/testkeyword.py index aa5b16b3..6f0b8446 100644 --- a/test/python/testscoring.py +++ b/test/python/testscoring/testkeyword.py @@ -1,5 +1,5 @@ """ -Scoring module tests +Keyword scoring tests """ import os @@ -12,9 +12,9 @@ from txtai.scoring import ScoringFactory, Scoring # pylint: disable=R0904 -class TestScoring(unittest.TestCase): +class TestKeyword(unittest.TestCase): """ - Scoring tests. + Sparse keyword scoring tests. """ @classmethod @@ -73,7 +73,7 @@ class TestScoring(unittest.TestCase): self.assertRaises(NotImplementedError, scoring.load, None) self.assertRaises(NotImplementedError, scoring.save, None) self.assertRaises(NotImplementedError, scoring.close) - self.assertRaises(NotImplementedError, scoring.hasterms) + self.assertRaises(NotImplementedError, scoring.issparse) self.assertRaises(NotImplementedError, scoring.isnormalized) @patch("sqlalchemy.orm.Query.params") @@ -109,8 +109,8 @@ class TestScoring(unittest.TestCase): scoring.delete([0]) self.assertEqual(scoring.count(), len(self.data) - 1) - # PGText is a normalized terms index - self.assertTrue(scoring.hasterms() and scoring.isnormalized()) + # PGText is a normalized sparse index + self.assertTrue(scoring.issparse() and scoring.isnormalized()) self.assertIsNone(scoring.weights("This is a test".split())) # Close scoring @@ -123,93 +123,6 @@ class TestScoring(unittest.TestCase): self.runTests("sif") - def testSparse(self): - """ - Test sparse vectors - """ - - # Models cache - models = {} - - # Test sparse scoring - scoring = ScoringFactory.create({"method": "sparse", "path": "sparse-encoder-testing/splade-bert-tiny-nq"}, models=models) - scoring.index((uid, {"text": text}, tags) for uid, text, tags in self.data) - - # Run search and validate correct result returned - index, _ = scoring.search("lottery ticket", 1)[0] - self.assertEqual(index, 4) - - # Run batch search - index, _ = scoring.batchsearch(["lottery ticket"], 1)[0][0] - self.assertEqual(index, 4) - - # Validate count - self.assertEqual(scoring.count(), len(self.data)) - - # Test delete - scoring.delete([4]) - self.assertEqual(scoring.count(), len(self.data) - 1) - - # Run search after delete - index, _ = scoring.search("lottery ticket", 1)[0] - self.assertEqual(index, 5) - - # Sparse vectors is a normalized terms index - self.assertTrue(scoring.hasterms() and scoring.isnormalized()) - self.assertIsNone(scoring.weights("This is a test".split())) - - # Close scoring - scoring.close() - - # Test model caching - scoring = ScoringFactory.create({"method": "sparse", "path": "sparse-encoder-testing/splade-bert-tiny-nq"}, models=models) - self.assertIsNotNone(scoring.model) - scoring.close() - - def testSparseEmpty(self): - """ - Test empty sparse vectors - """ - - scoring = ScoringFactory.create({"method": "sparse", "path": "sparse-encoder-testing/splade-bert-tiny-nq"}) - scoring.upsert((uid, {"text": text}, tags) for uid, text, tags in self.data) - self.assertEqual(scoring.count(), len(self.data)) - - @patch("torch.cuda.device_count") - def testSparseGPU(self, count): - """ - Test sparse vectors with GPU encoding - """ - - # Mock accelerator count - count.return_value = 2 - - # Test multiple gpus - scoring = ScoringFactory.create({"method": "sparse", "path": "sparse-encoder-testing/splade-bert-tiny-nq", "gpu": "all"}) - self.assertIsNotNone(scoring) - scoring.close() - - def testSparseIVFFlat(self): - """ - Test sparse vectors with IVFFlat clustering - """ - - # Expand dataset - data = self.data * 1000 - - # Test higher volume IVFFlat index with clustering - config = {"method": "sparse", "path": "sparse-encoder-testing/splade-bert-tiny-nq", "ivf": {"sample": 1.0}} - scoring = ScoringFactory.create(config) - scoring.index((uid, {"text": text}, tags) for uid, text, tags in data) - - # Recreate index - self.save(scoring, config, "scoring.sparse.index") - - # Run search and validate correct result returned - results = scoring.search("lottery ticket", 1) - self.assertGreater(len(results), 0) - scoring.close() - def testTFIDF(self): """ Test tfidf diff --git a/test/python/testscoring/testsparse.py b/test/python/testscoring/testsparse.py new file mode 100644 index 00000000..7471b212 --- /dev/null +++ b/test/python/testscoring/testsparse.py @@ -0,0 +1,136 @@ +""" +Sparse module tests +""" + +import os +import tempfile +import unittest + +from unittest.mock import patch + +from txtai.scoring import ScoringFactory + + +# pylint: disable=R0904 +class TestSparse(unittest.TestCase): + """ + Sparse vector scoring 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", + "wins wins wins", + "Make huge profits without work, earn up to $100,000 a day", + ] + + cls.data = [(uid, x, None) for uid, x in enumerate(cls.data)] + + def testGeneral(self): + """ + Test general sparse vector operations + """ + + # Models cache + models = {} + + # Test sparse scoring + scoring = ScoringFactory.create({"method": "sparse", "path": "sparse-encoder-testing/splade-bert-tiny-nq"}, models=models) + scoring.index((uid, {"text": text}, tags) for uid, text, tags in self.data) + + # Run search and validate correct result returned + index, _ = scoring.search("lottery ticket", 1)[0] + self.assertEqual(index, 4) + + # Run batch search + index, _ = scoring.batchsearch(["lottery ticket"], 1)[0][0] + self.assertEqual(index, 4) + + # Validate count + self.assertEqual(scoring.count(), len(self.data)) + + # Test delete + scoring.delete([4]) + self.assertEqual(scoring.count(), len(self.data) - 1) + + # Run search after delete + index, _ = scoring.search("lottery ticket", 1)[0] + self.assertEqual(index, 5) + + # Sparse vectors is a normalized sparse index + self.assertTrue(scoring.issparse() and scoring.isnormalized()) + self.assertIsNone(scoring.weights("This is a test".split())) + + # Close scoring + scoring.close() + + # Test model caching + scoring = ScoringFactory.create({"method": "sparse", "path": "sparse-encoder-testing/splade-bert-tiny-nq"}, models=models) + self.assertIsNotNone(scoring.model) + scoring.close() + + def testEmpty(self): + """ + Test empty sparse vectors + """ + + scoring = ScoringFactory.create({"method": "sparse", "path": "sparse-encoder-testing/splade-bert-tiny-nq"}) + scoring.upsert((uid, {"text": text}, tags) for uid, text, tags in self.data) + self.assertEqual(scoring.count(), len(self.data)) + + @patch("torch.cuda.device_count") + def testGPU(self, count): + """ + Test sparse vectors with GPU encoding + """ + + # Mock accelerator count + count.return_value = 2 + + # Test multiple gpus + scoring = ScoringFactory.create({"method": "sparse", "path": "sparse-encoder-testing/splade-bert-tiny-nq", "gpu": "all"}) + self.assertIsNotNone(scoring) + scoring.close() + + def testIVFFlat(self): + """ + Test sparse vectors with IVFFlat clustering + """ + + # Expand dataset + data = self.data * 1000 + + # Test higher volume IVFFlat index with clustering + config = { + "method": "sparse", + "vectormethod": "sentence-transformers", + "path": "sparse-encoder-testing/splade-bert-tiny-nq", + "ivfsparse": {"sample": 1.0}, + } + scoring = ScoringFactory.create(config) + scoring.index((uid, {"text": text}, tags) for uid, text, tags in data) + + # Generate temp file path + index = os.path.join(tempfile.gettempdir(), "scoring") + os.makedirs(index, exist_ok=True) + + # Save scoring instance + scoring.save(f"{index}/scoring.sparse.index") + + # Reload scoring instance + scoring = ScoringFactory.create(config) + scoring.load(f"{index}/scoring.sparse.index") + + # Run search and validate correct result returned + results = scoring.search("lottery ticket", 1) + self.assertGreater(len(results), 0) + scoring.close() diff --git a/test/python/testvectors/testdense/__init__.py b/test/python/testvectors/testdense/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/python/testvectors/testcustom.py b/test/python/testvectors/testdense/testcustom.py similarity index 100% rename from test/python/testvectors/testcustom.py rename to test/python/testvectors/testdense/testcustom.py diff --git a/test/python/testvectors/testexternal.py b/test/python/testvectors/testdense/testexternal.py similarity index 100% rename from test/python/testvectors/testexternal.py rename to test/python/testvectors/testdense/testexternal.py diff --git a/test/python/testvectors/testhuggingface.py b/test/python/testvectors/testdense/testhuggingface.py similarity index 100% rename from test/python/testvectors/testhuggingface.py rename to test/python/testvectors/testdense/testhuggingface.py diff --git a/test/python/testvectors/testlitellm.py b/test/python/testvectors/testdense/testlitellm.py similarity index 100% rename from test/python/testvectors/testlitellm.py rename to test/python/testvectors/testdense/testlitellm.py diff --git a/test/python/testvectors/testllama.py b/test/python/testvectors/testdense/testllama.py similarity index 100% rename from test/python/testvectors/testllama.py rename to test/python/testvectors/testdense/testllama.py diff --git a/test/python/testvectors/testm2v.py b/test/python/testvectors/testdense/testm2v.py similarity index 100% rename from test/python/testvectors/testm2v.py rename to test/python/testvectors/testdense/testm2v.py diff --git a/test/python/testvectors/testsbert.py b/test/python/testvectors/testdense/testsbert.py similarity index 100% rename from test/python/testvectors/testsbert.py rename to test/python/testvectors/testdense/testsbert.py diff --git a/test/python/testvectors/testvectors.py b/test/python/testvectors/testdense/testvectors.py similarity index 96% rename from test/python/testvectors/testvectors.py rename to test/python/testvectors/testdense/testvectors.py index d28a7612..e3cb0dc9 100644 --- a/test/python/testvectors/testvectors.py +++ b/test/python/testvectors/testdense/testvectors.py @@ -64,5 +64,5 @@ class TestVectors(unittest.TestCase): f.close() # Create the recovery instance with an empty checkpoint file - recovery = Recovery(checkpoint, "id") + recovery = Recovery(checkpoint, "id", np.load) self.assertIsNone(recovery()) diff --git a/test/python/testvectors/testwordvectors.py b/test/python/testvectors/testdense/testwordvectors.py similarity index 98% rename from test/python/testvectors/testwordvectors.py rename to test/python/testvectors/testdense/testwordvectors.py index 13eb35e3..7ee98e4b 100644 --- a/test/python/testvectors/testwordvectors.py +++ b/test/python/testvectors/testdense/testwordvectors.py @@ -12,7 +12,7 @@ import numpy as np from huggingface_hub.errors import HFValidationError from txtai.vectors import VectorsFactory -from txtai.vectors.words import create, transform +from txtai.vectors.dense.words import create, transform class TestWordVectors(unittest.TestCase): diff --git a/test/python/testvectors/testsparse/__init__.py b/test/python/testvectors/testsparse/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/python/testvectors/testsparse/testsbert.py b/test/python/testvectors/testsparse/testsbert.py new file mode 100644 index 00000000..f18af43f --- /dev/null +++ b/test/python/testvectors/testsparse/testsbert.py @@ -0,0 +1,32 @@ +""" +Sparse Sentence Transformers module tests +""" + +import os +import unittest + +from txtai.vectors import SparseVectorsFactory +from txtai.util import SparseArray + + +class TestSparseSTVectors(unittest.TestCase): + """ + SparseSTVectors tests + """ + + def testIndex(self): + """ + Test indexing with sentence-transformers vectors + """ + + model = SparseVectorsFactory.create({"method": "sentence-transformers", "path": "sparse-encoder-testing/splade-bert-tiny-nq"}) + ids, dimension, batches, stream = model.index([(0, "test", None)]) + + self.assertEqual(len(ids), 1) + self.assertEqual(dimension, 30522) + self.assertEqual(batches, 1) + self.assertIsNotNone(os.path.exists(stream)) + + # Test shape of serialized embeddings + with open(stream, "rb") as queue: + self.assertEqual(SparseArray().load(queue).shape, (1, 30522)) diff --git a/test/python/testvectors/testsparse/testvectors.py b/test/python/testvectors/testsparse/testvectors.py new file mode 100644 index 00000000..63e18ade --- /dev/null +++ b/test/python/testvectors/testsparse/testvectors.py @@ -0,0 +1,40 @@ +""" +Sparse Vectors module tests +""" + +import unittest + +from txtai.vectors import SparseVectors, SparseVectorsFactory + + +class TestSparseVectors(unittest.TestCase): + """ + Sparse Vectors tests. + """ + + def testCustom(self): + """ + Test custom sparse vectors instance + """ + + self.assertIsNotNone( + SparseVectorsFactory.create({"method": "txtai.vectors.SparseSTVectors", "path": "sparse-encoder-testing/splade-bert-tiny-nq"}) + ) + + def testNotSupported(self): + """ + Test exceptions for unsupported methods + """ + + vectors = SparseVectors(None, None, None) + + self.assertRaises(ValueError, vectors.truncate, None) + self.assertRaises(ValueError, vectors.quantize, None) + + def testNotFound(self): + """ + Test unresolvable vector backend + """ + + with self.assertRaises(ImportError): + SparseVectorsFactory.create({"method": "notfound.vectors"})