mseep-txtai 9.1.1__py3-none-any.whl

txtai/vectors/dense/factory.py

@@ -0,0 +1,121 @@
+"""
+Factory module
+"""
+
+from ...util import Resolver
+
+from .external import External
+from .huggingface import HFVectors
+from .litellm import LiteLLM
+from .llama import LlamaCpp
+from .m2v import Model2Vec
+from .sbert import STVectors
+from .words import WordVectors
+
+
+class VectorsFactory:
+    """
+    Methods to create dense vector models.
+    """
+
+    @staticmethod
+    def create(config, scoring=None, models=None):
+        """
+        Create a Vectors model instance.
+
+        Args:
+            config: vector configuration
+            scoring: scoring instance
+            models: models cache
+
+        Returns:
+            Vectors
+        """
+
+        # Determine vector method
+        method = VectorsFactory.method(config)
+
+        # External vectors
+        if method == "external":
+            return External(config, scoring, models)
+
+        # LiteLLM vectors
+        if method == "litellm":
+            return LiteLLM(config, scoring, models)
+
+        # llama.cpp vectors
+        if method == "llama.cpp":
+            return LlamaCpp(config, scoring, models)
+
+        # Model2vec vectors
+        if method == "model2vec":
+            return Model2Vec(config, scoring, models)
+
+        # Sentence Transformers vectors
+        if method == "sentence-transformers":
+            return STVectors(config, scoring, models) if config and config.get("path") else None
+
+        # Word vectors
+        if method == "words":
+            return WordVectors(config, scoring, models)
+
+        # Transformers vectors
+        if HFVectors.ismethod(method):
+            return HFVectors(config, scoring, models) if config and config.get("path") else None
+
+        # Resolve custom method
+        return VectorsFactory.resolve(method, config, scoring, models) if method else None
+
+    @staticmethod
+    def method(config):
+        """
+        Get or derive the vector method.
+
+        Args:
+            config: vector configuration
+
+        Returns:
+            vector method
+        """
+
+        # Determine vector method
+        method = config.get("method")
+        path = config.get("path")
+
+        # Infer method from path, if blank
+        if not method:
+            if path:
+                if LiteLLM.ismodel(path):
+                    method = "litellm"
+                elif LlamaCpp.ismodel(path):
+                    method = "llama.cpp"
+                elif Model2Vec.ismodel(path):
+                    method = "model2vec"
+                elif WordVectors.ismodel(path):
+                    method = "words"
+                else:
+                    method = "transformers"
+            elif config.get("transform"):
+                method = "external"
+
+        return method
+
+    @staticmethod
+    def resolve(backend, config, scoring, models):
+        """
+        Attempt to resolve a custom backend.
+
+        Args:
+            backend: backend class
+            config: vector configuration
+            scoring: scoring instance
+            models: models cache
+
+        Returns:
+            Vectors
+        """
+
+        try:
+            return Resolver()(backend)(config, scoring, models)
+        except Exception as e:
+            raise ImportError(f"Unable to resolve vectors backend: '{backend}'") from e

txtai/vectors/dense/huggingface.py

@@ -0,0 +1,44 @@
+"""
+Hugging Face module
+"""
+
+from ...models import Models, PoolingFactory
+
+from ..base import Vectors
+
+
+class HFVectors(Vectors):
+    """
+    Builds vectors using the Hugging Face transformers library.
+    """
+
+    @staticmethod
+    def ismethod(method):
+        """
+        Checks if this method uses local transformers-based models.
+
+        Args:
+            method: input method
+
+        Returns:
+            True if this is a local transformers-based model, False otherwise
+        """
+
+        return method in ("transformers", "pooling", "clspooling", "meanpooling")
+
+    def loadmodel(self, path):
+        # Build embeddings with transformers pooling
+        return PoolingFactory.create(
+            {
+                "method": self.config.get("method"),
+                "path": path,
+                "device": Models.deviceid(self.config.get("gpu", True)),
+                "tokenizer": self.config.get("tokenizer"),
+                "maxlength": self.config.get("maxlength"),
+                "modelargs": self.config.get("vectors", {}),
+            }
+        )
+
+    def encode(self, data, category=None):
+        # Encode data using vectors model
+        return self.model.encode(data, batch=self.encodebatch, category=category)

txtai/vectors/dense/litellm.py

@@ -0,0 +1,86 @@
+"""
+LiteLLM module
+"""
+
+import numpy as np
+
+from transformers.utils import cached_file
+
+# Conditional import
+try:
+    import litellm as api
+
+    LITELLM = True
+except ImportError:
+    LITELLM = False
+
+from ..base import Vectors
+
+
+class LiteLLM(Vectors):
+    """
+    Builds vectors using an external embeddings API via LiteLLM.
+    """
+
+    @staticmethod
+    def ismodel(path):
+        """
+        Checks if path is a LiteLLM model.
+
+        Args:
+            path: input path
+
+        Returns:
+            True if this is a LiteLLM model, False otherwise
+        """
+
+        # pylint: disable=W0702
+        if isinstance(path, str) and LITELLM:
+            debug = api.suppress_debug_info
+            try:
+                # Suppress debug messages for this test
+                api.suppress_debug_info = True
+                return api.get_llm_provider(path) and not LiteLLM.ishub(path)
+            except:
+                return False
+            finally:
+                # Restore debug info value to original value
+                api.suppress_debug_info = debug
+
+        return False
+
+    @staticmethod
+    def ishub(path):
+        """
+        Checks if path is available on the HF Hub.
+
+        Args:
+            input path
+
+        Returns:
+            True if this is a model on the HF Hub
+        """
+
+        # pylint: disable=W0702
+        try:
+            return cached_file(path_or_repo_id=path, filename="config.json") is not None if "/" in path else False
+        except:
+            return False
+
+    def __init__(self, config, scoring, models):
+        # Check before parent constructor since it calls loadmodel
+        if not LITELLM:
+            raise ImportError('LiteLLM is not available - install "vectors" extra to enable')
+
+        super().__init__(config, scoring, models)
+
+    def loadmodel(self, path):
+        return None
+
+    def encode(self, data, category=None):
+        # Call external embeddings API using LiteLLM
+        # Batching is handled server-side
+        response = api.embedding(model=self.config.get("path"), input=data, **self.config.get("vectors", {}))
+
+        # Read response into a NumPy array
+        return np.array([x["embedding"] for x in response.data], dtype=np.float32)

txtai/vectors/dense/llama.py

@@ -0,0 +1,84 @@
+"""
+Llama module
+"""
+
+import os
+
+import numpy as np
+
+from huggingface_hub import hf_hub_download
+
+# Conditional import
+try:
+    from llama_cpp import Llama
+
+    LLAMA_CPP = True
+except ImportError:
+    LLAMA_CPP = False
+
+from ..base import Vectors
+
+
+class LlamaCpp(Vectors):
+    """
+    Builds vectors using llama.cpp.
+    """
+
+    @staticmethod
+    def ismodel(path):
+        """
+        Checks if path is a llama.cpp model.
+
+        Args:
+            path: input path
+
+        Returns:
+            True if this is a llama.cpp model, False otherwise
+        """
+
+        return isinstance(path, str) and path.lower().endswith(".gguf")
+
+    def __init__(self, config, scoring, models):
+        # Check before parent constructor since it calls loadmodel
+        if not LLAMA_CPP:
+            raise ImportError('llama.cpp is not available - install "vectors" extra to enable')
+
+        super().__init__(config, scoring, models)
+
+    def loadmodel(self, path):
+        # Check if this is a local path, otherwise download from the HF Hub
+        path = path if os.path.exists(path) else self.download(path)
+
+        # Additional model arguments
+        modelargs = self.config.get("vectors", {})
+
+        # Default GPU layers if not already set
+        modelargs["n_gpu_layers"] = modelargs.get("n_gpu_layers", -1 if self.config.get("gpu", os.environ.get("LLAMA_NO_METAL") != "1") else 0)
+
+        # Create llama.cpp instance
+        return Llama(path, n_ctx=0, verbose=modelargs.pop("verbose", False), embedding=True, **modelargs)
+
+    def encode(self, data, category=None):
+        # Generate embeddings and return as a NumPy array
+        # llama-cpp-python has it's own batching built-in using n_batch parameter
+        return np.array(self.model.embed(data), dtype=np.float32)
+
+    def download(self, path):
+        """
+        Downloads path from the Hugging Face Hub.
+
+        Args:
+            path: full model path
+
+        Returns:
+            local cached model path
+        """
+
+        # Split into parts
+        parts = path.split("/")
+
+        # Calculate repo id split
+        repo = 2 if len(parts) > 2 else 1
+
+        # Download and cache file
+        return hf_hub_download(repo_id="/".join(parts[:repo]), filename="/".join(parts[repo:]))

txtai/vectors/dense/m2v.py

@@ -0,0 +1,67 @@
+"""
+Model2Vec module
+"""
+
+import json
+
+from huggingface_hub.errors import HFValidationError
+from transformers.utils import cached_file
+
+# Conditional import
+try:
+    from model2vec import StaticModel
+
+    MODEL2VEC = True
+except ImportError:
+    MODEL2VEC = False
+
+from ..base import Vectors
+
+
+class Model2Vec(Vectors):
+    """
+    Builds vectors using Model2Vec.
+    """
+
+    @staticmethod
+    def ismodel(path):
+        """
+        Checks if path is a Model2Vec model.
+
+        Args:
+            path: input path
+
+        Returns:
+            True if this is a Model2Vec model, False otherwise
+        """
+
+        try:
+            # Download file and parse JSON
+            path = cached_file(path_or_repo_id=path, filename="config.json")
+            if path:
+                with open(path, encoding="utf-8") as f:
+                    config = json.load(f)
+                    return config.get("model_type") == "model2vec"
+
+        # Ignore this error - invalid repo or directory
+        except (HFValidationError, OSError):
+            pass
+
+        return False
+
+    def __init__(self, config, scoring, models):
+        # Check before parent constructor since it calls loadmodel
+        if not MODEL2VEC:
+            raise ImportError('Model2Vec is not available - install "vectors" extra to enable')
+
+        super().__init__(config, scoring, models)
+
+    def loadmodel(self, path):
+        return StaticModel.from_pretrained(path)
+
+    def encode(self, data, category=None):
+        # Additional model arguments
+        modelargs = self.config.get("vectors", {})
+
+        # Encode data
+        return self.model.encode(data, batch_size=self.encodebatch, **modelargs)

txtai/vectors/dense/sbert.py

@@ -0,0 +1,92 @@
+"""
+Sentence Transformers module
+"""
+
+# Conditional import
+try:
+    from sentence_transformers import SentenceTransformer
+
+    SENTENCE_TRANSFORMERS = True
+except ImportError:
+    SENTENCE_TRANSFORMERS = False
+
+from ...models import Models
+
+from ..base import Vectors
+
+
+class STVectors(Vectors):
+    """
+    Builds 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')
+
+        # Pool parameter created here since loadmodel is called from parent constructor
+        self.pool = None
+
+        super().__init__(config, scoring, models)
+
+    def loadmodel(self, path):
+        # Get target device
+        gpu, pool = self.config.get("gpu", True), False
+
+        # Default mode uses a single GPU. Setting to all spawns a process per GPU.
+        if isinstance(gpu, str) and gpu == "all":
+            # Get number of accelerator devices available
+            devices = Models.acceleratorcount()
+
+            # Enable multiprocessing pooling only when multiple devices are available
+            gpu, pool = devices <= 1, devices > 1
+
+        # Tensor device id
+        deviceid = Models.deviceid(gpu)
+
+        # Additional model arguments
+        modelargs = self.config.get("vectors", {})
+
+        # Load sentence-transformers encoder
+        model = self.loadencoder(path, device=Models.device(deviceid), **modelargs)
+
+        # Start process pool for multiple GPUs
+        if pool:
+            self.pool = model.start_multi_process_pool()
+
+        # Return model
+        return model
+
+    def encode(self, data, category=None):
+        # Get encode method based on input category
+        encode = self.model.encode_query if category == "query" else self.model.encode_document if category == "data" else self.model.encode
+
+        # Additional encoding arguments
+        encodeargs = self.config.get("encodeargs", {})
+
+        # Encode with sentence transformers encoder
+        return encode(data, pool=self.pool, batch_size=self.encodebatch, **encodeargs)
+
+    def close(self):
+        # Close pool before model is closed in parent method
+        if self.pool:
+            self.model.stop_multi_process_pool(self.pool)
+            self.pool = None
+
+        super().close()
+
+    def loadencoder(self, path, device, **kwargs):
+        """
+        Loads the embeddings encoder model from path.
+
+        Args:
+            path: model path
+            device: tensor device
+            kwargs: additional keyword args
+
+        Returns:
+            embeddings encoder
+        """
+
+        return SentenceTransformer(path, device=device, **kwargs)

txtai/vectors/dense/words.py

@@ -0,0 +1,211 @@
+"""
+Word Vectors module
+"""
+
+import json
+import logging
+import os
+import tempfile
+
+from multiprocessing import Pool
+
+import numpy as np
+
+from huggingface_hub.errors import HFValidationError
+from transformers.utils import cached_file
+
+# Conditional import
+try:
+    from staticvectors import Database, StaticVectors
+
+    STATICVECTORS = True
+except ImportError:
+    STATICVECTORS = False
+
+from ...pipeline import Tokenizer
+
+from ..base import Vectors
+
+# Logging configuration
+logger = logging.getLogger(__name__)
+
+# Multiprocessing helper methods
+# pylint: disable=W0603
+PARAMETERS, VECTORS = None, None
+
+
+def create(config, scoring):
+    """
+    Multiprocessing helper method. Creates a global embeddings object to be accessed in a new subprocess.
+
+    Args:
+        config: vector configuration
+        scoring: scoring instance
+    """
+
+    global PARAMETERS
+    global VECTORS
+
+    # Store model parameters for lazy loading
+    PARAMETERS, VECTORS = (config, scoring, None), None
+
+
+def transform(document):
+    """
+    Multiprocessing helper method. Transforms document into an embeddings vector.
+
+    Args:
+        document: (id, data, tags)
+
+    Returns:
+        (id, embedding)
+    """
+
+    # Lazy load vectors model
+    global VECTORS
+    if not VECTORS:
+        VECTORS = WordVectors(*PARAMETERS)
+
+    return (document[0], VECTORS.transform(document))
+
+
+class WordVectors(Vectors):
+    """
+    Builds vectors using weighted word embeddings.
+    """
+
+    @staticmethod
+    def ismodel(path):
+        """
+        Checks if path is a WordVectors model.
+
+        Args:
+            path: input path
+
+        Returns:
+            True if this is a WordVectors model, False otherwise
+        """
+
+        # Check if this is a SQLite database
+        if WordVectors.isdatabase(path):
+            return True
+
+        try:
+            # Download file and parse JSON
+            path = cached_file(path_or_repo_id=path, filename="config.json")
+            if path:
+                with open(path, encoding="utf-8") as f:
+                    config = json.load(f)
+                    return config.get("model_type") == "staticvectors"
+
+        # Ignore this error - invalid repo or directory
+        except (HFValidationError, OSError):
+            pass
+
+        return False
+
+    @staticmethod
+    def isdatabase(path):
+        """
+        Checks if this is a SQLite database file which is the file format used for word vectors databases.
+
+        Args:
+            path: path to check
+
+        Returns:
+            True if this is a SQLite database
+        """
+
+        return isinstance(path, str) and STATICVECTORS and Database.isdatabase(path)
+
+    def __init__(self, config, scoring, models):
+        # Check before parent constructor since it calls loadmodel
+        if not STATICVECTORS:
+            raise ImportError('staticvectors is not available - install "vectors" extra to enable')
+
+        super().__init__(config, scoring, models)
+
+    def loadmodel(self, path):
+        return StaticVectors(path)
+
+    def encode(self, data, category=None):
+        # Iterate over each data element, tokenize (if necessary) and build an aggregated embeddings vector
+        embeddings = []
+        for tokens in data:
+            # Convert to tokens, if necessary. If tokenized list is empty, use input string.
+            if isinstance(tokens, str):
+                tokenlist = Tokenizer.tokenize(tokens)
+                tokens = tokenlist if tokenlist else [tokens]
+
+            # Generate weights for each vector using a scoring method
+            weights = self.scoring.weights(tokens) if self.scoring else None
+
+            # pylint: disable=E1133
+            if weights and [x for x in weights if x > 0]:
+                # Build weighted average embeddings vector. Create weights array as float32 to match embeddings precision.
+                embedding = np.average(self.lookup(tokens), weights=np.array(weights, dtype=np.float32), axis=0)
+            else:
+                # If no weights, use mean
+                embedding = np.mean(self.lookup(tokens), axis=0)
+
+            embeddings.append(embedding)
+
+        return np.array(embeddings, dtype=np.float32)
+
+    def index(self, documents, batchsize=500, checkpoint=None):
+        # Derive number of parallel processes
+        parallel = self.config.get("parallel", True)
+        parallel = os.cpu_count() if parallel and isinstance(parallel, bool) else int(parallel)
+
+        # Use default single process indexing logic
+        if not parallel:
+            return super().index(documents, batchsize)
+
+        # Customize indexing logic with multiprocessing pool to efficiently build vectors
+        ids, dimensions, batches, stream = [], None, 0, None
+
+        # Shared objects with Pool
+        args = (self.config, self.scoring)
+
+        # Convert all documents to embedding arrays, stream embeddings to disk to control memory usage
+        with Pool(parallel, initializer=create, initargs=args) as pool:
+            with tempfile.NamedTemporaryFile(mode="wb", suffix=".npy", delete=False) as output:
+                stream = output.name
+                embeddings = []
+                for uid, embedding in pool.imap(transform, documents, self.encodebatch):
+                    if not dimensions:
+                        # Set number of dimensions for embeddings
+                        dimensions = embedding.shape[0]
+
+                    ids.append(uid)
+                    embeddings.append(embedding)
+
+                    if len(embeddings) == batchsize:
+                        np.save(output, np.array(embeddings, dtype=np.float32), allow_pickle=False)
+                        batches += 1
+
+                        embeddings = []
+
+                # Final embeddings batch
+                if embeddings:
+                    np.save(output, np.array(embeddings, dtype=np.float32), allow_pickle=False)
+                    batches += 1
+
+        return (ids, dimensions, batches, stream)
+
+    def lookup(self, tokens):
+        """
+        Queries word vectors for given list of input tokens.
+
+        Args:
+            tokens: list of tokens to query
+
+        Returns:
+            word vectors array
+        """
+
+        return self.model.embeddings(tokens)
+
+    def tokens(self, data):
+        # Skip tokenization rules
+        return data