Merge branch 'main' into clip-export

* new: add gpu support, allow users to override providers

* fix: update poetry.lock

* fix: fix type hint for 3.8

* [readme] Remove similar work

* [README] Add GPU support for FastEmbed library

* [README]  Add device check

* fix: revert changes to pyproject and lock, update readme

* Update poetry.lock

* new: add type alias for providers, add explicit providers to embeddings

---------

Co-authored-by: Nirant Kasliwal <nirant.bits@gmail.com>
Co-authored-by: Nirant <NirantK@users.noreply.github.com>

.github/workflows/python-publish.yml

@@ -15,7 +15,6 @@ on:
     tags:
       - 'v*'           # Push events to every version tag
 
-
 jobs:
   deploy:
 

.github/workflows/python-tests.yml

@@ -2,7 +2,7 @@ name: Tests
 
 on:
   push:
-    branches: [ master, main ]
+    branches: [ master, main, gpu ]
   schedule:
     - cron: 0 0 * * *
   pull_request:

README.md

@@ -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,23 @@ 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.
+Make sure not to have the `fastembed` package installed, as it might interfere with the `fastembed-gpu` package.
+
+```bash
+pip install fastembed-gpu
+``` 
+
+```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 +73,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 +114,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.
+```

RELEASE.md

@@ -0,0 +1,41 @@
+# Releasing FastEmbed
+
+This is a guide how to release `fastembed` and `fastembed-gpu` packages.
+
+## How to
+
+1. Accumulate changes in the `main` branch.
+2. Bump the version in `pyproject.toml`
+
+3. Rebase the `gpu` branch on `main` and resolve conflicts if occurred:
+
+```bash
+git checkout gpu
+git rebase main
+git push origin gpu
+```
+
+4. Draft release notes
+5. Checkout to `main` and create a tag, e.g.:
+
+```bash
+git checkout main
+git tag -a v0.1.0 -m "Release v0.1.0"
+```
+
+6. Checkout `gpu` and create a tag, e.g.:
+
+```bash
+git checkout gpu
+git tag -a v0.1.0-gpu -m "Release v0.1.0"
+```
+
+7. Push tags:
+
+```bash
+git push --tags
+```
+
+8. Verify that both packages have been published successfully on PyPI. Try installing them and verify imports.
+9. Create a release on GitHub with the written release notes.
+

experiments/03_CLIP_TO_ONNX.ipynb

@@ -0,0 +1,323 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f7f8e989-cdc9-475e-918d-af20530fcfe6",
+   "metadata": {
+    "is_executing": true
+   },
+   "outputs": [],
+   "source": [
+    "!pip3 install -q torch transformers optimum pillow"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e8c9e276-58f2-45a4-af40-6d7bacc30eec",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "from pathlib import Path\n",
+    "from typing import Optional, Dict, Union, Tuple\n",
+    "\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from PIL import Image\n",
+    "from transformers import (\n",
+    "    CLIPVisionModelWithProjection,\n",
+    "    CLIPTextModelWithProjection,\n",
+    "    CLIPImageProcessor,\n",
+    "    CLIPTokenizerFast,\n",
+    ")\n",
+    "from transformers.models.clip.modeling_clip import (\n",
+    "    CLIPTextModelOutput,\n",
+    "    CLIPVisionModelOutput,\n",
+    "    CLIPModel,\n",
+    ")\n",
+    "from optimum.onnxruntime import ORTModelForCustomTasks\n",
+    "from optimum.exporters.onnx.model_configs import CLIPTextWithProjectionOnnxConfig, ViTOnnxConfig\n",
+    "from optimum.exporters.onnx import export_models"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bbceb40f-22cd-4d92-be6e-fe14f16f7bc2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_id = \"openai/clip-vit-base-patch32\"\n",
+    "output_dir = \"split-clip-onnx\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f96ff7fb-518e-405e-a7e0-46f836ffdec8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class CLIPVisionModelWithProjectionOnnxConfig(ViTOnnxConfig):\n",
+    "    @property\n",
+    "    def outputs(self) -> Dict[str, Dict[int, str]]:\n",
+    "        return {\n",
+    "            \"image_embeds\": {0: \"batch_size\"},\n",
+    "        }"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "18863f0b-6bd5-463f-bebc-38bf40d51b9c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class CLIPTextModelWithProjectionAndAttentionOnnxConfig(CLIPTextWithProjectionOnnxConfig):\n",
+    "    @property\n",
+    "    def inputs(self) -> Dict[str, Dict[int, str]]:\n",
+    "        return {\n",
+    "            \"input_ids\": {0: \"batch_size\", 1: \"sequence_length\"},\n",
+    "            \"attention_mask\": {0: \"batch_size\", 1: \"sequence_length\"},\n",
+    "        }"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5d37b16a-ec30-40e1-8404-0f0d51abfa76",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class CLIPTextModelWithProjectionNormalized(CLIPTextModelWithProjection):\n",
+    "    def forward(\n",
+    "        self,\n",
+    "        input_ids: Optional[torch.Tensor] = None,\n",
+    "        attention_mask: Optional[torch.Tensor] = None,\n",
+    "        position_ids: Optional[torch.Tensor] = None,\n",
+    "        output_attentions: Optional[bool] = None,\n",
+    "        output_hidden_states: Optional[bool] = None,\n",
+    "        return_dict: Optional[bool] = None,\n",
+    "    ) -> Union[Tuple, CLIPTextModelOutput]:\n",
+    "        text_outputs = super().forward(\n",
+    "            input_ids,\n",
+    "            attention_mask,\n",
+    "            position_ids,\n",
+    "            output_attentions,\n",
+    "            output_hidden_states,\n",
+    "            return_dict,\n",
+    "        )\n",
+    "        normalized_text_embeds = text_outputs.text_embeds / text_outputs.text_embeds.norm(\n",
+    "            p=2, dim=-1, keepdim=True\n",
+    "        )\n",
+    "        return CLIPTextModelOutput(\n",
+    "            text_embeds=normalized_text_embeds,\n",
+    "            last_hidden_state=text_outputs.last_hidden_state,\n",
+    "            hidden_states=text_outputs.hidden_states,\n",
+    "            attentions=text_outputs.attentions,\n",
+    "        )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2cd05d33-5f4f-4b36-aa2c-43fd97d5061d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class CLIPVisionModelWithProjectionNormalized(CLIPVisionModelWithProjection):\n",
+    "    def forward(\n",
+    "        self,\n",
+    "        pixel_values: Optional[torch.FloatTensor] = None,\n",
+    "        output_attentions: Optional[bool] = None,\n",
+    "        output_hidden_states: Optional[bool] = None,\n",
+    "        return_dict: Optional[bool] = None,\n",
+    "    ) -> Union[Tuple, CLIPVisionModelOutput]:\n",
+    "        vision_outputs = super().forward(pixel_values, return_dict)\n",
+    "        normalized_image_embeds = vision_outputs.image_embeds / vision_outputs.image_embeds.norm(\n",
+    "            p=2, dim=-1, keepdim=True\n",
+    "        )\n",
+    "        return CLIPVisionModelOutput(\n",
+    "            image_embeds=normalized_image_embeds,\n",
+    "            last_hidden_state=vision_outputs.last_hidden_state,\n",
+    "            hidden_states=vision_outputs.hidden_states,\n",
+    "            attentions=vision_outputs.attentions,\n",
+    "        )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fb5e5617-f8ce-4dcf-9148-68ddd91854c9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "text_model = CLIPTextModelWithProjectionNormalized.from_pretrained(model_id)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7f578131-c6bb-460e-8200-d0b7f0aa4135",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "vision_model = CLIPVisionModelWithProjectionNormalized.from_pretrained(model_id)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c57db080-642f-4b62-96ad-75af9c0ab277",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "text_config = CLIPTextModelWithProjectionAndAttentionOnnxConfig(text_model.config)\n",
+    "vision_config = CLIPVisionModelWithProjectionOnnxConfig(vision_model.config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cdb01aab-0dff-4fd5-9297-d16599274fdb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "text_model.config.save_pretrained(f\"./{output_dir}/text\")\n",
+    "vision_model.config.save_pretrained(f\"./{output_dir}/image\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "133bcd2c-57ae-4132-a691-3d129057f275",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "export_models(\n",
+    "    models_and_onnx_configs={\n",
+    "        \"text_model\": (text_model, text_config),\n",
+    "        \"vision_model\": (vision_model, vision_config),\n",
+    "    },\n",
+    "    output_dir=Path(f\"./{output_dir}\"),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fd9167b9-0d00-4a24-8f5e-41392d923b95",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "os.rename(f\"./{output_dir}/text_model.onnx\", f\"./{output_dir}/text/model.onnx\")\n",
+    "os.rename(f\"./{output_dir}/vision_model.onnx\", f\"./{output_dir}/image/model.onnx\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2281cc75-1027-4dbe-a7a5-86ee1dad4c3e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ort_vision_model = ORTModelForCustomTasks.from_pretrained(\n",
+    "    f\"./{output_dir}/image\", config=vision_config\n",
+    ")\n",
+    "image_processor = CLIPImageProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
+    "image_input = image_processor(images=Image.open(\"assets/image.jpeg\"), return_tensors=\"pt\")\n",
+    "\n",
+    "with torch.inference_mode():\n",
+    "    image_outputs = ort_vision_model(**image_input)\n",
+    "image_processor.save_pretrained(f\"./{output_dir}/image\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5ce61045-b1e7-4b62-a47b-4bcb27ac7288",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ort_text_model = ORTModelForCustomTasks.from_pretrained(f\"./{output_dir}/text\", config=text_config)\n",
+    "text_processor = CLIPTokenizerFast.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
+    "text_input = text_processor(\"What am I using?\", return_tensors=\"pt\")\n",
+    "\n",
+    "with torch.inference_mode():\n",
+    "    text_outputs = ort_text_model(**text_input)\n",
+    "text_processor.save_pretrained(f\"./{output_dir}/text\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ee1e3d13-6884-4aa3-a6ab-aaa41fc17134",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "clip_model = CLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
+    "inputs = {**text_input, **image_input}\n",
+    "clip_model.eval()\n",
+    "with torch.inference_mode():\n",
+    "    gt_output = clip_model(**inputs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ec07ce55-a0e7-42d2-b1f4-ba1370ab45b7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(np.allclose(gt_output.text_embeds.numpy(), text_outputs.text_embeds, atol=1e-6))\n",
+    "print(np.allclose(gt_output.image_embeds.numpy(), image_outputs.image_embeds, atol=1e-6))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e9b15597-3054-40f1-a080-fea19d378d88",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "token = \"<token>\"\n",
+    "# create_repo(repo_id='Qdrant/clip-ViT-B-32-vision', exist_ok=True, token=token)\n",
+    "# create_repo(repo_id='Qdrant/clip-ViT-B-32-text', exist_ok=True, token=token)\n",
+    "\n",
+    "ort_text_model.push_to_hub(\n",
+    "    save_directory=f\"./{output_dir}/text/\",\n",
+    "    repository_id=\"Qdrant/clip-ViT-B-32-text\",\n",
+    "    use_auth_token=token,\n",
+    ")\n",
+    "ort_vision_model.push_to_hub(\n",
+    "    save_directory=f\"./{output_dir}/image\",\n",
+    "    repository_id=\"Qdrant/clip-ViT-B-32-vision\",\n",
+    "    use_auth_token=token,\n",
+    ")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

fastembed/__init__.py

@@ -3,5 +3,10 @@ import importlib.metadata
 from fastembed.text import TextEmbedding
 from fastembed.sparse import SparseTextEmbedding, SparseEmbedding
 
-__version__ = importlib.metadata.version("fastembed")
+try:
+    version = importlib.metadata.version("fastembed")
+except importlib.metadata.PackageNotFoundError as _:
+    version = importlib.metadata.version("fastembed-gpu")
+
+__version__ = version
 __all__ = ["TextEmbedding", "SparseTextEmbedding", "SparseEmbedding"]

fastembed/common/__init__.py

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

fastembed/common/model_management.py

@@ -91,6 +91,7 @@ class ModelManagement:
         hf_source_repo: str,
         cache_dir: Optional[str] = None,
         extra_patterns: Optional[List[str]] = None,
+        **kwargs,
     ) -> str:
         """
         Downloads a model from HuggingFace Hub.
@@ -115,6 +116,7 @@ class ModelManagement:
             repo_id=hf_source_repo,
             allow_patterns=allow_patterns,
             cache_dir=cache_dir,
+            local_files_only=kwargs.get("local_files_only", False),
         )
 
     @classmethod
@@ -189,7 +191,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.
 
@@ -224,7 +226,10 @@ class ModelManagement:
             try:
                 return Path(
                     cls.download_files_from_huggingface(
-                        hf_source, cache_dir=str(cache_dir), extra_patterns=extra_patterns
+                        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:

fastembed/common/onnx_model.py

@@ -1,7 +1,19 @@
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Dict, Generic, Iterable, List, Optional, Tuple, Type, TypeVar, Union
+from typing import (
+    Any,
+    Dict,
+    Generic,
+    Iterable,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    Sequence,
+)
 
 import numpy as np
 import onnxruntime as ort
@@ -14,6 +26,8 @@ from fastembed.parallel_processor import ParallelWorkerPool, Worker
 # Holds type of the embedding result
 T = TypeVar("T")
 
+OnnxProvider = Union[str, Tuple[str, Dict[Any, Any]]]
+
 
 class OnnxModel(Generic[T]):
     @classmethod
@@ -39,11 +53,21 @@ class OnnxModel(Generic[T]):
         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"]
+
+        onnx_providers = ["CPUExecutionProvider"] if providers is None else list(providers)
+        available_providers = ort.get_available_providers()
+        for provider in onnx_providers:
+            # check providers available
+            provider_name = provider if isinstance(provider, str) else provider[0]
+            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

fastembed/sparse/sparse_embedding_base.py

@@ -32,6 +32,7 @@ class SparseTextEmbeddingBase(ModelManagement):
         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,

fastembed/sparse/sparse_text_embedding.py

@@ -1,5 +1,6 @@
-from typing import List, Type, Dict, Any, Union, Iterable, Optional
+from typing import List, Type, Dict, Any, Union, Iterable, Optional, Sequence
 
+from fastembed.common import OnnxProvider
 from fastembed.sparse.sparse_embedding_base import SparseTextEmbeddingBase, SparseEmbedding
 from fastembed.sparse.splade_pp import SpladePP
 
@@ -42,6 +43,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         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)
@@ -49,7 +51,9 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         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, **kwargs)
+                self.model = EMBEDDING_MODEL_TYPE(
+                    model_name, cache_dir, threads, providers=providers, **kwargs
+                )
                 return
 
         raise ValueError(

fastembed/sparse/splade_pp.py

@@ -1,8 +1,8 @@
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type, Sequence
 
 import numpy as np
 
-from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel
+from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxProvider
 from fastembed.common.utils import define_cache_dir
 from fastembed.sparse.sparse_embedding_base import SparseEmbedding, SparseTextEmbeddingBase
 
@@ -63,6 +63,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxModel[SparseEmbedding]):
         model_name: str,
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
         **kwargs,
     ):
         """
@@ -82,12 +83,15 @@ class SpladePP(SparseTextEmbeddingBase, OnnxModel[SparseEmbedding]):
         model_description = self._get_model_description(model_name)
         cache_dir = define_cache_dir(cache_dir)
 
-        model_dir = self.download_model(model_description, cache_dir)
+        model_dir = self.download_model(
+            model_description, cache_dir, local_files_only=self._local_files_only
+        )
 
         self.load_onnx_model(
             model_dir=model_dir,
             model_file=model_description["model_file"],
             threads=threads,
+            providers=providers,
         )
 
     def embed(

fastembed/text/onnx_embedding.py

@@ -1,8 +1,8 @@
-from typing import Dict, Optional, Tuple, Union, Iterable, Type, List, Any
+from typing import Dict, Optional, Tuple, Union, Iterable, Type, List, Any, Sequence
 
 import numpy as np
 
-from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker
+from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker, OnnxProvider
 from fastembed.common.models import normalize
 from fastembed.common.utils import define_cache_dir
 from fastembed.text.text_embedding_base import TextEmbeddingBase
@@ -211,6 +211,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxModel[np.ndarray]):
         model_name: str = "BAAI/bge-small-en-v1.5",
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
         **kwargs,
     ):
         """
@@ -229,12 +230,15 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxModel[np.ndarray]):
 
         model_description = self._get_model_description(model_name)
         cache_dir = define_cache_dir(cache_dir)
-        model_dir = self.download_model(model_description, cache_dir)
+        model_dir = self.download_model(
+            model_description, cache_dir, local_files_only=self._local_files_only
+        )
 
         self.load_onnx_model(
             model_dir=model_dir,
             model_file=model_description["model_file"],
             threads=threads,
+            providers=providers,
         )
 
     def embed(

fastembed/text/text_embedding.py

@@ -1,7 +1,8 @@
-from typing import Any, Dict, Iterable, List, Optional, Type, Union
+from typing import Any, Dict, Iterable, List, Optional, Type, Union, Sequence
 
 import numpy as np
 
+from fastembed.common import OnnxProvider
 from fastembed.text.e5_onnx_embedding import E5OnnxEmbedding
 from fastembed.text.jina_onnx_embedding import JinaOnnxEmbedding
 from fastembed.text.onnx_embedding import OnnxTextEmbedding
@@ -49,6 +50,7 @@ class TextEmbedding(TextEmbeddingBase):
         model_name: str = "BAAI/bge-small-en-v1.5",
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
         **kwargs,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
@@ -56,7 +58,9 @@ class TextEmbedding(TextEmbeddingBase):
         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, **kwargs)
+                self.model = EMBEDDING_MODEL_TYPE(
+                    model_name, cache_dir, threads, providers=providers, **kwargs
+                )
                 return
 
         raise ValueError(

fastembed/text/text_embedding_base.py

@@ -16,6 +16,7 @@ class TextEmbeddingBase(ModelManagement):
         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,

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed"
-version = "0.2.6"
+version = "0.2.7"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -16,7 +16,7 @@ onnx = "^1.15.0"
 onnxruntime = "^1.17.0"
 tqdm = "^4.66"
 requests = "^2.31"
-tokenizers = "^0.15.1"
+tokenizers = "^0.15"
 huggingface-hub = "^0.20"
 loguru = "^0.7.2"
 numpy = [