fastembed-bio 0.1.0__py3-none-any.whl

fastembed/late_interaction/token_embeddings.py

@@ -0,0 +1,83 @@
+from dataclasses import asdict
+from typing import Iterable, Any, Type
+
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import NumpyArray
+from fastembed.late_interaction.late_interaction_embedding_base import (
+    LateInteractionTextEmbeddingBase,
+)
+from fastembed.text.onnx_embedding import OnnxTextEmbedding
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
+
+
+supported_token_embeddings_models = [
+    DenseModelDescription(
+        model="jinaai/jina-embeddings-v2-small-en-tokens",
+        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",
+    ),
+]
+
+
+class TokenEmbeddingsModel(OnnxTextEmbedding, LateInteractionTextEmbeddingBase):
+    @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_token_embeddings_models
+
+    @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 _get_worker_class(cls) -> Type[TextEmbeddingWorker[NumpyArray]]:
+        return TokensEmbeddingWorker
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
+        # Size: (batch_size, sequence_length, hidden_size)
+        embeddings = output.model_output
+        # Size: (batch_size, sequence_length)
+        assert output.attention_mask is not None
+        masks = output.attention_mask
+
+        # For each document we only select those embeddings that are not masked out
+        for i in range(embeddings.shape[0]):
+            yield embeddings[i, masks[i] == 1]
+
+    def embed(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        yield from super().embed(documents, batch_size=batch_size, parallel=parallel, **kwargs)
+
+
+class TokensEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
+    def init_embedding(
+        self, model_name: str, cache_dir: str, **kwargs: Any
+    ) -> TokenEmbeddingsModel:
+        return TokenEmbeddingsModel(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/late_interaction_multimodal/__init__.py

@@ -0,0 +1,5 @@
+from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding import (
+    LateInteractionMultimodalEmbedding,
+)
+
+__all__ = ["LateInteractionMultimodalEmbedding"]

fastembed/late_interaction_multimodal/colmodernvbert.py

@@ -0,0 +1,532 @@
+import contextlib
+from typing import Any, Iterable, Type, Optional, Sequence
+import json
+
+import numpy as np
+from tokenizers import Encoding
+from PIL import Image
+
+from fastembed.common import ImageInput
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import NumpyArray, OnnxProvider
+from fastembed.common.utils import define_cache_dir, iter_batch
+from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import (
+    LateInteractionMultimodalEmbeddingBase,
+)
+from fastembed.late_interaction_multimodal.onnx_multimodal_model import (
+    OnnxMultimodalModel,
+    TextEmbeddingWorker,
+    ImageEmbeddingWorker,
+)
+
+supported_colmodernvbert_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="Qdrant/colmodernvbert",
+        dim=128,
+        description="The late-interaction version of ModernVBERT, CPU friendly, English, 2025.",
+        license="mit",
+        size_in_GB=1.0,
+        sources=ModelSource(hf="Qdrant/colmodernvbert"),
+        additional_files=["processor_config.json"],
+        model_file="model.onnx",
+    ),
+]
+
+
+class ColModernVBERT(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyArray]):
+    """
+    The ModernVBERT/colmodernvbert model implementation. This model uses
+    bidirectional attention, which proves to work better for retrieval.
+
+    See: https://huggingface.co/ModernVBERT/colmodernvbert
+    """
+
+    VISUAL_PROMPT_PREFIX = (
+        "<|begin_of_text|>User:<image>Describe the image.<end_of_utterance>\nAssistant:"
+    )
+    QUERY_AUGMENTATION_TOKEN = "<end_of_utterance>"
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[list[int]] = None,
+        lazy_load: bool = False,
+        device_id: Optional[int] = None,
+        specific_model_path: Optional[str] = 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 (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to False.
+            device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda=True`, 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.
+
+        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: Optional[int] = 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,
+        )
+        self.mask_token_id = None
+        self.pad_token_id = None
+        self.image_seq_len: Optional[int] = None
+        self.max_image_size: Optional[int] = None
+        self.image_size: Optional[int] = None
+
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    @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_colmodernvbert_models
+
+    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,
+        )
+
+        # Load image processing configuration
+        processor_config_path = self._model_dir / "processor_config.json"
+        with open(processor_config_path) as f:
+            processor_config = json.load(f)
+            self.image_seq_len = processor_config.get("image_seq_len", 64)
+
+        preprocessor_config_path = self._model_dir / "preprocessor_config.json"
+        with open(preprocessor_config_path) as f:
+            preprocessor_config = json.load(f)
+            self.max_image_size = preprocessor_config.get("max_image_size", {}).get(
+                "longest_edge", 512
+            )
+
+        # Load model configuration
+        config_path = self._model_dir / "config.json"
+        with open(config_path) as f:
+            model_config = json.load(f)
+            vision_config = model_config.get("vision_config", {})
+            self.image_size = vision_config.get("image_size", 512)
+
+    def _preprocess_onnx_text_input(
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        """
+        Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+
+        Returns:
+            Iterable[NumpyArray]: Post-processed output as NumPy arrays.
+        """
+        batch_size, seq_length = onnx_input["input_ids"].shape
+        empty_image_placeholder: NumpyArray = np.zeros(
+            (batch_size, seq_length, 3, self.image_size, self.image_size),
+            dtype=np.float32,  # type: ignore[type-var,arg-type,assignment]
+        )
+        onnx_input["pixel_values"] = empty_image_placeholder
+        return onnx_input
+
+    def _post_process_onnx_text_output(
+        self,
+        output: OnnxOutputContext,
+    ) -> Iterable[NumpyArray]:
+        """
+        Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+
+        Returns:
+            Iterable[NumpyArray]: Post-processed output as NumPy arrays.
+        """
+        return output.model_output
+
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
+        # Add query augmentation tokens (matching process_queries logic from colpali-engine)
+        augmented_queries = [doc + self.QUERY_AUGMENTATION_TOKEN * 10 for doc in documents]
+        encoded = self.tokenizer.encode_batch(augmented_queries)  # type: ignore[union-attr]
+        return encoded
+
+    def token_count(
+        self,
+        texts: str | Iterable[str],
+        batch_size: int = 1024,
+        include_extension: bool = False,
+        **kwargs: Any,
+    ) -> int:
+        if not hasattr(self, "model") or self.model is None:
+            self.load_onnx_model()  # loads the tokenizer as well
+        token_num = 0
+        texts = [texts] if isinstance(texts, str) else texts
+        assert self.tokenizer is not None
+        tokenize_func = self.tokenize if include_extension else self.tokenizer.encode_batch
+        for batch in iter_batch(texts, batch_size):
+            token_num += sum([sum(encoding.attention_mask) for encoding in tokenize_func(batch)])
+        return token_num
+
+    def onnx_embed_image(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutputContext:
+        with contextlib.ExitStack() as stack:
+            image_files = [
+                stack.enter_context(Image.open(image))
+                if not isinstance(image, Image.Image)
+                else image
+                for image in images
+            ]
+            assert self.processor is not None, "Processor is not initialized"
+            processed = self.processor(image_files)
+            encoded, attention_mask, metadata = self._process_nested_patches(processed)  # type: ignore[arg-type]
+
+        onnx_input = {"pixel_values": encoded, "attention_mask": attention_mask}
+        onnx_input = self._preprocess_onnx_image_input(onnx_input, **kwargs)
+        model_output = self.model.run(None, onnx_input)  # type: ignore[union-attr]
+
+        return OnnxOutputContext(
+            model_output=model_output[0],
+            attention_mask=attention_mask,  # type: ignore[arg-type]
+            metadata=metadata,
+        )
+
+    @staticmethod
+    def _process_nested_patches(
+        processed: list[list[NumpyArray]],
+    ) -> tuple[NumpyArray, NumpyArray, dict[str, Any]]:
+        """
+        Process nested image patches (from ImageSplitter).
+
+        Args:
+            processed: List of patch lists, one per image [[img1_patches], [img2_patches], ...]
+
+        Returns:
+            tuple: (encoded array, attention_mask, metadata)
+                - encoded: (batch_size, max_patches, C, H, W)
+                - attention_mask: (batch_size, max_patches) with 1 for real patches, 0 for padding
+                - metadata: Dict with 'patch_counts' key
+        """
+        patch_counts = [len(patches) for patches in processed]
+        max_patches = max(patch_counts)
+
+        # Get dimensions from first patch
+        channels, height, width = processed[0][0].shape
+        batch_size = len(processed)
+
+        # Create padded array
+        encoded = np.zeros(
+            (batch_size, max_patches, channels, height, width), dtype=processed[0][0].dtype
+        )
+
+        # Create attention mask (1 for real patches, 0 for padding)
+        attention_mask = np.zeros((batch_size, max_patches), dtype=np.int64)
+
+        # Fill in patches and attention mask
+        for i, patches in enumerate(processed):
+            for j, patch in enumerate(patches):
+                encoded[i, j] = patch
+                attention_mask[i, j] = 1
+
+        metadata = {"patch_counts": patch_counts}
+        return encoded, attention_mask, metadata  # type: ignore[return-value]
+
+    def _preprocess_onnx_image_input(
+        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        """
+        Add text input placeholders for image data, following Idefics3 processing logic.
+
+        Constructs input_ids dynamically based on the actual number of image patches,
+        using the same token expansion logic as Idefics3Processor.
+
+        Args:
+            onnx_input: Dict with 'pixel_values' (batch, num_patches, C, H, W)
+                        and 'attention_mask' (batch, num_patches) indicating real patches
+            **kwargs: Additional arguments
+
+        Returns:
+            Updated onnx_input with 'input_ids' and updated 'attention_mask' for token sequence
+        """
+        # The attention_mask in onnx_input has a shape of (batch_size, num_patches),
+        # and should be used to create an attention mask matching the input_ids shape.
+        patch_attention_mask = onnx_input["attention_mask"]
+        pixel_values = onnx_input["pixel_values"]
+
+        batch_size = pixel_values.shape[0]
+        batch_input_ids = []
+
+        # Build input_ids for each image based on its actual patch count
+        for i in range(batch_size):
+            # Count real patches (non-padded) from attention mask
+            patch_count = int(np.sum(patch_attention_mask[i]))
+
+            # Compute rows/cols from patch count
+            rows, cols = self._compute_rows_cols_from_patches(patch_count)
+
+            # Build input_ids for this image
+            input_ids = self._build_input_ids_for_image(rows, cols)
+            batch_input_ids.append(input_ids)
+
+        # Pad sequences to max length in batch
+        max_len = max(len(ids) for ids in batch_input_ids)
+
+        # Get padding config from tokenizer
+        padding_direction = self.tokenizer.padding["direction"]  # type: ignore[index,union-attr]
+        pad_token_id = self.tokenizer.padding["pad_id"]  # type: ignore[index,union-attr]
+
+        # Initialize with pad token
+        padded_input_ids = np.full((batch_size, max_len), pad_token_id, dtype=np.int64)
+        attention_mask = np.zeros((batch_size, max_len), dtype=np.int64)
+
+        for i, input_ids in enumerate(batch_input_ids):
+            seq_len = len(input_ids)
+            if padding_direction == "left":
+                # Left padding: place tokens at the END of the array
+                start_idx = max_len - seq_len
+                padded_input_ids[i, start_idx:] = input_ids
+                attention_mask[i, start_idx:] = 1
+            else:
+                # Right padding: place tokens at the START of the array
+                padded_input_ids[i, :seq_len] = input_ids
+                attention_mask[i, :seq_len] = 1
+
+        onnx_input["input_ids"] = padded_input_ids
+        # Update attention_mask with token-level data
+        onnx_input["attention_mask"] = attention_mask
+        return onnx_input
+
+    @staticmethod
+    def _compute_rows_cols_from_patches(patch_count: int) -> tuple[int, int]:
+        if patch_count <= 1:
+            return 0, 0
+
+        # Subtract 1 for the global image
+        grid_patches = patch_count - 1
+
+        # Find rows and cols (assume square or near-square grid)
+        rows = int(grid_patches**0.5)
+        cols = grid_patches // rows
+
+        # Verify the calculation
+        if rows * cols + 1 != patch_count:
+            # Handle non-square grids
+            for r in range(1, grid_patches + 1):
+                if grid_patches % r == 0:
+                    c = grid_patches // r
+                    if r * c + 1 == patch_count:
+                        return r, c
+            # Fallback: treat as unsplit
+            return 0, 0
+
+        return rows, cols
+
+    def _create_single_image_prompt_string(self) -> str:
+        return (
+            "<fake_token_around_image>"
+            + "<global-img>"
+            + "<image>" * self.image_seq_len  # type: ignore[operator]
+            + "<fake_token_around_image>"
+        )
+
+    def _create_split_image_prompt_string(self, rows: int, cols: int) -> str:
+        text_split_images = ""
+
+        # Add tokens for each patch in the grid
+        for n_h in range(rows):
+            for n_w in range(cols):
+                text_split_images += (
+                    "<fake_token_around_image>"
+                    + f"<row_{n_h + 1}_col_{n_w + 1}>"
+                    + "<image>" * self.image_seq_len  # type: ignore[operator]
+                )
+            text_split_images += "\n"
+
+        # Add global image at the end
+        text_split_images += (
+            "\n<fake_token_around_image>"
+            + "<global-img>"
+            + "<image>" * self.image_seq_len  # type: ignore[operator]
+            + "<fake_token_around_image>"
+        )
+
+        return text_split_images
+
+    def _build_input_ids_for_image(self, rows: int, cols: int) -> np.ndarray:
+        # Create the appropriate image prompt string
+        if rows == 0 and cols == 0:
+            image_prompt_tokens = self._create_single_image_prompt_string()
+        else:
+            image_prompt_tokens = self._create_split_image_prompt_string(rows, cols)
+
+        # Replace <image> in visual prompt with expanded tokens
+        # The visual prompt is: "<|begin_of_text|>User:<image>Describe the image.<end_of_utterance>\nAssistant:"
+        expanded_prompt = self.VISUAL_PROMPT_PREFIX.replace("<image>", image_prompt_tokens)
+
+        # Tokenize the complete prompt
+        encoded = self.tokenizer.encode(expanded_prompt)  # type: ignore[union-attr]
+
+        # Convert to numpy array
+        return np.array(encoded.ids, dtype=np.int64)
+
+    def _post_process_onnx_image_output(
+        self,
+        output: OnnxOutputContext,
+    ) -> Iterable[NumpyArray]:
+        """
+        Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+
+        Returns:
+            Iterable[NumpyArray]: Post-processed output as NumPy arrays.
+        """
+        assert self.model_description.dim is not None, "Model dim is not defined"
+        return output.model_output.reshape(
+            output.model_output.shape[0], -1, self.model_description.dim
+        )
+
+    def embed_text(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        """
+        Encode a list of documents into list of embeddings.
+
+        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,
+        )
+
+    def embed_image(
+        self,
+        images: ImageInput | Iterable[ImageInput],
+        batch_size: int = 16,
+        parallel: Optional[int] = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        """
+        Encode a list of images into list of embeddings.
+
+        Args:
+            images: Iterator of image paths or single image path 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_images(
+            model_name=self.model_name,
+            cache_dir=str(self.cache_dir),
+            images=images,
+            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_text_worker_class(cls) -> Type[TextEmbeddingWorker[NumpyArray]]:
+        return ColModernVBERTTextEmbeddingWorker
+
+    @classmethod
+    def _get_image_worker_class(cls) -> Type[ImageEmbeddingWorker[NumpyArray]]:
+        return ColModernVBERTImageEmbeddingWorker
+
+
+class ColModernVBERTTextEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> ColModernVBERT:
+        return ColModernVBERT(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )
+
+
+class ColModernVBERTImageEmbeddingWorker(ImageEmbeddingWorker[NumpyArray]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> ColModernVBERT:
+        return ColModernVBERT(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )