mseep-txtai 9.1.1__py3-none-any.whl

txtai/ann/dense/numpy.py

@@ -0,0 +1,164 @@
+"""
+NumPy module
+"""
+
+import numpy as np
+
+from ...serialize import SerializeFactory
+
+from ..base import ANN
+
+
+class NumPy(ANN):
+    """
+    Builds an ANN index backed by a NumPy array.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        # Array function definitions
+        self.all, self.cat, self.dot, self.zeros = np.all, np.concatenate, np.dot, np.zeros
+        self.argsort, self.xor, self.clip = np.argsort, np.bitwise_xor, np.clip
+
+        # Scalar quantization
+        quantize = self.config.get("quantize")
+        self.qbits = quantize if quantize and isinstance(quantize, int) and not isinstance(quantize, bool) else None
+
+    def load(self, path):
+        # Load array from file
+        try:
+            self.backend = self.tensor(np.load(path, allow_pickle=False))
+        except ValueError:
+            # Backwards compatible support for previously pickled data
+            self.backend = self.tensor(SerializeFactory.create("pickle").load(path))
+
+    def index(self, embeddings):
+        # Create index
+        self.backend = self.tensor(embeddings)
+
+        # Add id offset and index build metadata
+        self.config["offset"] = embeddings.shape[0]
+        self.metadata(self.settings())
+
+    def append(self, embeddings):
+        # Append new data to array
+        self.backend = self.cat((self.backend, self.tensor(embeddings)), axis=0)
+
+        # Update id offset and index metadata
+        self.config["offset"] += embeddings.shape[0]
+        self.metadata()
+
+    def delete(self, ids):
+        # Filter any index greater than size of array
+        ids = [x for x in ids if x < self.backend.shape[0]]
+
+        # Clear specified ids
+        self.backend[ids] = self.tensor(self.zeros((len(ids), self.backend.shape[1])))
+
+    def search(self, queries, limit):
+        if self.qbits:
+            # Calculate hamming score for integer vectors
+            scores = self.hammingscore(queries)
+        else:
+            # Dot product on normalized vectors is equal to cosine similarity
+            scores = self.dot(self.tensor(queries), self.backend.T)
+
+        # Get topn ids
+        ids = self.argsort(-scores)[:, :limit]
+
+        # Map results to [(id, score)]
+        results = []
+        for x, score in enumerate(scores):
+            # Add results
+            results.append(list(zip(ids[x].tolist(), score[ids[x]].tolist())))
+
+        return results
+
+    def count(self):
+        # Get count of non-zero rows (ignores deleted rows)
+        return self.backend[~self.all(self.backend == 0, axis=1)].shape[0]
+
+    def save(self, path):
+        # Save array to file. Use stream to prevent ".npy" suffix being added.
+        with open(path, "wb") as handle:
+            np.save(handle, self.numpy(self.backend), allow_pickle=False)
+
+    def tensor(self, array):
+        """
+        Handles backend-specific code such as loading to a GPU device.
+
+        Args:
+            array: data array
+
+        Returns:
+            array with backend-specific logic applied
+        """
+
+        return array
+
+    def numpy(self, array):
+        """
+        Handles backend-specific code to convert an array to numpy
+
+        Args:
+            array: data array
+
+        Returns:
+            numpy array
+        """
+
+        return array
+
+    def totype(self, array, dtype):
+        """
+        Casts array to dtype.
+
+        Args:
+            array: input array
+            dtype: dtype
+
+        Returns:
+            array cast as dtype
+        """
+
+        return np.int64(array) if dtype == np.int64 else array
+
+    def settings(self):
+        """
+        Returns settings for this array.
+
+        Returns:
+            dict
+        """
+
+        return {"numpy": np.__version__}
+
+    def hammingscore(self, queries):
+        """
+        Calculates a hamming distance score.
+
+        This is defined as:
+
+            score = 1.0 - (hamming distance / total number of bits)
+
+        Args:
+            queries: queries array
+
+        Returns:
+            scores
+        """
+
+        # Build table of number of bits for each distinct uint8 value
+        table = 1 << np.arange(8)
+        table = self.tensor(np.array([np.count_nonzero(x & table) for x in np.arange(256)]))
+
+        # Number of different bits
+        delta = self.xor(self.tensor(queries[:, None]), self.backend)
+
+        # Cast to long array
+        delta = self.totype(delta, np.int64)
+
+        # Calculate score as 1.0 - percentage of different bits
+        # Bound score from 0 to 1
+        return self.clip(1.0 - (table[delta].sum(axis=2) / (self.config["dimensions"] * 8)), 0.0, 1.0)

txtai/ann/dense/pgvector.py

@@ -0,0 +1,323 @@
+"""
+PGVector module
+"""
+
+import os
+
+import numpy as np
+
+# Conditional import
+try:
+    from pgvector.sqlalchemy import BIT, HALFVEC, VECTOR
+
+    from sqlalchemy import create_engine, delete, func, text, Column, Index, Integer, MetaData, StaticPool, Table
+    from sqlalchemy.orm import Session
+    from sqlalchemy.schema import CreateSchema
+
+    PGVECTOR = True
+except ImportError:
+    PGVECTOR = False
+
+from ..base import ANN
+
+
+# pylint: disable=R0904
+class PGVector(ANN):
+    """
+    Builds an ANN index backed by a Postgres database.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        if not PGVECTOR:
+            raise ImportError('PGVector is not available - install "ann" extra to enable')
+
+        # Database connection
+        self.engine, self.database, self.connection, self.table = None, None, None, None
+
+        # Scalar quantization
+        quantize = self.config.get("quantize")
+        self.qbits = quantize if quantize and isinstance(quantize, int) and not isinstance(quantize, bool) else None
+
+    def load(self, path):
+        # Initialize tables
+        self.initialize()
+
+    def index(self, embeddings):
+        # Initialize tables
+        self.initialize(recreate=True)
+
+        # Prepare embeddings and insert rows
+        self.database.execute(self.table.insert(), [{"indexid": x, "embedding": self.prepare(row)} for x, row in enumerate(embeddings)])
+
+        # Create index
+        self.createindex()
+
+        # Add id offset and index build metadata
+        self.config["offset"] = embeddings.shape[0]
+        self.metadata(self.settings())
+
+    def append(self, embeddings):
+        # Prepare embeddings and insert rows
+        self.database.execute(
+            self.table.insert(), [{"indexid": x + self.config["offset"], "embedding": self.prepare(row)} for x, row in enumerate(embeddings)]
+        )
+
+        # Update id offset and index metadata
+        self.config["offset"] += embeddings.shape[0]
+        self.metadata()
+
+    def delete(self, ids):
+        self.database.execute(delete(self.table).where(self.table.c["indexid"].in_(ids)))
+
+    def search(self, queries, limit):
+        results = []
+        for query in queries:
+            # Run query
+            query = self.database.query(self.table.c["indexid"], self.query(query)).order_by("score").limit(limit)
+
+            # Calculate and collect scores
+            results.append([(indexid, self.score(score)) for indexid, score in query])
+
+        return results
+
+    def count(self):
+        # pylint: disable=E1102
+        return self.database.query(func.count(self.table.c["indexid"])).scalar()
+
+    def save(self, path):
+        # Commit session and connection
+        self.database.commit()
+        self.connection.commit()
+
+    def close(self):
+        # Parent logic
+        super().close()
+
+        # Close database connection
+        if self.database:
+            self.database.close()
+            self.engine.dispose()
+
+    def initialize(self, recreate=False):
+        """
+        Initializes a new database session.
+
+        Args:
+            recreate: Recreates the database tables if True
+        """
+
+        # Connect to database
+        self.connect()
+
+        # Set the database schema
+        self.schema()
+
+        # Table name
+        table = self.setting("table", self.defaulttable())
+
+        # Create vectors table object
+        self.table = Table(table, MetaData(), Column("indexid", Integer, primary_key=True, autoincrement=False), Column("embedding", self.column()))
+
+        # Drop table, if necessary
+        if recreate:
+            self.table.drop(self.connection, checkfirst=True)
+
+        # Create table, if necessary
+        self.table.create(self.connection, checkfirst=True)
+
+    def createindex(self):
+        """
+        Creates a index with the current settings.
+        """
+
+        # Table name
+        table = self.setting("table", self.defaulttable())
+
+        # Create ANN index - inner product is equal to cosine similarity on normalized vectors
+        index = Index(
+            f"{table}-index",
+            self.table.c["embedding"],
+            postgresql_using="hnsw",
+            postgresql_with=self.settings(),
+            postgresql_ops={"embedding": self.operation()},
+        )
+
+        # Create or recreate index
+        index.drop(self.connection, checkfirst=True)
+        index.create(self.connection, checkfirst=True)
+
+    def connect(self):
+        """
+        Establishes a database connection. Cleans up any existing database connection first.
+        """
+
+        # Close existing connection
+        if self.database:
+            self.close()
+
+        # Create engine
+        self.engine = create_engine(self.url(), poolclass=StaticPool, echo=False)
+        self.connection = self.engine.connect()
+
+        # Start database session
+        self.database = Session(self.connection)
+
+        # Initialize pgvector extension
+        self.sqldialect(text("CREATE EXTENSION IF NOT EXISTS vector"))
+
+    def schema(self):
+        """
+        Sets the database schema, if available.
+        """
+
+        # Set default schema, if necessary
+        schema = self.setting("schema")
+        if schema:
+            with self.engine.begin():
+                self.sqldialect(CreateSchema(schema, if_not_exists=True))
+
+            self.sqldialect(text("SET search_path TO :schema,public"), {"schema": schema})
+
+    def settings(self):
+        """
+        Returns settings for this index.
+
+        Returns:
+            dict
+        """
+
+        return {"m": self.setting("m", 16), "ef_construction": self.setting("efconstruction", 200)}
+
+    def sqldialect(self, sql, parameters=None):
+        """
+        Executes a SQL statement based on the current SQL dialect.
+
+        Args:
+            sql: SQL to execute
+            parameters: optional bind parameters
+        """
+
+        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 for the current settings.
+
+        Returns:
+            embedding column definition
+        """
+
+        if self.qbits:
+            # If quantization is set, always return BIT vectors
+            return BIT(self.config["dimensions"] * 8)
+
+        if self.setting("precision") == "half":
+            # 16-bit HALF precision vectors
+            return HALFVEC(self.config["dimensions"])
+
+        # Default is full 32-bit FULL precision vectors
+        return VECTOR(self.config["dimensions"])
+
+    def operation(self):
+        """
+        Gets the index operation for the current settings.
+
+        Returns:
+            index operation
+        """
+
+        if self.qbits:
+            # If quantization is set, always return BIT vectors
+            return "bit_hamming_ops"
+
+        if self.setting("precision") == "half":
+            # 16-bit HALF precision vectors
+            return "halfvec_ip_ops"
+
+        # Default is full 32-bit FULL precision vectors
+        return "vector_ip_ops"
+
+    def prepare(self, data):
+        """
+        Prepares data for the embeddings column. This method returns a bit string for bit vectors and
+        the input data unmodified for float vectors.
+
+        Args:
+            data: input data
+
+        Returns:
+            data ready for the embeddings column
+        """
+
+        # Transform to a bit string when vector quantization is enabled
+        if self.qbits:
+            return "".join(np.where(np.unpackbits(data), "1", "0"))
+
+        # Return original data
+        return data
+
+    def query(self, query):
+        """
+        Creates a query statement from an input query. This method uses hamming distance for bit vectors and
+        the max_inner_product for float vectors.
+
+        Args:
+            query: input query
+
+        Returns:
+            query statement
+        """
+
+        # Prepare query embeddings
+        query = self.prepare(query)
+
+        # Bit vector query
+        if self.qbits:
+            return self.table.c["embedding"].hamming_distance(query).label("score")
+
+        # Float vector query
+        return self.table.c["embedding"].max_inner_product(query).label("score")
+
+    def score(self, score):
+        """
+        Calculates the index score from the input score. This method returns the hamming score
+        (1.0 - (hamming distance / total number of bits)) for bit vectors and the -score for
+        float vectors.
+
+        Args:
+            score: input score
+
+        Returns:
+            index score
+        """
+
+        # Calculate hamming score as 1.0 - (hamming distance / total number of bits)
+        # Bound score from 0 to 1
+        if self.qbits:
+            return min(max(0.0, 1.0 - (score / (self.config["dimensions"] * 8))), 1.0)
+
+        # pgvector returns negative inner product since Postgres only supports ASC order index scans on operators
+        return -score