bump version to 0.4.0

.github/workflows/python-tests.yml

@@ -1,5 +1,4 @@
 name: Tests
-run-name: Tests (gpu)
 
 on:
   push:
@@ -22,6 +21,8 @@ jobs:
           - '3.12.x'
         os:
           - ubuntu-latest
+          - macos-latest
+          - windows-latest
 
     runs-on: ${{ matrix.os }}
 
@@ -33,8 +34,6 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
-#      - name: Setup tmate session
-#        uses: mxschmitt/action-tmate@v3
       - name: Install dependencies
         run: |
           python -m pip install poetry

NOTICE

@@ -1,12 +0,0 @@
-Copyright 2024 Qdrant
-
-This product includes software developed by Qdrant
-
-This distribution includes the following Jina AI models, each with its respective license:
-- jinaai/jina-colbert-v2
-    - License: cc-by-nc-4.0
-
-These models are developed by Jina (https://jina.ai/) and are subject to Jina AI's licensing terms.
-
-Additional Notes:
-This project also includes third-party libraries with their respective licenses. Please refer to the documentation of each library for details regarding its usage and licensing terms.

docs/examples/Supported_Models.ipynb

@@ -2,14 +2,23 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 4,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2024-05-31T18:13:23.806907Z",
      "start_time": "2024-05-31T18:13:23.797078Z"
     }
    },
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "The autoreload extension is already loaded. To reload it, use:\n",
+      "  %reload_ext autoreload\n"
+     ]
+    }
+   ],
    "source": [
     "%load_ext autoreload\n",
     "%autoreload 2"
@@ -17,23 +26,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2024-05-31T18:14:31.147674Z",
      "start_time": "2024-05-31T18:14:31.134015Z"
     }
    },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/hossam/.pyenv/versions/.venv/lib/python3.10/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"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "import pandas as pd\n",
     "\n",
@@ -54,7 +54,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 5,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2024-05-31T18:13:25.863008Z",
@@ -86,7 +86,6 @@
        "      <th>model</th>\n",
        "      <th>dim</th>\n",
        "      <th>description</th>\n",
-       "      <th>license</th>\n",
        "      <th>size_in_GB</th>\n",
        "    </tr>\n",
        "  </thead>\n",
@@ -95,200 +94,175 @@
        "      <th>0</th>\n",
        "      <td>BAAI/bge-small-en-v1.5</td>\n",
        "      <td>384</td>\n",
-       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
-       "      <td>mit</td>\n",
+       "      <td>Fast and Default English model</td>\n",
        "      <td>0.067</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>1</th>\n",
        "      <td>BAAI/bge-small-zh-v1.5</td>\n",
        "      <td>512</td>\n",
-       "      <td>Text embeddings, Unimodal (text), Chinese, 512...</td>\n",
-       "      <td>mit</td>\n",
+       "      <td>Fast and recommended Chinese model</td>\n",
        "      <td>0.090</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>2</th>\n",
        "      <td>snowflake/snowflake-arctic-embed-xs</td>\n",
        "      <td>384</td>\n",
-       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
-       "      <td>apache-2.0</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>3</th>\n",
        "      <td>sentence-transformers/all-MiniLM-L6-v2</td>\n",
        "      <td>384</td>\n",
-       "      <td>Text embeddings, Unimodal (text), English, 256...</td>\n",
-       "      <td>apache-2.0</td>\n",
+       "      <td>Sentence Transformer model, MiniLM-L6-v2</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>Text embeddings, Unimodal (text), English, 819...</td>\n",
-       "      <td>apache-2.0</td>\n",
+       "      <td>English embedding model supporting 8192 sequen...</td>\n",
        "      <td>0.120</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>5</th>\n",
        "      <td>BAAI/bge-small-en</td>\n",
        "      <td>384</td>\n",
-       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
-       "      <td>mit</td>\n",
+       "      <td>Fast English model</td>\n",
        "      <td>0.130</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>6</th>\n",
        "      <td>snowflake/snowflake-arctic-embed-s</td>\n",
        "      <td>384</td>\n",
-       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
-       "      <td>apache-2.0</td>\n",
+       "      <td>Based on infloat/e5-small-unsupervised, does n...</td>\n",
        "      <td>0.130</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>7</th>\n",
        "      <td>nomic-ai/nomic-embed-text-v1.5-Q</td>\n",
        "      <td>768</td>\n",
-       "      <td>Text embeddings, Multimodal (text, image), Eng...</td>\n",
-       "      <td>apache-2.0</td>\n",
+       "      <td>Quantized 8192 context length english model</td>\n",
        "      <td>0.130</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>8</th>\n",
        "      <td>BAAI/bge-base-en-v1.5</td>\n",
        "      <td>768</td>\n",
-       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
-       "      <td>mit</td>\n",
+       "      <td>Base English model, v1.5</td>\n",
        "      <td>0.210</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>9</th>\n",
        "      <td>sentence-transformers/paraphrase-multilingual-...</td>\n",
        "      <td>384</td>\n",
-       "      <td>Text embeddings, Unimodal (text), Multilingual...</td>\n",
-       "      <td>apache-2.0</td>\n",
+       "      <td>Sentence Transformer model, paraphrase-multili...</td>\n",
        "      <td>0.220</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>10</th>\n",
        "      <td>Qdrant/clip-ViT-B-32-text</td>\n",
        "      <td>512</td>\n",
-       "      <td>Text embeddings, Multimodal (text&amp;image), Engl...</td>\n",
-       "      <td>mit</td>\n",
+       "      <td>CLIP text encoder</td>\n",
        "      <td>0.250</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>11</th>\n",
        "      <td>jinaai/jina-embeddings-v2-base-de</td>\n",
        "      <td>768</td>\n",
-       "      <td>Text embeddings, Unimodal (text), Multilingual...</td>\n",
-       "      <td>apache-2.0</td>\n",
+       "      <td>German embedding model supporting 8192 sequenc...</td>\n",
        "      <td>0.320</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>12</th>\n",
        "      <td>BAAI/bge-base-en</td>\n",
        "      <td>768</td>\n",
-       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
-       "      <td>mit</td>\n",
+       "      <td>Base English model</td>\n",
        "      <td>0.420</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>13</th>\n",
        "      <td>snowflake/snowflake-arctic-embed-m</td>\n",
        "      <td>768</td>\n",
-       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
-       "      <td>apache-2.0</td>\n",
+       "      <td>Based on intfloat/e5-base-unsupervised model, ...</td>\n",
        "      <td>0.430</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>14</th>\n",
        "      <td>nomic-ai/nomic-embed-text-v1.5</td>\n",
        "      <td>768</td>\n",
-       "      <td>Text embeddings, Multimodal (text, image), Eng...</td>\n",
-       "      <td>apache-2.0</td>\n",
+       "      <td>8192 context length english model</td>\n",
        "      <td>0.520</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>15</th>\n",
        "      <td>jinaai/jina-embeddings-v2-base-en</td>\n",
        "      <td>768</td>\n",
-       "      <td>Text embeddings, Unimodal (text), English, 819...</td>\n",
-       "      <td>apache-2.0</td>\n",
+       "      <td>English embedding model supporting 8192 sequen...</td>\n",
        "      <td>0.520</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>16</th>\n",
        "      <td>nomic-ai/nomic-embed-text-v1</td>\n",
        "      <td>768</td>\n",
-       "      <td>Text embeddings, Multimodal (text, image), Eng...</td>\n",
-       "      <td>apache-2.0</td>\n",
+       "      <td>8192 context length english model</td>\n",
        "      <td>0.520</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>17</th>\n",
        "      <td>snowflake/snowflake-arctic-embed-m-long</td>\n",
        "      <td>768</td>\n",
-       "      <td>Text embeddings, Unimodal (text), English, 204...</td>\n",
-       "      <td>apache-2.0</td>\n",
+       "      <td>Based on nomic-ai/nomic-embed-text-v1-unsuperv...</td>\n",
        "      <td>0.540</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>18</th>\n",
        "      <td>mixedbread-ai/mxbai-embed-large-v1</td>\n",
        "      <td>1024</td>\n",
-       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
-       "      <td>apache-2.0</td>\n",
+       "      <td>MixedBread Base sentence embedding model, does...</td>\n",
        "      <td>0.640</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>19</th>\n",
        "      <td>jinaai/jina-embeddings-v2-base-code</td>\n",
        "      <td>768</td>\n",
-       "      <td>Text embeddings, Unimodal (text), Multilingual...</td>\n",
-       "      <td>apache-2.0</td>\n",
+       "      <td>Source code embedding model supporting 8192 se...</td>\n",
        "      <td>0.640</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>20</th>\n",
        "      <td>sentence-transformers/paraphrase-multilingual-...</td>\n",
        "      <td>768</td>\n",
-       "      <td>Text embeddings, Unimodal (text), Multilingual...</td>\n",
-       "      <td>apache-2.0</td>\n",
+       "      <td>Sentence-transformers model for tasks like clu...</td>\n",
        "      <td>1.000</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>21</th>\n",
        "      <td>snowflake/snowflake-arctic-embed-l</td>\n",
        "      <td>1024</td>\n",
-       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
-       "      <td>apache-2.0</td>\n",
+       "      <td>Based on intfloat/e5-large-unsupervised, large...</td>\n",
        "      <td>1.020</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>22</th>\n",
        "      <td>thenlper/gte-large</td>\n",
        "      <td>1024</td>\n",
-       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
-       "      <td>mit</td>\n",
+       "      <td>Large general text embeddings model</td>\n",
        "      <td>1.200</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>23</th>\n",
        "      <td>BAAI/bge-large-en-v1.5</td>\n",
        "      <td>1024</td>\n",
-       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
-       "      <td>mit</td>\n",
+       "      <td>Large English model, v1.5</td>\n",
        "      <td>1.200</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>24</th>\n",
        "      <td>intfloat/multilingual-e5-large</td>\n",
        "      <td>1024</td>\n",
-       "      <td>Text embeddings, Unimodal (text), Multilingual...</td>\n",
-       "      <td>mit</td>\n",
+       "      <td>Multilingual model, e5-large. Recommend using ...</td>\n",
        "      <td>2.240</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
@@ -323,35 +297,35 @@
        "23                             BAAI/bge-large-en-v1.5  1024   \n",
        "24                     intfloat/multilingual-e5-large  1024   \n",
        "\n",
-       "                                          description     license  size_in_GB  \n",
-       "0   Text embeddings, Unimodal (text), English, 512...         mit       0.067  \n",
-       "1   Text embeddings, Unimodal (text), Chinese, 512...         mit       0.090  \n",
-       "2   Text embeddings, Unimodal (text), English, 512...  apache-2.0       0.090  \n",
-       "3   Text embeddings, Unimodal (text), English, 256...  apache-2.0       0.090  \n",
-       "4   Text embeddings, Unimodal (text), English, 819...  apache-2.0       0.120  \n",
-       "5   Text embeddings, Unimodal (text), English, 512...         mit       0.130  \n",
-       "6   Text embeddings, Unimodal (text), English, 512...  apache-2.0       0.130  \n",
-       "7   Text embeddings, Multimodal (text, image), Eng...  apache-2.0       0.130  \n",
-       "8   Text embeddings, Unimodal (text), English, 512...         mit       0.210  \n",
-       "9   Text embeddings, Unimodal (text), Multilingual...  apache-2.0       0.220  \n",
-       "10  Text embeddings, Multimodal (text&image), Engl...         mit       0.250  \n",
-       "11  Text embeddings, Unimodal (text), Multilingual...  apache-2.0       0.320  \n",
-       "12  Text embeddings, Unimodal (text), English, 512...         mit       0.420  \n",
-       "13  Text embeddings, Unimodal (text), English, 512...  apache-2.0       0.430  \n",
-       "14  Text embeddings, Multimodal (text, image), Eng...  apache-2.0       0.520  \n",
-       "15  Text embeddings, Unimodal (text), English, 819...  apache-2.0       0.520  \n",
-       "16  Text embeddings, Multimodal (text, image), Eng...  apache-2.0       0.520  \n",
-       "17  Text embeddings, Unimodal (text), English, 204...  apache-2.0       0.540  \n",
-       "18  Text embeddings, Unimodal (text), English, 512...  apache-2.0       0.640  \n",
-       "19  Text embeddings, Unimodal (text), Multilingual...  apache-2.0       0.640  \n",
-       "20  Text embeddings, Unimodal (text), Multilingual...  apache-2.0       1.000  \n",
-       "21  Text embeddings, Unimodal (text), English, 512...  apache-2.0       1.020  \n",
-       "22  Text embeddings, Unimodal (text), English, 512...         mit       1.200  \n",
-       "23  Text embeddings, Unimodal (text), English, 512...         mit       1.200  \n",
-       "24  Text embeddings, Unimodal (text), Multilingual...         mit       2.240  "
+       "                                          description  size_in_GB  \n",
+       "0                      Fast and Default English model       0.067  \n",
+       "1                  Fast and recommended Chinese model       0.090  \n",
+       "2   Based on all-MiniLM-L6-v2 model with only 22m ...       0.090  \n",
+       "3            Sentence Transformer model, MiniLM-L6-v2       0.090  \n",
+       "4   English embedding model supporting 8192 sequen...       0.120  \n",
+       "5                                  Fast English model       0.130  \n",
+       "6   Based on infloat/e5-small-unsupervised, does n...       0.130  \n",
+       "7         Quantized 8192 context length english model       0.130  \n",
+       "8                            Base English model, v1.5       0.210  \n",
+       "9   Sentence Transformer model, paraphrase-multili...       0.220  \n",
+       "10                                  CLIP text encoder       0.250  \n",
+       "11  German embedding model supporting 8192 sequenc...       0.320  \n",
+       "12                                 Base English model       0.420  \n",
+       "13  Based on intfloat/e5-base-unsupervised model, ...       0.430  \n",
+       "14                  8192 context length english model       0.520  \n",
+       "15  English embedding model supporting 8192 sequen...       0.520  \n",
+       "16                  8192 context length english model       0.520  \n",
+       "17  Based on nomic-ai/nomic-embed-text-v1-unsuperv...       0.540  \n",
+       "18  MixedBread Base sentence embedding model, does...       0.640  \n",
+       "19  Source code embedding model supporting 8192 se...       0.640  \n",
+       "20  Sentence-transformers model for tasks like clu...       1.000  \n",
+       "21  Based on intfloat/e5-large-unsupervised, large...       1.020  \n",
+       "22                Large general text embeddings model       1.200  \n",
+       "23                          Large English model, v1.5       1.200  \n",
+       "24  Multilingual model, e5-large. Recommend using ...       2.240  "
       ]
      },
-     "execution_count": 3,
+     "execution_count": 5,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -407,7 +381,6 @@
        "      <th>model</th>\n",
        "      <th>vocab_size</th>\n",
        "      <th>description</th>\n",
-       "      <th>license</th>\n",
        "      <th>size_in_GB</th>\n",
        "      <th>requires_idf</th>\n",
        "    </tr>\n",
@@ -418,7 +391,6 @@
        "      <td>Qdrant/bm25</td>\n",
        "      <td>NaN</td>\n",
        "      <td>BM25 as sparse embeddings meant to be used wit...</td>\n",
-       "      <td>apache-2.0</td>\n",
        "      <td>0.010</td>\n",
        "      <td>True</td>\n",
        "    </tr>\n",
@@ -427,25 +399,22 @@
        "      <td>Qdrant/bm42-all-minilm-l6-v2-attentions</td>\n",
        "      <td>30522.0</td>\n",
        "      <td>Light sparse embedding model, which assigns an...</td>\n",
-       "      <td>apache-2.0</td>\n",
        "      <td>0.090</td>\n",
        "      <td>True</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>2</th>\n",
-       "      <td>prithivida/Splade_PP_en_v1</td>\n",
+       "      <td>prithvida/Splade_PP_en_v1</td>\n",
        "      <td>30522.0</td>\n",
-       "      <td>Independent Implementation of SPLADE++ Model f...</td>\n",
-       "      <td>apache-2.0</td>\n",
+       "      <td>Misspelled version of the model. Retained for ...</td>\n",
        "      <td>0.532</td>\n",
        "      <td>NaN</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>3</th>\n",
-       "      <td>prithvida/Splade_PP_en_v1</td>\n",
+       "      <td>prithivida/Splade_PP_en_v1</td>\n",
        "      <td>30522.0</td>\n",
        "      <td>Independent Implementation of SPLADE++ Model f...</td>\n",
-       "      <td>apache-2.0</td>\n",
        "      <td>0.532</td>\n",
        "      <td>NaN</td>\n",
        "    </tr>\n",
@@ -457,20 +426,14 @@
        "                                     model  vocab_size  \\\n",
        "0                              Qdrant/bm25         NaN   \n",
        "1  Qdrant/bm42-all-minilm-l6-v2-attentions     30522.0   \n",
-       "2               prithivida/Splade_PP_en_v1     30522.0   \n",
-       "3                prithvida/Splade_PP_en_v1     30522.0   \n",
+       "2                prithvida/Splade_PP_en_v1     30522.0   \n",
+       "3               prithivida/Splade_PP_en_v1     30522.0   \n",
        "\n",
-       "                                         description     license  size_in_GB  \\\n",
-       "0  BM25 as sparse embeddings meant to be used wit...  apache-2.0       0.010   \n",
-       "1  Light sparse embedding model, which assigns an...  apache-2.0       0.090   \n",
-       "2  Independent Implementation of SPLADE++ Model f...  apache-2.0       0.532   \n",
-       "3  Independent Implementation of SPLADE++ Model f...  apache-2.0       0.532   \n",
-       "\n",
-       "  requires_idf  \n",
-       "0         True  \n",
-       "1         True  \n",
-       "2          NaN  \n",
-       "3          NaN  "
+       "                                         description  size_in_GB requires_idf  \n",
+       "0  BM25 as sparse embeddings meant to be used wit...       0.010         True  \n",
+       "1  Light sparse embedding model, which assigns an...       0.090         True  \n",
+       "2  Misspelled version of the model. Retained for ...       0.532          NaN  \n",
+       "3  Independent Implementation of SPLADE++ Model f...       0.532          NaN  "
       ]
      },
      "execution_count": 4,
@@ -480,7 +443,7 @@
    ],
    "source": [
     "(\n",
-    "    pd.DataFrame(SparseTextEmbedding.list_supported_models())\n",
+    "  pd.DataFrame(SparseTextEmbedding.list_supported_models())\n",
     "    .sort_values(\"size_in_GB\")\n",
     "    .drop(columns=[\"sources\", \"model_file\", \"additional_files\"])\n",
     "    .reset_index(drop=True)\n",
@@ -531,58 +494,24 @@
        "      <th>model</th>\n",
        "      <th>dim</th>\n",
        "      <th>description</th>\n",
-       "      <th>license</th>\n",
        "      <th>size_in_GB</th>\n",
-       "      <th>additional_files</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>0</th>\n",
-       "      <td>answerdotai/answerai-colbert-small-v1</td>\n",
-       "      <td>96</td>\n",
-       "      <td>Text embeddings, Unimodal (text), Multilingual...</td>\n",
-       "      <td>apache-2.0</td>\n",
-       "      <td>0.13</td>\n",
-       "      <td>NaN</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1</th>\n",
        "      <td>colbert-ir/colbertv2.0</td>\n",
        "      <td>128</td>\n",
        "      <td>Late interaction model</td>\n",
-       "      <td>mit</td>\n",
        "      <td>0.44</td>\n",
-       "      <td>NaN</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>2</th>\n",
-       "      <td>jinaai/jina-colbert-v2</td>\n",
-       "      <td>1024</td>\n",
-       "      <td>New model that expands capabilities of colbert...</td>\n",
-       "      <td>cc-by-nc-4.0</td>\n",
-       "      <td>2.24</td>\n",
-       "      <td>[onnx/model.onnx_data]</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "text/plain": [
-       "                                   model   dim  \\\n",
-       "0  answerdotai/answerai-colbert-small-v1    96   \n",
-       "1                 colbert-ir/colbertv2.0   128   \n",
-       "2                 jinaai/jina-colbert-v2  1024   \n",
-       "\n",
-       "                                         description       license  \\\n",
-       "0  Text embeddings, Unimodal (text), Multilingual...    apache-2.0   \n",
-       "1                             Late interaction model           mit   \n",
-       "2  New model that expands capabilities of colbert...  cc-by-nc-4.0   \n",
-       "\n",
-       "   size_in_GB        additional_files  \n",
-       "0        0.13                     NaN  \n",
-       "1        0.44                     NaN  \n",
-       "2        2.24  [onnx/model.onnx_data]  "
+       "                    model  dim             description  size_in_GB\n",
+       "0  colbert-ir/colbertv2.0  128  Late interaction model        0.44"
       ]
      },
      "execution_count": 5,
@@ -592,7 +521,7 @@
    ],
    "source": [
     "(\n",
-    "    pd.DataFrame(LateInteractionTextEmbedding.list_supported_models())\n",
+    "  pd.DataFrame(LateInteractionTextEmbedding.list_supported_models())\n",
     "    .sort_values(\"size_in_GB\")\n",
     "    .drop(columns=[\"sources\", \"model_file\"])\n",
     "    .reset_index(drop=True)\n",
@@ -643,7 +572,6 @@
        "      <th>model</th>\n",
        "      <th>dim</th>\n",
        "      <th>description</th>\n",
-       "      <th>license</th>\n",
        "      <th>size_in_GB</th>\n",
        "    </tr>\n",
        "  </thead>\n",
@@ -652,32 +580,28 @@
        "      <th>0</th>\n",
        "      <td>Qdrant/resnet50-onnx</td>\n",
        "      <td>2048</td>\n",
-       "      <td>Image embeddings, Unimodal (image), 2016 year</td>\n",
-       "      <td>apache-2.0</td>\n",
+       "      <td>ResNet-50 from `Deep Residual Learning for Ima...</td>\n",
        "      <td>0.10</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>1</th>\n",
        "      <td>Qdrant/clip-ViT-B-32-vision</td>\n",
        "      <td>512</td>\n",
-       "      <td>Image embeddings, Multimodal (text&amp;image), 202...</td>\n",
-       "      <td>mit</td>\n",
+       "      <td>CLIP vision encoder based on ViT-B/32</td>\n",
        "      <td>0.34</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>2</th>\n",
        "      <td>Qdrant/Unicom-ViT-B-32</td>\n",
        "      <td>512</td>\n",
-       "      <td>Image embeddings, Multimodal (text&amp;image), 202...</td>\n",
-       "      <td>apache-2.0</td>\n",
+       "      <td>Unicom Unicom-ViT-B-32 from open-metric-learning</td>\n",
        "      <td>0.48</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>3</th>\n",
        "      <td>Qdrant/Unicom-ViT-B-16</td>\n",
        "      <td>768</td>\n",
-       "      <td>Image embeddings (more detailed than Unicom-Vi...</td>\n",
-       "      <td>apache-2.0</td>\n",
+       "      <td>Unicom Unicom-ViT-B-16 from open-metric-learning</td>\n",
        "      <td>0.82</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
@@ -691,11 +615,11 @@
        "2       Qdrant/Unicom-ViT-B-32   512   \n",
        "3       Qdrant/Unicom-ViT-B-16   768   \n",
        "\n",
-       "                                         description     license  size_in_GB  \n",
-       "0      Image embeddings, Unimodal (image), 2016 year  apache-2.0        0.10  \n",
-       "1  Image embeddings, Multimodal (text&image), 202...         mit        0.34  \n",
-       "2  Image embeddings, Multimodal (text&image), 202...  apache-2.0        0.48  \n",
-       "3  Image embeddings (more detailed than Unicom-Vi...  apache-2.0        0.82  "
+       "                                         description  size_in_GB  \n",
+       "0  ResNet-50 from `Deep Residual Learning for Ima...        0.10  \n",
+       "1              CLIP vision encoder based on ViT-B/32        0.34  \n",
+       "2   Unicom Unicom-ViT-B-32 from open-metric-learning        0.48  \n",
+       "3   Unicom Unicom-ViT-B-16 from open-metric-learning        0.82  "
       ]
      },
      "execution_count": 6,
@@ -705,8 +629,7 @@
    ],
    "source": [
     "(\n",
-    "    pd.DataFrame(ImageEmbedding.list_supported_models())\n",
-    "    .sort_values(\"size_in_GB\")\n",
+    "  pd.DataFrame(ImageEmbedding.list_supported_models()).sort_values(\"size_in_GB\")\n",
     "    .drop(columns=[\"sources\", \"model_file\"])\n",
     "    .reset_index(drop=True)\n",
     ")"
@@ -715,7 +638,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": ".venv",
+   "display_name": "Python 3.8.18 ('base')",
    "language": "python",
    "name": "python3"
   },
@@ -729,9 +652,14 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.15"
+   "version": "3.11.8"
   },
-  "orig_nbformat": 4
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "c4a27af61e455bc18dcf16f5867a2ff0402fa12b01dd0f6ce3a79ae73ad15e91"
+   }
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2

fastembed/late_interaction/colbert.py

@@ -83,10 +83,9 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
         )
 
     def _tokenize_query(self, query: str) -> List[Encoding]:
-        # "@ " is added to a query to be replaced with a special query token
-        # make sure that "@ " is considered as a single token
-        query = f"@ {query}"
-        encoded = self.tokenizer.encode_batch([query])
+        # ". " is added to a query to be replaced with a special query token
+        query = [f". {query}"]
+        encoded = self.tokenizer.encode_batch(query)
         # colbert authors recommend to pad queries with [MASK] tokens for query augmentation to improve performance
         if len(encoded[0].ids) < self.MIN_QUERY_LENGTH:
             prev_padding = None
@@ -97,7 +96,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
                 pad_id=self.mask_token_id,
                 length=self.MIN_QUERY_LENGTH,
             )
-            encoded = self.tokenizer.encode_batch([query])
+            encoded = self.tokenizer.encode_batch(query)
             if prev_padding is None:
                 self.tokenizer.no_padding()
             else:
@@ -105,9 +104,8 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
         return encoded
 
     def _tokenize_documents(self, documents: List[str]) -> List[Encoding]:
-        # "@ " is added to a document to be replaced with a special document token
-        # make sure that "@ " is considered as a single token
-        documents = ["@ " + doc for doc in documents]
+        # ". " is added to a document to be replaced with a special document token
+        documents = [". " + doc for doc in documents]
         encoded = self.tokenizer.encode_batch(documents)
         return encoded
 
@@ -191,7 +189,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
             cuda=self.cuda,
             device_id=self.device_id,
         )
-        self.mask_token_id = self.special_token_to_id[self.MASK_TOKEN]
+        self.mask_token_id = self.special_token_to_id["[MASK]"]
         self.pad_token_id = self.tokenizer.padding["pad_id"]
         self.skip_list = {
             self.tokenizer.encode(symbol, add_special_tokens=False).ids[0]

fastembed/late_interaction/jina_colbert.py

@@ -1,62 +0,0 @@
-from typing import Any, Dict, List, Type
-
-import numpy as np
-
-from fastembed.late_interaction.colbert import Colbert
-from fastembed.text.onnx_text_model import TextEmbeddingWorker
-
-supported_jina_colbert_models = [
-    {
-        "model": "jinaai/jina-colbert-v2",
-        "dim": 128,
-        "description": "New model that expands capabilities of colbert-v1 with multilingual and context length of 8192, 2024 year",
-        "license": "cc-by-nc-4.0",
-        "size_in_GB": 2.24,
-        "sources": {
-            "hf": "jinaai/jina-colbert-v2",
-        },
-        "model_file": "onnx/model.onnx",
-        "additional_files": ["onnx/model.onnx_data"],
-    },
-]
-
-
-class JinaColbert(Colbert):
-    QUERY_MARKER_TOKEN_ID = 250002
-    DOCUMENT_MARKER_TOKEN_ID = 250003
-    MIN_QUERY_LENGTH = 32
-    MASK_TOKEN = "<mask>"
-
-    @classmethod
-    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
-        return JinaColbertEmbeddingWorker
-
-    @classmethod
-    def list_supported_models(cls) -> List[Dict[str, Any]]:
-        """Lists the supported models.
-
-        Returns:
-            List[Dict[str, Any]]: A list of dictionaries containing the model information.
-        """
-        return supported_jina_colbert_models
-
-    def _preprocess_onnx_input(
-        self, onnx_input: Dict[str, np.ndarray], is_doc: bool = True
-    ) -> Dict[str, np.ndarray]:
-        if is_doc:
-            onnx_input["input_ids"][:, 1] = self.DOCUMENT_MARKER_TOKEN_ID
-        else:
-            onnx_input["input_ids"][:, 1] = self.QUERY_MARKER_TOKEN_ID
-            # the attention mask for jina-colbert-v2 is always 1 in queries
-            onnx_input["attention_mask"][:] = 1
-        return onnx_input
-
-
-class JinaColbertEmbeddingWorker(TextEmbeddingWorker):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> JinaColbert:
-        return JinaColbert(
-            model_name=model_name,
-            cache_dir=cache_dir,
-            threads=1,
-            **kwargs,
-        )

fastembed/late_interaction/late_interaction_text_embedding.py

@@ -4,14 +4,15 @@ import numpy as np
 
 from fastembed.common import OnnxProvider
 from fastembed.late_interaction.colbert import Colbert
-from fastembed.late_interaction.jina_colbert import JinaColbert
 from fastembed.late_interaction.late_interaction_embedding_base import (
     LateInteractionTextEmbeddingBase,
 )
 
 
 class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
-    EMBEDDINGS_REGISTRY: List[Type[LateInteractionTextEmbeddingBase]] = [Colbert, JinaColbert]
+    EMBEDDINGS_REGISTRY: List[Type[LateInteractionTextEmbeddingBase]] = [
+        Colbert,
+    ]
 
     @classmethod
     def list_supported_models(cls) -> List[Dict[str, Any]]:

fastembed/sparse/bm25.py

@@ -6,7 +6,8 @@ from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Union
 
 import mmh3
 import numpy as np
-from py_rust_stemmers import SnowballStemmer
+from snowballstemmer import stemmer as get_stemmer
+
 from fastembed.common.utils import (
     define_cache_dir,
     iter_batch,
@@ -129,7 +130,7 @@ class Bm25(SparseTextEmbeddingBase):
         self.punctuation = set(get_all_punctuation())
         self.stopwords = set(self._load_stopwords(self._model_dir, self.language))
 
-        self.stemmer = SnowballStemmer(language)
+        self.stemmer = get_stemmer(language)
         self.tokenizer = SimpleTokenizer
 
     @classmethod
@@ -234,7 +235,7 @@ class Bm25(SparseTextEmbeddingBase):
             if len(token) > self.token_max_length:
                 continue
 
-            stemmed_token = self.stemmer.stem_word(token.lower())
+            stemmed_token = self.stemmer.stemWord(token.lower())
 
             if stemmed_token:
                 stemmed_tokens.append(stemmed_token)

fastembed/sparse/bm42.py

@@ -5,7 +5,7 @@ from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type, U
 
 import mmh3
 import numpy as np
-from py_rust_stemmers import SnowballStemmer
+from snowballstemmer import stemmer as get_stemmer
 
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
@@ -119,8 +119,8 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         self.special_tokens = set()
         self.special_tokens_ids = set()
         self.punctuation = set(string.punctuation)
-        self.stopwords = set(self._load_stopwords(self._model_dir))
-        self.stemmer = SnowballStemmer(MODEL_TO_LANGUAGE[model_name])
+        self.stopwords = set()
+        self.stemmer = get_stemmer(MODEL_TO_LANGUAGE[model_name])
         self.alpha = alpha
 
         if not self.lazy_load:
@@ -152,7 +152,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     def _stem_pair_tokens(self, tokens: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:
         result = []
         for token, value in tokens:
-            processed_token = self.stemmer.stem_word(token)
+            processed_token = self.stemmer.stemWord(token)
             result.append((processed_token, value))
         return result
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
-name = "fastembed-gpu"
-version = "0.4.2"
+name = "fastembed"
+version = "0.4.0"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -13,7 +13,7 @@ keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-trans
 [tool.poetry.dependencies]
 python = ">=3.8.0,<3.13"
 onnx = "^1.15.0"
-onnxruntime-gpu = ">=1.17.0,<1.20.0"
+onnxruntime = "^1.17.0"
 tqdm = "^4.66"
 requests = "^2.31"
 tokenizers = ">=0.15,<1.0"
@@ -24,8 +24,9 @@ numpy = [
     { version = ">=1.26, <2", python = ">=3.12" }
 ]
 pillow = "^10.3.0"
+snowballstemmer = "^2.2.0"
+PyStemmer = "^2.2.0"
 mmh3 = "^4.1.0"
-py-rust-stemmers = "^0.1.0"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.4.2"

tests/test_attention_embeddings.py

@@ -1,10 +1,10 @@
 import os
+import shutil
 
 import numpy as np
 import pytest
 
 from fastembed import SparseTextEmbedding
-from tests.utils import delete_model_cache
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"])
@@ -67,7 +67,7 @@ def test_attention_embeddings(model_name):
         assert len(result.indices) == 2
 
     if is_ci:
-        delete_model_cache(model.model._model_dir)
+        shutil.rmtree(model.model._model_dir)
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"])
@@ -92,7 +92,7 @@ def test_parallel_processing(model_name):
         assert np.allclose(emb_1.values, emb_3.values)
 
     if is_ci:
-        delete_model_cache(model.model._model_dir)
+        shutil.rmtree(model.model._model_dir)
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/bm25"])
@@ -118,7 +118,7 @@ def test_multilanguage(model_name):
     assert embeddings[1].indices.shape == (4,)
 
     if is_ci:
-        delete_model_cache(model.model._model_dir)
+        shutil.rmtree(model.model._model_dir)
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/bm25"])
@@ -141,7 +141,7 @@ def test_special_characters(model_name):
         assert embeddings[idx].indices.shape == (shape,)
 
     if is_ci:
-        delete_model_cache(model.model._model_dir)
+        shutil.rmtree(model.model._model_dir)
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions"])

tests/test_image_onnx_embeddings.py

@@ -1,4 +1,5 @@
 import os
+import shutil
 from io import BytesIO
 
 import numpy as np
@@ -8,7 +9,6 @@ from PIL import Image
 
 from fastembed import ImageEmbedding
 from tests.config import TEST_MISC_DIR
-from tests.utils import delete_model_cache
 
 CANONICAL_VECTOR_VALUES = {
     "Qdrant/clip-ViT-B-32-vision": np.array([-0.0098, 0.0128, -0.0274, 0.002, -0.0059]),
@@ -54,7 +54,7 @@ def test_embedding():
         assert np.allclose(embeddings[1], embeddings[2]), model_desc["model"]
 
         if is_ci:
-            delete_model_cache(model.model._model_dir)
+            shutil.rmtree(model.model._model_dir)
 
 
 @pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
@@ -74,7 +74,7 @@ def test_batch_embedding(n_dims, model_name):
 
     assert embeddings.shape == (len(test_images) * n_images, n_dims)
     if is_ci:
-        delete_model_cache(model.model._model_dir)
+        shutil.rmtree(model.model._model_dir)
 
 
 @pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
@@ -102,12 +102,11 @@ def test_parallel_processing(n_dims, model_name):
     assert np.allclose(embeddings, embeddings_2, atol=1e-3)
     assert np.allclose(embeddings, embeddings_3, atol=1e-3)
     if is_ci:
-        delete_model_cache(model.model._model_dir)
+        shutil.rmtree(model.model._model_dir)
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/clip-ViT-B-32-vision"])
 def test_lazy_load(model_name):
-    is_ci = os.getenv("CI")
     model = ImageEmbedding(model_name=model_name, lazy_load=True)
     assert not hasattr(model.model, "model")
     images = [
@@ -116,5 +115,3 @@ def test_lazy_load(model_name):
     ]
     list(model.embed(images))
     assert hasattr(model.model, "model")
-    if is_ci:
-        delete_model_cache(model.model._model_dir)

tests/test_late_interaction_embeddings.py

@@ -1,4 +1,5 @@
 import os
+import shutil
 
 import pytest
 import numpy as np
@@ -6,7 +7,6 @@ import numpy as np
 from fastembed.late_interaction.late_interaction_text_embedding import (
     LateInteractionTextEmbedding,
 )
-from tests.utils import delete_model_cache
 
 # vectors are abridged and rounded for brevity
 CANONICAL_COLUMN_VALUES = {
@@ -28,15 +28,6 @@ CANONICAL_COLUMN_VALUES = {
             [-0.07281, 0.04633, -0.04711, 0.00762, -0.07374],
         ]
     ),
-    "jinaai/jina-colbert-v2": np.array(
-        [
-            [0.0742, 0.0591, -0.2403, -0.1774, 0.02],
-            [0.1318, 0.0882, -0.1138, -0.2066, 0.146],
-            [-0.0183, -0.1354, -0.0139, -0.1079, -0.051],
-            [0.0003, -0.1184, -0.07, -0.0479, -0.0649],
-            [0.0766, 0.0452, -0.2343, -0.183, 0.0058],
-        ]
-    ),
 }
 
 CANONICAL_QUERY_VALUES = {
@@ -112,42 +103,6 @@ CANONICAL_QUERY_VALUES = {
             [-0.03473, 0.04792, -0.07033, 0.02196, -0.05314],
         ]
     ),
-    "jinaai/jina-colbert-v2": np.array(
-        [
-            [0.0477, 0.0255, -0.2224, -0.1085, -0.03],
-            [0.0206, -0.0845, -0.0075, -0.1712, 0.0156],
-            [-0.0056, -0.0957, -0.0147, -0.1277, -0.0225],
-            [0.0486, -0.0499, -0.1609, 0.0194, 0.0274],
-            [0.0481, 0.0253, -0.2278, -0.1126, -0.0294],
-            [0.0599, -0.0678, -0.0956, -0.0757, 0.0236],
-            [0.0592, -0.0862, -0.0621, -0.1084, 0.0155],
-            [0.0874, -0.0714, -0.0772, -0.1414, 0.037],
-            [0.1009, -0.0552, -0.0669, -0.163, 0.0493],
-            [0.1135, -0.047, -0.0576, -0.1699, 0.0538],
-            [0.1228, -0.0428, -0.0507, -0.1725, 0.0562],
-            [0.1291, -0.0388, -0.042, -0.1753, 0.0569],
-            [0.1365, -0.0337, -0.0326, -0.1786, 0.0574],
-            [0.1439, -0.026, -0.024, -0.1831, 0.0574],
-            [0.1527, -0.0099, -0.0179, -0.1874, 0.057],
-            [0.1555, 0.0186, -0.023, -0.1801, 0.0539],
-            [0.1389, 0.054, -0.0345, -0.1636, 0.0429],
-            [0.1058, 0.0862, -0.0418, -0.1455, 0.0222],
-            [0.0713, 0.1061, -0.0438, -0.1288, 0.0002],
-            [0.0453, 0.1143, -0.0457, -0.1119, -0.019],
-            [0.0346, 0.1131, -0.0487, -0.0952, -0.0338],
-            [0.0355, 0.1073, -0.0493, -0.0823, -0.0438],
-            [0.0424, 0.1041, -0.0459, -0.0761, -0.048],
-            [0.048, 0.102, -0.0421, -0.0718, -0.0477],
-            [0.0474, 0.0989, -0.0413, -0.0654, -0.0431],
-            [0.0434, 0.095, -0.0415, -0.0589, -0.0345],
-            [0.0408, 0.0897, -0.0405, -0.0554, -0.0197],
-            [0.0433, 0.0811, -0.0407, -0.0545, 0.0055],
-            [0.0514, 0.0629, -0.0446, -0.0549, 0.0368],
-            [0.058, 0.048, -0.0527, -0.0607, 0.0568],
-            [0.0561, 0.0447, -0.0661, -0.0702, 0.0764],
-            [0.0204, -0.0856, -0.0386, -0.1232, -0.0332],
-        ]
-    ),
 }
 
 docs = ["Hello World"]
@@ -164,10 +119,10 @@ def test_batch_embedding():
 
         for value in result:
             token_num, abridged_dim = expected_result.shape
-            assert np.allclose(value[:, :abridged_dim], expected_result, atol=2e-3)
+            assert np.allclose(value[:, :abridged_dim], expected_result, atol=10e-4)
 
         if is_ci:
-            delete_model_cache(model.model._model_dir)
+            shutil.rmtree(model.model._model_dir)
 
 
 def test_single_embedding():
@@ -179,10 +134,10 @@ def test_single_embedding():
         model = LateInteractionTextEmbedding(model_name=model_name)
         result = next(iter(model.embed(docs_to_embed, batch_size=6)))
         token_num, abridged_dim = expected_result.shape
-        assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
+        assert np.allclose(result[:, :abridged_dim], expected_result, atol=10e-4)
 
         if is_ci:
-            delete_model_cache(model.model._model_dir)
+            shutil.rmtree(model.model._model_dir)
 
 
 def test_single_embedding_query():
@@ -194,10 +149,10 @@ def test_single_embedding_query():
         model = LateInteractionTextEmbedding(model_name=model_name)
         result = next(iter(model.query_embed(queries_to_embed)))
         token_num, abridged_dim = expected_result.shape
-        assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
+        assert np.allclose(result[:, :abridged_dim], expected_result, atol=10e-4)
 
         if is_ci:
-            delete_model_cache(model.model._model_dir)
+            shutil.rmtree(model.model._model_dir)
 
 
 def test_parallel_processing():
@@ -219,7 +174,7 @@ def test_parallel_processing():
     assert np.allclose(embeddings, embeddings_3, atol=1e-3)
 
     if is_ci:
-        delete_model_cache(model.model._model_dir)
+        shutil.rmtree(model.model._model_dir)
 
 
 @pytest.mark.parametrize(
@@ -227,8 +182,6 @@ def test_parallel_processing():
     ["colbert-ir/colbertv2.0"],
 )
 def test_lazy_load(model_name):
-    is_ci = os.getenv("CI")
-
     model = LateInteractionTextEmbedding(model_name=model_name, lazy_load=True)
     assert not hasattr(model.model, "model")
 
@@ -241,6 +194,3 @@ def test_lazy_load(model_name):
 
     model = LateInteractionTextEmbedding(model_name=model_name, lazy_load=True)
     list(model.passage_embed(docs))
-
-    if is_ci:
-        delete_model_cache(model.model._model_dir)

tests/test_sparse_embeddings.py

@@ -1,11 +1,11 @@
 import os
+import shutil
 
 import pytest
 import numpy as np
 
 from fastembed.sparse.bm25 import Bm25
 from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
-from tests.utils import delete_model_cache
 
 CANONICAL_COLUMN_VALUES = {
     "prithvida/Splade_PP_en_v1": {
@@ -61,7 +61,7 @@ def test_batch_embedding():
         for i, value in enumerate(result.values):
             assert pytest.approx(value, abs=0.001) == expected_result["values"][i]
         if is_ci:
-            delete_model_cache(model.model._model_dir)
+            shutil.rmtree(model.model._model_dir)
 
 
 def test_single_embedding():
@@ -77,7 +77,7 @@ def test_single_embedding():
             for i, value in enumerate(result.values):
                 assert pytest.approx(value, abs=0.001) == expected_result["values"][i]
         if is_ci:
-            delete_model_cache(model.model._model_dir)
+            shutil.rmtree(model.model._model_dir)
 
 
 def test_parallel_processing():
@@ -107,7 +107,7 @@ def test_parallel_processing():
         assert np.allclose(sparse_embedding.values, sparse_embedding_all.values, atol=1e-3)
 
     if is_ci:
-        delete_model_cache(model.model._model_dir)
+        shutil.rmtree(model.model._model_dir)
 
 
 @pytest.fixture
@@ -116,7 +116,7 @@ def bm25_instance():
     model = Bm25("Qdrant/bm25", language="english")
     yield model
     if ci:
-        delete_model_cache(model._model_dir)
+        shutil.rmtree(model._model_dir)
 
 
 def test_stem_with_stopwords_and_punctuation(bm25_instance):
@@ -150,13 +150,11 @@ def test_stem_case_insensitive_stopwords(bm25_instance):
     expected = ["quick", "brown", "fox", "test", "sentenc"]
     assert result == expected, f"Expected {expected}, but got {result}"
 
-
 @pytest.mark.parametrize(
     "model_name",
     ["prithivida/Splade_PP_en_v1"],
 )
 def test_lazy_load(model_name):
-    is_ci = os.getenv("CI")
     model = SparseTextEmbedding(model_name=model_name, lazy_load=True)
     assert not hasattr(model.model, "model")
 
@@ -169,6 +167,3 @@ def test_lazy_load(model_name):
 
     model = SparseTextEmbedding(model_name=model_name, lazy_load=True)
     list(model.passage_embed(docs))
-
-    if is_ci:
-        delete_model_cache(model.model._model_dir)

tests/test_text_cross_encoder.py

@@ -2,9 +2,9 @@ import os
 
 import numpy as np
 import pytest
+import shutil
 
 from fastembed.rerank.cross_encoder import TextCrossEncoder
-from tests.utils import delete_model_cache
 
 CANONICAL_SCORE_VALUES = {
     "Xenova/ms-marco-MiniLM-L-6-v2": np.array([8.500708, -2.541011]),
@@ -32,7 +32,7 @@ def test_rerank():
             scores, canonical_scores, atol=1e-3
         ), f"Model: {model_name}, Scores: {scores}, Expected: {canonical_scores}"
         if is_ci:
-            delete_model_cache(model.model._model_dir)
+            shutil.rmtree(model.model._model_dir)
 
 
 @pytest.mark.parametrize(
@@ -55,7 +55,7 @@ def test_batch_rerank(model_name):
         scores, canonical_scores, atol=1e-3
     ), f"Model: {model_name}, Scores: {scores}, Expected: {canonical_scores}"
     if is_ci:
-        delete_model_cache(model.model._model_dir)
+        shutil.rmtree(model.model._model_dir)
 
 
 @pytest.mark.parametrize(
@@ -63,13 +63,9 @@ def test_batch_rerank(model_name):
     ["Xenova/ms-marco-MiniLM-L-6-v2"],
 )
 def test_lazy_load(model_name):
-    is_ci = os.getenv("CI")
     model = TextCrossEncoder(model_name=model_name, lazy_load=True)
     assert not hasattr(model.model, "model")
     query = "What is the capital of France?"
     documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
     list(model.rerank(query, documents))
     assert hasattr(model.model, "model")
-
-    if is_ci:
-        delete_model_cache(model.model._model_dir)

tests/test_text_onnx_embeddings.py

@@ -1,10 +1,10 @@
 import os
+import shutil
 
 import numpy as np
 import pytest
 
 from fastembed.text.text_embedding import TextEmbedding
-from tests.utils import delete_model_cache
 
 CANONICAL_VECTOR_VALUES = {
     "BAAI/bge-small-en": np.array([-0.0232, -0.0255, 0.0174, -0.0639, -0.0006]),
@@ -85,7 +85,7 @@ def test_embedding():
             embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
         ), model_desc["model"]
         if is_ci:
-            delete_model_cache(model.model._model_dir)
+            shutil.rmtree(model.model._model_dir)
 
 
 @pytest.mark.parametrize(
@@ -102,7 +102,7 @@ def test_batch_embedding(n_dims, model_name):
 
     assert embeddings.shape == (200, n_dims)
     if is_ci:
-        delete_model_cache(model.model._model_dir)
+        shutil.rmtree(model.model._model_dir)
 
 
 @pytest.mark.parametrize(
@@ -128,7 +128,7 @@ def test_parallel_processing(n_dims, model_name):
     assert np.allclose(embeddings, embeddings_3, atol=1e-3)
 
     if is_ci:
-        delete_model_cache(model.model._model_dir)
+        shutil.rmtree(model.model._model_dir)
 
 
 @pytest.mark.parametrize(
@@ -136,7 +136,6 @@ def test_parallel_processing(n_dims, model_name):
     ["BAAI/bge-small-en-v1.5"],
 )
 def test_lazy_load(model_name):
-    is_ci = os.getenv("CI")
     model = TextEmbedding(model_name=model_name, lazy_load=True)
     assert not hasattr(model.model, "model")
     docs = ["hello world", "flag embedding"]
@@ -148,6 +147,3 @@ def test_lazy_load(model_name):
 
     model = TextEmbedding(model_name=model_name, lazy_load=True)
     list(model.passage_embed(docs))
-
-    if is_ci:
-        delete_model_cache(model.model._model_dir)

tests/utils.py

@@ -1,24 +0,0 @@
-import shutil
-
-from pathlib import Path
-from typing import Union
-
-
-def delete_model_cache(model_dir: Union[str, Path]) -> None:
-    """Delete the model cache directory.
-
-    If a model was downloaded from the HuggingFace model hub, then _model_dir is the dir to snapshots, removing
-    it won't help to release the memory, because data is in blobs directory.
-    If a model was downloaded from GCS, then we can just remove model_dir
-
-    Args:
-        model_dir (Union[str, Path]): The path to the model cache directory.
-    """
-    if isinstance(model_dir, str):
-        model_dir = Path(model_dir)
-
-    if model_dir.parent.parent.name.startswith("models--"):
-        model_dir = model_dir.parent.parent
-
-    if model_dir.exists():
-        shutil.rmtree(model_dir)