fastembed-bio 0.1.0__py3-none-any.whl

fastembed/sparse/splade_pp.py

@@ -0,0 +1,196 @@
+from typing import Any, Iterable, Sequence, Type
+
+import numpy as np
+from fastembed.common import OnnxProvider
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import Device
+from fastembed.common.utils import define_cache_dir
+from fastembed.sparse.sparse_embedding_base import (
+    SparseEmbedding,
+    SparseTextEmbeddingBase,
+)
+from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
+from fastembed.common.model_description import SparseModelDescription, ModelSource
+
+supported_splade_models: list[SparseModelDescription] = [
+    SparseModelDescription(
+        model="prithivida/Splade_PP_en_v1",
+        vocab_size=30522,
+        description="Independent Implementation of SPLADE++ Model for English.",
+        license="apache-2.0",
+        size_in_GB=0.532,
+        sources=ModelSource(hf="Qdrant/Splade_PP_en_v1"),
+        model_file="model.onnx",
+    ),
+    SparseModelDescription(
+        model="prithvida/Splade_PP_en_v1",
+        vocab_size=30522,
+        description="Independent Implementation of SPLADE++ Model for English.",
+        license="apache-2.0",
+        size_in_GB=0.532,
+        sources=ModelSource(hf="Qdrant/Splade_PP_en_v1"),
+        model_file="model.onnx",
+    ),
+]
+
+
+class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[SparseEmbedding]:
+        if output.attention_mask is None:
+            raise ValueError("attention_mask must be provided for document post-processing")
+
+        relu_log = np.log(1 + np.maximum(output.model_output, 0))
+
+        weighted_log = relu_log * np.expand_dims(output.attention_mask, axis=-1)
+
+        scores = np.max(weighted_log, axis=1)
+
+        # Score matrix of shape (batch_size, vocab_size)
+        # Most of the values are 0, only a few are non-zero
+        for row_scores in scores:
+            indices = row_scores.nonzero()[0]
+            scores = row_scores[indices]
+            yield SparseEmbedding(values=scores, indices=indices)
+
+    def token_count(
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
+    ) -> int:
+        return self._token_count(texts, batch_size=batch_size, **kwargs)
+
+    @classmethod
+    def _list_supported_models(cls) -> list[SparseModelDescription]:
+        """Lists the supported models.
+
+        Returns:
+            list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information.
+        """
+        return supported_splade_models
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        lazy_load: bool = False,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
+        **kwargs: Any,
+    ):
+        """
+        Args:
+            model_name (str): The name of the model to use.
+            cache_dir (str, optional): The path to the cache directory.
+                                       Can be set using the `FASTEMBED_CACHE_PATH` env variable.
+                                       Defaults to `fastembed_cache` in the system's temp directory.
+            threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
+            providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
+                Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.
+            device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
+            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
+                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+            specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
+
+        Raises:
+            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+        """
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+        self.providers = providers
+        self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
+
+        # List of device ids, that can be used for data parallel processing in workers
+        self.device_ids = device_ids
+        self.cuda = cuda
+
+        # This device_id will be used if we need to load model in current process
+        self.device_id: int | None = None
+        if device_id is not None:
+            self.device_id = device_id
+        elif self.device_ids is not None:
+            self.device_id = self.device_ids[0]
+
+        self.model_description = self._get_model_description(model_name)
+        self.cache_dir = str(define_cache_dir(cache_dir))
+
+        self._specific_model_path = specific_model_path
+        self._model_dir = self.download_model(
+            self.model_description,
+            self.cache_dir,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+        )
+
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    def load_onnx_model(self) -> None:
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description.model_file,
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
+        )
+
+    def embed(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[SparseEmbedding]:
+        """
+        Encode a list of documents into list of embeddings.
+        We use mean pooling with attention so that the model can handle variable-length inputs.
+
+        Args:
+            documents: Iterator of documents or single document to embed
+            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
+            parallel:
+                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
+                If 0, use all available cores.
+                If None, don't use data-parallel processing, use default onnxruntime threading instead.
+
+        Returns:
+            List of embeddings, one per document
+        """
+        yield from self._embed_documents(
+            model_name=self.model_name,
+            cache_dir=str(self.cache_dir),
+            documents=documents,
+            batch_size=batch_size,
+            parallel=parallel,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
+            **kwargs,
+        )
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker[SparseEmbedding]]:
+        return SpladePPEmbeddingWorker
+
+
+class SpladePPEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> SpladePP:
+        return SpladePP(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/sparse/utils/minicoil_encoder.py

@@ -0,0 +1,146 @@
+"""
+Pure numpy implementation of encoder model for a single word.
+
+This model is not trainable, and should only be used for inference.
+"""
+
+import numpy as np
+from fastembed.common.types import NumpyArray
+
+
+class Encoder:
+    """
+    Encoder(768, 4, 10000)
+
+    Will look like this:
+
+
+                                         Per-word
+                                         Encoder Matrix
+     ┌─────────────────────┐
+     │ Token Embedding(768)├──────┐      (10k, 768, 4)
+     └─────────────────────┘      │         ┌─────────┐
+                                  │         │         │
+     ┌─────────────────────┐      │       ┌─┴───────┐ │
+     │                     │      │       │         │ │
+     └─────────────────────┘      │     ┌─┴───────┐ │ │      ┌─────────┐
+                                  └────►│         │ │ ├─────►│Tanh     │
+     ┌─────────────────────┐            │         │ │ │      └─────────┘
+     │                     │            │         │ ├─┘
+     └─────────────────────┘            │         ├─┘
+                                        │         │
+     ┌─────────────────────┐            └─────────┘
+     │                     │
+     └─────────────────────┘
+
+     Final linear transformation is accompanied by a non-linear activation function: Tanh.
+
+     Tanh is used to ensure that the output is in the range [-1, 1].
+     It would be easier to visually interpret the output of the model, assuming that each dimension
+     would need to encode a type of semantic cluster.
+    """
+
+    def __init__(
+        self,
+        weights: NumpyArray,
+    ):
+        self.weights = weights
+        self.vocab_size, self.input_dim, self.output_dim = weights.shape
+
+        self.encoder_weights: NumpyArray = weights
+
+        # Activation function
+        self.activation = np.tanh
+
+    @staticmethod
+    def convert_vocab_ids(vocab_ids: NumpyArray) -> NumpyArray:
+        """
+        Convert vocab_ids of shape (batch_size, seq_len) into (batch_size, seq_len, 2)
+        by appending batch_id alongside each vocab_id.
+        """
+        batch_size, seq_len = vocab_ids.shape
+        batch_ids = np.arange(batch_size, dtype=vocab_ids.dtype).reshape(batch_size, 1)
+        batch_ids = np.repeat(batch_ids, seq_len, axis=1)
+        # Stack vocab_ids and batch_ids along the last dimension
+        combined: NumpyArray = np.stack((vocab_ids, batch_ids), axis=2).astype(np.int32)
+        return combined
+
+    @classmethod
+    def avg_by_vocab_ids(
+        cls, vocab_ids: NumpyArray, embeddings: NumpyArray
+    ) -> tuple[NumpyArray, NumpyArray]:
+        """
+        Takes:
+            vocab_ids: (batch_size, seq_len) int array
+            embeddings: (batch_size, seq_len, input_dim) float array
+
+        Returns:
+            unique_flattened_vocab_ids: (total_unique, 2) array of [vocab_id, batch_id]
+            unique_flattened_embeddings: (total_unique, input_dim) averaged embeddings
+        """
+        input_dim = embeddings.shape[2]
+
+        # Flatten vocab_ids and embeddings
+        # flattened_vocab_ids: (batch_size*seq_len, 2)
+        flattened_vocab_ids = cls.convert_vocab_ids(vocab_ids).reshape(-1, 2)
+
+        # flattened_embeddings: (batch_size*seq_len, input_dim)
+        flattened_embeddings = embeddings.reshape(-1, input_dim)
+
+        # Find unique (vocab_id, batch_id) pairs
+        unique_flattened_vocab_ids, inverse_indices = np.unique(
+            flattened_vocab_ids, axis=0, return_inverse=True
+        )
+
+        # Prepare arrays to accumulate sums
+        unique_count = unique_flattened_vocab_ids.shape[0]
+        unique_flattened_embeddings = np.zeros((unique_count, input_dim), dtype=np.float32)
+        unique_flattened_count = np.zeros(unique_count, dtype=np.int32)
+
+        # Use np.add.at to accumulate sums based on inverse indices
+        np.add.at(unique_flattened_embeddings, inverse_indices, flattened_embeddings)
+        np.add.at(unique_flattened_count, inverse_indices, 1)
+
+        # Compute averages
+        unique_flattened_embeddings /= unique_flattened_count[:, None]
+
+        return unique_flattened_vocab_ids.astype(np.int32), unique_flattened_embeddings.astype(
+            np.float32
+        )
+
+    def forward(
+        self, vocab_ids: NumpyArray, embeddings: NumpyArray
+    ) -> tuple[NumpyArray, NumpyArray]:
+        """
+        Args:
+            vocab_ids: (batch_size, seq_len) int array
+            embeddings: (batch_size, seq_len, input_dim) float array
+
+        Returns:
+            unique_flattened_vocab_ids_and_batch_ids: (total_unique, 2)
+            unique_flattened_encoded: (total_unique, output_dim)
+        """
+        # Average embeddings for duplicate vocab_ids
+        unique_flattened_vocab_ids_and_batch_ids, unique_flattened_embeddings = (
+            self.avg_by_vocab_ids(vocab_ids, embeddings)
+        )
+
+        # Select the encoder weights for each unique vocab_id
+        unique_flattened_vocab_ids = unique_flattened_vocab_ids_and_batch_ids[:, 0].astype(
+            np.int32
+        )
+
+        # unique_encoder_weights: (total_unique, input_dim, output_dim)
+        unique_encoder_weights = self.encoder_weights[unique_flattened_vocab_ids]
+
+        # Compute linear transform: (total_unique, output_dim)
+        # Using Einstein summation for matrix multiplication:
+        # 'bi,bio->bo' means: for each "b" (batch element), multiply embeddings (b,i) by weights (b,i,o) -> (b,o)
+        unique_flattened_encoded = np.einsum(
+            "bi,bio->bo", unique_flattened_embeddings, unique_encoder_weights
+        )
+
+        # Apply Tanh activation and ensure float32 type
+        unique_flattened_encoded = self.activation(unique_flattened_encoded).astype(np.float32)
+
+        return unique_flattened_vocab_ids_and_batch_ids.astype(np.int32), unique_flattened_encoded

fastembed/sparse/utils/sparse_vectors_converter.py

@@ -0,0 +1,244 @@
+import copy
+from dataclasses import dataclass
+
+import mmh3
+import numpy as np
+from py_rust_stemmers import SnowballStemmer
+
+from fastembed.common.utils import get_all_punctuation, remove_non_alphanumeric
+from fastembed.sparse.sparse_embedding_base import SparseEmbedding
+
+GAP = 32000
+INT32_MAX = 2**31 - 1
+
+
+@dataclass
+class WordEmbedding:
+    word: str
+    forms: list[str]
+    count: int
+    word_id: int
+    embedding: list[float]
+
+
+class SparseVectorConverter:
+    def __init__(
+        self,
+        stopwords: set[str],
+        stemmer: SnowballStemmer,
+        k: float = 1.2,
+        b: float = 0.75,
+        avg_len: float = 150.0,
+    ):
+        punctuation = set(get_all_punctuation())
+        special_tokens = {"[CLS]", "[SEP]", "[PAD]", "[UNK]", "[MASK]"}
+
+        self.stemmer = stemmer
+        self.unwanted_tokens = punctuation | special_tokens | stopwords
+
+        self.k = k
+        self.b = b
+        self.avg_len = avg_len
+
+    @classmethod
+    def unkn_word_token_id(
+        cls, word: str, shift: int
+    ) -> int:  # 2-3 words can collide in 1 index with this mapping, not considering mm3 collisions
+        token_hash = abs(mmh3.hash(word))
+
+        range_size = INT32_MAX - shift
+        remapped_hash = shift + (token_hash % range_size)
+
+        return remapped_hash
+
+    def bm25_tf(self, num_occurrences: int, sentence_len: int) -> float:
+        res = num_occurrences * (self.k + 1)
+        res /= num_occurrences + self.k * (1 - self.b + self.b * sentence_len / self.avg_len)
+        return res
+
+    @classmethod
+    def normalize_vector(cls, vector: list[float]) -> list[float]:
+        norm = sum([x**2 for x in vector]) ** 0.5
+        if norm < 1e-8:
+            return vector
+        return [x / norm for x in vector]
+
+    def clean_words(
+        self, sentence_embedding: dict[str, WordEmbedding], token_max_length: int = 40
+    ) -> dict[str, WordEmbedding]:
+        """
+        Clean miniCOIL-produced sentence_embedding, as unknown to the miniCOIL's stemmer tokens should fully resemble
+        our BM25 token representation.
+
+        sentence_embedding = {"9°": {"word": "9°", "word_id": -1, "count": 2, "embedding": [1], "forms": ["9°"]},
+                "9": {"word": "9", "word_id": -1, "count": 2, "embedding": [1], "forms": ["9"]},
+                "bat": {"word": "bat", "word_id": 2, "count": 3, "embedding": [0.2, 0.1, -0.2, -0.2], "forms": ["bats", "bat"]},
+                "9°9": {"word": "9°9", "word_id": -1, "count": 1, "embedding": [1], "forms": ["9°9"]},
+                "screech": {"word": "screech", "word_id": -1, "count": 1, "embedding": [1], "forms": ["screech"]},
+                "screeched": {"word": "screeched", "word_id": -1, "count": 1, "embedding": [1], "forms": ["screeched"]}
+                }
+        cleaned_embedding_ground_truth = {
+                "9": {"word": "9", "word_id": -1, "count": 6, "embedding": [1], "forms": ["9°", "9", "9°9", "9°9"]},
+                "bat": {"word": "bat", "word_id": 2, "count": 3, "embedding": [0.2, 0.1, -0.2, -0.2], "forms": ["bats", "bat"]},
+                "screech": {"word": "screech", "word_id": -1, "count": 2, "embedding": [1], "forms": ["screech", "screeched"]}
+                }
+        """
+
+        new_sentence_embedding: dict[str, WordEmbedding] = {}
+
+        for word, embedding in sentence_embedding.items():
+            # embedding = {
+            #     "word": "vector",
+            #     "forms": ["vector", "vectors"],
+            #     "count": 2,
+            #     "word_id": 1231,
+            #     "embedding": [0.1, 0.2, 0.3, 0.4]
+            # }
+            if embedding.word_id > 0:
+                # Known word, no need to clean
+                new_sentence_embedding[word] = embedding
+            else:
+                # Unknown word
+                if word in self.unwanted_tokens:
+                    continue
+
+                # Example complex word split:
+                # word = `word^vec`
+                word_cleaned = remove_non_alphanumeric(word).strip()
+                # word_cleaned = `word vec`
+
+                if len(word_cleaned) > 0:
+                    # Subwords: ['word', 'vec']
+                    for subword in word_cleaned.split():
+                        stemmed_subword: str = self.stemmer.stem_word(subword)
+                        if (
+                            len(stemmed_subword) <= token_max_length
+                            and stemmed_subword not in self.unwanted_tokens
+                        ):
+                            if stemmed_subword not in new_sentence_embedding:
+                                new_sentence_embedding[stemmed_subword] = copy.deepcopy(embedding)
+                                new_sentence_embedding[stemmed_subword].word = stemmed_subword
+                            else:
+                                new_sentence_embedding[stemmed_subword].count += embedding.count
+                                new_sentence_embedding[stemmed_subword].forms += embedding.forms
+
+        return new_sentence_embedding
+
+    def embedding_to_vector(
+        self,
+        sentence_embedding: dict[str, WordEmbedding],
+        embedding_size: int,
+        vocab_size: int,
+    ) -> SparseEmbedding:
+        """
+        Convert miniCOIL sentence embedding to Qdrant sparse vector
+
+        Example input:
+
+        ```
+        {
+            "vector": WordEmbedding({ // Vocabulary word, encoded with miniCOIL normally
+                "word": "vector",
+                "forms": ["vector", "vectors"],
+                "count": 2,
+                "word_id": 1231,
+                "embedding": [0.1, 0.2, 0.3, 0.4]
+            }),
+            "axiotic": WordEmbedding({ // Out-of-vocabulary word, fallback to BM25
+                "word": "axiotic",
+                "forms": ["axiotics"],
+                "count": 1,
+                "word_id": -1,
+            })
+        }
+        ```
+
+        """
+
+        indices: list[int] = []
+        values: list[float] = []
+
+        # Example:
+        # vocab_size = 10000
+        # embedding_size = 4
+        # GAP = 32000
+        #
+        # We want to start random words section from the bucket, that is guaranteed to not
+        # include any vocab words.
+        # We need (vocab_size * embedding_size) slots for vocab words.
+        # Therefore we need (vocab_size * embedding_size) // GAP + 1 buckets for vocab words.
+        # Therefore, we can start random words from bucket (vocab_size * embedding_size) // GAP + 1 + 1
+
+        # ID at which the scope of OOV words starts
+        unknown_words_shift = ((vocab_size * embedding_size) // GAP + 2) * GAP
+        sentence_embedding_cleaned = self.clean_words(sentence_embedding)
+
+        # Calculate sentence length after cleaning
+        sentence_len = 0
+        for embedding in sentence_embedding_cleaned.values():
+            sentence_len += embedding.count
+
+        for embedding in sentence_embedding_cleaned.values():
+            word_id = embedding.word_id
+            num_occurrences = embedding.count
+            tf = self.bm25_tf(num_occurrences, sentence_len)
+            if (
+                word_id > 0
+            ):  # miniCOIL starts with ID 1, we generally won't have word_id == 0 (UNK), as we don't add
+                # these words to sentence_embedding
+                embedding_values = embedding.embedding
+                normalized_embedding = self.normalize_vector(embedding_values)
+
+                for val_id, value in enumerate(normalized_embedding):
+                    indices.append(
+                        word_id * embedding_size + val_id
+                    )  # since miniCOIL IDs start with 1
+                    values.append(value * tf)
+            else:
+                indices.append(self.unkn_word_token_id(embedding.word, unknown_words_shift))
+                values.append(tf)
+
+        return SparseEmbedding(
+            indices=np.array(indices, dtype=np.int32),
+            values=np.array(values, dtype=np.float32),
+        )
+
+    def embedding_to_vector_query(
+        self,
+        sentence_embedding: dict[str, WordEmbedding],
+        embedding_size: int,
+        vocab_size: int,
+    ) -> SparseEmbedding:
+        """
+        Same as `embedding_to_vector`, but no TF
+        """
+
+        indices: list[int] = []
+        values: list[float] = []
+
+        # ID at which the scope of OOV words starts
+        unknown_words_shift = ((vocab_size * embedding_size) // GAP + 2) * GAP
+
+        sentence_embedding_cleaned = self.clean_words(sentence_embedding)
+
+        for embedding in sentence_embedding_cleaned.values():
+            word_id = embedding.word_id
+            tf = 1.0
+
+            if word_id >= 0:  # miniCOIL starts with ID 1
+                embedding_values = embedding.embedding
+                normalized_embedding = self.normalize_vector(embedding_values)
+
+                for val_id, value in enumerate(normalized_embedding):
+                    indices.append(
+                        word_id * embedding_size + val_id
+                    )  # since miniCOIL IDs start with 1
+                    values.append(value * tf)
+            else:
+                indices.append(self.unkn_word_token_id(embedding.word, unknown_words_shift))
+                values.append(tf)
+
+        return SparseEmbedding(
+            indices=np.array(indices, dtype=np.int32),
+            values=np.array(values, dtype=np.float32),
+        )