codekingpro/portable-devtools

1from typing import Any, Iterable, Type2 3from fastembed.common.types import NumpyArray4from fastembed.common.onnx_model import OnnxOutputContext5from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker6from fastembed.common.model_description import DenseModelDescription, ModelSource7 8supported_clip_models: list[DenseModelDescription] = [9    DenseModelDescription(10        model="Qdrant/clip-ViT-B-32-text",11        dim=512,12        description=(13            "Text embeddings, Multimodal (text&image), English, 77 input tokens truncation, "14            "Prefixes for queries/documents: not necessary, 2021 year"15        ),16        license="mit",17        size_in_GB=0.25,18        sources=ModelSource(hf="Qdrant/clip-ViT-B-32-text"),19        model_file="model.onnx",20    ),21]22 23 24class CLIPOnnxEmbedding(OnnxTextEmbedding):25    @classmethod26    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:27        return CLIPEmbeddingWorker28 29    @classmethod30    def _list_supported_models(cls) -> list[DenseModelDescription]:31        """Lists the supported models.32 33        Returns:34            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.35        """36        return supported_clip_models37 38    def _post_process_onnx_output(39        self, output: OnnxOutputContext, **kwargs: Any40    ) -> Iterable[NumpyArray]:41        return output.model_output42 43 44class CLIPEmbeddingWorker(OnnxTextEmbeddingWorker):45    def init_embedding(46        self,47        model_name: str,48        cache_dir: str,49        **kwargs: Any,50    ) -> OnnxTextEmbedding:51        return CLIPOnnxEmbedding(52            model_name=model_name,53            cache_dir=cache_dir,54            threads=1,55            **kwargs,56        )57