mseep-txtai 9.1.1__py3-none-any.whl

txtai/embeddings/search/base.py

@@ -0,0 +1,344 @@
+"""
+Search module
+"""
+
+import logging
+
+from .errors import IndexNotFoundError
+from .scan import Scan
+
+# Logging configuration
+logger = logging.getLogger(__name__)
+
+
+class Search:
+    """
+    Executes a batch search action. A search can be both index and/or database driven.
+    """
+
+    def __init__(self, embeddings, indexids=False, indexonly=False):
+        """
+        Creates a new search action.
+
+        Args:
+            embeddings: embeddings instance
+            indexids: searches return indexids when True, otherwise run standard search
+            indexonly: always runs an index search even when a database is available
+        """
+
+        self.embeddings = embeddings
+        self.indexids = indexids or indexonly
+        self.indexonly = indexonly
+
+        # Alias embeddings attributes
+        self.ann = embeddings.ann
+        self.batchtransform = embeddings.batchtransform
+        self.database = embeddings.database
+        self.ids = embeddings.ids
+        self.indexes = embeddings.indexes
+        self.graph = embeddings.graph
+        self.query = embeddings.query
+        self.scoring = embeddings.scoring if embeddings.issparse() else None
+
+    def __call__(self, queries, limit=None, weights=None, index=None, parameters=None):
+        """
+        Executes a batch search for queries. This method will run either an index search or an index + database search
+        depending on if a database is available.
+
+        Args:
+            queries: list of queries
+            limit: maximum results
+            weights: hybrid score weights
+            index: index name
+            parameters: list of dicts of named parameters to bind to placeholders
+
+        Returns:
+            list of (id, score) per query for index search
+            list of dict per query for an index + database search
+            list of graph results for a graph index search
+        """
+
+        # Default input parameters
+        limit = limit if limit else 3
+        weights = weights if weights is not None else 0.5
+
+        # Return empty results if there is no database and indexes
+        if not self.ann and not self.scoring and not self.indexes and not self.database:
+            return [[]] * len(queries)
+
+        # Default index name if only subindexes set
+        if not index and not self.ann and not self.scoring and self.indexes:
+            index = self.indexes.default()
+
+        # Graph search
+        if self.graph and self.graph.isquery(queries):
+            return self.graphsearch(queries, limit, weights, index)
+
+        # Database search
+        if not self.indexonly and self.database:
+            return self.dbsearch(queries, limit, weights, index, parameters)
+
+        # Default vector index query (sparse, dense or hybrid)
+        return self.search(queries, limit, weights, index)
+
+    def search(self, queries, limit, weights, index):
+        """
+        Executes an index search. When only a sparse index is enabled, this is a a keyword search. When only
+        a dense index is enabled, this is an ann search. When both are enabled, this is a hybrid search.
+
+        This method will also query subindexes, if available.
+
+        Args:
+            queries: list of queries
+            limit: maximum results
+            weights: hybrid score weights
+            index: index name
+
+        Returns:
+            list of (id, score) per query
+        """
+
+        # Run against specified subindex
+        if index:
+            return self.subindex(queries, limit, weights, index)
+
+        # Run against base indexes
+        hybrid = self.ann and self.scoring
+        dense = self.dense(queries, limit * 10 if hybrid else limit) if self.ann else None
+        sparse = self.sparse(queries, limit * 10 if hybrid else limit) if self.scoring else None
+
+        # Combine scores together
+        if hybrid:
+            # Create weights array if single number passed
+            if isinstance(weights, (int, float)):
+                weights = [weights, 1 - weights]
+
+            # Create weighted scores
+            results = []
+            for vectors in zip(dense, sparse):
+                uids = {}
+                for v, scores in enumerate(vectors):
+                    for r, (uid, score) in enumerate(scores if weights[v] > 0 else []):
+                        # Initialize score
+                        if uid not in uids:
+                            uids[uid] = 0.0
+
+                        # Create hybrid score
+                        #  - Convex Combination when sparse scores are normalized
+                        #  - Reciprocal Rank Fusion (RRF) when sparse scores aren't normalized
+                        if self.scoring.isnormalized():
+                            uids[uid] += score * weights[v]
+                        else:
+                            uids[uid] += (1.0 / (r + 1)) * weights[v]
+
+                results.append(sorted(uids.items(), key=lambda x: x[1], reverse=True)[:limit])
+
+            return results
+
+        # Raise an error if when no indexes are available
+        if not sparse and not dense:
+            raise IndexNotFoundError("No indexes available")
+
+        # Return single query results
+        return dense if dense else sparse
+
+    def subindex(self, queries, limit, weights, index):
+        """
+        Executes a subindex search.
+
+        Args:
+            queries: list of queries
+            limit: maximum results
+            weights: hybrid score weights
+            index: index name
+
+        Returns:
+            list of (id, score) per query
+        """
+
+        # Check that index exists
+        if not self.indexes or index not in self.indexes:
+            raise IndexNotFoundError(f"Index '{index}' not found")
+
+        # Run subindex search
+        results = self.indexes[index].batchsearch(queries, limit, weights)
+        return self.resolve(results)
+
+    def dense(self, queries, limit):
+        """
+        Executes an dense vector search with an approximate nearest neighbor index.
+
+        Args:
+            queries: list of queries
+            limit: maximum results
+
+        Returns:
+            list of (id, score) per query
+        """
+
+        # Convert queries to embedding vectors
+        embeddings = self.batchtransform((None, query, None) for query in queries)
+
+        # Search approximate nearest neighbor index
+        results = self.ann.search(embeddings, limit)
+
+        # Require scores to be greater than 0
+        results = [[(i, score) for i, score in r if score > 0] for r in results]
+
+        return self.resolve(results)
+
+    def sparse(self, queries, limit):
+        """
+        Executes a sparse vector search with a sparse keyword or sparse vector index.
+
+        Args:
+            queries: list of queries
+            limit: maximum results
+
+        Returns:
+            list of (id, score) per query
+        """
+
+        # Search sparse index
+        results = self.scoring.batchsearch(queries, limit)
+
+        # Require scores to be greater than 0
+        results = [[(i, score) for i, score in r if score > 0] for r in results]
+
+        return self.resolve(results)
+
+    def resolve(self, results):
+        """
+        Resolves index ids. This is only executed when content is disabled.
+
+        Args:
+            results: results
+
+        Returns:
+            results with resolved ids
+        """
+
+        # Map indexids to ids if embeddings ids are available
+        if not self.indexids and self.ids:
+            return [[(self.ids[i], score) for i, score in r] for r in results]
+
+        return results
+
+    def dbsearch(self, queries, limit, weights, index, parameters):
+        """
+        Executes an index + database search.
+
+        Args:
+            queries: list of queries
+            limit: maximum results
+            weights: default hybrid score weights
+            index: default index name
+            parameters: list of dicts of named parameters to bind to placeholders
+
+        Returns:
+            list of dict per query
+        """
+
+        # Parse queries
+        queries = self.parse(queries)
+
+        # Override limit with query limit, if applicable
+        limit = max(limit, self.limit(queries))
+
+        # Bulk index scan
+        scan = Scan(self.search, limit, weights, index)(queries, parameters)
+
+        # Combine index search results with database search results
+        results = []
+        for x, query in enumerate(queries):
+            # Run the database query, get matching bulk searches for current query
+            result = self.database.search(
+                query, [r for y, r in scan if x == y], limit, parameters[x] if parameters and parameters[x] else None, self.indexids
+            )
+            results.append(result)
+
+        return results
+
+    def parse(self, queries):
+        """
+        Parses a list of database queries.
+
+        Args:
+            queries: list of queries
+
+        Returns:
+            parsed queries
+        """
+
+        # Parsed queries
+        parsed = []
+
+        for query in queries:
+            # Parse query
+            parse = self.database.parse(query)
+
+            # Transform query if SQL not parsed and reparse
+            if self.query and "select" not in parse:
+                # Generate query
+                query = self.query(query)
+                logger.debug(query)
+
+                # Reparse query
+                parse = self.database.parse(query)
+
+            parsed.append(parse)
+
+        return parsed
+
+    def limit(self, queries):
+        """
+        Parses the largest LIMIT clause from queries.
+
+        Args:
+            queries: list of queries
+
+        Returns:
+            largest limit number or 0 if not found
+        """
+
+        # Override limit with largest limit from database queries
+        qlimit = 0
+        for query in queries:
+            # Parse out qlimit
+            l = query.get("limit")
+            if l and l.isdigit():
+                l = int(l)
+
+            qlimit = l if l and l > qlimit else qlimit
+
+        return qlimit
+
+    def graphsearch(self, queries, limit, weights, index):
+        """
+        Executes an index + graph search.
+
+        Args:
+            queries: list of queries
+            limit: maximum results
+            weights: default hybrid score weights
+            index: default index name
+
+        Returns:
+            graph search results
+        """
+
+        # Parse queries
+        queries = [self.graph.parse(query) for query in queries]
+
+        # Override limit with query limit, if applicable
+        limit = max(limit, self.limit(queries))
+
+        # Bulk index scan
+        scan = Scan(self.search, limit, weights, index)(queries, None)
+
+        # Combine index search results with database search results
+        for x, query in enumerate(queries):
+            # Add search results to query
+            query["results"] = [r for y, r in scan if x == y]
+
+        return self.graph.batchsearch(queries, limit, self.indexids)

txtai/embeddings/search/errors.py

@@ -0,0 +1,9 @@
+"""
+Errors module
+"""
+
+
+class IndexNotFoundError(Exception):
+    """
+    Raised when an embeddings query fails to locate an index
+    """

txtai/embeddings/search/explain.py

@@ -0,0 +1,120 @@
+"""
+Explain module
+"""
+
+import numpy as np
+
+
+class Explain:
+    """
+    Explains the importance of each token in an input text element for a query. This method creates n permutations of the input text, where n
+    is the number of tokens in the input text. This effectively masks each token to determine its importance to the query.
+    """
+
+    def __init__(self, embeddings):
+        """
+        Creates a new explain action.
+
+        Args:
+            embeddings: embeddings instance
+        """
+
+        self.embeddings = embeddings
+        self.content = embeddings.config.get("content")
+
+        # Alias embeddings attributes
+        self.database = embeddings.database
+
+    def __call__(self, queries, texts, limit):
+        """
+        Explains the importance of each input token in text for a list of queries.
+
+        Args:
+            query: input queries
+            texts: optional list of (text|list of tokens), otherwise runs search queries
+            limit: optional limit if texts is None
+
+        Returns:
+            list of dict per input text per query where a higher token scores represents higher importance relative to the query
+        """
+
+        # Construct texts elements per query
+        texts = self.texts(queries, texts, limit)
+
+        # Explain each query-texts combination
+        return [self.explain(query, texts[x]) for x, query in enumerate(queries)]
+
+    def texts(self, queries, texts, limit):
+        """
+        Constructs lists of dict for each input query.
+
+        Args:
+            queries: input queries
+            texts: optional list of texts
+            limit: optional limit if texts is None
+
+        Returns:
+            lists of dict for each input query
+        """
+
+        # Calculate similarity scores per query if texts present
+        if texts:
+            results = []
+            for scores in self.embeddings.batchsimilarity(queries, texts):
+                results.append([{"id": uid, "text": texts[uid], "score": score} for uid, score in scores])
+
+            return results
+
+        # Query for results if texts is None and content is enabled
+        return self.embeddings.batchsearch(queries, limit) if self.content else [[]] * len(queries)
+
+    def explain(self, query, texts):
+        """
+        Explains the importance of each input token in text for a list of queries.
+
+        Args:
+            query: input query
+            texts: list of text
+
+        Returns:
+            list of {"id": value, "text": value, "score": value, "tokens": value} covering each input text element
+        """
+
+        # Explain results
+        results = []
+
+        # Parse out similar clauses, if necessary
+        if self.database:
+            # Parse query
+            query = self.database.parse(query)
+
+            # Extract query from similar clause
+            query = " ".join([" ".join(clause) for clause in query["similar"]]) if "similar" in query else None
+
+        # Return original texts if query, text or score not present
+        if not query or not texts or "score" not in texts[0] or "text" not in texts[0]:
+            return texts
+
+        # Calculate result per input text element
+        for result in texts:
+            text = result["text"]
+            tokens = text if isinstance(text, list) else text.split()
+
+            # Create permutations of input text, masking each token to determine importance
+            permutations = []
+            for i in range(len(tokens)):
+                data = tokens.copy()
+                data.pop(i)
+                permutations.append([" ".join(data)])
+
+            # Calculate similarity for each input text permutation and get score delta as importance
+            scores = [(i, result["score"] - np.abs(s)) for i, s in self.embeddings.similarity(query, permutations)]
+
+            # Append tokens to result
+            result["tokens"] = [(tokens[i], score) for i, score in sorted(scores, key=lambda x: x[0])]
+
+            # Add data sorted in index order
+            results.append(result)
+
+        # Sort score descending and return
+        return sorted(results, key=lambda x: x["score"], reverse=True)

txtai/embeddings/search/ids.py

@@ -0,0 +1,61 @@
+"""
+Ids module
+"""
+
+
+class Ids:
+    """
+    Resolves internal ids for lists of ids.
+    """
+
+    def __init__(self, embeddings):
+        """
+        Create a new ids action.
+
+        Args:
+            embeddings: embeddings instance
+        """
+
+        self.database = embeddings.database
+        self.ids = embeddings.ids
+
+    def __call__(self, ids):
+        """
+        Resolve internal ids.
+
+        Args:
+            ids: ids
+
+        Returns:
+            internal ids
+        """
+
+        # Resolve ids using database if available, otherwise fallback to embeddings ids
+        results = self.database.ids(ids) if self.database else self.scan(ids)
+
+        # Create dict of id: [iids] given there is a one to many relationship
+        ids = {}
+        for iid, uid in results:
+            if uid not in ids:
+                ids[uid] = []
+            ids[uid].append(iid)
+
+        return ids
+
+    def scan(self, ids):
+        """
+        Scans embeddings ids array for matches when content is disabled.
+
+        Args:
+            ids: search ids
+
+        Returns:
+            internal ids
+        """
+
+        # Find existing ids
+        indices = []
+        for uid in ids:
+            indices.extend([(index, value) for index, value in enumerate(self.ids) if uid == value])
+
+        return indices

txtai/embeddings/search/query.py

@@ -0,0 +1,69 @@
+"""
+Query module
+"""
+
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, T5ForConditionalGeneration
+
+
+class Query:
+    """
+    Query translation model.
+    """
+
+    def __init__(self, path, prefix=None, maxlength=512):
+        """
+        Creates a query translation model.
+
+        Args:
+            path: path to query model
+            prefix: text prefix
+            maxlength: max sequence length to generate
+        """
+
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(path)
+
+        # Default prefix if not provided for T5 models
+        if not prefix and isinstance(self.model, T5ForConditionalGeneration):
+            prefix = "translate English to SQL: "
+
+        self.prefix = prefix
+        self.maxlength = maxlength
+
+    def __call__(self, query):
+        """
+        Runs query translation model.
+
+        Args:
+            query: input query
+
+        Returns:
+            transformed query
+        """
+
+        # Add prefix, if necessary
+        if self.prefix:
+            query = f"{self.prefix}{query}"
+
+        # Tokenize and generate text using model
+        features = self.tokenizer([query], return_tensors="pt")
+        output = self.model.generate(input_ids=features["input_ids"], attention_mask=features["attention_mask"], max_length=self.maxlength)
+
+        # Decode tokens to text
+        result = self.tokenizer.decode(output[0], skip_special_tokens=True)
+
+        # Clean and return generated text
+        return self.clean(result)
+
+    def clean(self, text):
+        """
+        Applies a series of rules to clean generated text.
+
+        Args:
+            text: input text
+
+        Returns:
+            clean text
+        """
+
+        return text.replace("$=", "<=")