codekingpro/portable-devtools

1from typing import Any, Type2 3from fastembed.common.types import NumpyArray4from fastembed.late_interaction.colbert import Colbert, ColbertEmbeddingWorker5from fastembed.common.model_description import DenseModelDescription, ModelSource6 7supported_jina_colbert_models: list[DenseModelDescription] = [8    DenseModelDescription(9        model="jinaai/jina-colbert-v2",10        dim=128,11        description="New model that expands capabilities of colbert-v1 with multilingual and context length of 8192, 2024 year",12        license="cc-by-nc-4.0",13        size_in_GB=2.24,14        sources=ModelSource(hf="jinaai/jina-colbert-v2"),15        model_file="onnx/model.onnx",16        additional_files=["onnx/model.onnx_data"],17    )18]19 20 21class JinaColbert(Colbert):22    QUERY_MARKER_TOKEN_ID = 25000223    DOCUMENT_MARKER_TOKEN_ID = 25000324    MIN_QUERY_LENGTH = 31  # it's 32, we add one additional special token in the beginning25    MASK_TOKEN = "<mask>"26 27    @classmethod28    def _get_worker_class(cls) -> Type[ColbertEmbeddingWorker]:29        return JinaColbertEmbeddingWorker30 31    @classmethod32    def _list_supported_models(cls) -> list[DenseModelDescription]:33        """Lists the supported models.34 35        Returns:36            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.37        """38        return supported_jina_colbert_models39 40    def _preprocess_onnx_input(41        self, onnx_input: dict[str, NumpyArray], is_doc: bool = True, **kwargs: Any42    ) -> dict[str, NumpyArray]:43        onnx_input = super()._preprocess_onnx_input(onnx_input, is_doc)44 45        # the attention mask for jina-colbert-v2 is always 1 in queries46        if not is_doc:47            onnx_input["attention_mask"][:] = 148        return onnx_input49 50 51class JinaColbertEmbeddingWorker(ColbertEmbeddingWorker):52    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> JinaColbert:53        return JinaColbert(54            model_name=model_name,55            cache_dir=cache_dir,56            threads=1,57            **kwargs,58        )59