mseep-txtai 9.1.1__py3-none-any.whl

txtai/vectors/recovery.py

@@ -0,0 +1,57 @@
+"""
+Recovery module
+"""
+
+import os
+import shutil
+
+
+class Recovery:
+    """
+    Vector embeddings recovery. This class handles streaming embeddings from a vector checkpoint file.
+    """
+
+    def __init__(self, checkpoint, vectorsid, load):
+        """
+        Creates a Recovery instance.
+
+        Args:
+            checkpoint: checkpoint directory
+            vectorsid: vectors uid for current configuration
+            load: load embeddings method
+        """
+
+        self.spool, self.path, self.load = None, None, load
+
+        # Get unique file id
+        path = f"{checkpoint}/{vectorsid}"
+        if os.path.exists(path):
+            # Generate recovery path
+            self.path = f"{checkpoint}/recovery"
+
+            # Copy current checkpoint to recovery
+            shutil.copyfile(path, self.path)
+
+            # Open file an return
+            # pylint: disable=R1732
+            self.spool = open(self.path, "rb")
+
+    def __call__(self):
+        """
+        Reads and returns the next batch of embeddings.
+
+        Returns
+            batch of embeddings
+        """
+
+        try:
+            return self.load(self.spool) if self.spool else None
+        except EOFError:
+            # End of spool file, cleanup
+            self.spool.close()
+            os.remove(self.path)
+
+            # Clear parameters
+            self.spool, self.path = None, None
+
+            return None

txtai/vectors/sparse/__init__.py

@@ -0,0 +1,7 @@
+"""
+Sparse vectors imports
+"""
+
+from .base import SparseVectors
+from .factory import SparseVectorsFactory
+from .sbert import SparseSTVectors

txtai/vectors/sparse/base.py

@@ -0,0 +1,90 @@
+"""
+SparseVectors module
+"""
+
+# Conditional import
+try:
+    from scipy.sparse import csr_matrix, vstack
+    from sklearn.preprocessing import normalize
+    from sklearn.utils.extmath import safe_sparse_dot
+
+    SPARSE = True
+except ImportError:
+    SPARSE = False
+
+from ...util import SparseArray
+from ..base import Vectors
+
+
+# pylint: disable=W0223
+class SparseVectors(Vectors):
+    """
+    Base class for sparse vector models. Vector models transform input content into sparse arrays.
+    """
+
+    def __init__(self, config, scoring, models):
+        # Check before parent constructor since it calls loadmodel
+        if not SPARSE:
+            raise ImportError('SparseVectors is not available - install "vectors" extra to enable')
+
+        super().__init__(config, scoring, models)
+
+        # Get normalization setting
+        self.isnormalize = self.config.get("normalize", self.defaultnormalize()) if self.config else None
+
+    def encode(self, data, category=None):
+        # Encode data to embeddings
+        embeddings = super().encode(data, category)
+
+        # Get sparse torch vector attributes
+        embeddings = embeddings.cpu().coalesce()
+        indices = embeddings.indices().numpy()
+        values = embeddings.values().numpy()
+
+        # Return as SciPy CSR Matrix
+        return csr_matrix((values, indices), shape=embeddings.size())
+
+    def vectors(self, documents, batchsize=500, checkpoint=None, buffer=None, dtype=None):
+        # Run indexing
+        ids, dimensions, batches, stream = self.index(documents, batchsize, checkpoint)
+
+        # Rebuild sparse array
+        embeddings = None
+        with open(stream, "rb") as queue:
+            for _ in range(batches):
+                # Read in array batch
+                data = self.loadembeddings(queue)
+                embeddings = vstack((embeddings, data)) if embeddings is not None else data
+
+        # Return sparse array
+        return (ids, dimensions, embeddings)
+
+    def dot(self, queries, data):
+        return safe_sparse_dot(queries, data.T, dense_output=True).tolist()
+
+    def loadembeddings(self, f):
+        return SparseArray().load(f)
+
+    def saveembeddings(self, f, embeddings):
+        SparseArray().save(f, embeddings)
+
+    def truncate(self, embeddings):
+        raise ValueError("Truncate is not supported for sparse vectors")
+
+    def normalize(self, embeddings):
+        # Optionally normalize embeddings using method that supports sparse vectors
+        return normalize(embeddings, copy=False) if self.isnormalize else embeddings
+
+    def quantize(self, embeddings):
+        raise ValueError("Quantize is not supported for sparse vectors")
+
+    def defaultnormalize(self):
+        """
+        Returns the default normalization setting.
+
+        Returns:
+            default normalization setting
+        """
+
+        # Sparse vector embeddings typically perform better as unnormalized
+        return False

txtai/vectors/sparse/factory.py

@@ -0,0 +1,55 @@
+"""
+Factory module
+"""
+
+from ...util import Resolver
+
+from .sbert import SparseSTVectors
+
+
+class SparseVectorsFactory:
+    """
+    Methods to create sparse vector models.
+    """
+
+    @staticmethod
+    def create(config, models=None):
+        """
+        Create a Vectors model instance.
+
+        Args:
+            config: vector configuration
+            models: models cache
+
+        Returns:
+            Vectors
+        """
+
+        # Get vector method
+        method = config.get("method", "sentence-transformers")
+
+        # Sentence Transformers vectors
+        if method == "sentence-transformers":
+            return SparseSTVectors(config, None, models) if config and config.get("path") else None
+
+        # Resolve custom method
+        return SparseVectorsFactory.resolve(method, config, models) if method else None
+
+    @staticmethod
+    def resolve(backend, config, models):
+        """
+        Attempt to resolve a custom backend.
+
+        Args:
+            backend: backend class
+            config: vector configuration
+            models: models cache
+
+        Returns:
+            Vectors
+        """
+
+        try:
+            return Resolver()(backend)(config, None, models)
+        except Exception as e:
+            raise ImportError(f"Unable to resolve sparse vectors backend: '{backend}'") from e

txtai/vectors/sparse/sbert.py

@@ -0,0 +1,34 @@
+"""
+Sparse Sentence Transformers module
+"""
+
+# Conditional import
+try:
+    from sentence_transformers import SparseEncoder
+
+    SENTENCE_TRANSFORMERS = True
+except ImportError:
+    SENTENCE_TRANSFORMERS = False
+
+from ..dense.sbert import STVectors
+from .base import SparseVectors
+
+
+class SparseSTVectors(SparseVectors, STVectors):
+    """
+    Builds sparse vectors using sentence-transformers (aka SBERT).
+    """
+
+    def __init__(self, config, scoring, models):
+        # Check before parent constructor since it calls loadmodel
+        if not SENTENCE_TRANSFORMERS:
+            raise ImportError('sentence-transformers is not available - install "vectors" extra to enable')
+
+        super().__init__(config, scoring, models)
+
+    def loadencoder(self, path, device, **kwargs):
+        return SparseEncoder(path, device=device, **kwargs)
+
+    def defaultnormalize(self):
+        # Enable normalization by default if similarity function is cosine
+        return self.model and self.model.similarity_fn_name == "cosine"

txtai/version.py

@@ -0,0 +1,6 @@
+"""
+Version strings
+"""
+
+# Current version tag
+__version__ = "9.0.0"

txtai/workflow/__init__.py

@@ -0,0 +1,8 @@
+"""
+Workflow imports
+"""
+
+from .base import Workflow
+from .execute import Execute
+from .factory import WorkflowFactory
+from .task import *

txtai/workflow/base.py

@@ -0,0 +1,184 @@
+"""
+Workflow module
+"""
+
+import logging
+import time
+import traceback
+
+from datetime import datetime
+
+# Conditional import
+try:
+    from croniter import croniter
+
+    CRONITER = True
+except ImportError:
+    CRONITER = False
+
+from .execute import Execute
+
+# Logging configuration
+logger = logging.getLogger(__name__)
+
+
+class Workflow:
+    """
+    Base class for all workflows.
+    """
+
+    def __init__(self, tasks, batch=100, workers=None, name=None, stream=None):
+        """
+        Creates a new workflow. Workflows are lists of tasks to execute.
+
+        Args:
+            tasks: list of workflow tasks
+            batch: how many items to process at a time, defaults to 100
+            workers: number of concurrent workers
+            name: workflow name
+            stream: workflow stream processor
+        """
+
+        self.tasks = tasks
+        self.batch = batch
+        self.workers = workers
+        self.name = name
+        self.stream = stream
+
+        # Set default number of executor workers to max number of actions in a task
+        self.workers = max(len(task.action) for task in self.tasks) if not self.workers else self.workers
+
+    def __call__(self, elements):
+        """
+        Executes a workflow for input elements. This method returns a generator that yields transformed
+        data elements.
+
+        Args:
+            elements: iterable data elements
+
+        Returns:
+            generator that yields transformed data elements
+        """
+
+        # Create execute instance for this run
+        with Execute(self.workers) as executor:
+            # Run task initializers
+            self.initialize()
+
+            # Process elements with stream processor, if available
+            elements = self.stream(elements) if self.stream else elements
+
+            # Process elements in batches
+            for batch in self.chunk(elements):
+                yield from self.process(batch, executor)
+
+            # Run task finalizers
+            self.finalize()
+
+    def schedule(self, cron, elements, iterations=None):
+        """
+        Schedules a workflow using a cron expression and elements.
+
+        Args:
+            cron: cron expression
+            elements: iterable data elements passed to workflow each call
+            iterations: number of times to run workflow, defaults to run indefinitely
+        """
+
+        # Check that croniter is installed
+        if not CRONITER:
+            raise ImportError('Workflow scheduling is not available - install "workflow" extra to enable')
+
+        logger.info("'%s' scheduler started with schedule %s", self.name, cron)
+
+        maxiterations = iterations
+        while iterations is None or iterations > 0:
+            # Schedule using localtime
+            schedule = croniter(cron, datetime.now().astimezone()).get_next(datetime)
+            logger.info("'%s' next run scheduled for %s", self.name, schedule.isoformat())
+            time.sleep(schedule.timestamp() - time.time())
+
+            # Run workflow
+            # pylint: disable=W0703
+            try:
+                for _ in self(elements):
+                    pass
+            except Exception:
+                logger.error(traceback.format_exc())
+
+            # Decrement iterations remaining, if necessary
+            if iterations is not None:
+                iterations -= 1
+
+        logger.info("'%s' max iterations (%d) reached", self.name, maxiterations)
+
+    def initialize(self):
+        """
+        Runs task initializer methods (if any) before processing data.
+        """
+
+        # Run task initializers
+        for task in self.tasks:
+            if task.initialize:
+                task.initialize()
+
+    def chunk(self, elements):
+        """
+        Splits elements into batches. This method efficiently processes both fixed size inputs and
+        dynamically generated inputs.
+
+        Args:
+            elements: iterable data elements
+
+        Returns:
+            evenly sized batches with the last batch having the remaining elements
+        """
+
+        # Build batches by slicing elements, more efficient for fixed sized inputs
+        if hasattr(elements, "__len__") and hasattr(elements, "__getitem__"):
+            for x in range(0, len(elements), self.batch):
+                yield elements[x : x + self.batch]
+
+        # Build batches by iterating over elements when inputs are dynamically generated (i.e. generators)
+        else:
+            batch = []
+            for x in elements:
+                batch.append(x)
+
+                if len(batch) == self.batch:
+                    yield batch
+                    batch = []
+
+            # Final batch
+            if batch:
+                yield batch
+
+    def process(self, elements, executor):
+        """
+        Processes a batch of data elements.
+
+        Args:
+            elements: iterable data elements
+            executor: execute instance, enables concurrent task actions
+
+        Returns:
+            transformed data elements
+        """
+
+        # Run elements through each task
+        for x, task in enumerate(self.tasks):
+            logger.debug("Running Task #%d", x)
+            elements = task(elements, executor)
+
+        # Yield results processed by all tasks
+        yield from elements
+
+    def finalize(self):
+        """
+        Runs task finalizer methods (if any) after all data processed.
+        """
+
+        # Run task finalizers
+        for task in self.tasks:
+            if task.finalize:
+                task.finalize()

txtai/workflow/execute.py

@@ -0,0 +1,99 @@
+"""
+Execute module
+"""
+
+from multiprocessing.pool import Pool, ThreadPool
+
+import torch.multiprocessing
+
+
+class Execute:
+    """
+    Supports sequential, multithreading and multiprocessing based execution of tasks.
+    """
+
+    def __init__(self, workers=None):
+        """
+        Creates a new execute instance. Functions can be executed sequentially, in a thread pool
+        or in a process pool. Once created, the thread and/or process pool will stay open until the
+        close method is called.
+
+        Args:
+            workers: number of workers for thread/process pools
+        """
+
+        # Number of workers to use in thread/process pools
+        self.workers = workers
+
+        self.thread = None
+        self.process = None
+
+    def __del__(self):
+        self.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, etype, value, traceback):
+        self.close()
+
+    def run(self, method, function, args):
+        """
+        Runs multiple calls of function for each tuple in args. The method parameter controls if the calls are
+        sequential (method = None), multithreaded (method = "thread") or with multiprocessing (method="process").
+
+        Args:
+            method: run method - "thread" for multithreading, "process" for multiprocessing, otherwise runs sequentially
+            function: function to run
+            args: list of tuples with arguments to each call
+        """
+
+        # Concurrent processing
+        if method and len(args) > 1:
+            pool = self.pool(method)
+            if pool:
+                return pool.starmap(function, args, 1)
+
+        # Sequential processing
+        return [function(*arg) for arg in args]
+
+    def pool(self, method):
+        """
+        Gets a handle to a concurrent processing pool. This method will create the pool if it doesn't already exist.
+
+        Args:
+            method: pool type - "thread" or "process"
+
+        Returns:
+            concurrent processing pool or None if no pool of that type available
+        """
+
+        if method == "thread":
+            if not self.thread:
+                self.thread = ThreadPool(self.workers)
+
+            return self.thread
+
+        if method == "process":
+            if not self.process:
+                # Importing torch.multiprocessing will register torch shared memory serialization for cuda
+                self.process = Pool(self.workers, context=torch.multiprocessing.get_context("spawn"))
+
+            return self.process
+
+        return None
+
+    def close(self):
+        """
+        Closes concurrent processing pools.
+        """
+
+        if hasattr(self, "thread") and self.thread:
+            self.thread.close()
+            self.thread.join()
+            self.thread = None
+
+        if hasattr(self, "process") and self.process:
+            self.process.close()
+            self.process.join()
+            self.process = None

txtai/workflow/factory.py

@@ -0,0 +1,42 @@
+"""
+Workflow factory module
+"""
+
+from .base import Workflow
+from .task import TaskFactory
+
+
+class WorkflowFactory:
+    """
+    Workflow factory. Creates new Workflow instances.
+    """
+
+    @staticmethod
+    def create(config, name):
+        """
+        Creates a new Workflow instance.
+
+        Args:
+            config: Workflow configuration
+            name: Workflow name
+
+        Returns:
+            Workflow
+        """
+
+        # Resolve workflow tasks
+        tasks = []
+        for tconfig in config["tasks"]:
+            task = tconfig.pop("task") if "task" in tconfig else ""
+            tasks.append(TaskFactory.create(tconfig, task))
+
+        config["tasks"] = tasks
+
+        if "stream" in config:
+            sconfig = config["stream"]
+            task = sconfig.pop("task") if "task" in sconfig else "stream"
+
+            config["stream"] = TaskFactory.create(sconfig, task)
+
+        # Create workflow
+        return Workflow(**config, name=name)

txtai/workflow/task/__init__.py

@@ -0,0 +1,18 @@
+"""
+Task imports
+"""
+
+from .base import Task
+from .console import ConsoleTask
+from .export import ExportTask
+from .factory import TaskFactory
+from .file import FileTask
+from .image import ImageTask
+from .retrieve import RetrieveTask
+from .service import ServiceTask
+from .storage import StorageTask
+from .stream import StreamTask
+from .template import *
+from .template import RagTask as ExtractorTask
+from .url import UrlTask
+from .workflow import WorkflowTask