codekingpro/portable-devtools

1from pathlib import Path2 3from typing import Any, Sequence, Iterable, Type4 5import numpy as np6from numpy.typing import NDArray7from py_rust_stemmers import SnowballStemmer8from tokenizers import Tokenizer9 10from fastembed.common.model_description import SparseModelDescription, ModelSource11from fastembed.common.onnx_model import OnnxOutputContext12from fastembed.common import OnnxProvider13from fastembed.common.types import Device14from fastembed.common.utils import define_cache_dir15from fastembed.sparse.sparse_embedding_base import (16    SparseEmbedding,17    SparseTextEmbeddingBase,18)19from fastembed.sparse.utils.minicoil_encoder import Encoder20from fastembed.sparse.utils.sparse_vectors_converter import SparseVectorConverter, WordEmbedding21from fastembed.sparse.utils.vocab_resolver import VocabResolver, VocabTokenizer22from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker23 24 25MINICOIL_MODEL_FILE = "minicoil.triplet.model.npy"26MINICOIL_VOCAB_FILE = "minicoil.triplet.model.vocab"27STOPWORDS_FILE = "stopwords.txt"28 29 30supported_minicoil_models: list[SparseModelDescription] = [31    SparseModelDescription(32        model="Qdrant/minicoil-v1",33        vocab_size=19125,34        description="Sparse embedding model, that resolves semantic meaning of the words, "35        "while keeping exact keyword match behavior. "36        "Based on jinaai/jina-embeddings-v2-small-en-tokens",37        license="apache-2.0",38        size_in_GB=0.09,39        sources=ModelSource(hf="Qdrant/minicoil-v1"),40        model_file="onnx/model.onnx",41        additional_files=[42            STOPWORDS_FILE,43            MINICOIL_MODEL_FILE,44            MINICOIL_VOCAB_FILE,45        ],46        requires_idf=True,47    ),48]49 50_MODEL_TO_LANGUAGE = {51    "Qdrant/minicoil-v1": "english",52}53MODEL_TO_LANGUAGE = {54    model_name.lower(): language for model_name, language in _MODEL_TO_LANGUAGE.items()55}56 57 58def get_language_by_model_name(model_name: str) -> str:59    return MODEL_TO_LANGUAGE[model_name.lower()]60 61 62class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):63    """64        MiniCOIL is a sparse embedding model, that resolves semantic meaning of the words,65        while keeping exact keyword match behavior.66 67        Each vocabulary token is converted into 4d component of a sparse vector, which is then weighted by the token frequency in the corpus.68        If the token is not found in the corpus, it is treated exactly like in BM25.69    `70        The model is based on `jinaai/jina-embeddings-v2-small-en-tokens`71    """72 73    def __init__(74        self,75        model_name: str,76        cache_dir: str | None = None,77        threads: int | None = None,78        providers: Sequence[OnnxProvider] | None = None,79        k: float = 1.2,80        b: float = 0.75,81        avg_len: float = 150.0,82        cuda: bool | Device = Device.AUTO,83        device_ids: list[int] | None = None,84        lazy_load: bool = False,85        device_id: int | None = None,86        specific_model_path: str | None = None,87        **kwargs: Any,88    ):89        """90        Args:91            model_name (str): The name of the model to use.92            cache_dir (str, optional): The path to the cache directory.93                                       Can be set using the `FASTEMBED_CACHE_PATH` env variable.94                                       Defaults to `fastembed_cache` in the system's temp directory.95            threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.96            providers (Optional[Sequence[OnnxProvider]], optional): The providers to use for onnxruntime.97            k (float, optional): The k parameter in the BM25 formula. Defines the saturation of the term frequency.98                I.e. defines how fast the moment when additional terms stop to increase the score. Defaults to 1.2.99            b (float, optional): The b parameter in the BM25 formula. Defines the importance of the document length.100                Defaults to 0.75.101            avg_len (float, optional): The average length of the documents in the corpus. Defaults to 150.0.102            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`103                Defaults to Device.AUTO.104            device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in105                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive106                with `providers`. Defaults to None.107            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.108                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.109            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.110            specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else111 112        Raises:113            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.114        """115 116        super().__init__(model_name, cache_dir, threads, **kwargs)117        self.providers = providers118        self.lazy_load = lazy_load119        self.device_ids = device_ids120        self.cuda = cuda121        self.device_id = device_id122        self._extra_session_options = self._select_exposed_session_options(kwargs)123 124        self.k = k125        self.b = b126        self.avg_len = avg_len127 128        # Initialize class attributes129        self.tokenizer: Tokenizer | None = None130        self.invert_vocab: dict[int, str] = {}131        self.special_tokens: set[str] = set()132        self.special_tokens_ids: set[int] = set()133        self.stopwords: set[str] = set()134        self.vocab_resolver: VocabResolver | None = None135        self.encoder: Encoder | None = None136        self.output_dim: int | None = None137        self.sparse_vector_converter: SparseVectorConverter | None = None138 139        self.model_description = self._get_model_description(model_name)140        self.cache_dir = str(define_cache_dir(cache_dir))141        self._specific_model_path = specific_model_path142        self._model_dir = self.download_model(143            self.model_description,144            self.cache_dir,145            local_files_only=self._local_files_only,146            specific_model_path=self._specific_model_path,147        )148 149        if not self.lazy_load:150            self.load_onnx_model()151 152    def load_onnx_model(self) -> None:153        self._load_onnx_model(154            model_dir=self._model_dir,155            model_file=self.model_description.model_file,156            threads=self.threads,157            providers=self.providers,158            cuda=self.cuda,159            device_id=self.device_id,160            extra_session_options=self._extra_session_options,161        )162 163        assert self.tokenizer is not None164 165        for token, idx in self.tokenizer.get_vocab().items():  # type: ignore[union-attr]166            self.invert_vocab[idx] = token167        self.special_tokens = set(self.special_token_to_id.keys())168        self.special_tokens_ids = set(self.special_token_to_id.values())169        self.stopwords = set(self._load_stopwords(self._model_dir))170 171        stemmer = SnowballStemmer(get_language_by_model_name(self.model_name))172 173        self.vocab_resolver = VocabResolver(174            tokenizer=VocabTokenizer(self.tokenizer),175            stopwords=self.stopwords,176            stemmer=stemmer,177        )178        self.vocab_resolver.load_json_vocab(str(self._model_dir / MINICOIL_VOCAB_FILE))179 180        weights = np.load(str(self._model_dir / MINICOIL_MODEL_FILE), mmap_mode="r")181        self.encoder = Encoder(weights)182        self.output_dim = self.encoder.output_dim183 184        self.sparse_vector_converter = SparseVectorConverter(185            stopwords=self.stopwords,186            stemmer=stemmer,187            k=self.k,188            b=self.b,189            avg_len=self.avg_len,190        )191 192    def token_count(193        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any194    ) -> int:195        return self._token_count(texts, batch_size=batch_size, **kwargs)196 197    def embed(198        self,199        documents: str | Iterable[str],200        batch_size: int = 256,201        parallel: int | None = None,202        **kwargs: Any,203    ) -> Iterable[SparseEmbedding]:204        """205        Encode a list of documents into list of embeddings.206        We use mean pooling with attention so that the model can handle variable-length inputs.207 208        Args:209            documents: Iterator of documents or single document to embed210            batch_size: Batch size for encoding -- higher values will use more memory, but be faster211            parallel:212                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.213                If 0, use all available cores.214                If None, don't use data-parallel processing, use default onnxruntime threading instead.215 216        Returns:217            List of embeddings, one per document218        """219        yield from self._embed_documents(220            model_name=self.model_name,221            cache_dir=str(self.cache_dir),222            documents=documents,223            batch_size=batch_size,224            parallel=parallel,225            providers=self.providers,226            cuda=self.cuda,227            device_ids=self.device_ids,228            k=self.k,229            b=self.b,230            avg_len=self.avg_len,231            is_query=False,232            local_files_only=self._local_files_only,233            specific_model_path=self._specific_model_path,234            extra_session_options=self._extra_session_options,235            **kwargs,236        )237 238    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:239        """240        Encode a list of queries into list of embeddings.241        """242        yield from self._embed_documents(243            model_name=self.model_name,244            cache_dir=str(self.cache_dir),245            documents=query,246            providers=self.providers,247            cuda=self.cuda,248            device_ids=self.device_ids,249            k=self.k,250            b=self.b,251            avg_len=self.avg_len,252            is_query=True,253            local_files_only=self._local_files_only,254            specific_model_path=self._specific_model_path,255            **kwargs,256        )257 258    @classmethod259    def _load_stopwords(cls, model_dir: Path) -> list[str]:260        stopwords_path = model_dir / STOPWORDS_FILE261        if not stopwords_path.exists():262            return []263 264        with open(stopwords_path, "r") as f:265            return f.read().splitlines()266 267    @classmethod268    def _list_supported_models(cls) -> list[SparseModelDescription]:269        """Lists the supported models.270 271        Returns:272            list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information.273        """274        return supported_minicoil_models275 276    def _post_process_onnx_output(277        self, output: OnnxOutputContext, is_query: bool = False, **kwargs: Any278    ) -> Iterable[SparseEmbedding]:279        if output.input_ids is None:280            raise ValueError("input_ids must be provided for document post-processing")281 282        assert self.vocab_resolver is not None283        assert self.encoder is not None284        assert self.sparse_vector_converter is not None285 286        # Size: (batch_size, sequence_length, hidden_size)287        embeddings = output.model_output288        # Size: (batch_size, sequence_length)289        assert output.attention_mask is not None290        masks = output.attention_mask291 292        vocab_size = self.vocab_resolver.vocab_size()293        embedding_size = self.encoder.output_dim294 295        # For each document we only select those embeddings that are not masked out296 297        for i in range(embeddings.shape[0]):298            # Size: (sequence_length, hidden_size)299            token_embeddings = embeddings[i, masks[i] == 1]300 301            # Size: (sequence_length)302            token_ids: NDArray[np.int64] = output.input_ids[i, masks[i] == 1]303 304            word_ids_array, counts, oov, forms = self.vocab_resolver.resolve_tokens(token_ids)305 306            # Size: (1, words)307            word_ids_array_expanded: NDArray[np.int64] = np.expand_dims(word_ids_array, axis=0)308 309            # Size: (1, words, embedding_size)310            token_embeddings_array: NDArray[np.float32] = np.expand_dims(token_embeddings, axis=0)311 312            assert word_ids_array_expanded.shape[1] == token_embeddings_array.shape[1]313 314            # Size of word_ids_mapping: (unique_words, 2) - [vocab_id, batch_id]315            # Size of embeddings: (unique_words, embedding_size)316            ids_mapping, minicoil_embeddings = self.encoder.forward(317                word_ids_array_expanded, token_embeddings_array318            )319 320            # Size of counts: (unique_words)321            words_ids: list[int] = ids_mapping[:, 0].tolist()  # type: ignore[assignment]322 323            sentence_result: dict[str, WordEmbedding] = {}324 325            words = [self.vocab_resolver.lookup_word(word_id) for word_id in words_ids]326 327            for word, word_id, emb in zip(words, words_ids, minicoil_embeddings.tolist()):  # type: ignore[arg-type]328                if word_id == 0:329                    continue330 331                sentence_result[word] = WordEmbedding(332                    word=word,333                    forms=forms[word],334                    count=int(counts[word_id]),335                    word_id=int(word_id),336                    embedding=emb,  # type: ignore[arg-type]337                )338 339            for oov_word, count in oov.items():340                # {341                #     "word": oov_word,342                #     "forms": [oov_word],343                #     "count": int(count),344                #     "word_id": -1,345                #     "embedding": [1]346                # }347                sentence_result[oov_word] = WordEmbedding(348                    word=oov_word, forms=[oov_word], count=int(count), word_id=-1, embedding=[1]349                )350 351            if not is_query:352                yield self.sparse_vector_converter.embedding_to_vector(353                    sentence_result, vocab_size=vocab_size, embedding_size=embedding_size354                )355            else:356                yield self.sparse_vector_converter.embedding_to_vector_query(357                    sentence_result, vocab_size=vocab_size, embedding_size=embedding_size358                )359 360    @classmethod361    def _get_worker_class(cls) -> Type["MiniCoilTextEmbeddingWorker"]:362        return MiniCoilTextEmbeddingWorker363 364 365class MiniCoilTextEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):366    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> MiniCOIL:367        return MiniCOIL(368            model_name=model_name,369            cache_dir=cache_dir,370            threads=1,371            **kwargs,372        )373