mseep-txtai 9.1.1__py3-none-any.whl

txtai/models/pooling/base.py

@@ -0,0 +1,141 @@
+"""
+Pooling module
+"""
+
+import numpy as np
+import torch
+
+from torch import nn
+
+from ..models import Models
+
+
+class Pooling(nn.Module):
+    """
+    Builds pooled vectors usings outputs from a transformers model.
+    """
+
+    def __init__(self, path, device, tokenizer=None, maxlength=None, modelargs=None):
+        """
+        Creates a new Pooling model.
+
+        Args:
+            path: path to model, accepts Hugging Face model hub id or local path
+            device: tensor device id
+            tokenizer: optional path to tokenizer
+            maxlength: max sequence length
+            modelargs: additional model arguments
+        """
+
+        super().__init__()
+
+        self.model = Models.load(path, modelargs=modelargs)
+        self.tokenizer = Models.tokenizer(tokenizer if tokenizer else path)
+        self.device = Models.device(device)
+
+        # Detect unbounded tokenizer typically found in older models
+        Models.checklength(self.model, self.tokenizer)
+
+        # Set max length
+        self.maxlength = maxlength if maxlength else self.tokenizer.model_max_length if self.tokenizer.model_max_length != int(1e30) else None
+
+        # Move to device
+        self.to(self.device)
+
+    def encode(self, documents, batch=32, category=None):
+        """
+        Builds an array of pooled embeddings for documents.
+
+        Args:
+            documents: list of documents used to build embeddings
+            batch: model batch size
+            category: embeddings category (query or data)
+
+        Returns:
+            pooled embeddings
+        """
+
+        # Split documents into batches and process
+        results = []
+
+        # Apply pre encoding transformation logic
+        documents = self.preencode(documents, category)
+
+        # Sort document indices from largest to smallest to enable efficient batching
+        # This performance tweak matches logic in sentence-transformers
+        lengths = np.argsort([-len(x) if x else 0 for x in documents])
+        documents = [documents[x] for x in lengths]
+
+        for chunk in self.chunk(documents, batch):
+            # Tokenize input
+            inputs = self.tokenizer(chunk, padding=True, truncation="longest_first", return_tensors="pt", max_length=self.maxlength)
+
+            # Move inputs to device
+            inputs = inputs.to(self.device)
+
+            # Run inputs through model
+            with torch.no_grad():
+                outputs = self.forward(**inputs)
+
+            # Add batch result
+            results.extend(outputs.cpu().to(torch.float32).numpy())
+
+        # Apply post encoding transformation logic
+        results = self.postencode(results, category)
+
+        # Restore original order and return array
+        return np.asarray([results[x] for x in np.argsort(lengths)])
+
+    def chunk(self, texts, size):
+        """
+        Splits texts into separate batch sizes specified by size.
+
+        Args:
+            texts: text elements
+            size: batch size
+
+        Returns:
+            list of evenly sized batches with the last batch having the remaining elements
+        """
+
+        return [texts[x : x + size] for x in range(0, len(texts), size)]
+
+    def forward(self, **inputs):
+        """
+        Runs inputs through transformers model and returns outputs.
+
+        Args:
+            inputs: model inputs
+
+        Returns:
+            model outputs
+        """
+
+        return self.model(**inputs)[0]
+
+    # pylint: disable=W0613
+    def preencode(self, documents, category):
+        """
+        Applies pre encoding transformation logic.
+
+        Args:
+            documents: list of documents used to build embeddings
+            category: embeddings category (query or data)
+        """
+
+        return documents
+
+    # pylint: disable=W0613
+    def postencode(self, results, category):
+        """
+        Applies post encoding transformation logic.
+
+        Args:
+            results: list of results
+            category: embeddings category (query or data)
+
+        Returns:
+            results with transformation logic applied
+        """
+
+        return results

txtai/models/pooling/cls.py

@@ -0,0 +1,28 @@
+"""
+CLS module
+"""
+
+from .base import Pooling
+
+
+class ClsPooling(Pooling):
+    """
+    Builds CLS pooled vectors using outputs from a transformers model.
+    """
+
+    def forward(self, **inputs):
+        """
+        Runs CLS pooling on token embeddings.
+
+        Args:
+            inputs: model inputs
+
+        Returns:
+            CLS pooled embeddings using output token embeddings (i.e. last hidden state)
+        """
+
+        # Run through transformers model
+        tokens = super().forward(**inputs)
+
+        # CLS token pooling
+        return tokens[:, 0]

txtai/models/pooling/factory.py

@@ -0,0 +1,144 @@
+"""
+Factory module
+"""
+
+import json
+import os
+
+from huggingface_hub.errors import HFValidationError
+from transformers.utils import cached_file
+
+from .base import Pooling
+from .cls import ClsPooling
+from .late import LatePooling
+from .mean import MeanPooling
+
+
+class PoolingFactory:
+    """
+    Method to create pooling models.
+    """
+
+    @staticmethod
+    def create(config):
+        """
+        Create a Pooling model.
+
+        Args:
+            config: pooling configuration
+
+        Returns:
+            Pooling
+        """
+
+        # Unpack parameters
+        method, path, device, tokenizer, maxlength, modelargs = [
+            config.get(x) for x in ["method", "path", "device", "tokenizer", "maxlength", "modelargs"]
+        ]
+
+        # Derive maxlength, if applicable
+        maxlength = PoolingFactory.maxlength(path) if isinstance(maxlength, bool) and maxlength else maxlength
+
+        # Default pooling returns hidden state
+        if isinstance(path, bytes) or (isinstance(path, str) and os.path.isfile(path)) or method == "pooling":
+            return Pooling(path, device, tokenizer, maxlength, modelargs)
+
+        # Derive pooling method if it's not specified and path is a string
+        if (not method or method not in ("clspooling", "meanpooling", "latepooling")) and isinstance(path, str):
+            method = PoolingFactory.method(path)
+
+        # Check for cls pooling
+        if method == "clspooling":
+            return ClsPooling(path, device, tokenizer, maxlength, modelargs)
+
+        # Check for late pooling
+        if method == "latepooling":
+            return LatePooling(path, device, tokenizer, maxlength, modelargs)
+
+        # Default to mean pooling
+        return MeanPooling(path, device, tokenizer, maxlength, modelargs)
+
+    @staticmethod
+    def method(path):
+        """
+        Determines the pooling method using the sentence transformers pooling config.
+
+        Args:
+            path: model path
+
+        Returns:
+            pooling method
+        """
+
+        # Default method
+        method = "meanpooling"
+
+        # Load 1_Pooling/config.json file
+        config = PoolingFactory.load(path, "1_Pooling/config.json")
+
+        # Set to CLS pooling if it's enabled and mean pooling is disabled
+        if config and config["pooling_mode_cls_token"] and not config["pooling_mode_mean_tokens"]:
+            method = "clspooling"
+
+        # Check for late interaction pooling
+        if not config:
+            # Load 1_Dense/config.json
+            config = PoolingFactory.load(path, "1_Dense/config.json")
+            if config:
+                method = "latepooling"
+
+            # Load config.json and check architecture
+            else:
+                config = PoolingFactory.load(path, "config.json")
+                if config and "HF_ColBERT" in config.get("architectures", []):
+                    method = "latepooling"
+
+        return method
+
+    @staticmethod
+    def maxlength(path):
+        """
+        Reads the max_seq_length parameter from sentence transformers config.
+
+        Args:
+            path: model path
+
+        Returns:
+            max sequence length
+        """
+
+        # Default length is unset
+        maxlength = None
+
+        # Read max_seq_length from sentence_bert_config.json
+        config = PoolingFactory.load(path, "sentence_bert_config.json")
+        maxlength = config.get("max_seq_length") if config else maxlength
+
+        return maxlength
+
+    @staticmethod
+    def load(path, name):
+        """
+        Loads a JSON config file from the Hugging Face Hub.
+
+        Args:
+            path: model path
+            name: file to load
+
+        Returns:
+            config
+        """
+
+        # Download file and parse JSON
+        config = None
+        try:
+            path = cached_file(path_or_repo_id=path, filename=name)
+            if path:
+                with open(path, encoding="utf-8") as f:
+                    config = json.load(f)
+
+        # Ignore this error - invalid repo or directory
+        except (HFValidationError, OSError):
+            pass
+
+        return config

txtai/models/pooling/late.py

@@ -0,0 +1,173 @@
+"""
+Late module
+"""
+
+import json
+
+import numpy as np
+import torch
+
+from huggingface_hub.errors import HFValidationError
+from safetensors import safe_open
+from torch import nn
+from transformers.utils import cached_file
+
+from .base import Pooling
+from .muvera import Muvera
+
+
+class LatePooling(Pooling):
+    """
+    Builds late pooled vectors using outputs from a transformers model.
+    """
+
+    def __init__(self, path, device, tokenizer=None, maxlength=None, modelargs=None):
+        # Check if fixed dimensional encoder is enabled
+        modelargs = modelargs.copy() if modelargs else {}
+        muvera = modelargs.pop("muvera", {})
+        self.encoder = Muvera(**muvera) if muvera is not None else None
+
+        # Call parent initialization
+        super().__init__(path, device, tokenizer, maxlength, modelargs)
+
+        # Get linear weights path
+        config = self.load(path, "1_Dense/config.json")
+        if config:
+            # PyLate weights format
+            name = "1_Dense/model.safetensors"
+        else:
+            # Stanford weights format
+            name = "model.safetensors"
+
+        # Read model settings
+        self.qprefix, self.qlength, self.dprefix, self.dlength = self.settings(path, config)
+
+        # Load linear layer
+        path = cached_file(path_or_repo_id=path, filename=name)
+        with safe_open(filename=path, framework="pt") as f:
+            weights = f.get_tensor("linear.weight")
+
+            # Load weights into linear layer
+            self.linear = nn.Linear(weights.shape[1], weights.shape[0], bias=False, device=self.device, dtype=weights.dtype)
+            with torch.no_grad():
+                self.linear.weight.copy_(weights)
+
+    def forward(self, **inputs):
+        """
+        Runs late pooling on token embeddings.
+
+        Args:
+            inputs: model inputs
+
+        Returns:
+            Late pooled embeddings using output token embeddings (i.e. last hidden state)
+        """
+
+        # Run through transformers model
+        tokens = super().forward(**inputs)
+
+        # Run through final linear layer and return
+        return self.linear(tokens)
+
+    def preencode(self, documents, category):
+        """
+        Apply prefixes and lengths to data.
+
+        Args:
+            documents: list of documents used to build embeddings
+            category: embeddings category (query or data)
+        """
+
+        results = []
+
+        # Apply prefix
+        for text in documents:
+            prefix = self.qprefix if category == "query" else self.dprefix
+            if prefix:
+                text = f"{prefix}{text}"
+
+            results.append(text)
+
+        # Set maxlength
+        maxlength = self.qlength if category == "query" else self.dlength
+        if maxlength:
+            self.maxlength = maxlength
+
+        return results
+
+    def postencode(self, results, category):
+        """
+        Normalizes and pads results.
+
+        Args:
+            results: input results
+
+        Returns:
+            normalized results with padding
+        """
+
+        length = 0
+        for vectors in results:
+            # Get max length
+            if vectors.shape[0] > length:
+                length = vectors.shape[0]
+
+            # Normalize vectors
+            vectors /= np.linalg.norm(vectors, axis=1)[:, np.newaxis]
+
+        # Pad values
+        data = []
+        for vectors in results:
+            data.append(np.pad(vectors, [(0, length - vectors.shape[0]), (0, 0)]))
+
+        # Build NumPy array
+        data = np.asarray(data)
+
+        # Apply fixed dimesional encoder, if necessary
+        return self.encoder(data, category) if self.encoder else data
+
+    def settings(self, path, config):
+        """
+        Reads model settings.
+
+        Args:
+            path: model path
+            config: PyLate model format if provided, otherwise read from Stanford format
+        """
+
+        if config:
+            # PyLate format
+            config = self.load(path, "config_sentence_transformers.json")
+            params = ["query_prefix", "query_length", "document_prefix", "document_length"]
+        else:
+            # Stanford format
+            config = self.load(path, "artifact.metadata")
+            params = ["query_token_id", "query_maxlen", "doc_token_id", "doc_maxlen"]
+
+        return [config.get(p) for p in params]
+
+    def load(self, path, name):
+        """
+        Loads a JSON config file from the Hugging Face Hub.
+
+        Args:
+            path: model path
+            name: file to load
+
+        Returns:
+            config
+        """
+
+        # Download file and parse JSON
+        config = None
+        try:
+            path = cached_file(path_or_repo_id=path, filename=name)
+            if path:
+                with open(path, encoding="utf-8") as f:
+                    config = json.load(f)
+
+        # Ignore this error - invalid repo or directory
+        except (HFValidationError, OSError):
+            pass
+
+        return config

txtai/models/pooling/mean.py

@@ -0,0 +1,33 @@
+"""
+Mean module
+"""
+
+import torch
+
+from .base import Pooling
+
+
+class MeanPooling(Pooling):
+    """
+    Builds mean pooled vectors usings outputs from a transformers model.
+    """
+
+    def forward(self, **inputs):
+        """
+        Runs mean pooling on token embeddings taking the input mask into account.
+
+        Args:
+            inputs: model inputs
+
+        Returns:
+            mean pooled embeddings using output token embeddings (i.e. last hidden state)
+        """
+
+        # Run through transformers model
+        tokens = super().forward(**inputs)
+        mask = inputs["attention_mask"]
+
+        # Mean pooling
+        # pylint: disable=E1101
+        mask = mask.unsqueeze(-1).expand(tokens.size()).float()
+        return torch.sum(tokens * mask, 1) / torch.clamp(mask.sum(1), min=1e-9)

txtai/models/pooling/muvera.py

@@ -0,0 +1,164 @@
+"""
+Muvera module
+"""
+
+import numpy as np
+
+
+class Muvera:
+    """
+    Implements the MUVERA (Multi-Vector Retrieval via Fixed Dimensional Encodings) algorithm. This reduces late interaction multi-vector
+    outputs to a single fixed vector.
+
+    The size of the output vectors are set using the following parameters
+
+        output dimensions = repetitions * 2^hashes * projected
+
+    For example, the default parameters create vectors with the following output dimensions.
+
+        output dimensions = 20 * 2^5 * 16 = 10240
+
+    This code is based on the following:
+      - Paper: https://arxiv.org/abs/2405.19504
+      - GitHub: https://github.com/google/graph-mining/tree/main/sketching/point_cloud
+      - Python port of the original C++ code: https://github.com/sigridjineth/muvera-py
+    """
+
+    def __init__(self, repetitions=20, hashes=5, projection=16, seed=42):
+        """
+        Creates a Muvera instance.
+
+        Args:
+            repetitions: number of iterations
+            hashes: number of simhash partitions as 2^hashes
+            projection: dimensionality reduction, uses an identity projection when set to None
+            seed: random seed
+        """
+
+        # Number of repetitions
+        self.repetitions = repetitions
+
+        # Number of simhash projections
+        self.hashes = hashes
+
+        # Optional number of projected dimensions
+        self.projection = projection
+
+        # Seed
+        self.seed = seed
+
+    def __call__(self, data, category):
+        """
+        Transforms a list of multi-vector collections into single fixed vector outputs.
+
+        Args:
+            data: array of multi-vector vectors
+            category: embeddings category (query or data)
+        """
+
+        # Get stats
+        dimension, length = data[0].shape[1], len(data)
+
+        # Determine projection dimension
+        identity = not self.projection
+        projection = dimension if identity else self.projection
+
+        # Number of simhash partitions
+        partitions = 2**self.hashes
+
+        # Document tracking
+        lengths = np.array([len(doc) for doc in data], dtype=np.int32)
+        total = np.sum(lengths)
+        documents = np.repeat(np.arange(length), lengths)
+
+        # Stack all vectors
+        points = np.vstack(data).astype(np.float32)
+
+        # Output vectors
+        size = self.repetitions * partitions * projection
+        vectors = np.zeros((length, size), dtype=np.float32)
+
+        # Process each repetition
+        for number in range(self.repetitions):
+            seed = self.seed + number
+
+            # Calculate the simhash
+            sketches = points @ self.random(dimension, self.hashes, seed)
+
+            # Dimensionality reduction, if necessary
+            projected = points if identity else (points @ self.reducer(dimension, projection, seed))
+
+            # Get partition indices
+            bits = (sketches > 0).astype(np.uint32)
+            indices = np.zeros(total, dtype=np.uint32)
+
+            # Calculate vector indices
+            for x in range(self.hashes):
+                indices = (indices << 1) + (bits[:, x] ^ (indices & 1))
+
+            # Initialize storage
+            fdesum = np.zeros((length * partitions * projection,), dtype=np.float32)
+            counts = np.zeros((length, partitions), dtype=np.int32)
+
+            # Count vectors per partition per document
+            np.add.at(counts, (documents, indices), 1)
+
+            # Aggregate vectors using flattened indexing for efficiency
+            part = documents * partitions + indices
+            base = part * projection
+
+            for d in range(projection):
+                flat = base + d
+                np.add.at(fdesum, flat, projected[:, d])
+
+            # Reshape for easier manipulation
+            # pylint: disable=E1121
+            fdesum = fdesum.reshape(length, partitions, projection)
+
+            # Convert sums to averages for data category
+            if category == "data":
+                # Safe division (avoid divide by zero)
+                counts = counts[:, :, np.newaxis]
+                np.divide(fdesum, counts, out=fdesum, where=counts > 0)
+
+            # Save results
+            start = number * partitions * projection
+            vectors[:, start : start + partitions * projection] = fdesum.reshape(length, -1)
+
+        return vectors
+
+    def random(self, dimension, projection, seed):
+        """
+        Generates a random matrix for simhash projections.
+
+        Args:
+            dimensions: number of dimensions for input vectors
+            projections: number of projection dimensions
+            seed: random seed
+
+        Returns:
+            random matrix for simhash projections
+        """
+
+        rng = np.random.default_rng(seed)
+        return rng.normal(loc=0.0, scale=1.0, size=(dimension, projection)).astype(np.float32)
+
+    def reducer(self, dimension, projection, seed):
+        """
+        Generates a random matrix for dimensionality reduction using the AMS sketch algorithm.
+
+        Args:
+            dimension: number of input dimensions
+            projected: number of dimensions to project inputs to
+
+        Returns:
+            Dimensionality reduced matrix
+        """
+
+        rng = np.random.default_rng(seed)
+        out = np.zeros((dimension, projection), dtype=np.float32)
+        indices = rng.integers(0, projection, size=dimension)
+        signs = rng.choice([-1.0, 1.0], size=dimension)
+        out[np.arange(dimension), indices] = signs
+
+        return out