fastembed-bio 0.1.0__py3-none-any.whl

fastembed/sparse/utils/tokenizer.py

@@ -0,0 +1,120 @@
+# This code is a modified copy of the `NLTKWordTokenizer` class from `NLTK` library.
+
+import re
+
+
+class SimpleTokenizer:
+    @staticmethod
+    def tokenize(text: str) -> list[str]:
+        text = re.sub(r"[^\w]", " ", text.lower())
+        text = re.sub(r"\s+", " ", text)
+
+        return text.strip().split()
+
+
+class WordTokenizer:
+    """The tokenizer is "destructive" such that the regexes applied will munge the
+    input string to a state beyond re-construction.
+    """
+
+    # Starting quotes.
+    STARTING_QUOTES = [
+        (re.compile("([«“‘„]|[`]+)", re.U), r" \1 "),
+        (re.compile(r"^\""), r"``"),
+        (re.compile(r"(``)"), r" \1 "),
+        (re.compile(r"([ \(\[{<])(\"|\'{2})"), r"\1 `` "),
+        (re.compile(r"(?i)(\')(?!re|ve|ll|m|t|s|d|n)(\w)\b", re.U), r"\1 \2"),
+    ]
+
+    # Ending quotes.
+    ENDING_QUOTES = [
+        (re.compile("([»”’])", re.U), r" \1 "),
+        (re.compile(r"''"), " '' "),
+        (re.compile(r'"'), " '' "),
+        (re.compile(r"([^' ])('[sS]|'[mM]|'[dD]|') "), r"\1 \2 "),
+        (re.compile(r"([^' ])('ll|'LL|'re|'RE|'ve|'VE|n't|N'T) "), r"\1 \2 "),
+    ]
+
+    # Punctuation.
+    PUNCTUATION = [
+        (re.compile(r'([^\.])(\.)([\]\)}>"\'' "»”’ " r"]*)\s*$", re.U), r"\1 \2 \3 "),
+        (re.compile(r"([:,])([^\d])"), r" \1 \2"),
+        (re.compile(r"([:,])$"), r" \1 "),
+        (
+            re.compile(r"\.{2,}", re.U),
+            r" \g<0> ",
+        ),
+        (re.compile(r"[;@#$%&]"), r" \g<0> "),
+        (
+            re.compile(r'([^\.])(\.)([\]\)}>"\']*)\s*$'),
+            r"\1 \2\3 ",
+        ),  # Handles the final period.
+        (re.compile(r"[?!]"), r" \g<0> "),
+        (re.compile(r"([^'])' "), r"\1 ' "),
+        (
+            re.compile(r"[*]", re.U),
+            r" \g<0> ",
+        ),
+    ]
+
+    # Pads parentheses
+    PARENS_BRACKETS = (re.compile(r"[\]\[\(\)\{\}\<\>]"), r" \g<0> ")
+    DOUBLE_DASHES = (re.compile(r"--"), r" -- ")
+
+    # List of contractions adapted from Robert MacIntyre's tokenizer.
+    CONTRACTIONS2 = [
+        re.compile(pattern)
+        for pattern in (
+            r"(?i)\b(can)(?#X)(not)\b",
+            r"(?i)\b(d)(?#X)('ye)\b",
+            r"(?i)\b(gim)(?#X)(me)\b",
+            r"(?i)\b(gon)(?#X)(na)\b",
+            r"(?i)\b(got)(?#X)(ta)\b",
+            r"(?i)\b(lem)(?#X)(me)\b",
+            r"(?i)\b(more)(?#X)('n)\b",
+            r"(?i)\b(wan)(?#X)(na)(?=\s)",
+        )
+    ]
+    CONTRACTIONS3 = [
+        re.compile(pattern) for pattern in (r"(?i) ('t)(?#X)(is)\b", r"(?i) ('t)(?#X)(was)\b")
+    ]
+
+    @classmethod
+    def tokenize(cls, text: str) -> list[str]:
+        """Return a tokenized copy of `text`.
+
+        >>> s = '''Good muffins cost $3.88 (roughly 3,36 euros)\nin New York.'''
+        >>> WordTokenizer().tokenize(s)
+        ['Good', 'muffins', 'cost', '$', '3.88', '(', 'roughly', '3,36', 'euros', ')', 'in', 'New', 'York', '.']
+
+        Args:
+            text: The text to be tokenized.
+
+        Returns:
+            A list of tokens.
+        """
+        for regexp, substitution in cls.STARTING_QUOTES:
+            text = regexp.sub(substitution, text)
+
+        for regexp, substitution in cls.PUNCTUATION:
+            text = regexp.sub(substitution, text)
+
+        # Handles parentheses.
+        regexp, substitution = cls.PARENS_BRACKETS
+        text = regexp.sub(substitution, text)
+
+        # Handles double dash.
+        regexp, substitution = cls.DOUBLE_DASHES
+        text = regexp.sub(substitution, text)
+
+        # add extra space to make things easier
+        text = " " + text + " "
+
+        for regexp, substitution in cls.ENDING_QUOTES:
+            text = regexp.sub(substitution, text)
+
+        for regexp in cls.CONTRACTIONS2:
+            text = regexp.sub(r" \1 \2 ", text)
+        for regexp in cls.CONTRACTIONS3:
+            text = regexp.sub(r" \1 \2 ", text)
+        return text.split()

fastembed/sparse/utils/vocab_resolver.py

@@ -0,0 +1,202 @@
+from collections import defaultdict
+from typing import Iterable
+
+from py_rust_stemmers import SnowballStemmer
+import numpy as np
+from tokenizers import Tokenizer
+from numpy.typing import NDArray
+
+from fastembed.common.types import NumpyArray
+
+
+class VocabTokenizerBase:
+    def tokenize(self, sentence: str) -> NumpyArray:
+        raise NotImplementedError()
+
+    def convert_ids_to_tokens(self, token_ids: NumpyArray) -> list[str]:
+        raise NotImplementedError()
+
+
+class VocabTokenizer(VocabTokenizerBase):
+    def __init__(self, tokenizer: Tokenizer):
+        self.tokenizer = tokenizer
+
+    def tokenize(self, sentence: str) -> NumpyArray:
+        return np.array(self.tokenizer.encode(sentence).ids)
+
+    def convert_ids_to_tokens(self, token_ids: NumpyArray) -> list[str]:
+        return [self.tokenizer.id_to_token(token_id) for token_id in token_ids]
+
+
+class VocabResolver:
+    def __init__(self, tokenizer: VocabTokenizerBase, stopwords: set[str], stemmer: SnowballStemmer):
+        # Word to id mapping
+        self.vocab: dict[str, int] = {}
+        # Id to word mapping
+        self.words: list[str] = []
+        # Lemma to word mapping
+        self.stem_mapping: dict[str, str] = {}
+        self.tokenizer: VocabTokenizerBase = tokenizer
+        self.stemmer = stemmer
+        self.stopwords: set[str] = stopwords
+
+    def tokenize(self, sentence: str) -> NumpyArray:
+        return self.tokenizer.tokenize(sentence)
+
+    def lookup_word(self, word_id: int) -> str:
+        if word_id == 0:
+            return "UNK"
+        return self.words[word_id - 1]
+
+    def convert_ids_to_tokens(self, token_ids: NumpyArray) -> list[str]:
+        return self.tokenizer.convert_ids_to_tokens(token_ids)
+
+    def vocab_size(self) -> int:
+        # We need +1 for UNK token
+        return len(self.vocab) + 1
+
+    def save_vocab(self, path: str) -> None:
+        with open(path, "w") as f:
+            for word in self.words:
+                f.write(word + "\n")
+
+    def save_json_vocab(self, path: str) -> None:
+        import json
+
+        with open(path, "w") as f:
+            json.dump({"vocab": self.words, "stem_mapping": self.stem_mapping}, f, indent=2)
+
+    def load_json_vocab(self, path: str) -> None:
+        import json
+
+        with open(path, "r") as f:
+            data = json.load(f)
+            self.words = data["vocab"]
+            self.vocab = {word: idx + 1 for idx, word in enumerate(self.words)}
+            self.stem_mapping = data["stem_mapping"]
+
+    def add_word(self, word: str) -> None:
+        if word not in self.vocab:
+            self.vocab[word] = len(self.vocab) + 1
+            self.words.append(word)
+            stem = self.stemmer.stem_word(word)
+            if stem not in self.stem_mapping:
+                self.stem_mapping[stem] = word
+            else:
+                existing_word = self.stem_mapping[stem]
+                if len(existing_word) > len(word):
+                    # Prefer shorter words for the same stem
+                    # Example: "swim" is preferred over "swimming"
+                    self.stem_mapping[stem] = word
+
+    def load_vocab(self, path: str) -> None:
+        with open(path, "r") as f:
+            for line in f:
+                self.add_word(line.strip())
+
+    @classmethod
+    def _reconstruct_bpe(
+        cls, bpe_tokens: Iterable[tuple[int, str]]
+    ) -> list[tuple[str, list[int]]]:
+        result: list[tuple[str, list[int]]] = []
+        acc: str = ""
+        acc_idx: list[int] = []
+
+        continuing_subword_prefix = "##"
+        continuing_subword_prefix_len = len(continuing_subword_prefix)
+
+        for idx, token in bpe_tokens:
+            if token.startswith(continuing_subword_prefix):
+                acc += token[continuing_subword_prefix_len:]
+                acc_idx.append(idx)
+            else:
+                if acc:
+                    result.append((acc, acc_idx))
+                    acc_idx = []
+                acc = token
+                acc_idx.append(idx)
+
+        if acc:
+            result.append((acc, acc_idx))
+        return result
+
+    def resolve_tokens(
+        self, token_ids: NDArray[np.int64]
+    ) -> tuple[NDArray[np.int64], dict[int, int], dict[str, int], dict[str, list[str]]]:
+        """
+        Mark known tokens (including composed tokens) with vocab ids.
+
+        Args:
+            token_ids: (seq_len) - list of ids of tokens
+                Example:
+                    [
+                        101,  3897, 19332, 12718, 23348,
+                        1010,  1996,  7151,  2296, 4845,
+                        2359,  2005,  4234,  1010,  4332,
+                        2871,  3191,  2062, 102
+                    ]
+
+            returns:
+                - token_ids with vocab ids
+                    [
+                        0,  151, 151, 0, 0,
+                        912,  0,  0,  0, 332,
+                        332,  332,  0,  7121,  191,
+                        0,  0,  332, 0
+                    ]
+                - counts of each token
+                    {
+                        151: 1,
+                        332: 3,
+                        7121: 1,
+                        191: 1,
+                        912: 1
+                    }
+                - oov counts of each token
+                    {
+                        "the": 1,
+                        "a": 1,
+                        "[CLS]": 1,
+                        "[SEP]": 1,
+                        ...
+                    }
+                - forms of each token
+                    {
+                        "hello": ["hello"],
+                        "world": ["worlds", "world", "worlding"],
+                    }
+
+        """
+        tokens = self.convert_ids_to_tokens(token_ids)
+        tokens_mapping = self._reconstruct_bpe(enumerate(tokens))
+
+        counts: dict[int, int] = defaultdict(int)
+        oov_count: dict[str, int] = defaultdict(int)
+
+        forms: dict[str, list[str]] = defaultdict(list)
+
+        for token, mapped_token_ids in tokens_mapping:
+            vocab_id = 0
+            if token in self.stopwords:
+                vocab_id = 0
+            elif token in self.vocab:
+                vocab_id = self.vocab[token]
+                forms[token].append(token)
+            elif token in self.stem_mapping:
+                vocab_id = self.vocab[self.stem_mapping[token]]
+                forms[self.stem_mapping[token]].append(token)
+            else:
+                stem = self.stemmer.stem_word(token)
+                if stem in self.stem_mapping:
+                    vocab_id = self.vocab[self.stem_mapping[stem]]
+                    forms[self.stem_mapping[stem]].append(token)
+
+            for token_id in mapped_token_ids:
+                token_ids[token_id] = vocab_id
+
+            if vocab_id == 0:
+                oov_count[token] += 1
+            else:
+                counts[vocab_id] += 1
+        return token_ids, counts, oov_count, forms
+

fastembed/text/__init__.py

@@ -0,0 +1,3 @@
+from fastembed.text.text_embedding import TextEmbedding
+
+__all__ = ["TextEmbedding"]

fastembed/text/clip_embedding.py

@@ -0,0 +1,56 @@
+from typing import Any, Iterable, Type
+
+from fastembed.common.types import NumpyArray
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+
+supported_clip_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="Qdrant/clip-ViT-B-32-text",
+        dim=512,
+        description=(
+            "Text embeddings, Multimodal (text&image), English, 77 input tokens truncation, "
+            "Prefixes for queries/documents: not necessary, 2021 year"
+        ),
+        license="mit",
+        size_in_GB=0.25,
+        sources=ModelSource(hf="Qdrant/clip-ViT-B-32-text"),
+        model_file="model.onnx",
+    ),
+]
+
+
+class CLIPOnnxEmbedding(OnnxTextEmbedding):
+    @classmethod
+    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
+        return CLIPEmbeddingWorker
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        """Lists the supported models.
+
+        Returns:
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+        """
+        return supported_clip_models
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
+        return output.model_output
+
+
+class CLIPEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ) -> OnnxTextEmbedding:
+        return CLIPOnnxEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/text/custom_text_embedding.py

@@ -0,0 +1,97 @@
+from typing import Sequence, Any, Iterable
+from dataclasses import dataclass
+
+import numpy as np
+from numpy.typing import NDArray
+
+from fastembed.common import OnnxProvider
+from fastembed.common.model_description import (
+    PoolingType,
+    DenseModelDescription,
+)
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import NumpyArray, Device
+from fastembed.common.utils import normalize, mean_pooling
+from fastembed.text.onnx_embedding import OnnxTextEmbedding
+
+
+@dataclass(frozen=True)
+class PostprocessingConfig:
+    pooling: PoolingType
+    normalization: bool
+
+
+class CustomTextEmbedding(OnnxTextEmbedding):
+    SUPPORTED_MODELS: list[DenseModelDescription] = []
+    POSTPROCESSING_MAPPING: dict[str, PostprocessingConfig] = {}
+
+    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,
+    ):
+        super().__init__(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=threads,
+            providers=providers,
+            cuda=cuda,
+            device_ids=device_ids,
+            lazy_load=lazy_load,
+            device_id=device_id,
+            specific_model_path=specific_model_path,
+            **kwargs,
+        )
+        self._pooling = self.POSTPROCESSING_MAPPING[model_name].pooling
+        self._normalization = self.POSTPROCESSING_MAPPING[model_name].normalization
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        return cls.SUPPORTED_MODELS
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
+        return self._normalize(self._pool(output.model_output, output.attention_mask))
+
+    def _pool(
+        self, embeddings: NumpyArray, attention_mask: NDArray[np.int64] | None = None
+    ) -> NumpyArray:
+        if self._pooling == PoolingType.CLS:
+            return embeddings[:, 0]
+
+        if self._pooling == PoolingType.MEAN:
+            if attention_mask is None:
+                raise ValueError("attention_mask must be provided for mean pooling")
+            return mean_pooling(embeddings, attention_mask)
+
+        if self._pooling == PoolingType.DISABLED:
+            return embeddings
+
+        raise ValueError(
+            f"Unsupported pooling type {self._pooling}. "
+            f"Supported types are: {PoolingType.CLS}, {PoolingType.MEAN}, {PoolingType.DISABLED}."
+        )
+
+    def _normalize(self, embeddings: NumpyArray) -> NumpyArray:
+        return normalize(embeddings) if self._normalization else embeddings
+
+    @classmethod
+    def add_model(
+        cls,
+        model_description: DenseModelDescription,
+        pooling: PoolingType,
+        normalization: bool,
+    ) -> None:
+        cls.SUPPORTED_MODELS.append(model_description)
+        cls.POSTPROCESSING_MAPPING[model_description.model] = PostprocessingConfig(
+            pooling=pooling, normalization=normalization
+        )

fastembed/text/multitask_embedding.py

@@ -0,0 +1,109 @@
+from enum import Enum
+from typing import Any, Type, Iterable
+
+import numpy as np
+
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import NumpyArray
+from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
+from fastembed.text.onnx_embedding import OnnxTextEmbeddingWorker
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+
+supported_multitask_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="jinaai/jina-embeddings-v3",
+        dim=1024,
+        tasks={
+            "retrieval.query": 0,
+            "retrieval.passage": 1,
+            "separation": 2,
+            "classification": 3,
+            "text-matching": 4,
+        },
+        description=(
+            "Multi-task unimodal (text) embedding model, multi-lingual (~100), "
+            "1024 tokens truncation, and 8192 sequence length. Prefixes for queries/documents: not necessary, 2024 year."
+        ),
+        license="cc-by-nc-4.0",
+        size_in_GB=2.29,
+        sources=ModelSource(hf="jinaai/jina-embeddings-v3"),
+        model_file="onnx/model.onnx",
+        additional_files=["onnx/model.onnx_data"],
+    ),
+]
+
+
+class Task(int, Enum):
+    RETRIEVAL_QUERY = 0
+    RETRIEVAL_PASSAGE = 1
+    SEPARATION = 2
+    CLASSIFICATION = 3
+    TEXT_MATCHING = 4
+
+
+class JinaEmbeddingV3(PooledNormalizedEmbedding):
+    PASSAGE_TASK = Task.RETRIEVAL_PASSAGE
+    QUERY_TASK = Task.RETRIEVAL_QUERY
+
+    def __init__(self, *args: Any, task_id: int | None = None, **kwargs: Any):
+        super().__init__(*args, **kwargs)
+        self.default_task_id: Task | int = task_id if task_id is not None else self.PASSAGE_TASK
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
+        return JinaEmbeddingV3Worker
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        return supported_multitask_models
+
+    def _preprocess_onnx_input(
+        self,
+        onnx_input: dict[str, NumpyArray],
+        task_id: int | Task | None = None,
+        **kwargs: Any,
+    ) -> dict[str, NumpyArray]:
+        if task_id is None:
+            raise ValueError(f"task_id must be provided for JinaEmbeddingV3, got <{task_id}>")
+        onnx_input["task_id"] = np.array(task_id, dtype=np.int64)
+        return onnx_input
+
+    def embed(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        task_id: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        task_id = (
+            task_id if task_id is not None else self.default_task_id
+        )  # required for multiprocessing
+        yield from super().embed(documents, batch_size, parallel, task_id=task_id, **kwargs)
+
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+        yield from super().embed(query, task_id=self.QUERY_TASK, **kwargs)
+
+    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+        yield from super().embed(texts, task_id=self.PASSAGE_TASK, **kwargs)
+
+
+class JinaEmbeddingV3Worker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ) -> JinaEmbeddingV3:
+        return JinaEmbeddingV3(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )
+
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, OnnxOutputContext]]:
+        self.model: JinaEmbeddingV3  # mypy complaints `self.model` does not have `default_task_id`
+        for idx, batch in items:
+            onnx_output = self.model.onnx_embed(batch, task_id=self.model.default_task_id)
+            yield idx, onnx_output