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/ISSUE_TEMPLATE/bug-report.yml

@@ -15,13 +15,23 @@ body:
       value: "A bug happened!"
     validations:
       required: true
+  - type: textarea
+    id: Python version
+    attributes:
+      label: What Python version are you on? e.g. python --version
+      description: Also tell us, what package manager are you using e.g. conda, pip, poetry?
+      placeholder: Python3.10
+    validations:
+      required: true
   - type: dropdown
     id: version
     attributes:
       label: Version
       description: What version of FastEmbed are you running? python -c "import fastembed; print(fastembed.__version__)". If you're not on the latest, please upgrade and see if the problem persists.
       options:
-        - 0.2.4 (Latest)
+        - 0.2.6 (Latest)
+        - 0.2.5
+        - 0.2.4 
         - 0.2.3 
         - 0.2.2
         - 0.2.1
@@ -43,4 +53,4 @@ body:
     attributes:
       label: Relevant stack traces and/or logs
       description: Please copy and paste any relevant raised exceptions. This will be automatically formatted into code, so no need for backticks.
-      render: shell
+      render: shell

.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,9 @@ name: Tests
 
 on:
   push:
-    branches: [ master, main ]
+    branches: [ master, main, gpu ]
+  schedule:
+    - cron: 0 0 * * *
   pull_request:
 
 env:
@@ -29,9 +31,9 @@ jobs:
     name: Python ${{ matrix.python-version }} on ${{ matrix.os }} test
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - name: Set up Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
@@ -39,8 +41,14 @@ jobs:
           python -m pip install poetry
           poetry config virtualenvs.create false
           poetry install --no-interaction --no-ansi --without docs
-      - name: Run tests
-        run: |
-          export IS_UBUNTU_CI=$(test "${{ matrix.os }}" = "ubuntu-latest" && echo "true" || echo "false")
-          pytest
-        shell: bash
+
+      - name: Install Test Dependencies
+        run: pip install pytest pytest-md pytest-emoji
+
+      - name: Run pytest
+        uses: pavelzw/pytest-action@v2
+        with:
+            verbose: true
+            emoji: true
+            job-summary: true
+            report-title: 'FastEmbed Test Report'

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,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.
+

docs/examples/Supported_Models.ipynb

@@ -17,7 +17,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -35,7 +35,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
@@ -89,90 +89,125 @@
        "    </tr>\n",
        "    <tr>\n",
        "      <th>3</th>\n",
+       "      <td>snowflake/snowflake-arctic-embed-xs</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Based on all-MiniLM-L6-v2 model with only 22m ...</td>\n",
+       "      <td>0.090</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
        "      <td>jinaai/jina-embeddings-v2-small-en</td>\n",
        "      <td>512</td>\n",
        "      <td>English embedding model supporting 8192 sequen...</td>\n",
        "      <td>0.120</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>4</th>\n",
+       "      <th>5</th>\n",
+       "      <td>snowflake/snowflake-arctic-embed-s</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Based on infloat/e5-small-unsupervised, does n...</td>\n",
+       "      <td>0.130</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6</th>\n",
        "      <td>BAAI/bge-small-en</td>\n",
        "      <td>384</td>\n",
        "      <td>Fast English model</td>\n",
        "      <td>0.130</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>5</th>\n",
+       "      <th>7</th>\n",
        "      <td>BAAI/bge-base-en-v1.5</td>\n",
        "      <td>768</td>\n",
        "      <td>Base English model, v1.5</td>\n",
        "      <td>0.210</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>6</th>\n",
+       "      <th>8</th>\n",
        "      <td>sentence-transformers/paraphrase-multilingual-...</td>\n",
        "      <td>384</td>\n",
        "      <td>Sentence Transformer model, paraphrase-multili...</td>\n",
        "      <td>0.220</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>7</th>\n",
+       "      <th>9</th>\n",
        "      <td>BAAI/bge-base-en</td>\n",
        "      <td>768</td>\n",
        "      <td>Base English model</td>\n",
        "      <td>0.420</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>8</th>\n",
-       "      <td>nomic-ai/nomic-embed-text-v1</td>\n",
-       "      <td>768</td>\n",
-       "      <td>8192 context length english model</td>\n",
-       "      <td>0.520</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>9</th>\n",
-       "      <td>nomic-ai/nomic-embed-text-v1.5</td>\n",
-       "      <td>768</td>\n",
-       "      <td>8192 context length english model</td>\n",
-       "      <td>0.520</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
        "      <th>10</th>\n",
+       "      <td>snowflake/snowflake-arctic-embed-m</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Based on intfloat/e5-base-unsupervised model, ...</td>\n",
+       "      <td>0.430</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>11</th>\n",
        "      <td>jinaai/jina-embeddings-v2-base-en</td>\n",
        "      <td>768</td>\n",
        "      <td>English embedding model supporting 8192 sequen...</td>\n",
        "      <td>0.520</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>11</th>\n",
+       "      <th>12</th>\n",
+       "      <td>nomic-ai/nomic-embed-text-v1</td>\n",
+       "      <td>768</td>\n",
+       "      <td>8192 context length english model</td>\n",
+       "      <td>0.520</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>13</th>\n",
+       "      <td>nomic-ai/nomic-embed-text-v1.5</td>\n",
+       "      <td>768</td>\n",
+       "      <td>8192 context length english model</td>\n",
+       "      <td>0.520</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>14</th>\n",
+       "      <td>snowflake/snowflake-arctic-embed-m-long</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Based on nomic-ai/nomic-embed-text-v1-unsuperv...</td>\n",
+       "      <td>0.540</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>15</th>\n",
        "      <td>mixedbread-ai/mxbai-embed-large-v1</td>\n",
        "      <td>1024</td>\n",
        "      <td>MixedBread Base sentence embedding model, does...</td>\n",
        "      <td>0.640</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>12</th>\n",
+       "      <th>16</th>\n",
        "      <td>sentence-transformers/paraphrase-multilingual-...</td>\n",
        "      <td>768</td>\n",
        "      <td>Sentence-transformers model for tasks like clu...</td>\n",
        "      <td>1.000</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>13</th>\n",
+       "      <th>17</th>\n",
+       "      <td>snowflake/snowflake-arctic-embed-l</td>\n",
+       "      <td>1024</td>\n",
+       "      <td>Based on intfloat/e5-large-unsupervised, large...</td>\n",
+       "      <td>1.020</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>18</th>\n",
        "      <td>BAAI/bge-large-en-v1.5</td>\n",
        "      <td>1024</td>\n",
        "      <td>Large English model, v1.5</td>\n",
        "      <td>1.200</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>14</th>\n",
+       "      <th>19</th>\n",
        "      <td>thenlper/gte-large</td>\n",
        "      <td>1024</td>\n",
        "      <td>Large general text embeddings model</td>\n",
        "      <td>1.200</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>15</th>\n",
+       "      <th>20</th>\n",
        "      <td>intfloat/multilingual-e5-large</td>\n",
        "      <td>1024</td>\n",
        "      <td>Multilingual model, e5-large. Recommend using ...</td>\n",
@@ -187,40 +222,50 @@
        "0                              BAAI/bge-small-en-v1.5   384   \n",
        "1                              BAAI/bge-small-zh-v1.5   512   \n",
        "2              sentence-transformers/all-MiniLM-L6-v2   384   \n",
-       "3                  jinaai/jina-embeddings-v2-small-en   512   \n",
-       "4                                   BAAI/bge-small-en   384   \n",
-       "5                               BAAI/bge-base-en-v1.5   768   \n",
-       "6   sentence-transformers/paraphrase-multilingual-...   384   \n",
-       "7                                    BAAI/bge-base-en   768   \n",
-       "8                        nomic-ai/nomic-embed-text-v1   768   \n",
-       "9                      nomic-ai/nomic-embed-text-v1.5   768   \n",
-       "10                  jinaai/jina-embeddings-v2-base-en   768   \n",
-       "11                 mixedbread-ai/mxbai-embed-large-v1  1024   \n",
-       "12  sentence-transformers/paraphrase-multilingual-...   768   \n",
-       "13                             BAAI/bge-large-en-v1.5  1024   \n",
-       "14                                 thenlper/gte-large  1024   \n",
-       "15                     intfloat/multilingual-e5-large  1024   \n",
+       "3                 snowflake/snowflake-arctic-embed-xs   384   \n",
+       "4                  jinaai/jina-embeddings-v2-small-en   512   \n",
+       "5                  snowflake/snowflake-arctic-embed-s   384   \n",
+       "6                                   BAAI/bge-small-en   384   \n",
+       "7                               BAAI/bge-base-en-v1.5   768   \n",
+       "8   sentence-transformers/paraphrase-multilingual-...   384   \n",
+       "9                                    BAAI/bge-base-en   768   \n",
+       "10                 snowflake/snowflake-arctic-embed-m   768   \n",
+       "11                  jinaai/jina-embeddings-v2-base-en   768   \n",
+       "12                       nomic-ai/nomic-embed-text-v1   768   \n",
+       "13                     nomic-ai/nomic-embed-text-v1.5   768   \n",
+       "14            snowflake/snowflake-arctic-embed-m-long   768   \n",
+       "15                 mixedbread-ai/mxbai-embed-large-v1  1024   \n",
+       "16  sentence-transformers/paraphrase-multilingual-...   768   \n",
+       "17                 snowflake/snowflake-arctic-embed-l  1024   \n",
+       "18                             BAAI/bge-large-en-v1.5  1024   \n",
+       "19                                 thenlper/gte-large  1024   \n",
+       "20                     intfloat/multilingual-e5-large  1024   \n",
        "\n",
        "                                          description  size_in_GB  \n",
        "0                      Fast and Default English model       0.067  \n",
        "1                  Fast and recommended Chinese model       0.090  \n",
        "2            Sentence Transformer model, MiniLM-L6-v2       0.090  \n",
-       "3   English embedding model supporting 8192 sequen...       0.120  \n",
-       "4                                  Fast English model       0.130  \n",
-       "5                            Base English model, v1.5       0.210  \n",
-       "6   Sentence Transformer model, paraphrase-multili...       0.220  \n",
-       "7                                  Base English model       0.420  \n",
-       "8                   8192 context length english model       0.520  \n",
-       "9                   8192 context length english model       0.520  \n",
-       "10  English embedding model supporting 8192 sequen...       0.520  \n",
-       "11  MixedBread Base sentence embedding model, does...       0.640  \n",
-       "12  Sentence-transformers model for tasks like clu...       1.000  \n",
-       "13                          Large English model, v1.5       1.200  \n",
-       "14                Large general text embeddings model       1.200  \n",
-       "15  Multilingual model, e5-large. Recommend using ...       2.240  "
+       "3   Based on all-MiniLM-L6-v2 model with only 22m ...       0.090  \n",
+       "4   English embedding model supporting 8192 sequen...       0.120  \n",
+       "5   Based on infloat/e5-small-unsupervised, does n...       0.130  \n",
+       "6                                  Fast English model       0.130  \n",
+       "7                            Base English model, v1.5       0.210  \n",
+       "8   Sentence Transformer model, paraphrase-multili...       0.220  \n",
+       "9                                  Base English model       0.420  \n",
+       "10  Based on intfloat/e5-base-unsupervised model, ...       0.430  \n",
+       "11  English embedding model supporting 8192 sequen...       0.520  \n",
+       "12                  8192 context length english model       0.520  \n",
+       "13                  8192 context length english model       0.520  \n",
+       "14  Based on nomic-ai/nomic-embed-text-v1-unsuperv...       0.540  \n",
+       "15  MixedBread Base sentence embedding model, does...       0.640  \n",
+       "16  Sentence-transformers model for tasks like clu...       1.000  \n",
+       "17  Based on intfloat/e5-large-unsupervised, large...       1.020  \n",
+       "18                          Large English model, v1.5       1.200  \n",
+       "19                Large general text embeddings model       1.200  \n",
+       "20  Multilingual model, e5-large. Recommend using ...       2.240  "
       ]
      },
-     "execution_count": 6,
+     "execution_count": 2,
      "metadata": {},
      "output_type": "execute_result"
     }

docs/qdrant/Binary_Quantization_with_Qdrant.ipynb

@@ -33,8 +33,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
+   "execution_count": 1,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:33:03.013948Z",
+     "start_time": "2024-04-01T16:33:01.019043Z"
+    }
+   },
    "outputs": [],
    "source": [
     "!pip install qdrant-client pandas dataset --quiet --upgrade"
@@ -42,12 +47,34 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
+   "execution_count": 2,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:33:03.914729Z",
+     "start_time": "2024-04-01T16:33:03.015394Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/joein/work/qdrant/fastembed/venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
    "source": [
+    "import os\n",
+    "import random\n",
+    "import time\n",
+    "\n",
+    "import numpy as np\n",
     "import pandas as pd\n",
-    "from qdrant_client import QdrantClient, models"
+    "from qdrant_client import QdrantClient, models\n",
+    "\n",
+    "random.seed(37)\n",
+    "np.random.seed(37)"
    ]
   },
   {
@@ -61,41 +88,59 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
+   "execution_count": 3,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:33:09.085853Z",
+     "start_time": "2024-04-01T16:33:03.912688Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "100000"
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "import datasets\n",
     "\n",
     "dataset = datasets.load_dataset(\n",
     "    \"Qdrant/dbpedia-entities-openai3-text-embedding-3-small-1536-100K\", split=\"train\"\n",
-    ")"
+    ")\n",
+    "len(dataset)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "len(dataset)\n",
-    "# dataset[0]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
+   "execution_count": 4,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:33:09.176212Z",
+     "start_time": "2024-04-01T16:33:09.084550Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "True"
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "client = QdrantClient(\n",
-    "    timeout=600,\n",
     "    prefer_grpc=True,\n",
     ")\n",
     "\n",
     "collection_name = \"binary-quantization\"\n",
     "client.recreate_collection(\n",
-    "    collection_name=f\"{collection_name}\",\n",
+    "    collection_name=collection_name,\n",
     "    vectors_config=models.VectorParams(\n",
     "        size=1536,\n",
     "        distance=models.Distance.DOT,\n",
@@ -109,29 +154,48 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
+   "execution_count": 5,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:34:13.383986Z",
+     "start_time": "2024-04-01T16:33:09.175725Z"
+    }
+   },
    "outputs": [],
    "source": [
-    "import os\n",
+    "def iter_dataset(dataset):\n",
+    "    for point in dataset:\n",
+    "        yield point[\"openai\"], {\"text\": point[\"text\"]}\n",
     "\n",
     "\n",
-    "bs = 1000\n",
-    "for i in range(0, len(dataset), bs):\n",
-    "    client.upload_collection(\n",
-    "        collection_name=collection_name,\n",
-    "        ids=range(i, i + bs),\n",
-    "        vectors=dataset[i : i + bs][\"openai\"],\n",
-    "        payload=[{\"text\": x} for x in dataset[i : i + bs][\"text\"]],\n",
-    "        parallel=max(1, (os.cpu_count() // 2)),\n",
-    "    )"
+    "vectors, payload = zip(*iter_dataset(dataset))\n",
+    "client.upload_collection(\n",
+    "    collection_name=collection_name,\n",
+    "    vectors=vectors,\n",
+    "    payload=payload,\n",
+    "    parallel=max(1, (os.cpu_count() // 2)),\n",
+    ")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
+   "execution_count": 6,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:34:13.390886Z",
+     "start_time": "2024-04-01T16:34:13.385961Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "{'status': <CollectionStatus.YELLOW: 'yellow'>,\n 'optimizer_status': <OptimizersStatusOneOf.OK: 'ok'>,\n 'vectors_count': 116640,\n 'indexed_vectors_count': 43520,\n 'points_count': 116640,\n 'segments_count': 6,\n 'config': {'params': {'vectors': {'size': 1536,\n    'distance': <Distance.DOT: 'Dot'>,\n    'hnsw_config': None,\n    'quantization_config': None,\n    'on_disk': True},\n   'shard_number': 1,\n   'sharding_method': None,\n   'replication_factor': 1,\n   'write_consistency_factor': 1,\n   'read_fan_out_factor': None,\n   'on_disk_payload': True,\n   'sparse_vectors': None},\n  'hnsw_config': {'m': 16,\n   'ef_construct': 100,\n   'full_scan_threshold': 10000,\n   'max_indexing_threads': 0,\n   'on_disk': False,\n   'payload_m': None},\n  'optimizer_config': {'deleted_threshold': 0.2,\n   'vacuum_min_vector_number': 1000,\n   'default_segment_number': 0,\n   'max_segment_size': None,\n   'memmap_threshold': None,\n   'indexing_threshold': 20000,\n   'flush_interval_sec': 5,\n   'max_optimization_threads': None},\n  'wal_config': {'wal_capacity_mb': 32, 'wal_segments_ahead': 0},\n  'quantization_config': {'binary': {'always_ram': True}}},\n 'payload_schema': {}}"
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "collection_info = client.get_collection(collection_name=f\"{collection_name}\")\n",
     "collection_info.dict()"
@@ -150,30 +214,41 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
+   "execution_count": 7,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:34:13.453626Z",
+     "start_time": "2024-04-01T16:34:13.391567Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "[89391,\n 79659,\n 12006,\n 80978,\n 87219,\n 97885,\n 83155,\n 67504,\n 4645,\n 82711,\n 48395,\n 57375,\n 69208,\n 14136,\n 89515,\n 59880,\n 78730,\n 36952,\n 49620,\n 96486,\n 55473,\n 58179,\n 18926,\n 6489,\n 11931,\n 54146,\n 9850,\n 71259,\n 37825,\n 47331,\n 84964,\n 92399,\n 56669,\n 77042,\n 73744,\n 47993,\n 83780,\n 92429,\n 75114,\n 4463,\n 69030,\n 81185,\n 27950,\n 66217,\n 54652,\n 8260,\n 1151,\n 993,\n 85954,\n 66863,\n 47303,\n 8992,\n 92688,\n 76030,\n 29472,\n 3077,\n 42454,\n 46120,\n 69140,\n 20877,\n 2844,\n 95423,\n 1770,\n 28568,\n 96448,\n 94227,\n 40837,\n 91684,\n 29785,\n 66936,\n 85121,\n 39546,\n 81910,\n 5514,\n 37068,\n 35731,\n 93990,\n 26685,\n 63076,\n 18762,\n 27922,\n 34916,\n 80976,\n 83189,\n 6328,\n 57508,\n 58860,\n 13758,\n 72976,\n 85030,\n 332,\n 34963,\n 85009,\n 31344,\n 11560,\n 58108,\n 85163,\n 17064,\n 44712,\n 45962]"
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
-    "import random\n",
-    "from random import randint\n",
-    "\n",
-    "random.seed(37)\n",
-    "\n",
-    "query_indices = [randint(0, len(dataset)) for _ in range(100)]\n",
+    "query_indices = random.sample(range(len(dataset)), 100)\n",
     "query_dataset = dataset[query_indices]\n",
     "query_indices"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
+   "execution_count": 8,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:34:13.453928Z",
+     "start_time": "2024-04-01T16:34:13.452405Z"
+    }
+   },
    "outputs": [],
    "source": [
     "## Add Gaussian noise to any vector\n",
-    "import numpy as np\n",
-    "\n",
-    "np.random.seed(37)\n",
     "\n",
     "\n",
     "def add_noise(vector, noise=0.05):\n",
@@ -182,76 +257,98 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
+   "execution_count": 9,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:34:13.457839Z",
+     "start_time": "2024-04-01T16:34:13.455431Z"
+    }
+   },
    "outputs": [],
    "source": [
-    "import time\n",
-    "\n",
-    "\n",
     "def correct(results, text):\n",
-    "    result_texts = [x.payload[\"text\"] for x in results]\n",
-    "    return text in result_texts\n",
+    "    return text in [x.payload[\"text\"] for x in results]\n",
     "\n",
     "\n",
     "def count_correct(query_dataset, limit=1, oversampling=1, rescore=False):\n",
     "    correct_results = 0\n",
-    "    for qv, text in zip(query_dataset[\"openai\"], query_dataset[\"text\"]):\n",
+    "    for query_vector, text in zip(query_dataset[\"openai\"], query_dataset[\"text\"]):\n",
     "        results = client.search(\n",
     "            collection_name=collection_name,\n",
-    "            query_vector=add_noise(np.array(qv)),\n",
+    "            query_vector=add_noise(np.array(query_vector)),\n",
     "            limit=limit,\n",
     "            search_params=models.SearchParams(\n",
     "                quantization=models.QuantizationSearchParams(\n",
-    "                    ignore=False,\n",
     "                    rescore=rescore,\n",
     "                    oversampling=oversampling,\n",
     "                )\n",
     "            ),\n",
     "        )\n",
     "        correct_results += correct(results, text)\n",
-    "    return correct_results\n",
-    "\n",
-    "\n",
-    "limit_grid = [1, 3, 5, 10, 20, 50]\n",
-    "# limit_grid = [1, 3, 5]\n",
-    "oversampling_grid = [1.0, 1.5, 2.0, 3.0, 5.0]\n",
-    "# oversampling_grid = [1.0, 1.5, 2.0]\n",
-    "rescore_grid = [False, True]\n",
-    "results = []\n",
-    "for limit in limit_grid:\n",
-    "    for oversampling in oversampling_grid:\n",
-    "        for rescore in rescore_grid:\n",
-    "            # print(f\"limit={limit}, oversampling={oversampling}, rescore={rescore}\")\n",
-    "            start = time.time()\n",
-    "            correct_results = count_correct(\n",
-    "                query_dataset, limit=limit, oversampling=oversampling, rescore=rescore\n",
-    "            )\n",
-    "            end = time.time()\n",
-    "            results.append(\n",
-    "                {\n",
-    "                    \"limit\": limit,\n",
-    "                    \"oversampling\": oversampling,\n",
-    "                    \"rescore\": rescore,\n",
-    "                    \"correct\": correct_results,\n",
-    "                    \"total queries\": len(query_dataset[\"text\"]),\n",
-    "                    \"time\": end - start,\n",
-    "                }\n",
-    "            )\n",
-    "\n",
-    "results_df = pd.DataFrame(results)\n",
-    "results_df"
+    "    return correct_results"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
+   "execution_count": 18,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:40:48.347002Z",
+     "start_time": "2024-04-01T16:40:42.228551Z"
+    },
+    "collapsed": false
+   },
    "outputs": [],
    "source": [
-    "df = results_df.copy()\n",
-    "df[\"candidates\"] = df[\"oversampling\"] * df[\"limit\"]\n",
-    "df[[\"candidates\", \"rescore\", \"time\"]]\n",
+    "limit_grid = [1, 3, 10, 20, 50]\n",
+    "oversampling_grid = [1.0, 3.0, 5.0]\n",
+    "rescore_grid = [False, True]\n",
+    "results = []\n",
+    "\n",
+    "for limit in limit_grid:\n",
+    "    for oversampling in oversampling_grid:\n",
+    "        for rescore in rescore_grid:\n",
+    "            start = time.perf_counter()\n",
+    "            correct_results = count_correct(\n",
+    "                query_dataset, limit=limit, oversampling=oversampling, rescore=rescore\n",
+    "            )\n",
+    "            end = time.perf_counter()\n",
+    "            results.append(\n",
+    "                {\n",
+    "                    \"limit\": limit,\n",
+    "                    \"oversampling\": oversampling,\n",
+    "                    \"candidates\": int(oversampling * limit),\n",
+    "                    \"rescore\": rescore,\n",
+    "                    \"accuracy\": correct_results / 100,\n",
+    "                    \"total queries\": len(query_dataset[\"text\"]),\n",
+    "                    \"time\": end - start,\n",
+    "                }\n",
+    "            )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:41:55.445405Z",
+     "start_time": "2024-04-01T16:41:55.442687Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>candidates</th>\n      <th>rescore</th>\n      <th>accuracy</th>\n      <th>time</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>1</td>\n      <td>False</td>\n      <td>0.90</td>\n      <td>0.221826</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>1</td>\n      <td>True</td>\n      <td>0.91</td>\n      <td>0.134167</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>3</td>\n      <td>False</td>\n      <td>0.88</td>\n      <td>0.115299</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>3</td>\n      <td>True</td>\n      <td>0.97</td>\n      <td>0.209320</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>5</td>\n      <td>False</td>\n      <td>0.84</td>\n      <td>0.154485</td>\n    </tr>\n    <tr>\n      <th>5</th>\n      <td>5</td>\n      <td>True</td>\n      <td>0.91</td>\n      <td>0.124424</td>\n    </tr>\n    <tr>\n      <th>6</th>\n      <td>3</td>\n      <td>False</td>\n      <td>0.99</td>\n      <td>0.121695</td>\n    </tr>\n    <tr>\n      <th>7</th>\n      <td>3</td>\n      <td>True</td>\n      <td>0.96</td>\n      <td>0.123257</td>\n    </tr>\n    <tr>\n      <th>8</th>\n      <td>9</td>\n      <td>False</td>\n      <td>0.94</td>\n      <td>0.119629</td>\n    </tr>\n    <tr>\n      <th>9</th>\n      <td>9</td>\n      <td>True</td>\n      <td>0.98</td>\n      <td>0.119372</td>\n    </tr>\n    <tr>\n      <th>10</th>\n      <td>15</td>\n      <td>False</td>\n      <td>0.90</td>\n      <td>0.121621</td>\n    </tr>\n    <tr>\n      <th>11</th>\n      <td>15</td>\n      <td>True</td>\n      <td>0.97</td>\n      <td>0.125466</td>\n    </tr>\n    <tr>\n      <th>12</th>\n      <td>10</td>\n      <td>False</td>\n      <td>0.93</td>\n      <td>0.135910</td>\n    </tr>\n    <tr>\n      <th>13</th>\n      <td>10</td>\n      <td>True</td>\n      <td>0.95</td>\n      <td>0.138135</td>\n    </tr>\n    <tr>\n      <th>14</th>\n      <td>30</td>\n      <td>False</td>\n      <td>0.94</td>\n      <td>0.177928</td>\n    </tr>\n    <tr>\n      <th>15</th>\n      <td>30</td>\n      <td>True</td>\n      <td>0.98</td>\n      <td>0.254588</td>\n    </tr>\n    <tr>\n      <th>16</th>\n      <td>50</td>\n      <td>False</td>\n      <td>0.94</td>\n      <td>0.268659</td>\n    </tr>\n    <tr>\n      <th>17</th>\n      <td>50</td>\n      <td>True</td>\n      <td>0.96</td>\n      <td>0.269792</td>\n    </tr>\n    <tr>\n      <th>18</th>\n      <td>20</td>\n      <td>False</td>\n      <td>0.96</td>\n      <td>0.249941</td>\n    </tr>\n    <tr>\n      <th>19</th>\n      <td>20</td>\n      <td>True</td>\n      <td>0.96</td>\n      <td>0.247138</td>\n    </tr>\n    <tr>\n      <th>20</th>\n      <td>60</td>\n      <td>False</td>\n      <td>0.97</td>\n      <td>0.251301</td>\n    </tr>\n    <tr>\n      <th>21</th>\n      <td>60</td>\n      <td>True</td>\n      <td>0.98</td>\n      <td>0.256504</td>\n    </tr>\n    <tr>\n      <th>22</th>\n      <td>100</td>\n      <td>False</td>\n      <td>0.98</td>\n      <td>0.270049</td>\n    </tr>\n    <tr>\n      <th>23</th>\n      <td>100</td>\n      <td>True</td>\n      <td>0.97</td>\n      <td>0.248972</td>\n    </tr>\n    <tr>\n      <th>24</th>\n      <td>50</td>\n      <td>False</td>\n      <td>0.97</td>\n      <td>0.306356</td>\n    </tr>\n    <tr>\n      <th>25</th>\n      <td>50</td>\n      <td>True</td>\n      <td>0.98</td>\n      <td>0.257544</td>\n    </tr>\n    <tr>\n      <th>26</th>\n      <td>150</td>\n      <td>False</td>\n      <td>0.98</td>\n      <td>0.238811</td>\n    </tr>\n    <tr>\n      <th>27</th>\n      <td>150</td>\n      <td>True</td>\n      <td>0.99</td>\n      <td>0.263939</td>\n    </tr>\n    <tr>\n      <th>28</th>\n      <td>250</td>\n      <td>False</td>\n      <td>0.99</td>\n      <td>0.256558</td>\n    </tr>\n    <tr>\n      <th>29</th>\n      <td>250</td>\n      <td>True</td>\n      <td>1.00</td>\n      <td>0.335823</td>\n    </tr>\n  </tbody>\n</table>\n</div>",
+      "text/plain": "    candidates  rescore  accuracy      time\n0            1    False      0.90  0.221826\n1            1     True      0.91  0.134167\n2            3    False      0.88  0.115299\n3            3     True      0.97  0.209320\n4            5    False      0.84  0.154485\n5            5     True      0.91  0.124424\n6            3    False      0.99  0.121695\n7            3     True      0.96  0.123257\n8            9    False      0.94  0.119629\n9            9     True      0.98  0.119372\n10          15    False      0.90  0.121621\n11          15     True      0.97  0.125466\n12          10    False      0.93  0.135910\n13          10     True      0.95  0.138135\n14          30    False      0.94  0.177928\n15          30     True      0.98  0.254588\n16          50    False      0.94  0.268659\n17          50     True      0.96  0.269792\n18          20    False      0.96  0.249941\n19          20     True      0.96  0.247138\n20          60    False      0.97  0.251301\n21          60     True      0.98  0.256504\n22         100    False      0.98  0.270049\n23         100     True      0.97  0.248972\n24          50    False      0.97  0.306356\n25          50     True      0.98  0.257544\n26         150    False      0.98  0.238811\n27         150     True      0.99  0.263939\n28         250    False      0.99  0.256558\n29         250     True      1.00  0.335823"
+     },
+     "execution_count": 22,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "df = pd.DataFrame(results)\n",
+    "df[[\"candidates\", \"rescore\", \"accuracy\", \"time\"]]\n",
     "# df.to_csv(\"candidates-rescore-time.csv\", index=False)"
    ]
   }

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

@@ -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,36 @@ 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",
+        ]
+        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 +173,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 +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.
 
@@ -219,9 +220,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/common/onnx_model.py

@@ -1,19 +1,33 @@
 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
 
-from fastembed.common.model_management import locate_model_file
 from fastembed.common.models import load_tokenizer
 from fastembed.common.utils import iter_batch
 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
@@ -34,11 +48,26 @@ class OnnxModel(Generic[T]):
         """
         return onnx_input
 
-    def load_onnx_model(self, model_dir: Path, threads: Optional[int], max_length: int) -> None:
-        model_path = locate_model_file(model_dir, ["model.onnx", "model_optimized.onnx"])
+    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"]
+
+        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
@@ -47,7 +76,7 @@ class OnnxModel(Generic[T]):
             so.intra_op_num_threads = threads
             so.inter_op_num_threads = threads
 
-        self.tokenizer = load_tokenizer(model_dir=model_dir, max_length=max_length)
+        self.tokenizer = load_tokenizer(model_dir=model_dir)
         self.model = ort.InferenceSession(
             str(model_path), providers=onnx_providers, sess_options=so
         )

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
 
@@ -15,6 +15,7 @@ supported_splade_models = [
         "sources": {
             "hf": "Qdrant/SPLADE_PP_en_v1",
         },
+        "model_file": "model.onnx",
     },
     {
         "model": "prithivida/Splade_PP_en_v1",
@@ -24,6 +25,7 @@ supported_splade_models = [
         "sources": {
             "hf": "Qdrant/SPLADE_PP_en_v1",
         },
+        "model_file": "model.onnx",
     },
 ]
 
@@ -61,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,
     ):
         """
@@ -77,14 +80,19 @@ class SpladePP(SparseTextEmbeddingBase, OnnxModel[SparseEmbedding]):
 
         super().__init__(model_name, cache_dir, threads, **kwargs)
 
-        self.model_name = model_name
-        self._model_description = self._get_model_description(model_name)
+        model_description = self._get_model_description(model_name)
+        cache_dir = define_cache_dir(cache_dir)
 
-        self._cache_dir = define_cache_dir(cache_dir)
-        self._model_dir = self.download_model(self._model_description, self._cache_dir)
-        self._max_length = 512
+        model_dir = self.download_model(
+            model_description, cache_dir, local_files_only=self._local_files_only
+        )
 
-        self.load_onnx_model(self._model_dir, self.threads, self._max_length)
+        self.load_onnx_model(
+            model_dir=model_dir,
+            model_file=model_description["model_file"],
+            threads=threads,
+            providers=providers,
+        )
 
     def embed(
         self,
@@ -110,7 +118,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxModel[SparseEmbedding]):
         """
         yield from self._embed_documents(
             model_name=self.model_name,
-            cache_dir=str(self._cache_dir),
+            cache_dir=str(self.cache_dir),
             documents=documents,
             batch_size=batch_size,
             parallel=parallel,

fastembed/text/e5_onnx_embedding.py

@@ -15,6 +15,8 @@ supported_multilingual_e5_models = [
             "url": "https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
             "hf": "qdrant/multilingual-e5-large-onnx",
         },
+        "model_file": "model.onnx",
+        "additional_files": ["model.onnx_data"],
     },
     {
         "model": "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
@@ -24,6 +26,7 @@ supported_multilingual_e5_models = [
         "sources": {
             "hf": "xenova/paraphrase-multilingual-mpnet-base-v2",
         },
+        "model_file": "onnx/model.onnx",
     },
 ]
 

fastembed/text/jina_onnx_embedding.py

@@ -13,6 +13,7 @@ supported_jina_models = [
         "description": "English embedding model supporting 8192 sequence length",
         "size_in_GB": 0.52,
         "sources": {"hf": "xenova/jina-embeddings-v2-base-en"},
+        "model_file": "onnx/model.onnx",
     },
     {
         "model": "jinaai/jina-embeddings-v2-small-en",
@@ -20,6 +21,7 @@ supported_jina_models = [
         "description": "English embedding model supporting 8192 sequence length",
         "size_in_GB": 0.12,
         "sources": {"hf": "xenova/jina-embeddings-v2-small-en"},
+        "model_file": "onnx/model.onnx",
     },
 ]
 

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
@@ -16,6 +16,7 @@ supported_onnx_models = [
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz",
         },
+        "model_file": "model_optimized.onnx",
     },
     {
         "model": "BAAI/bge-base-en-v1.5",
@@ -26,6 +27,7 @@ supported_onnx_models = [
             "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
             "hf": "qdrant/bge-base-en-v1.5-onnx-q",
         },
+        "model_file": "model_optimized.onnx",
     },
     {
         "model": "BAAI/bge-large-en-v1.5",
@@ -35,6 +37,7 @@ supported_onnx_models = [
         "sources": {
             "hf": "qdrant/bge-large-en-v1.5-onnx",
         },
+        "model_file": "model.onnx",
     },
     {
         "model": "BAAI/bge-small-en",
@@ -44,18 +47,8 @@ supported_onnx_models = [
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
         },
+        "model_file": "model_optimized.onnx",
     },
-    # {
-    #     "model": "BAAI/bge-small-en",
-    #     "dim": 384,
-    #     "description": "Fast English model",
-    #     "size_in_GB": 0.2,
-    #     "hf_sources": [],
-    #     "compressed_url_sources": [
-    #         "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-en.tar.gz",
-    #         "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz"
-    #     ]
-    # },
     {
         "model": "BAAI/bge-small-en-v1.5",
         "dim": 384,
@@ -64,6 +57,7 @@ supported_onnx_models = [
         "sources": {
             "hf": "qdrant/bge-small-en-v1.5-onnx-q",
         },
+        "model_file": "model_optimized.onnx",
     },
     {
         "model": "BAAI/bge-small-zh-v1.5",
@@ -73,6 +67,7 @@ supported_onnx_models = [
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
         },
+        "model_file": "model_optimized.onnx",
     },
     {
         "model": "sentence-transformers/all-MiniLM-L6-v2",
@@ -83,6 +78,7 @@ supported_onnx_models = [
             "url": "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz",
             "hf": "qdrant/all-MiniLM-L6-v2-onnx",
         },
+        "model_file": "model.onnx",
     },
     {
         "model": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
@@ -92,6 +88,7 @@ supported_onnx_models = [
         "sources": {
             "hf": "qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q",
         },
+        "model_file": "model_optimized.onnx",
     },
     {
         "model": "nomic-ai/nomic-embed-text-v1",
@@ -101,6 +98,7 @@ supported_onnx_models = [
         "sources": {
             "hf": "nomic-ai/nomic-embed-text-v1",
         },
+        "model_file": "onnx/model.onnx",
     },
     {
         "model": "nomic-ai/nomic-embed-text-v1.5",
@@ -110,6 +108,17 @@ supported_onnx_models = [
         "sources": {
             "hf": "nomic-ai/nomic-embed-text-v1.5",
         },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "nomic-ai/nomic-embed-text-v1.5-Q",
+        "dim": 768,
+        "description": "Quantized 8192 context length english model",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1.5",
+        },
+        "model_file": "onnx/model_quantized.onnx",
     },
     {
         "model": "thenlper/gte-large",
@@ -119,20 +128,8 @@ supported_onnx_models = [
         "sources": {
             "hf": "qdrant/gte-large-onnx",
         },
+        "model_file": "model.onnx",
     },
-    # {
-    #     "model": "sentence-transformers/all-MiniLM-L6-v2",
-    #     "dim": 384,
-    #     "description": "Sentence Transformer model, MiniLM-L6-v2",
-    #     "size_in_GB": 0.09,
-    #     "hf_sources": [
-    #         "qdrant/all-MiniLM-L6-v2-onnx"
-    #     ],
-    #     "compressed_url_sources": [
-    #         "https://storage.googleapis.com/qdrant-fastembed/fast-all-MiniLM-L6-v2.tar.gz",
-    #         "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz"
-    #     ]
-    # }
     {
         "model": "mixedbread-ai/mxbai-embed-large-v1",
         "dim": 1024,
@@ -141,6 +138,57 @@ supported_onnx_models = [
         "sources": {
             "hf": "mixedbread-ai/mxbai-embed-large-v1",
         },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-xs",
+        "dim": 384,
+        "description": "Based on all-MiniLM-L6-v2 model with only 22m parameters, ideal for latency/TCO budgets.",
+        "size_in_GB": 0.09,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-xs",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-s",
+        "dim": 384,
+        "description": "Based on infloat/e5-small-unsupervised, does not trade off retrieval accuracy for its small size.",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-s",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-m",
+        "dim": 768,
+        "description": "Based on intfloat/e5-base-unsupervised model, provides the best retrieval without slowing down inference.",
+        "size_in_GB": 0.43,
+        "sources": {
+            "hf": "Snowflake/snowflake-arctic-embed-m",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-m-long",
+        "dim": 768,
+        "description": "Based on nomic-ai/nomic-embed-text-v1-unsupervised model, 8192 context-length model",
+        "size_in_GB": 0.54,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-m-long",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-l",
+        "dim": 1024,
+        "description": "Based on intfloat/e5-large-unsupervised, large model for most accurate retrieval.",
+        "size_in_GB": 1.02,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-l",
+        },
+        "model_file": "onnx/model.onnx",
     },
 ]
 
@@ -163,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,
     ):
         """
@@ -179,14 +228,18 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxModel[np.ndarray]):
 
         super().__init__(model_name, cache_dir, threads, **kwargs)
 
-        self.model_name = model_name
-        self._model_description = self._get_model_description(model_name)
+        model_description = self._get_model_description(model_name)
+        cache_dir = define_cache_dir(cache_dir)
+        model_dir = self.download_model(
+            model_description, cache_dir, local_files_only=self._local_files_only
+        )
 
-        self._cache_dir = define_cache_dir(cache_dir)
-        self._model_dir = self.download_model(self._model_description, self._cache_dir)
-        self._max_length = 512
-
-        self.load_onnx_model(self._model_dir, self.threads, self._max_length)
+        self.load_onnx_model(
+            model_dir=model_dir,
+            model_file=model_description["model_file"],
+            threads=threads,
+            providers=providers,
+        )
 
     def embed(
         self,
@@ -212,7 +265,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxModel[np.ndarray]):
         """
         yield from self._embed_documents(
             model_name=self.model_name,
-            cache_dir=str(self._cache_dir),
+            cache_dir=str(self.cache_dir),
             documents=documents,
             batch_size=batch_size,
             parallel=parallel,

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 = [

tests/test_text_onnx_embeddings.py

@@ -26,16 +26,26 @@ CANONICAL_VECTOR_VALUES = {
     "nomic-ai/nomic-embed-text-v1.5": np.array(
         [-1.6531514e-02, 8.5380634e-05, -1.8171231e-01, -3.9333291e-03, 1.2763254e-02]
     ),
+    "nomic-ai/nomic-embed-text-v1.5-Q": np.array(
+        [-0.01554983,  0.0129992 , -0.17909265, -0.01062993,  0.00512859]
+    ),
     "thenlper/gte-large": np.array([-0.01920587, 0.00113156, -0.00708992, -0.00632304, -0.04025577]),
     "mixedbread-ai/mxbai-embed-large-v1": np.array([0.02295546, 0.03196154, 0.016512, -0.04031524, -0.0219634]),
+    "snowflake/snowflake-arctic-embed-xs": np.array([0.0092, 0.0619, 0.0196, 0.009, -0.0114]),
+    "snowflake/snowflake-arctic-embed-s": np.array([-0.0416, -0.0867, 0.0209, 0.0554, -0.0272]),
+    "snowflake/snowflake-arctic-embed-m": np.array([-0.0329, 0.0364, 0.0481, 0.0016, 0.0328]),
+    "snowflake/snowflake-arctic-embed-m-long": np.array(
+        [0.0080, -0.0266, -0.0335, 0.0282, 0.0143]
+    ),
+    "snowflake/snowflake-arctic-embed-l": np.array([0.0189, -0.0673, 0.0183, 0.0124, 0.0146]),
 }
 
 
 def test_embedding():
-    is_ubuntu_ci = os.getenv("IS_UBUNTU_CI")
+    is_ci = os.getenv("CI")
 
     for model_desc in TextEmbedding.list_supported_models():
-        if is_ubuntu_ci == "false" and model_desc["size_in_GB"] > 1:
+        if not is_ci and model_desc["size_in_GB"] > 1:
             continue
 
         dim = model_desc["dim"]