mseep-txtai 9.1.1__py3-none-any.whl

txtai/embeddings/index/__init__.py

@@ -0,0 +1,14 @@
+"""
+Index imports
+"""
+
+from .action import Action
+from .autoid import AutoId
+from .configuration import Configuration
+from .documents import Documents
+from .functions import Functions
+from .indexes import Indexes
+from .indexids import IndexIds
+from .reducer import Reducer
+from .stream import Stream
+from .transform import Transform

txtai/embeddings/index/action.py

@@ -0,0 +1,15 @@
+"""
+Action module
+"""
+
+from enum import Enum
+
+
+class Action(Enum):
+    """
+    Index action types
+    """
+
+    INDEX = 1
+    UPSERT = 2
+    REINDEX = 3

txtai/embeddings/index/autoid.py

@@ -0,0 +1,92 @@
+"""
+AutoId module
+"""
+
+import inspect
+import uuid
+
+
+class AutoId:
+    """
+    Generates unique ids.
+    """
+
+    def __init__(self, method=None):
+        """
+        Creates a unique id generator.
+
+        Args:
+            method: generation method - supports int sequence (default) or UUID function
+        """
+
+        # Initialize variables
+        self.method, self.function, self.value = None, None, None
+
+        # Set id generation method
+        if not method or isinstance(method, int):
+            # Incrementing sequence (default)
+            self.method = self.sequence
+            self.value = method if method else 0
+        else:
+            # UUID generation function
+            self.method = self.uuid
+            self.function = getattr(uuid, method)
+
+        # Check if signature takes a namespace argument (deterministic)
+        args = inspect.getfullargspec(self.function).args if self.function else []
+        self.deterministic = "namespace" in args
+
+    def __call__(self, data=None):
+        """
+        Generates a unique id.
+
+        Args:
+            data: optional data to use for deterministic algorithms (i.e. uuid3, uuid5)
+
+        Returns:
+            unique id
+        """
+
+        return self.method(data)
+
+    # pylint: disable=W0613
+    def sequence(self, data):
+        """
+        Gets and increments sequence.
+
+        Args:
+            data: not used
+
+        Returns:
+            current sequence value
+        """
+
+        # Get and increment sequence
+        value = self.value
+        self.value += 1
+
+        return value
+
+    def uuid(self, data):
+        """
+        Generates a UUID and return as a string.
+
+        Args:
+            data: used with determistic algorithms (uuid3, uuid5)
+
+        Returns:
+            UUID string
+        """
+
+        uid = self.function(uuid.NAMESPACE_DNS, str(data)) if self.deterministic else self.function()
+        return str(uid)
+
+    def current(self):
+        """
+        Get the current sequence value. Only applicable for sequence ids, will be None for UUID methods.
+
+        Returns:
+            current sequence value
+        """
+
+        return self.value

txtai/embeddings/index/configuration.py

@@ -0,0 +1,71 @@
+"""
+Configuration module
+"""
+
+import json
+import os
+
+from ...serialize import SerializeFactory
+
+
+class Configuration:
+    """
+    Loads and saves index configuration.
+    """
+
+    def load(self, path):
+        """
+        Loads index configuration. This method supports both config.json and config pickle files.
+
+        Args:
+            path: path to directory
+
+        Returns:
+            dict
+        """
+
+        # Configuration
+        config = None
+
+        # Determine if config is json or pickle
+        jsonconfig = os.path.exists(f"{path}/config.json")
+
+        # Set config file name
+        name = "config.json" if jsonconfig else "config"
+
+        # Load configuration
+        with open(f"{path}/{name}", "r" if jsonconfig else "rb", encoding="utf-8" if jsonconfig else None) as handle:
+            # Load JSON, also backwards-compatible with pickle configuration
+            config = json.load(handle) if jsonconfig else SerializeFactory.create("pickle").loadstream(handle)
+
+        # Add format parameter
+        config["format"] = "json" if jsonconfig else "pickle"
+
+        return config
+
+    def save(self, config, path):
+        """
+        Saves index configuration. This method defaults to JSON and falls back to pickle.
+
+        Args:
+            config: configuration to save
+            path: path to directory
+
+        Returns:
+            dict
+        """
+
+        # Default to JSON config
+        jsonconfig = config.get("format", "json") == "json"
+
+        # Set config file name
+        name = "config.json" if jsonconfig else "config"
+
+        # Write configuration
+        with open(f"{path}/{name}", "w" if jsonconfig else "wb", encoding="utf-8" if jsonconfig else None) as handle:
+            if jsonconfig:
+                # Write config as JSON
+                json.dump(config, handle, default=str, indent=2)
+            else:
+                # Backwards compatible method to save pickle configuration
+                SerializeFactory.create("pickle").savestream(config, handle)

txtai/embeddings/index/documents.py

@@ -0,0 +1,86 @@
+"""
+Documents module
+"""
+
+import os
+import tempfile
+
+from ...serialize import SerializeFactory
+
+
+class Documents:
+    """
+    Streams documents to temporary storage. Allows queuing large volumes of content for later indexing.
+    """
+
+    def __init__(self):
+        """
+        Creates a new documents stream.
+        """
+
+        self.documents = None
+        self.batch = 0
+        self.size = 0
+
+        # Pickle serialization - local temporary data
+        self.serializer = SerializeFactory.create("pickle", allowpickle=True)
+
+    def __len__(self):
+        """
+        Returns total number of queued documents.
+        """
+
+        return self.size
+
+    def __iter__(self):
+        """
+        Streams all queued documents.
+        """
+
+        # Close streaming file
+        self.documents.close()
+
+        # Open stream file
+        with open(self.documents.name, "rb") as queue:
+            # Read each batch
+            for _ in range(self.batch):
+                documents = self.serializer.loadstream(queue)
+
+                # Yield each document
+                yield from documents
+
+    def add(self, documents):
+        """
+        Adds a batch of documents for indexing.
+
+        Args:
+            documents: list of (id, data, tag) tuples
+
+        Returns:
+            documents
+        """
+
+        # Create documents file if not already open
+        # pylint: disable=R1732
+        if not self.documents:
+            self.documents = tempfile.NamedTemporaryFile(mode="wb", suffix=".docs", delete=False)
+
+        # Add batch
+        self.serializer.savestream(documents, self.documents)
+        self.batch += 1
+        self.size += len(documents)
+
+        return documents
+
+    def close(self):
+        """
+        Closes and resets this instance. New sets of documents can be added with additional calls to add.
+        """
+
+        # Cleanup stream file
+        os.remove(self.documents.name)
+
+        # Reset document parameters
+        self.documents = None
+        self.batch = 0
+        self.size = 0

txtai/embeddings/index/functions.py

@@ -0,0 +1,155 @@
+"""
+Functions module
+"""
+
+from types import FunctionType, MethodType
+
+
+class Functions:
+    """
+    Resolves function configuration to function references.
+    """
+
+    def __init__(self, embeddings):
+        """
+        Creates a new function resolver.
+
+        Args:
+            embeddings: embeddings instance
+        """
+
+        self.embeddings = embeddings
+
+        # Handle to all reference objects
+        self.references = None
+
+    def __call__(self, config):
+        """
+        Resolves a list of functions to function references.
+
+        Args:
+            config: configuration
+
+        Returns:
+            list of function references
+        """
+
+        # Initialize stored references array
+        self.references = []
+
+        # Resolve callable functions
+        functions = []
+        for fn in config["functions"]:
+            if isinstance(fn, dict):
+                fn = fn.copy()
+                fn["function"] = self.function(fn["function"])
+            else:
+                fn = self.function(fn)
+            functions.append(fn)
+
+        return functions
+
+    def reset(self):
+        """
+        Clears all resolved references.
+        """
+
+        if self.references:
+            for reference in self.references:
+                reference.reset()
+
+    def function(self, function):
+        """
+        Resolves function configuration. If function is a string, it's split on '.' and each part
+        is separately resolved to an object, attribute or function. Each part is resolved upon the
+        first invocation of the function. Otherwise, the input is returned.
+
+        Args:
+            function: function configuration
+
+        Returns:
+            function reference
+        """
+
+        if isinstance(function, str):
+            parts = function.split(".")
+
+            if hasattr(self.embeddings, parts[0]):
+                m = Reference(self.embeddings, parts[0])
+                self.references.append(m)
+            else:
+                module = ".".join(parts[:-1])
+                m = __import__(module)
+
+            for comp in parts[1:]:
+                m = Reference(m, comp)
+                self.references.append(m)
+
+            return m
+
+        return function
+
+
+class Reference:
+    """
+    Stores a reference to an object attribute. This attribute is resolved by invoking the __call__ method.
+    This allows for functions to be independent of the initialization order of an embeddings instance.
+    """
+
+    def __init__(self, obj, attribute):
+        """
+        Create a new reference.
+
+        Args:
+            obj: object handle
+            attribute: attribute name
+        """
+
+        # Object handle and attribute
+        self.obj = obj
+        self.attribute = attribute
+
+        # Keep a handle to the original inputs
+        self.inputs = (obj, attribute)
+
+        # True if the object and attribute have been resolved
+        self.resolved = False
+
+        # True if the attribute is a function
+        self.function = None
+
+    def __call__(self, *args):
+        """
+        Resolves an object attribute reference. If the attribute is a function, the function is executed.
+        Otherwise, the object attribute value is returned.
+
+        Args:
+            args: list of function arguments to the object attribute, when attribute is a function
+
+        Returns:
+            object attribute function result or object attribute value
+        """
+
+        # Resolve nested function arguments, if necessary
+        if not self.resolved:
+            self.obj = self.obj() if isinstance(self.obj, Reference) else self.obj
+            self.attribute = self.attribute() if isinstance(self.attribute, Reference) else self.attribute
+            self.resolved = True
+
+        # Lookup attribute
+        attribute = getattr(self.obj, self.attribute)
+
+        # Determine if attribute is a function
+        if self.function is None:
+            self.function = isinstance(attribute, (FunctionType, MethodType)) or (hasattr(attribute, "__call__") and args)
+
+        # If attribute is a function, execute and return, otherwise return attribute
+        return attribute(*args) if self.function else attribute
+
+    def reset(self):
+        """
+        Clears resolved references.
+        """
+
+        self.obj, self.attribute = self.inputs
+        self.resolved = False

txtai/embeddings/index/indexes.py

@@ -0,0 +1,199 @@
+"""
+Indexes module
+"""
+
+import os
+
+from .documents import Documents
+
+
+class Indexes:
+    """
+    Manages a collection of subindexes for an embeddings instance.
+    """
+
+    def __init__(self, embeddings, indexes):
+        """
+        Creates a new indexes instance.
+
+        Args:
+            embeddings: embeddings instance
+            indexes: dict of subindexes to add
+        """
+
+        self.embeddings = embeddings
+        self.indexes = indexes
+
+        self.documents = None
+        self.checkpoint = 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 instance
+        self.indexing = embeddings.model or embeddings.scoring
+
+    def __contains__(self, name):
+        """
+        Returns True if name is in this instance, False otherwise.
+
+        Returns:
+            True if name is in this instance, False otherwise
+        """
+
+        return name in self.indexes
+
+    def __getitem__(self, name):
+        """
+        Looks up an index by name.
+
+        Args:
+            name: index name
+
+        Returns:
+            index
+        """
+
+        return self.indexes[name]
+
+    def __getattr__(self, name):
+        """
+        Looks up an index by attribute name.
+
+        Args:
+            name: index name
+
+        Returns:
+            index
+        """
+
+        try:
+            return self.indexes[name]
+        except Exception as e:
+            raise AttributeError(e) from e
+
+    def default(self):
+        """
+        Gets the default/first index.
+
+        Returns:
+            default index
+        """
+
+        return list(self.indexes.keys())[0]
+
+    def findmodel(self, index=None):
+        """
+        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].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, 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 = []
+        for _, document, _ in documents:
+            # Add to documents collection if text or object field is set
+            parent = document
+            if isinstance(parent, dict):
+                parent = parent.get(self.text, document.get(self.object))
+
+            # Add if field is available or top-level indexing is disabled
+            if parent is not None or not self.indexing:
+                batch.append((index, document, None))
+                index += 1
+
+        # Add filtered documents batch
+        self.documents.add(batch)
+
+    def delete(self, ids):
+        """
+        Deletes ids from each subindex.
+
+        Args:
+            ids: list of ids to delete
+        """
+
+        for index in self.indexes.values():
+            index.delete(ids)
+
+    def index(self):
+        """
+        Builds each subindex.
+        """
+
+        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):
+        """
+        Runs upsert for each subindex.
+        """
+
+        for index in self.indexes.values():
+            index.upsert(self.documents)
+
+        # Reset document stream
+        self.documents.close()
+        self.documents = None
+
+    def load(self, path):
+        """
+        Loads each subindex from path.
+
+        Args:
+            path: directory path to load subindexes
+        """
+
+        for name, index in self.indexes.items():
+            # Load subindex if it exists, subindexes aren't required to have data
+            directory = os.path.join(path, name)
+            if index.exists(directory):
+                index.load(directory)
+
+    def save(self, path):
+        """
+        Saves each subindex to path.
+
+        Args:
+            path: directory path to save subindexes
+        """
+
+        for name, index in self.indexes.items():
+            index.save(os.path.join(path, name))
+
+    def close(self):
+        """
+        Close and free resources used by this instance.
+        """
+
+        for index in self.indexes.values():
+            index.close()

txtai/embeddings/index/indexids.py

@@ -0,0 +1,60 @@
+"""
+IndexIds module
+"""
+
+from ...serialize import Serializer
+
+
+class IndexIds:
+    """
+    Stores index ids when content is disabled.
+    """
+
+    def __init__(self, embeddings, ids=None):
+        """
+        Creates an IndexIds instance.
+
+        Args:
+            embeddings: embeddings instance
+            ids: ids to store
+        """
+
+        self.config = embeddings.config
+        self.ids = ids
+
+    def __iter__(self):
+        yield from self.ids
+
+    def __getitem__(self, index):
+        return self.ids[index]
+
+    def __setitem__(self, index, value):
+        self.ids[index] = value
+
+    def __add__(self, ids):
+        return self.ids + ids
+
+    def load(self, path):
+        """
+        Loads IndexIds from path.
+
+        Args:
+            path: path to load
+        """
+
+        if "ids" in self.config:
+            # Legacy ids format
+            self.ids = self.config.pop("ids")
+        else:
+            # Standard ids format
+            self.ids = Serializer.load(path)
+
+    def save(self, path):
+        """
+        Saves IndexIds to path.
+
+        Args:
+            path: path to save
+        """
+
+        Serializer.save(self.ids, path)