codekingpro/portable-devtools

1from typing import Iterable, Any2 3from fastembed.common.model_description import DenseModelDescription4from fastembed.common.types import NumpyArray5from fastembed.common.model_management import ModelManagement6from fastembed.common.types import ImageInput7 8 9class ImageEmbeddingBase(ModelManagement[DenseModelDescription]):10    def __init__(11        self,12        model_name: str,13        cache_dir: str | None = None,14        threads: int | None = None,15        **kwargs: Any,16    ):17        self.model_name = model_name18        self.cache_dir = cache_dir19        self.threads = threads20        self._local_files_only = kwargs.pop("local_files_only", False)21        self._embedding_size: int | None = None22 23    def embed(24        self,25        images: ImageInput | Iterable[ImageInput],26        batch_size: int = 16,27        parallel: int | None = None,28        **kwargs: Any,29    ) -> Iterable[NumpyArray]:30        """31        Embeds a list of images into a list of embeddings.32 33        Args:34            images: The list of image paths to preprocess and embed.35            batch_size: Batch size for encoding36            parallel:37                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.38                If 0, use all available cores.39                If None, don't use data-parallel processing, use default onnxruntime threading instead.40            **kwargs: Additional keyword argument to pass to the embed method.41 42        Yields:43            Iterable[NdArray]: The embeddings.44        """45        raise NotImplementedError()46 47    @classmethod48    def get_embedding_size(cls, model_name: str) -> int:49        """Returns embedding size of the chosen model."""50        raise NotImplementedError("Subclasses must implement this method")51 52    @property53    def embedding_size(self) -> int:54        """Returns embedding size for the current model"""55        raise NotImplementedError("Subclasses must implement this method")56