fastembed 0.2.6__tar.gz → 0.3.0__tar.gz

{fastembed-0.2.6 → fastembed-0.3.0}/PKG-INFO

@@ -1,12 +1,12 @@
 Metadata-Version: 2.1
 Name: fastembed
-Version: 0.2.6
+Version: 0.3.0
 Summary: Fast, light, accurate library built for retrieval embedding generation
 Home-page: https://github.com/qdrant/fastembed
 License: Apache License
 Keywords: vector,embedding,neural,search,qdrant,sentence-transformers
-Author: NirantK
-Author-email: nirant.bits@gmail.com
+Author: Qdrant Team
+Author-email: info@qdrant.tech
 Requires-Python: >=3.8.0,<3.13
 Classifier: License :: Other/Proprietary License
 Classifier: Programming Language :: Python :: 3
@@ -15,14 +15,18 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: huggingface-hub (>=0.20,<0.21)
+Requires-Dist: PyStemmer (>=2.2.0,<3.0.0)
+Requires-Dist: huggingface-hub (>=0.20,<1.0)
 Requires-Dist: loguru (>=0.7.2,<0.8.0)
+Requires-Dist: mmh3 (>=4.0,<5.0)
 Requires-Dist: numpy (>=1.21) ; python_version < "3.12"
 Requires-Dist: numpy (>=1.26) ; python_version >= "3.12"
 Requires-Dist: onnx (>=1.15.0,<2.0.0)
 Requires-Dist: onnxruntime (>=1.17.0,<2.0.0)
+Requires-Dist: pillow (>=10.3.0,<11.0.0)
 Requires-Dist: requests (>=2.31,<3.0)
-Requires-Dist: tokenizers (>=0.15.1,<0.16.0)
+Requires-Dist: snowballstemmer (>=2.2.0,<3.0.0)
+Requires-Dist: tokenizers (>=0.15,<1.0)
 Requires-Dist: tqdm (>=4.66,<5.0)
 Project-URL: Repository, https://github.com/qdrant/fastembed
 Description-Content-Type: text/markdown
@@ -31,7 +35,7 @@ Description-Content-Type: text/markdown
 
 FastEmbed is a lightweight, fast, Python library built for embedding generation. We [support popular text models](https://qdrant.github.io/fastembed/examples/Supported_Models/). Please [open a GitHub issue](https://github.com/qdrant/fastembed/issues/new) if you want us to add a new model.
 
-The default text embedding (`TextEmbedding`) model is Flag Embedding, presented in the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. It supports "query" and "passage" prefixes for the input text. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
+The default text embedding (`TextEmbedding`) model is Flag Embedding, presented in the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. It supports "query" and "passage" prefixes for the input text. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/qdrant/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/qdrant/Usage_With_Qdrant/).
 
 ## 📈 Why FastEmbed?
 
@@ -43,12 +47,18 @@ The default text embedding (`TextEmbedding`) model is Flag Embedding, presented
 
 ## 🚀 Installation
 
-To install the FastEmbed library, pip works:
+To install the FastEmbed library, pip works best. You can install it with or without GPU support:
 
 ```bash
 pip install fastembed
 ```
 
+### ⚡️ With GPU
+
+```bash
+pip install fastembed-gpu
+```
+
 ## 📖 Quickstart
 
 ```python
@@ -71,6 +81,28 @@ embeddings_list = list(embedding_model.embed(documents))
 len(embeddings_list[0]) # Vector of 384 dimensions
 ```
 
+### ⚡️ FastEmbed on a GPU
+
+FastEmbed supports running on GPU devices.
+It requires installation of the `fastembed-gpu` package.
+
+```bash
+pip install fastembed-gpu
+```
+
+Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for the detailed instructions and CUDA 12.x support.
+
+```python
+from fastembed import TextEmbedding
+
+embedding_model = TextEmbedding(
+    model_name="BAAI/bge-small-en-v1.5", 
+    providers=["CUDAExecutionProvider"]
+)
+print("The model BAAI/bge-small-en-v1.5 is ready to use on a GPU.")
+
+```
+
 ## Usage with Qdrant
 
 Installation with Qdrant Client in Python:
@@ -79,7 +111,13 @@ Installation with Qdrant Client in Python:
 pip install qdrant-client[fastembed]
 ```
 
-You might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh.
+or 
+
+```bash
+pip install qdrant-client[fastembed-gpu]
+```
+
+You might have to use quotes ```pip install 'qdrant-client[fastembed]'``` on zsh.
 
 ```python
 from qdrant_client import QdrantClient
@@ -115,8 +153,3 @@ search_result = client.query(
 )
 print(search_result)
 ```
-
-#### Similar Work
-
-Ilyas M. wrote about using [FlagEmbeddings with Optimum](https://twitter.com/IlysMoutawwakil/status/1705215192425288017) over CUDA.
-

{fastembed-0.2.6 → fastembed-0.3.0}/README.md

@@ -2,7 +2,7 @@
 
 FastEmbed is a lightweight, fast, Python library built for embedding generation. We [support popular text models](https://qdrant.github.io/fastembed/examples/Supported_Models/). Please [open a GitHub issue](https://github.com/qdrant/fastembed/issues/new) if you want us to add a new model.
 
-The default text embedding (`TextEmbedding`) model is Flag Embedding, presented in the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. It supports "query" and "passage" prefixes for the input text. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
+The default text embedding (`TextEmbedding`) model is Flag Embedding, presented in the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. It supports "query" and "passage" prefixes for the input text. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/qdrant/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/qdrant/Usage_With_Qdrant/).
 
 ## 📈 Why FastEmbed?
 
@@ -14,12 +14,18 @@ The default text embedding (`TextEmbedding`) model is Flag Embedding, presented
 
 ## 🚀 Installation
 
-To install the FastEmbed library, pip works:
+To install the FastEmbed library, pip works best. You can install it with or without GPU support:
 
 ```bash
 pip install fastembed
 ```
 
+### ⚡️ With GPU
+
+```bash
+pip install fastembed-gpu
+```
+
 ## 📖 Quickstart
 
 ```python
@@ -42,6 +48,28 @@ embeddings_list = list(embedding_model.embed(documents))
 len(embeddings_list[0]) # Vector of 384 dimensions
 ```
 
+### ⚡️ FastEmbed on a GPU
+
+FastEmbed supports running on GPU devices.
+It requires installation of the `fastembed-gpu` package.
+
+```bash
+pip install fastembed-gpu
+```
+
+Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for the detailed instructions and CUDA 12.x support.
+
+```python
+from fastembed import TextEmbedding
+
+embedding_model = TextEmbedding(
+    model_name="BAAI/bge-small-en-v1.5", 
+    providers=["CUDAExecutionProvider"]
+)
+print("The model BAAI/bge-small-en-v1.5 is ready to use on a GPU.")
+
+```
+
 ## Usage with Qdrant
 
 Installation with Qdrant Client in Python:
@@ -50,7 +78,13 @@ Installation with Qdrant Client in Python:
 pip install qdrant-client[fastembed]
 ```
 
-You might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh.
+or 
+
+```bash
+pip install qdrant-client[fastembed-gpu]
+```
+
+You might have to use quotes ```pip install 'qdrant-client[fastembed]'``` on zsh.
 
 ```python
 from qdrant_client import QdrantClient
@@ -85,8 +119,4 @@ search_result = client.query(
     query_text="This is a query document"
 )
 print(search_result)
-```
-
-#### Similar Work
-
-Ilyas M. wrote about using [FlagEmbeddings with Optimum](https://twitter.com/IlysMoutawwakil/status/1705215192425288017) over CUDA.
+```

fastembed-0.3.0/fastembed/__init__.py

@@ -0,0 +1,20 @@
+import importlib.metadata
+
+from fastembed.image import ImageEmbedding
+from fastembed.text import TextEmbedding
+from fastembed.sparse import SparseTextEmbedding, SparseEmbedding
+from fastembed.late_interaction import LateInteractionTextEmbedding
+
+try:
+    version = importlib.metadata.version("fastembed")
+except importlib.metadata.PackageNotFoundError as _:
+    version = importlib.metadata.version("fastembed-gpu")
+
+__version__ = version
+__all__ = [
+    "TextEmbedding",
+    "SparseTextEmbedding",
+    "SparseEmbedding",
+    "ImageEmbedding",
+    "LateInteractionTextEmbedding",
+]

fastembed-0.3.0/fastembed/common/__init__.py

@@ -0,0 +1,3 @@
+from fastembed.common.types import OnnxProvider, ImageInput, PathInput
+
+__all__ = ["OnnxProvider", "ImageInput", "PathInput"]

{fastembed-0.2.6 → fastembed-0.3.0}/fastembed/common/model_management.py

@@ -11,23 +11,6 @@ from tqdm import tqdm
 from loguru import logger
 
 
-def locate_model_file(model_dir: Path, file_names: List[str]) -> Path:
-    """
-    Find model path for both TransformerJS style `onnx`  subdirectory structure and direct model weights structure used
-    by Optimum and Qdrant
-    """
-    if not model_dir.is_dir():
-        raise ValueError(f"Provided model path '{model_dir}' is not a directory.")
-
-    for file_name in file_names:
-        file_paths = [path for path in model_dir.rglob(file_name) if path.is_file()]
-
-        if file_paths:
-            return file_paths[0]
-
-    raise ValueError(f"Could not find either of {', '.join(file_names)} in {model_dir}")
-
-
 class ModelManagement:
     @classmethod
     def list_supported_models(cls) -> List[Dict[str, Any]]:
@@ -104,21 +87,37 @@ class ModelManagement:
 
     @classmethod
     def download_files_from_huggingface(
-        cls, hf_source_repo: str, cache_dir: Optional[str] = None
+        cls,
+        hf_source_repo: str,
+        cache_dir: Optional[str] = None,
+        extra_patterns: Optional[List[str]] = None,
+        **kwargs,
     ) -> str:
         """
         Downloads a model from HuggingFace Hub.
         Args:
             hf_source_repo (str): Name of the model on HuggingFace Hub, e.g. "qdrant/all-MiniLM-L6-v2-onnx".
             cache_dir (Optional[str]): The path to the cache directory.
+            extra_patterns (Optional[List[str]]): extra patterns to allow in the snapshot download, typically
+                includes the required model files.
         Returns:
             Path: The path to the model directory.
         """
+        allow_patterns = [
+            "config.json",
+            "tokenizer.json",
+            "tokenizer_config.json",
+            "special_tokens_map.json",
+            "preprocessor_config.json",
+        ]
+        if extra_patterns is not None:
+            allow_patterns.extend(extra_patterns)
 
         return snapshot_download(
             repo_id=hf_source_repo,
-            ignore_patterns=["model.safetensors", "pytorch_model.bin"],
+            allow_patterns=allow_patterns,
             cache_dir=cache_dir,
+            local_files_only=kwargs.get("local_files_only", False),
         )
 
     @classmethod
@@ -175,6 +174,9 @@ class ModelManagement:
 
         model_tar_gz = Path(cache_dir) / f"{fast_model_name}.tar.gz"
 
+        if model_tar_gz.exists():
+            model_tar_gz.unlink()
+
         cls.download_file_from_gcs(
             source_url,
             output_path=str(model_tar_gz),
@@ -190,7 +192,7 @@ class ModelManagement:
         return model_dir
 
     @classmethod
-    def download_model(cls, model: Dict[str, Any], cache_dir: Path) -> Path:
+    def download_model(cls, model: Dict[str, Any], cache_dir: Path, **kwargs) -> Path:
         """
         Downloads a model from HuggingFace Hub or Google Cloud Storage.
 
@@ -219,9 +221,17 @@ class ModelManagement:
         url_source = model.get("sources", {}).get("url")
 
         if hf_source:
+            extra_patterns = [model["model_file"]]
+            extra_patterns.extend(model.get("additional_files", []))
+
             try:
                 return Path(
-                    cls.download_files_from_huggingface(hf_source, cache_dir=str(cache_dir))
+                    cls.download_files_from_huggingface(
+                        hf_source,
+                        cache_dir=str(cache_dir),
+                        extra_patterns=extra_patterns,
+                        local_files_only=kwargs.get("local_files_only", False),
+                    )
                 )
             except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
                 logger.error(

fastembed-0.3.0/fastembed/common/onnx_model.py

@@ -0,0 +1,113 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, Generic, Iterable, Optional, Tuple, Type, TypeVar, Sequence
+import warnings
+
+import numpy as np
+import onnxruntime as ort
+
+from fastembed.common.types import OnnxProvider
+from fastembed.parallel_processor import Worker
+
+
+# Holds type of the embedding result
+T = TypeVar("T")
+
+
+@dataclass
+class OnnxOutputContext:
+    model_output: np.ndarray
+    attention_mask: Optional[np.ndarray] = None
+    input_ids: Optional[np.ndarray] = None
+
+
+class OnnxModel(Generic[T]):
+    @classmethod
+    def _get_worker_class(cls) -> Type["EmbeddingWorker"]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def __init__(self) -> None:
+        self.model = None
+        self.tokenizer = None
+
+    def _preprocess_onnx_input(
+        self, onnx_input: Dict[str, np.ndarray], **kwargs
+    ) -> Dict[str, np.ndarray]:
+        """
+        Preprocess the onnx input.
+        """
+        return onnx_input
+
+    def load_onnx_model(
+        self,
+        model_dir: Path,
+        model_file: str,
+        threads: Optional[int],
+        providers: Optional[Sequence[OnnxProvider]] = None,
+    ) -> None:
+        model_path = model_dir / model_file
+        # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
+
+        onnx_providers = ["CPUExecutionProvider"] if providers is None else list(providers)
+        available_providers = ort.get_available_providers()
+        requested_provider_names = []
+        for provider in onnx_providers:
+            # check providers available
+            provider_name = provider if isinstance(provider, str) else provider[0]
+            requested_provider_names.append(provider_name)
+            if provider_name not in available_providers:
+                raise ValueError(
+                    f"Provider {provider_name} is not available. Available providers: {available_providers}"
+                )
+
+        so = ort.SessionOptions()
+        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+        if threads is not None:
+            so.intra_op_num_threads = threads
+            so.inter_op_num_threads = threads
+
+        self.model = ort.InferenceSession(
+            str(model_path), providers=onnx_providers, sess_options=so
+        )
+        if "CUDAExecutionProvider" in requested_provider_names:
+            current_providers = self.model.get_providers()
+            if "CUDAExecutionProvider" not in current_providers:
+                warnings.warn(
+                    f"Attempt to set CUDAExecutionProvider failed. Current providers: {current_providers}."
+                    "If you are using CUDA 12.x, install onnxruntime-gpu via "
+                    "`pip install onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/`",
+                    RuntimeWarning,
+                )
+
+    def onnx_embed(self, *args, **kwargs) -> OnnxOutputContext:
+        raise NotImplementedError("Subclasses must implement this method")
+
+
+class EmbeddingWorker(Worker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+    ) -> OnnxModel:
+        raise NotImplementedError()
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str,
+    ):
+        self.model = self.init_embedding(model_name, cache_dir)
+
+    @classmethod
+    def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker":
+        return cls(
+            model_name=model_name,
+            cache_dir=cache_dir,
+        )
+
+    def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
+        raise NotImplementedError("Subclasses must implement this method")

fastembed-0.2.6/fastembed/common/models.py → fastembed-0.3.0/fastembed/common/preprocessor_utils.py

@@ -1,11 +1,24 @@
 import json
 from pathlib import Path
+from typing import Tuple
 
-import numpy as np
 from tokenizers import Tokenizer, AddedToken
 
+from fastembed.image.transform.operators import Compose
 
-def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tokenizer:
+
+def load_special_tokens(model_dir: Path) -> dict:
+    tokens_map_path = model_dir / "special_tokens_map.json"
+    if not tokens_map_path.exists():
+        raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
+
+    with open(str(tokens_map_path)) as tokens_map_file:
+        tokens_map = json.load(tokens_map_file)
+
+    return tokens_map
+
+
+def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tuple[Tokenizer, dict]:
     config_path = model_dir / "config.json"
     if not config_path.exists():
         raise ValueError(f"Could not find config.json in {model_dir}")
@@ -18,18 +31,13 @@ def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tokenizer:
     if not tokenizer_config_path.exists():
         raise ValueError(f"Could not find tokenizer_config.json in {model_dir}")
 
-    tokens_map_path = model_dir / "special_tokens_map.json"
-    if not tokens_map_path.exists():
-        raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
-
     with open(str(config_path)) as config_file:
         config = json.load(config_file)
 
     with open(str(tokenizer_config_path)) as tokenizer_config_file:
         tokenizer_config = json.load(tokenizer_config_file)
 
-    with open(str(tokens_map_path)) as tokens_map_file:
-        tokens_map = json.load(tokens_map_file)
+    tokens_map = load_special_tokens(model_dir)
 
     tokenizer = Tokenizer.from_file(str(tokenizer_path))
     tokenizer.enable_truncation(max_length=min(tokenizer_config["model_max_length"], max_length))
@@ -43,12 +51,24 @@ def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tokenizer:
         elif isinstance(token, dict):
             tokenizer.add_special_tokens([AddedToken(**token)])
 
-    return tokenizer
+    special_token_to_id = {}
+
+    for token in tokens_map.values():
+        if isinstance(token, str):
+            special_token_to_id[token] = tokenizer.token_to_id(token)
+        elif isinstance(token, dict):
+            token_str = token.get("content", "")
+            special_token_to_id[token_str] = tokenizer.token_to_id(token_str)
+
+    return tokenizer, special_token_to_id
+
 
+def load_preprocessor(model_dir: Path) -> Compose:
+    preprocessor_config_path = model_dir / "preprocessor_config.json"
+    if not preprocessor_config_path.exists():
+        raise ValueError(f"Could not find preprocessor_config.json in {model_dir}")
 
-def normalize(input_array, p=2, dim=1, eps=1e-12) -> np.ndarray:
-    # Calculate the Lp norm along the specified dimension
-    norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
-    norm = np.maximum(norm, eps)  # Avoid division by zero
-    normalized_array = input_array / norm
-    return normalized_array
+    with open(str(preprocessor_config_path)) as preprocessor_config_file:
+        preprocessor_config = json.load(preprocessor_config_file)
+        transforms = Compose.from_config(preprocessor_config)
+    return transforms

fastembed-0.3.0/fastembed/common/types.py

@@ -0,0 +1,14 @@
+import os
+import sys
+from typing import Union, Iterable, Tuple, Dict, Any
+
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias
+
+
+PathInput: TypeAlias = Union[str, os.PathLike]
+ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput]]
+
+OnnxProvider: TypeAlias = Union[str, Tuple[str, Dict[Any, Any]]]

{fastembed-0.2.6 → fastembed-0.3.0}/fastembed/common/utils.py

@@ -4,6 +4,16 @@ from itertools import islice
 from pathlib import Path
 from typing import Union, Iterable, Generator, Optional
 
+import numpy as np
+
+
+def normalize(input_array, p=2, dim=1, eps=1e-12) -> np.ndarray:
+    # Calculate the Lp norm along the specified dimension
+    norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
+    norm = np.maximum(norm, eps)  # Avoid division by zero
+    normalized_array = input_array / norm
+    return normalized_array
+
 
 def iter_batch(iterable: Union[Iterable, Generator], size: int) -> Iterable:
     """

fastembed-0.3.0/fastembed/image/__init__.py

@@ -0,0 +1,4 @@
+from fastembed.image.image_embedding import ImageEmbedding
+
+
+__all__ = ["ImageEmbedding"]

fastembed-0.3.0/fastembed/image/image_embedding.py

@@ -0,0 +1,87 @@
+from typing import Any, Dict, Iterable, List, Optional, Type, Sequence
+
+import numpy as np
+
+from fastembed.common import ImageInput, OnnxProvider
+from fastembed.image.image_embedding_base import ImageEmbeddingBase
+from fastembed.image.onnx_embedding import OnnxImageEmbedding
+
+
+class ImageEmbedding(ImageEmbeddingBase):
+    EMBEDDINGS_REGISTRY: List[Type[ImageEmbeddingBase]] = [OnnxImageEmbedding]
+
+    @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.
+
+            Example:
+                ```
+                [
+                    {
+                        "model": "Qdrant/clip-ViT-B-32-vision",
+                        "dim": 512,
+                        "description": "CLIP vision encoder based on ViT-B/32",
+                        "size_in_GB": 0.33,
+                        "sources": {
+                            "hf": "Qdrant/clip-ViT-B-32-vision",
+                        },
+                        "model_file": "model.onnx",
+                    }
+                ]
+                ```
+        """
+        result = []
+        for embedding in cls.EMBEDDINGS_REGISTRY:
+            result.extend(embedding.list_supported_models())
+        return result
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        **kwargs,
+    ):
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+
+        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, cache_dir, threads, providers=providers, **kwargs
+                )
+                return
+
+        raise ValueError(
+            f"Model {model_name} is not supported in TextEmbedding."
+            "Please check the supported models using `TextEmbedding.list_supported_models()`"
+        )
+
+    def embed(
+        self,
+        images: ImageInput,
+        batch_size: int = 16,
+        parallel: Optional[int] = None,
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
+        """
+        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:
+            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.model.embed(images, batch_size, parallel, **kwargs)

fastembed-0.3.0/fastembed/image/image_embedding_base.py

@@ -0,0 +1,39 @@
+from typing import Iterable, Optional
+
+import numpy as np
+
+from fastembed.common.model_management import ModelManagement
+from fastembed.common.types import ImageInput
+
+
+class ImageEmbeddingBase(ModelManagement):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        **kwargs,
+    ):
+        self.model_name = model_name
+        self.cache_dir = cache_dir
+        self.threads = threads
+        self._local_files_only = kwargs.pop("local_files_only", False)
+
+    def embed(
+        self,
+        images: ImageInput,
+        batch_size: int = 16,
+        parallel: Optional[int] = None,
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
+        """
+        Embeds a list of images into a list of embeddings.
+
+        Args:
+            images - The list of image paths to preprocess and embed.
+            **kwargs: Additional keyword argument to pass to the embed method.
+
+        Yields:
+            Iterable[np.ndarray]: The embeddings.
+        """
+        raise NotImplementedError()