fastembed-bio 0.1.0__py3-none-any.whl

fastembed/text/pooled_normalized_embedding.py

@@ -0,0 +1,164 @@
+from typing import Any, Iterable, Type
+
+
+from fastembed.common.types import NumpyArray
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.utils import normalize
+from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
+from fastembed.text.pooled_embedding import PooledEmbedding
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+
+supported_pooled_normalized_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="sentence-transformers/all-MiniLM-L6-v2",
+        dim=384,
+        description=(
+            "Text embeddings, Unimodal (text), English, 256 input tokens truncation, "
+            "Prefixes for queries/documents: not necessary, 2021 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=0.09,
+        sources=ModelSource(
+            url="https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz",
+            hf="qdrant/all-MiniLM-L6-v2-onnx",
+            _deprecated_tar_struct=True,
+        ),
+        model_file="model.onnx",
+    ),
+    DenseModelDescription(
+        model="jinaai/jina-embeddings-v2-base-en",
+        dim=768,
+        description=(
+            "Text embeddings, Unimodal (text), English, 8192 input tokens truncation, "
+            "Prefixes for queries/documents: not necessary, 2023 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=0.52,
+        sources=ModelSource(hf="xenova/jina-embeddings-v2-base-en"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="jinaai/jina-embeddings-v2-small-en",
+        dim=512,
+        description=(
+            "Text embeddings, Unimodal (text), English, 8192 input tokens truncation, "
+            "Prefixes for queries/documents: not necessary, 2023 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=0.12,
+        sources=ModelSource(hf="xenova/jina-embeddings-v2-small-en"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="jinaai/jina-embeddings-v2-base-de",
+        dim=768,
+        description=(
+            "Text embeddings, Unimodal (text), Multilingual (German, English), 8192 input tokens truncation, "
+            "Prefixes for queries/documents: not necessary, 2024 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=0.32,
+        sources=ModelSource(hf="jinaai/jina-embeddings-v2-base-de"),
+        model_file="onnx/model_fp16.onnx",
+    ),
+    DenseModelDescription(
+        model="jinaai/jina-embeddings-v2-base-code",
+        dim=768,
+        description=(
+            "Text embeddings, Unimodal (text), Multilingual (English, 30 programming languages), "
+            "8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=0.64,
+        sources=ModelSource(hf="jinaai/jina-embeddings-v2-base-code"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="jinaai/jina-embeddings-v2-base-zh",
+        dim=768,
+        description=(
+            "Text embeddings, Unimodal (text), supports mixed Chinese-English input text, "
+            "8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=0.64,
+        sources=ModelSource(hf="jinaai/jina-embeddings-v2-base-zh"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="jinaai/jina-embeddings-v2-base-es",
+        dim=768,
+        description=(
+            "Text embeddings, Unimodal (text), supports mixed Spanish-English input text, "
+            "8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=0.64,
+        sources=ModelSource(hf="jinaai/jina-embeddings-v2-base-es"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="thenlper/gte-base",
+        dim=768,
+        description=(
+            "General text embeddings, Unimodal (text), supports English only input text, "
+            "512 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year."
+        ),
+        license="mit",
+        size_in_GB=0.44,
+        sources=ModelSource(hf="thenlper/gte-base"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="thenlper/gte-large",
+        dim=1024,
+        description=(
+            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
+            "Prefixes for queries/documents: not necessary, 2023 year."
+        ),
+        license="mit",
+        size_in_GB=1.20,
+        sources=ModelSource(hf="qdrant/gte-large-onnx"),
+        model_file="model.onnx",
+    ),
+]
+
+
+class PooledNormalizedEmbedding(PooledEmbedding):
+    @classmethod
+    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
+        return PooledNormalizedEmbeddingWorker
+
+    @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_pooled_normalized_models
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
+        if output.attention_mask is None:
+            raise ValueError("attention_mask must be provided for document post-processing")
+
+        embeddings = output.model_output
+        attn_mask = output.attention_mask
+        return normalize(self.mean_pooling(embeddings, attn_mask))
+
+
+class PooledNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ) -> OnnxTextEmbedding:
+        return PooledNormalizedEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/text/text_embedding.py

@@ -0,0 +1,228 @@
+import warnings
+from typing import Any, Iterable, Sequence, Type
+from dataclasses import asdict
+
+from fastembed.common.types import NumpyArray, OnnxProvider, Device
+from fastembed.text.clip_embedding import CLIPOnnxEmbedding
+from fastembed.text.custom_text_embedding import CustomTextEmbedding
+from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
+from fastembed.text.pooled_embedding import PooledEmbedding
+from fastembed.text.multitask_embedding import JinaEmbeddingV3
+from fastembed.text.onnx_embedding import OnnxTextEmbedding
+from fastembed.text.text_embedding_base import TextEmbeddingBase
+from fastembed.common.model_description import DenseModelDescription, ModelSource, PoolingType
+
+
+class TextEmbedding(TextEmbeddingBase):
+    EMBEDDINGS_REGISTRY: list[Type[TextEmbeddingBase]] = [
+        OnnxTextEmbedding,
+        CLIPOnnxEmbedding,
+        PooledNormalizedEmbedding,
+        PooledEmbedding,
+        JinaEmbeddingV3,
+        CustomTextEmbedding,
+    ]
+
+    @classmethod
+    def list_supported_models(cls) -> list[dict[str, Any]]:
+        """Lists the supported models.
+
+        Returns:
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
+        """
+        return [asdict(model) for model in cls._list_supported_models()]
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        result: list[DenseModelDescription] = []
+        for embedding in cls.EMBEDDINGS_REGISTRY:
+            result.extend(embedding._list_supported_models())
+        return result
+
+    @classmethod
+    def add_custom_model(
+        cls,
+        model: str,
+        pooling: PoolingType,
+        normalization: bool,
+        sources: ModelSource,
+        dim: int,
+        model_file: str = "onnx/model.onnx",
+        description: str = "",
+        license: str = "",
+        size_in_gb: float = 0.0,
+        additional_files: list[str] | None = None,
+    ) -> None:
+        registered_models = cls._list_supported_models()
+        for registered_model in registered_models:
+            if model.lower() == registered_model.model.lower():
+                raise ValueError(
+                    f"Model {model} is already registered in TextEmbedding, if you still want to add this model, "
+                    f"please use another model name"
+                )
+
+        CustomTextEmbedding.add_model(
+            DenseModelDescription(
+                model=model,
+                sources=sources,
+                dim=dim,
+                model_file=model_file,
+                description=description,
+                license=license,
+                size_in_GB=size_in_gb,
+                additional_files=additional_files or [],
+            ),
+            pooling=pooling,
+            normalization=normalization,
+        )
+
+    def __init__(
+        self,
+        model_name: str = "BAAI/bge-small-en-v1.5",
+        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,
+        **kwargs: Any,
+    ):
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+        if model_name.lower() == "nomic-ai/nomic-embed-text-v1.5-Q".lower():
+            warnings.warn(
+                "The model 'nomic-ai/nomic-embed-text-v1.5-Q' has been updated on HuggingFace. Please review "
+                "the latest documentation on HF and release notes to ensure compatibility with your workflow. ",
+                UserWarning,
+                stacklevel=2,
+            )
+        if model_name.lower() in {
+            "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2".lower(),
+            "thenlper/gte-large".lower(),
+            "intfloat/multilingual-e5-large".lower(),
+            "sentence-transformers/paraphrase-multilingual-mpnet-base-v2".lower(),
+        }:
+            warnings.warn(
+                f"The model {model_name} now uses mean pooling instead of CLS embedding. "
+                f"In order to preserve the previous behaviour, consider either pinning fastembed version to 0.5.1 or "
+                "using `add_custom_model` functionality.",
+                UserWarning,
+                stacklevel=2,
+            )
+        for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
+            supported_models = EMBEDDING_MODEL_TYPE._list_supported_models()
+            if any(model_name.lower() == model.model.lower() for model in supported_models):
+                self.model = EMBEDDING_MODEL_TYPE(
+                    model_name=model_name,
+                    cache_dir=cache_dir,
+                    threads=threads,
+                    providers=providers,
+                    cuda=cuda,
+                    device_ids=device_ids,
+                    lazy_load=lazy_load,
+                    **kwargs,
+                )
+                return
+
+        raise ValueError(
+            f"Model {model_name} is not supported in TextEmbedding. "
+            "Please check the supported models using `TextEmbedding.list_supported_models()`"
+        )
+
+    @property
+    def embedding_size(self) -> int:
+        """Get the embedding size of the current model"""
+        if self._embedding_size is None:
+            self._embedding_size = self.get_embedding_size(self.model_name)
+        return self._embedding_size
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Get the embedding size of the passed model
+
+        Args:
+            model_name (str): The name of the model to get embedding size for.
+
+        Returns:
+            int: The size of the embedding.
+
+        Raises:
+            ValueError: If the model name is not found in the supported models.
+        """
+        descriptions = cls._list_supported_models()
+        embedding_size: int | None = None
+        for description in descriptions:
+            if description.model.lower() == model_name.lower():
+                embedding_size = description.dim
+                break
+        if embedding_size is None:
+            model_names = [description.model for description in descriptions]
+            raise ValueError(
+                f"Embedding size for model {model_name} was None. "
+                f"Available model names: {model_names}"
+            )
+        return embedding_size
+
+    def embed(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        """
+        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.model.embed(documents, batch_size, parallel, **kwargs)
+
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+        """
+        Embeds queries
+
+        Args:
+            query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
+
+        Returns:
+            Iterable[NumpyArray]: The embeddings.
+        """
+        # This is model-specific, so that different models can have specialized implementations
+        yield from self.model.query_embed(query, **kwargs)
+
+    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+        """
+        Embeds a list of text passages into a list of embeddings.
+
+        Args:
+            texts (Iterable[str]): The list of texts to embed.
+            **kwargs: Additional keyword argument to pass to the embed method.
+
+        Yields:
+            Iterable[SparseEmbedding]: The sparse embeddings.
+        """
+        # This is model-specific, so that different models can have specialized implementations
+        yield from self.model.passage_embed(texts, **kwargs)
+
+    def token_count(
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
+    ) -> int:
+        """Returns the number of tokens in the texts.
+
+        Args:
+            texts (str | Iterable[str]): The list of texts to embed.
+            batch_size (int): Batch size for encoding
+
+        Returns:
+            int: Sum of number of tokens in the texts.
+        """
+        return self.model.token_count(texts, batch_size=batch_size, **kwargs)

fastembed/text/text_embedding_base.py

@@ -0,0 +1,75 @@
+from typing import Iterable, Any
+
+from fastembed.common.model_description import DenseModelDescription
+from fastembed.common.types import NumpyArray
+from fastembed.common.model_management import ModelManagement
+
+
+class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        **kwargs: Any,
+    ):
+        self.model_name = model_name
+        self.cache_dir = cache_dir
+        self.threads = threads
+        self._local_files_only = kwargs.pop("local_files_only", False)
+        self._embedding_size: int | None = None
+
+    def embed(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        raise NotImplementedError()
+
+    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+        """
+        Embeds a list of text passages into a list of embeddings.
+
+        Args:
+            texts (Iterable[str]): The list of texts to embed.
+            **kwargs: Additional keyword argument to pass to the embed method.
+
+        Yields:
+            Iterable[NumpyArray]: The embeddings.
+        """
+
+        # This is model-specific, so that different models can have specialized implementations
+        yield from self.embed(texts, **kwargs)
+
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+        """
+        Embeds queries
+
+        Args:
+            query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
+
+        Returns:
+            Iterable[NumpyArray]: The embeddings.
+        """
+
+        # This is model-specific, so that different models can have specialized implementations
+        if isinstance(query, str):
+            yield from self.embed([query], **kwargs)
+        else:
+            yield from self.embed(query, **kwargs)
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Returns embedding size of the passed model."""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @property
+    def embedding_size(self) -> int:
+        """Returns embedding size for the current model"""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def token_count(self, texts: str | Iterable[str], **kwargs: Any) -> int:
+        """Returns the number of tokens in the texts."""
+        raise NotImplementedError("Subclasses must implement this method")