mseep-txtai 9.1.1__py3-none-any.whl

txtai/scoring/factory.py

@@ -0,0 +1,95 @@
+"""
+Factory module
+"""
+
+from ..util import Resolver
+
+from .bm25 import BM25
+from .pgtext import PGText
+from .sif import SIF
+from .sparse import Sparse
+from .tfidf import TFIDF
+
+
+class ScoringFactory:
+    """
+    Methods to create Scoring indexes.
+    """
+
+    @staticmethod
+    def create(config, models=None):
+        """
+        Factory method to construct a Scoring instance.
+
+        Args:
+            config: scoring configuration parameters
+            models: models cache
+
+        Returns:
+            Scoring
+        """
+
+        # Scoring instance
+        scoring = None
+
+        # Support string and dict configuration
+        if isinstance(config, str):
+            config = {"method": config}
+
+        # Get scoring method
+        method = config.get("method", "bm25")
+
+        if method == "bm25":
+            scoring = BM25(config)
+        elif method == "pgtext":
+            scoring = PGText(config)
+        elif method == "sif":
+            scoring = SIF(config)
+        elif method == "sparse":
+            scoring = Sparse(config, models)
+        elif method == "tfidf":
+            scoring = TFIDF(config)
+        else:
+            # Resolve custom method
+            scoring = ScoringFactory.resolve(method, config)
+
+        # Store config back
+        config["method"] = method
+
+        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):
+        """
+        Attempt to resolve a custom backend.
+
+        Args:
+            backend: backend class
+            config: index configuration parameters
+
+        Returns:
+            Scoring
+        """
+
+        try:
+            return Resolver()(backend)(config)
+        except Exception as e:
+            raise ImportError(f"Unable to resolve scoring backend: '{backend}'") from e

txtai/scoring/pgtext.py

@@ -0,0 +1,181 @@
+"""
+PGText module
+"""
+
+import os
+
+# Conditional import
+try:
+    from sqlalchemy import create_engine, desc, delete, func, text
+    from sqlalchemy import Column, Computed, Index, Integer, MetaData, StaticPool, Table, Text
+    from sqlalchemy.dialects.postgresql import TSVECTOR
+    from sqlalchemy.orm import Session
+    from sqlalchemy.schema import CreateSchema
+
+    PGTEXT = True
+except ImportError:
+    PGTEXT = False
+
+from .base import Scoring
+
+
+class PGText(Scoring):
+    """
+    Postgres full text search (FTS) based scoring.
+    """
+
+    def __init__(self, config=None):
+        super().__init__(config)
+
+        if not PGTEXT:
+            raise ImportError('PGText is not available - install "scoring" extra to enable')
+
+        # Database connection
+        self.engine, self.database, self.connection, self.table = None, None, None, None
+
+        # Language
+        self.language = self.config.get("language", "english")
+
+    def insert(self, documents, index=None, checkpoint=None):
+        # Initialize tables
+        self.initialize(recreate=True)
+
+        # Collection of rows to insert
+        rows = []
+
+        # Collect rows
+        for uid, document, _ in documents:
+            # Extract text, if necessary
+            if isinstance(document, dict):
+                document = document.get(self.text, document.get(self.object))
+
+            if document is not None:
+                # If index is passed, use indexid, otherwise use id
+                uid = index if index is not None else uid
+
+                # Add row if the data type is accepted
+                if isinstance(document, (str, list)):
+                    rows.append((uid, " ".join(document) if isinstance(document, list) else document))
+
+                # Increment index
+                index = index + 1 if index is not None else None
+
+        # Insert rows
+        self.database.execute(self.table.insert(), [{"indexid": x, "text": text} for x, text in rows])
+
+    def delete(self, ids):
+        self.database.execute(delete(self.table).where(self.table.c["indexid"].in_(ids)))
+
+    def weights(self, tokens):
+        # Not supported
+        return None
+
+    def search(self, query, limit=3):
+        # Run query
+        query = (
+            self.database.query(self.table.c["indexid"], text("ts_rank(vector, plainto_tsquery(:language, :query)) rank"))
+            .order_by(desc(text("rank")))
+            .limit(limit)
+            .params({"language": self.language, "query": query})
+        )
+
+        return [(uid, score) for uid, score in query if score > 1e-5]
+
+    def batchsearch(self, queries, limit=3, threads=True):
+        return [self.search(query, limit) for query in queries]
+
+    def count(self):
+        # pylint: disable=E1102
+        return self.database.query(func.count(self.table.c["indexid"])).scalar()
+
+    def load(self, path):
+        # Reset database to original checkpoint
+        if self.database:
+            self.database.rollback()
+            self.connection.rollback()
+
+        # Initialize tables
+        self.initialize()
+
+    def save(self, path):
+        # Commit session and connection
+        if self.database:
+            self.database.commit()
+            self.connection.commit()
+
+    def close(self):
+        if self.database:
+            self.database.close()
+            self.engine.dispose()
+
+    def issparse(self):
+        return True
+
+    def isnormalized(self):
+        return True
+
+    def initialize(self, recreate=False):
+        """
+        Initializes a new database session.
+
+        Args:
+            recreate: Recreates the database tables if True
+        """
+
+        if not self.database:
+            # Create engine, connection and session
+            self.engine = create_engine(self.config.get("url", os.environ.get("SCORING_URL")), poolclass=StaticPool, echo=False)
+            self.connection = self.engine.connect()
+            self.database = Session(self.connection)
+
+            # Set default schema, if necessary
+            schema = self.config.get("schema")
+            if schema:
+                with self.engine.begin():
+                    self.sqldialect(CreateSchema(schema, if_not_exists=True))
+
+                self.sqldialect(text("SET search_path TO :schema"), {"schema": schema})
+
+            # Table name
+            table = self.config.get("table", "scoring")
+
+            # Create vectors table
+            self.table = Table(
+                table,
+                MetaData(),
+                Column("indexid", Integer, primary_key=True, autoincrement=False),
+                Column("text", Text),
+                (
+                    Column("vector", TSVECTOR, Computed(f"to_tsvector('{self.language}', text)", persisted=True))
+                    if self.engine.dialect.name == "postgresql"
+                    else Column("vector", Integer)
+                ),
+            )
+
+            # Create text index
+            index = Index(
+                f"{table}-index",
+                self.table.c["vector"],
+                postgresql_using="gin",
+            )
+
+            # Drop and recreate table
+            if recreate:
+                self.table.drop(self.connection, checkfirst=True)
+                index.drop(self.connection, checkfirst=True)
+
+            # Create table and index
+            self.table.create(self.connection, checkfirst=True)
+            index.create(self.connection, checkfirst=True)
+
+    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)

txtai/scoring/sif.py

@@ -0,0 +1,32 @@
+"""
+SIF module
+"""
+
+import numpy as np
+
+from .tfidf import TFIDF
+
+
+class SIF(TFIDF):
+    """
+    Smooth Inverse Frequency (SIF) scoring.
+    """
+
+    def __init__(self, config=None):
+        super().__init__(config)
+
+        # SIF configurable parameters
+        self.a = self.config.get("a", 1e-3)
+
+    def computefreq(self, tokens):
+        # Default method computes frequency for a single entry
+        # SIF uses word frequencies across entire index
+        return {token: self.wordfreq[token] for token in tokens}
+
+    def score(self, freq, idf, length):
+        # Set freq to word frequencies across entire index when freq and idf shape don't match
+        if isinstance(freq, np.ndarray) and freq.shape != np.array(idf).shape:
+            freq.fill(freq.sum())
+
+        # Calculate SIF score
+        return self.a / (self.a + freq / self.tokens)

txtai/scoring/sparse.py

@@ -0,0 +1,218 @@
+"""
+Sparse module
+"""
+
+from queue import Queue
+from threading import Thread
+
+from ..ann import SparseANNFactory
+from ..vectors import SparseVectorsFactory
+
+from .base import Scoring
+
+
+class Sparse(Scoring):
+    """
+    Sparse vector scoring.
+    """
+
+    # End of stream message
+    COMPLETE = 1
+
+    def __init__(self, config=None, models=None):
+        super().__init__(config)
+
+        # Vector configuration
+        mapping = {"vectormethod": "method", "vectornormalize": "normalize"}
+        config = {k: v for k, v in config.items() if k not in mapping.values()}
+        for k, v in mapping.items():
+            if k in config:
+                config[v] = config[k]
+
+        # Load the SparseVectors model
+        self.model = SparseVectorsFactory.create(config, models)
+
+        # Normalize search outputs if vectors are not normalized already
+        # A float can also be provided to set the normalization factor (defaults to 30.0)
+        self.isnormalize = self.config.get("normalize", True)
+
+        # Sparse ANN
+        self.ann = None
+
+        # Encoding processing parameters
+        self.batch = self.config.get("batch", 1024)
+        self.thread, self.queue, self.data = None, None, None
+
+    def insert(self, documents, index=None, checkpoint=None):
+        # Start processing thread, if necessary
+        self.start(checkpoint)
+
+        data = []
+        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((uid, " ".join(document) if isinstance(document, list) else document, tags))
+
+        # Add batch of data
+        self.queue.put(data)
+
+    def delete(self, ids):
+        self.ann.delete(ids)
+
+    def index(self, documents=None):
+        # Insert documents, if provided
+        if documents:
+            self.insert(documents)
+
+        # 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)
+
+        # 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()
+
+    def weights(self, tokens):
+        # Not supported
+        return None
+
+    def search(self, query, limit=3):
+        return self.batchsearch([query], limit)[0]
+
+    def batchsearch(self, queries, limit=3, threads=True):
+        # Convert queries to embedding vectors
+        embeddings = self.model.batchtransform((None, query, None) for query in queries)
+
+        # Run ANN search
+        scores = self.ann.search(embeddings, limit)
+
+        # Normalize scores if normalization IS enabled AND vector normalization IS NOT enabled
+        return self.normalize(embeddings, scores) if self.isnormalize and not self.model.isnormalize else scores
+
+    def count(self):
+        return self.ann.count()
+
+    def load(self, path):
+        self.ann = SparseANNFactory.create(self.config)
+        self.ann.load(path)
+
+    def save(self, path):
+        # Save Sparse ANN
+        if self.ann:
+            self.ann.save(path)
+
+    def close(self):
+        # Close Sparse ANN
+        if self.ann:
+            self.ann.close()
+
+        # Clear parameters
+        self.model, self.ann, self.thread, self.queue = None, None, None, None
+
+    def issparse(self):
+        return True
+
+    def isnormalized(self):
+        return self.isnormalize or self.model.isnormalize
+
+    def start(self, checkpoint):
+        """
+        Starts an encoding processing thread.
+
+        Args:
+            checkpoint: checkpoint directory
+        """
+
+        if not self.thread:
+            self.queue = Queue(5)
+            self.thread = Thread(target=self.encode, args=(checkpoint,))
+            self.thread.start()
+
+    def stop(self):
+        """
+        Stops an encoding processing thread. Return processed results.
+
+        Returns:
+            results
+        """
+
+        results = None
+        if self.thread:
+            # Send EOS message
+            self.queue.put(Sparse.COMPLETE)
+
+            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()
+
+    def normalize(self, queries, scores):
+        """
+        Normalize query result using the max query score.
+
+        Args:
+            queries: query vectors
+            scores: query results
+
+        Returns:
+            normalized query results
+        """
+
+        # Get normalize scale factor
+        scale = 30.0 if isinstance(self.isnormalize, bool) else self.isnormalize
+
+        # Normalize scores using max scores
+        maxscores = self.model.dot(queries, queries)
+
+        # Normalize results and return
+        results = []
+        for x, result in enumerate(scores):
+            maxscore = max(maxscores[x][x] / scale, scale)
+            maxscore = max(maxscore, result[0][1]) if result else maxscore
+
+            results.append([(uid, score / maxscore) for uid, score in result])
+
+        return results