Refactor of ann and vectors packages for sparse vectors. Closes #930. Closes #934. Closes #935. Closes #936. Closes #937.

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

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

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 *

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

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

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

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.

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):

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):

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.

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):

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

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

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)

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)

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.

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):
         """

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:

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

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

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.

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):
         """

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):

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()

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]
 

src/python/txtai/util/__init__.py

@@ -3,4 +3,5 @@ Utility imports
 """
 
 from .resolver import Resolver
+from .sparsearray import SparseArray
 from .template import TemplateFormatter

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)

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 *

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.

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

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):

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

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)

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", {}))

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)

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", {})
 

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)

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)

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()

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

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")

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

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)

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")

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)

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]])

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

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

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()

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())

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):

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))

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"})