mseep-txtai 9.1.1__py3-none-any.whl

txtai/embeddings/index/reducer.py

@@ -0,0 +1,104 @@
+"""
+Reducer module
+"""
+
+from zipfile import BadZipFile
+
+# Conditionally import dimensionality reduction libraries as they aren't installed by default
+try:
+    import skops.io as sio
+
+    from sklearn.decomposition import TruncatedSVD
+
+    REDUCER = True
+except ImportError:
+    REDUCER = False
+
+from ...serialize import SerializeFactory
+
+
+class Reducer:
+    """
+    LSA dimensionality reduction model
+    """
+
+    def __init__(self, embeddings=None, components=None):
+        """
+        Creates a dimensionality reduction model.
+
+        Args:
+            embeddings: input embeddings matrix
+            components: number of model components
+        """
+
+        if not REDUCER:
+            raise ImportError('Dimensionality reduction is not available - install "vectors" extra to enable')
+
+        self.model = self.build(embeddings, components) if embeddings is not None and components else None
+
+    def __call__(self, embeddings):
+        """
+        Applies a dimensionality reduction model to embeddings, removed the top n principal components. Operation applied
+        directly on array.
+
+        Args:
+            embeddings: input embeddings matrix
+        """
+
+        pc = self.model.components_
+        factor = embeddings.dot(pc.transpose())
+
+        # Apply LSA model
+        # Calculation is different if n_components = 1
+        if pc.shape[0] == 1:
+            embeddings -= factor * pc
+        elif len(embeddings.shape) > 1:
+            # Apply model on a row-wise basis to limit memory usage
+            for x in range(embeddings.shape[0]):
+                embeddings[x] -= factor[x].dot(pc)
+        else:
+            # Single embedding
+            embeddings -= factor.dot(pc)
+
+    def build(self, embeddings, components):
+        """
+        Builds a LSA model. This model is used to remove the principal component within embeddings. This helps to
+        smooth out noisy embeddings (common words with less value).
+
+        Args:
+            embeddings: input embeddings matrix
+            components: number of model components
+
+        Returns:
+            LSA model
+        """
+
+        model = TruncatedSVD(n_components=components, random_state=0)
+        model.fit(embeddings)
+
+        return model
+
+    def load(self, path):
+        """
+        Loads a Reducer object from path.
+
+        Args:
+            path: directory path to load model
+        """
+
+        # Dimensionality reduction
+        try:
+            self.model = sio.load(path)
+        except (BadZipFile, KeyError):
+            # Backwards compatible support for pickled models
+            self.model = SerializeFactory.create("pickle").load(path)
+
+    def save(self, path):
+        """
+        Saves a Reducer object to path.
+
+        Args:
+            path: directory path to save model
+        """
+
+        sio.dump(self.model, path)

txtai/embeddings/index/stream.py

@@ -0,0 +1,67 @@
+"""
+Stream module
+"""
+
+from .autoid import AutoId
+from .transform import Action
+
+
+class Stream:
+    """
+    Yields input document as standard (id, data, tags) tuples.
+    """
+
+    def __init__(self, embeddings, action=None):
+        """
+        Create a new stream.
+
+        Args:
+            embeddings: embeddings instance
+            action: optional index action
+        """
+
+        self.embeddings = embeddings
+        self.action = action
+
+        # Alias embeddings attributes
+        self.config = embeddings.config
+
+        # Get config parameters
+        self.offset = self.config.get("offset", 0) if action == Action.UPSERT else 0
+        autoid = self.config.get("autoid", self.offset)
+
+        # Create autoid generator, reset int sequence if this isn't an UPSERT
+        autoid = 0 if isinstance(autoid, int) and action != Action.UPSERT else autoid
+        self.autoid = AutoId(autoid)
+
+    def __call__(self, documents):
+        """
+        Yield (id, data, tags) tuples from a stream of documents.
+
+        Args:
+            documents: input documents
+        """
+
+        # Iterate over documents and yield standard (id, data, tag) tuples
+        for document in documents:
+            if isinstance(document, dict):
+                # Create (id, data, tags) tuple from dictionary
+                document = document.get("id"), document, document.get("tags")
+            elif isinstance(document, tuple):
+                # Create (id, data, tags) tuple
+                document = document if len(document) >= 3 else (document[0], document[1], None)
+            else:
+                # Create (id, data, tags) tuple with empty fields
+                document = None, document, None
+
+            # Set autoid if the action is set
+            if self.action and document[0] is None:
+                document = (self.autoid(document[1]), document[1], document[2])
+
+            # Yield (id, data, tags) tuple
+            yield document
+
+        # Save autoid sequence if used
+        current = self.autoid.current()
+        if self.action and current:
+            self.config["autoid"] = current

txtai/embeddings/index/transform.py

@@ -0,0 +1,205 @@
+"""
+Transform module
+"""
+
+import numpy as np
+
+from .action import Action
+
+
+class Transform:
+    """
+    Executes a transform. Processes a stream of documents, loads batches into enabled data stores and vectorizes documents.
+    """
+
+    def __init__(self, embeddings, action, checkpoint=None):
+        """
+        Creates a new transform.
+
+        Args:
+            embeddings: embeddings instance
+            action: index action
+            checkpoint: optional checkpoint directory, enables indexing restart
+        """
+
+        self.embeddings = embeddings
+        self.action = action
+        self.checkpoint = checkpoint
+
+        # Alias embeddings attributes
+        self.config = embeddings.config
+        self.delete = embeddings.delete
+        self.model = embeddings.model
+        self.database = embeddings.database
+        self.graph = embeddings.graph
+        self.indexes = embeddings.indexes
+        self.scoring = embeddings.scoring if embeddings.issparse() else None
+
+        # Get config parameters
+        self.offset = embeddings.config.get("offset", 0) if action == Action.UPSERT else 0
+        self.batch = embeddings.config.get("batch", 1024)
+
+        # Scalar quantization
+        quantize = embeddings.config.get("quantize")
+        self.qbits = quantize if isinstance(quantize, int) and not isinstance(quantize, bool) else None
+
+        # Transform columns
+        columns = embeddings.config.get("columns", {})
+        self.text = columns.get("text", "text")
+        self.object = columns.get("object", "object")
+
+        # Check if top-level indexing is enabled for this embeddings
+        self.indexing = embeddings.model or embeddings.scoring
+
+        # List of deleted ids with this action
+        self.deletes = set()
+
+    def __call__(self, documents, buffer):
+        """
+        Processes an iterable collection of documents, handles any iterable including generators.
+
+        This method loads a stream of documents into enabled data stores and vectorizes documents into an embeddings array.
+
+        Args:
+            documents: iterable of (id, data, tags)
+            buffer: file path used for memmap buffer
+
+        Returns:
+            (document ids, dimensions, embeddings)
+        """
+
+        # Return parameters
+        ids, dimensions, embeddings = None, None, None
+
+        if self.model:
+            ids, dimensions, embeddings = self.vectors(documents, buffer)
+        else:
+            ids = self.ids(documents)
+
+        return (ids, dimensions, embeddings)
+
+    def vectors(self, documents, buffer):
+        """
+        Runs a vectors transform operation when dense indexing is enabled.
+
+        Args:
+            documents: iterable of (id, data, tags)
+            buffer: file path used for memmap buffer
+
+        Returns:
+            (document ids, dimensions, embeddings)
+        """
+
+        # Determine dtype
+        dtype = np.uint8 if self.qbits else np.float32
+
+        # Transform documents into vectors
+        return self.model.vectors(self.stream(documents), self.batch, self.checkpoint, buffer, dtype)
+
+    def ids(self, documents):
+        """
+        Runs an ids transform operation when dense indexing is disabled.
+
+        Args:
+            documents: iterable of (id, data, tags)
+
+        Returns:
+            document ids
+        """
+
+        # Consume stream and build extract ids
+        ids = []
+        for uid, _, _ in self.stream(documents):
+            ids.append(uid)
+
+        # Save offset when dense indexing is disabled
+        self.config["offset"] = self.offset
+
+        return ids
+
+    def stream(self, documents):
+        """
+        This method does two things:
+
+        1. Filter and yield data to vectorize
+        2. Batch and load original documents into enabled data stores (database, graph, scoring)
+
+        Documents are yielded for vectorization if one of the following is True:
+            - dict with a text or object field
+            - not a dict
+
+        Otherwise, documents are only batched and inserted into data stores
+
+        Args:
+            documents: iterable collection (id, data, tags)
+        """
+
+        # Batch and index offset. Index offset increments by count of documents streamed for vectorization
+        batch, offset = [], 0
+
+        # Iterate and process documents stream
+        for document in documents:
+            if isinstance(document[1], dict):
+                # Set text field to uid when top-level indexing is disabled and text empty
+                if not self.indexing and not document[1].get(self.text):
+                    document[1][self.text] = str(document[0])
+
+                if self.text in document[1]:
+                    yield (document[0], document[1][self.text], document[2])
+                    offset += 1
+                elif self.object in document[1]:
+                    yield (document[0], document[1][self.object], document[2])
+                    offset += 1
+            else:
+                yield document
+                offset += 1
+
+            # Batch document
+            batch.append(document)
+            if len(batch) == self.batch:
+                self.load(batch, offset)
+                batch, offset = [], 0
+
+        # Final batch
+        if batch:
+            self.load(batch, offset)
+
+    def load(self, batch, offset):
+        """
+        Loads a document batch. This method deletes existing ids from an embeddings index and
+        loads into enabled data stores (database, graph, scoring).
+
+        Args:
+            batch: list of (id, data, tags)
+            offset: index offset for batch
+        """
+
+        # Delete from embeddings index first (which deletes from underlying indexes and datastores) if this is an upsert
+        if self.action == Action.UPSERT:
+            # Get list of ids not yet seen and deleted
+            deletes = [uid for uid, _, _ in batch if uid not in self.deletes]
+            if deletes:
+                # Execute delete
+                self.delete(deletes)
+
+                # Save deleted ids as a delete must only occur once per action
+                self.deletes.update(deletes)
+
+        # Load batch into database except if this is a reindex
+        if self.database and self.action != Action.REINDEX:
+            self.database.insert(batch, self.offset)
+
+        # Load batch into scoring
+        if self.scoring:
+            self.scoring.insert(batch, self.offset, self.checkpoint)
+
+        # Load batch into subindex documents stream
+        if self.indexes:
+            self.indexes.insert(batch, self.offset, self.checkpoint)
+
+        # Load batch into graph
+        if self.graph:
+            self.graph.insert(batch, self.offset)
+
+        # Increment offset
+        self.offset += offset

txtai/embeddings/search/__init__.py

@@ -0,0 +1,11 @@
+"""
+Search imports
+"""
+
+from .base import Search
+from .errors import *
+from .explain import Explain
+from .ids import Ids
+from .query import Query
+from .scan import Scan
+from .terms import Terms