debug: try out with markers

fix: fix type hint for 3.8
fix: update poetry.lock
2026-09-24 06:57:51 -05:00 · 2024-05-01 10:51:32 +02:00 · 2024-04-26 21:45:24 +02:00 · 2024-04-26 21:45:24 +02:00 · 2024-04-26 21:45:24 +02:00 · 2024-04-26 21:45:02 +02:00
14 changed files with 1044 additions and 745 deletions
@@ -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
@@ -3,6 +3,8 @@ name: Tests
 on:
  push:
    branches: [ master, main ]
+  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'
@@ -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"
    }
@@ -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)"
   ]
  }
@@ -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,20 +87,33 @@ 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,
    ) -> 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,
        )

@@ -175,6 +171,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),
@@ -219,9 +218,14 @@ 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
+                    )
                )
            except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
                logger.error(
@@ -1,16 +1,28 @@
 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")

@@ -34,11 +46,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[Union[str, Tuple[str, Dict[Any, Any]]]]] = 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 +74,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
        )
@@ -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",
    },
 ]

@@ -77,14 +79,16 @@ 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)

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

    def embed(
        self,
@@ -110,7 +114,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,
@@ -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",
    },
 ]

@@ -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",
    },
 ]

@@ -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",
    },
 ]

@@ -179,14 +227,15 @@ 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)

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

    def embed(
        self,
@@ -212,7 +261,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,
@@ -12,8 +12,16 @@ keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-trans

 [tool.poetry.dependencies]
 python = ">=3.8.0,<3.13"
-onnx = "^1.15.0"
-onnxruntime = "^1.17.0"
+onnx = [
+    {version = "^1.15.0", optional = true, markers = "extra != 'gpu'"}
+]
+onnxruntime = [
+    {version = "^1.17.0", optional = true, markers = "extra != 'gpu'"}
+]
+onnxruntime-gpu = [
+    { version = "^1.17.0", optional = true, python = "<3.13", markers = "extra == 'gpu'" }
+]
+
 tqdm = "^4.66"
 requests = "^2.31"
 tokenizers = "^0.15.1"
@@ -37,6 +45,8 @@ pillow = "^10.2.0"
 cairosvg = "^2.7.1"
 mknotebooks = "^0.8.0"

+[tool.poetry.extras]
+gpu = ["onnxruntime-gpu"]

 [build-system]
 requires = ["poetry-core"]
@@ -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"]
Author	SHA1	Message	Date
George Panchuk	d9e85fe670	debug: try out with markers	2024-05-01 10:51:32 +02:00
George Panchuk	383e12c760	fix: fix type hint for 3.8	2024-04-26 21:45:24 +02:00
George Panchuk	0d6801d960	fix: update poetry.lock	2024-04-26 21:45:24 +02:00
George Panchuk	bff1b67637	new: add gpu support, allow users to override providers	2024-04-26 21:45:24 +02:00
George	5cde2898bc	new: remove slurm environment variables (#213 )	2024-04-26 21:45:02 +02:00
AnushandGeorge	ab7a99a748	feat: Quantized models (#201 ) * feat: Quantized models * refactor: use model_file for GCS * refactoring: refactor model downloading (#209) * refactoring: refactor model downloading * refactor: update docstring Co-authored-by: Anush <anushshetty90@gmail.com> * Update fastembed/common/model_management.py Co-authored-by: George <george.panchuk@qdrant.tech> * fix: model_file for Snowflake models --------- Co-authored-by: George <george.panchuk@qdrant.tech>	2024-04-26 17:27:16 +02:00
Anush	466886a317	ci: Schedule python-tests.yml (#211 ) * ci: Schedule python-tests.yml * ci: use emojis * ci: Bump action versions python-tests.yml * ci: python-tests.yml	2024-04-26 10:28:37 +05:30
Anush	cc4112d859	feat: Snowflake models (#207 ) * feat: Snowflake models * Added snowflake/snowflake-arctic-embed-m * docs: snowflake/snowflake-arctic-embed-m	2024-04-19 19:00:10 +05:30
Nirant	864217a7d9	Re-run of identical hardware and generate graphs (#205 )	2024-04-18 11:13:18 +05:30
Andrew Green	9bad44368e	Workaround for running on SLURM (#198 ) * Workaround for running on SLURM onnxruntime would usually get the number of threads from OMP_NUM_THREADS, but that isn't set on SLURM which handles the number of threads differently. This addition tries to figure out if we're running under SLURM, and if so sets the session options accordingly using SLURM environment variables instead. Tested working with latest versions of onnxruntime and fastembed on slurm 23.02.07 * onnxruntime requires number of threads be an integer Caused by me having mis-matched version from my machine and the slurm cluster :( * `os.getenv` returns None for unset environment variables, fix logic Instead of empty string, which I thought it did	2024-04-15 19:14:11 +05:30
dependabot[bot]andNirant	a5cab7a31a	build(deps): bump idna from 3.6 to 3.7 (#195 ) Bumps [idna](https://github.com/kjd/idna) from 3.6 to 3.7. - [Release notes](https://github.com/kjd/idna/releases) - [Changelog](https://github.com/kjd/idna/blob/master/HISTORY.rst) - [Commits](https://github.com/kjd/idna/compare/v3.6...v3.7) --- updated-dependencies: - dependency-name: idna dependency-type: indirect ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Nirant <NirantK@users.noreply.github.com>	2024-04-15 06:04:54 +05:30
Anush	e55c145924	chore: Exclude unused model repo files (#196 ) * chore: Exclude unused model files * fix: blob pattern	2024-04-12 20:15:02 +05:30
dependabot[bot]andNirant	ad02f60eea	build(deps-dev): bump pillow from 10.2.0 to 10.3.0 (#186 ) Bumps [pillow](https://github.com/python-pillow/Pillow) from 10.2.0 to 10.3.0. - [Release notes](https://github.com/python-pillow/Pillow/releases) - [Changelog](https://github.com/python-pillow/Pillow/blob/main/CHANGES.rst) - [Commits](https://github.com/python-pillow/Pillow/compare/10.2.0...10.3.0) --- updated-dependencies: - dependency-name: pillow dependency-type: direct:development ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Nirant <NirantK@users.noreply.github.com>	2024-04-08 16:21:32 +05:30
NirantandGeorge	2fcec07f1b	Add Python version to the ISSUE_TEMPLATE (#188 ) * Add Python version * Update .github/ISSUE_TEMPLATE/bug-report.yml Co-authored-by: George <george.panchuk@qdrant.tech> --------- Co-authored-by: George <george.panchuk@qdrant.tech>	2024-04-08 16:05:29 +05:30
George	f340a73a3e	refactoring: update binary quantization notebook (#180 )	2024-04-02 14:45:27 +05:30
Nirant	335f673f3a	Update bug-report.yml	2024-04-02 13:51:51 +05:30
Nirant	ee7ba0a536	Update bug-report.yml	2024-04-02 13:51:27 +05:30
George	912e95e5c8	fix: remove model archive if model extraction was not finished correctly (#179 )	2024-04-01 20:11:32 +02:00