new: gpu package (#224)

This causes an incompatibility with spaCy because no matching
versions of numpy can be found when using both libraries at once.

It looks like the tests pass again without the numpy<2 requirement
so there is no reason to keep this requirement anymore.

.github/ISSUE_TEMPLATE/bug-report.yml

@@ -1,6 +1,6 @@
-name: Bug/New Model Request
-description: File a bug report/Request a new Model
-title: "[Bug/Model Request]: "
+name: Bug
+description: File a bug report
+title: "[Bug]: "
 body:
   - type: markdown
     attributes:
@@ -10,11 +10,22 @@ body:
     id: what-happened
     attributes:
       label: What happened?
-      description: Also tell us, what did you expect to happen?
-      placeholder: Tell us what you see!
-      value: "A bug happened!"
+      description: Describe the error you encountered.
+      placeholder: <Description>
     validations:
       required: true
+  - type: textarea
+    id: expected
+    attributes:
+      label: What is the expected behaviour?
+      description: Describe the way you expected the code to behave.
+      placeholder: <Description>
+  - type: textarea
+    id: code-snippet
+    attributes:
+      label: A minimal reproducible example
+      description: It would really help us to fix the problem if you could provide a code snippet that reproduces the issue.
+      placeholder: <Code snippet>
   - type: textarea
     id: python-version
     attributes:
@@ -23,21 +34,12 @@ body:
       placeholder: Python3.10
     validations:
       required: true
-  - type: dropdown
+  - type: textarea
     id: version
     attributes:
-      label: Version
+      label: FastEmbed 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.7 (Latest)
-        - 0.2.6
-        - 0.2.5
-        - 0.2.4 
-        - 0.2.3 
-        - 0.2.2
-        - 0.2.1
-        - 0.1.x 
-      default: 0
+      placeholder: v0.4.2
     validations:
       required: true
   - type: dropdown

.github/ISSUE_TEMPLATE/config.yaml

@@ -1,4 +1,4 @@
-blank_issues_enabled: false
+blank_issues_enabled: true
 contact_links:
   - name: GitHub Community Support
     url: https://github.com/qdrant/fastembed/discussions

.github/ISSUE_TEMPLATE/feature-request.yml

@@ -0,0 +1,22 @@
+name: Feature
+description: New functionality request
+title: "[Feature]: "
+body:
+  - type: markdown
+    attributes:
+      value: |
+        Thanks for taking the time to fill out this report!
+  - type: textarea
+    id: feature-description
+    attributes:
+      label: What feature would you like to request?
+      description: Please provide the description of the feature you would like to request.
+      placeholder: <Description>
+    validations:
+      required: true
+  - type: textarea
+    id: additional-info
+    attributes:
+      label: Is there any additional information you would like to provide?
+      description: Please provide any additional information that you think might be useful.
+      placeholder: <Info>

.github/ISSUE_TEMPLATE/model-request.yml

@@ -0,0 +1,22 @@
+name: Model
+description: Request a new model
+title: "[Model]: "
+body:
+  - type: markdown
+    attributes:
+      value: |
+        Thanks for taking the time to fill out this report!
+  - type: textarea
+    id: model-name
+    attributes:
+      label: Which model would you like to support?
+      description: Please provide the name of the model you would like to see supported.
+      placeholder: Link to the model (e.g. on HuggingFace)
+    validations:
+      required: true
+  - type: textarea
+    id: motivation
+    attributes:
+      label: What are the main advantages of this model?
+      description: Please describe the main advantages of this model comparing to the existing ones and provide links to benchmarks if there are any.
+      placeholder: <Description>

.github/PULL_REQUEST_TEMPLATE.md

@@ -0,0 +1,19 @@
+### All Submissions:
+
+* [ ] Have you followed the guidelines in our Contributing document?
+* [ ] Have you checked to ensure there aren't other open [Pull Requests](../../../pulls) for the same update/change?
+
+<!-- You can erase any parts of this template not applicable to your Pull Request. -->
+
+### New Feature Submissions:
+
+* [ ] Does your submission pass the existing tests?
+* [ ] Have you added tests for your feature?
+* [ ] Have you installed `pre-commit` with `pip3 install pre-commit` and set up hooks with `pre-commit install`?
+
+### New models submission:
+
+* [ ] Have you added an explanation of why it's important to include this model?
+* [ ] Have you added tests for the new model? Were canonical values for tests computed via the original model?
+* [ ] Have you added the code snippet for how canonical values were computed?
+* [ ] Have you successfully ran tests with your changes locally?

.github/workflows/python-tests.yml

@@ -1,4 +1,5 @@
 name: Tests
+run-name: Tests (gpu)
 
 on:
   push:
@@ -21,8 +22,6 @@ jobs:
           - '3.12.x'
         os:
           - ubuntu-latest
-          - macos-latest
-          - windows-latest
 
     runs-on: ${{ matrix.os }}
 
@@ -42,4 +41,4 @@ jobs:
 
       - name: Run pytest
         run: |
-          poetry run pytest
+          poetry run pytest 

NOTICE

@@ -5,6 +5,8 @@ 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
+- jinaai/jina-reranker-v2-base-multilingual
+    - License: cc-by-nc-4.0
 
 These models are developed by Jina (https://jina.ai/) and are subject to Jina AI's licensing terms.
 

README.md

@@ -137,6 +137,21 @@ embeddings = list(model.embed(images))
 # ]
 ```
 
+### 🔄 Rerankers
+```python
+from typing import List
+from fastembed.rerank.cross_encoder import TextCrossEncoder
+
+query = "Who is maintaining Qdrant?"
+documents: List[str] = [
+    "This is built to be faster and lighter than other embedding libraries e.g. Transformers, Sentence-Transformers, etc.",
+    "fastembed is supported by and maintained by Qdrant.",
+]
+encoder = TextCrossEncoder(model_name="Xenova/ms-marco-MiniLM-L-6-v2")
+scores = list(encoder.rerank(query, documents))
+
+# [-11.48061752319336, 5.472434997558594]
+```
 
 ## ⚡️ FastEmbed on a GPU
 

docs/examples/Supported_Models.ipynb

@@ -2,38 +2,36 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-05-31T18:13:23.806907Z",
-     "start_time": "2024-05-31T18:13:23.797078Z"
+     "end_time": "2024-11-13T09:01:03.324551Z",
+     "start_time": "2024-11-13T09:01:03.234711Z"
     }
    },
-   "outputs": [],
    "source": [
     "%load_ext autoreload\n",
     "%autoreload 2"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-05-31T18:14:31.147674Z",
-     "start_time": "2024-05-31T18:14:31.134015Z"
-    }
-   },
+   ],
    "outputs": [
     {
-     "name": "stderr",
+     "name": "stdout",
      "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"
+      "The autoreload extension is already loaded. To reload it, use:\n",
+      "  %reload_ext autoreload\n"
      ]
     }
    ],
+   "execution_count": 10
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-11-13T09:01:04.505772Z",
+     "start_time": "2024-11-13T09:01:04.493296Z"
+    }
+   },
    "source": [
     "import pandas as pd\n",
     "\n",
@@ -42,8 +40,11 @@
     "    TextEmbedding,\n",
     "    LateInteractionTextEmbedding,\n",
     "    ImageEmbedding,\n",
-    ")"
-   ]
+    ")\n",
+    "from fastembed.rerank.cross_encoder import TextCrossEncoder"
+   ],
+   "outputs": [],
+   "execution_count": 11
   },
   {
    "cell_type": "markdown",
@@ -54,16 +55,79 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-05-31T18:13:25.863008Z",
-     "start_time": "2024-05-31T18:13:25.837795Z"
+     "end_time": "2024-11-13T09:01:05.812271Z",
+     "start_time": "2024-11-13T09:01:05.795846Z"
     }
    },
+   "source": [
+    "supported_models = (\n",
+    "    pd.DataFrame(TextEmbedding.list_supported_models())\n",
+    "    .sort_values(\"size_in_GB\")\n",
+    "    .drop(columns=[\"sources\", \"model_file\", \"additional_files\"])\n",
+    "    .reset_index(drop=True)\n",
+    ")\n",
+    "supported_models"
+   ],
    "outputs": [
     {
      "data": {
+      "text/plain": [
+       "                                                model   dim  \\\n",
+       "0                              BAAI/bge-small-en-v1.5   384   \n",
+       "1                              BAAI/bge-small-zh-v1.5   512   \n",
+       "2                 snowflake/snowflake-arctic-embed-xs   384   \n",
+       "3              sentence-transformers/all-MiniLM-L6-v2   384   \n",
+       "4                  jinaai/jina-embeddings-v2-small-en   512   \n",
+       "5                                   BAAI/bge-small-en   384   \n",
+       "6                  snowflake/snowflake-arctic-embed-s   384   \n",
+       "7                    nomic-ai/nomic-embed-text-v1.5-Q   768   \n",
+       "8                               BAAI/bge-base-en-v1.5   768   \n",
+       "9   sentence-transformers/paraphrase-multilingual-...   384   \n",
+       "10                          Qdrant/clip-ViT-B-32-text   512   \n",
+       "11                  jinaai/jina-embeddings-v2-base-de   768   \n",
+       "12                                   BAAI/bge-base-en   768   \n",
+       "13                 snowflake/snowflake-arctic-embed-m   768   \n",
+       "14                     nomic-ai/nomic-embed-text-v1.5   768   \n",
+       "15                  jinaai/jina-embeddings-v2-base-en   768   \n",
+       "16                       nomic-ai/nomic-embed-text-v1   768   \n",
+       "17            snowflake/snowflake-arctic-embed-m-long   768   \n",
+       "18                 mixedbread-ai/mxbai-embed-large-v1  1024   \n",
+       "19                jinaai/jina-embeddings-v2-base-code   768   \n",
+       "20  sentence-transformers/paraphrase-multilingual-...   768   \n",
+       "21                 snowflake/snowflake-arctic-embed-l  1024   \n",
+       "22                                 thenlper/gte-large  1024   \n",
+       "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  "
+      ],
       "text/html": [
        "<div>\n",
        "<style scoped>\n",
@@ -294,77 +358,14 @@
        "  </tbody>\n",
        "</table>\n",
        "</div>"
-      ],
-      "text/plain": [
-       "                                                model   dim  \\\n",
-       "0                              BAAI/bge-small-en-v1.5   384   \n",
-       "1                              BAAI/bge-small-zh-v1.5   512   \n",
-       "2                 snowflake/snowflake-arctic-embed-xs   384   \n",
-       "3              sentence-transformers/all-MiniLM-L6-v2   384   \n",
-       "4                  jinaai/jina-embeddings-v2-small-en   512   \n",
-       "5                                   BAAI/bge-small-en   384   \n",
-       "6                  snowflake/snowflake-arctic-embed-s   384   \n",
-       "7                    nomic-ai/nomic-embed-text-v1.5-Q   768   \n",
-       "8                               BAAI/bge-base-en-v1.5   768   \n",
-       "9   sentence-transformers/paraphrase-multilingual-...   384   \n",
-       "10                          Qdrant/clip-ViT-B-32-text   512   \n",
-       "11                  jinaai/jina-embeddings-v2-base-de   768   \n",
-       "12                                   BAAI/bge-base-en   768   \n",
-       "13                 snowflake/snowflake-arctic-embed-m   768   \n",
-       "14                     nomic-ai/nomic-embed-text-v1.5   768   \n",
-       "15                  jinaai/jina-embeddings-v2-base-en   768   \n",
-       "16                       nomic-ai/nomic-embed-text-v1   768   \n",
-       "17            snowflake/snowflake-arctic-embed-m-long   768   \n",
-       "18                 mixedbread-ai/mxbai-embed-large-v1  1024   \n",
-       "19                jinaai/jina-embeddings-v2-base-code   768   \n",
-       "20  sentence-transformers/paraphrase-multilingual-...   768   \n",
-       "21                 snowflake/snowflake-arctic-embed-l  1024   \n",
-       "22                                 thenlper/gte-large  1024   \n",
-       "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  "
       ]
      },
-     "execution_count": 3,
+     "execution_count": 12,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
-   "source": [
-    "supported_models = (\n",
-    "    pd.DataFrame(TextEmbedding.list_supported_models())\n",
-    "    .sort_values(\"size_in_GB\")\n",
-    "    .drop(columns=[\"sources\", \"model_file\", \"additional_files\"])\n",
-    "    .reset_index(drop=True)\n",
-    ")\n",
-    "supported_models"
-   ]
+   "execution_count": 12
   },
   {
    "cell_type": "markdown",
@@ -375,16 +376,42 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-05-31T18:13:27.124747Z",
-     "start_time": "2024-05-31T18:13:27.096212Z"
+     "end_time": "2024-11-13T09:01:07.038954Z",
+     "start_time": "2024-11-13T09:01:07.019656Z"
     }
    },
+   "source": [
+    "(\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",
+    ")"
+   ],
    "outputs": [
     {
      "data": {
+      "text/plain": [
+       "                                     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",
+       "\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  "
+      ],
       "text/html": [
        "<div>\n",
        "<style scoped>\n",
@@ -452,40 +479,14 @@
        "  </tbody>\n",
        "</table>\n",
        "</div>"
-      ],
-      "text/plain": [
-       "                                     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",
-       "\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  "
       ]
      },
-     "execution_count": 4,
+     "execution_count": 13,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
-   "source": [
-    "(\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",
-    ")"
-   ]
+   "execution_count": 13
   },
   {
    "cell_type": "markdown",
@@ -498,17 +499,40 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
    "metadata": {
+    "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-05-31T18:14:34.370252Z",
-     "start_time": "2024-05-31T18:14:34.354270Z"
-    },
-    "collapsed": false
+     "end_time": "2024-11-13T09:01:08.074442Z",
+     "start_time": "2024-11-13T09:01:08.056138Z"
+    }
    },
+   "source": [
+    "(\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",
+    ")"
+   ],
    "outputs": [
     {
      "data": {
+      "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  128   \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]  "
+      ],
       "text/html": [
        "<div>\n",
        "<style scoped>\n",
@@ -558,7 +582,7 @@
        "    <tr>\n",
        "      <th>2</th>\n",
        "      <td>jinaai/jina-colbert-v2</td>\n",
-       "      <td>1024</td>\n",
+       "      <td>128</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",
@@ -567,37 +591,14 @@
        "  </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]  "
       ]
      },
-     "execution_count": 5,
+     "execution_count": 14,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
-   "source": [
-    "(\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",
-    ")"
-   ]
+   "execution_count": 14
   },
   {
    "cell_type": "markdown",
@@ -610,17 +611,37 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
    "metadata": {
+    "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-05-31T18:14:42.501881Z",
-     "start_time": "2024-05-31T18:14:42.484726Z"
-    },
-    "collapsed": false
+     "end_time": "2024-11-13T09:01:09.171647Z",
+     "start_time": "2024-11-13T09:01:09.150940Z"
+    }
    },
+   "source": [
+    "(\n",
+    "    pd.DataFrame(ImageEmbedding.list_supported_models())\n",
+    "    .sort_values(\"size_in_GB\")\n",
+    "    .drop(columns=[\"sources\", \"model_file\"])\n",
+    "    .reset_index(drop=True)\n",
+    ")"
+   ],
    "outputs": [
     {
      "data": {
+      "text/plain": [
+       "                         model   dim  \\\n",
+       "0         Qdrant/resnet50-onnx  2048   \n",
+       "1  Qdrant/clip-ViT-B-32-vision   512   \n",
+       "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  "
+      ],
       "text/html": [
        "<div>\n",
        "<style scoped>\n",
@@ -683,39 +704,149 @@
        "  </tbody>\n",
        "</table>\n",
        "</div>"
-      ],
-      "text/plain": [
-       "                         model   dim  \\\n",
-       "0         Qdrant/resnet50-onnx  2048   \n",
-       "1  Qdrant/clip-ViT-B-32-vision   512   \n",
-       "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  "
       ]
      },
-     "execution_count": 6,
+     "execution_count": 15,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
+   "execution_count": 15
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Supported Rerank Cross Encoder Models"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-11-13T09:01:10.313943Z",
+     "start_time": "2024-11-13T09:01:10.298428Z"
+    }
+   },
    "source": [
     "(\n",
-    "    pd.DataFrame(ImageEmbedding.list_supported_models())\n",
+    "    pd.DataFrame(TextCrossEncoder.list_supported_models())\n",
     "    .sort_values(\"size_in_GB\")\n",
     "    .drop(columns=[\"sources\", \"model_file\"])\n",
     "    .reset_index(drop=True)\n",
     ")"
-   ]
+   ],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "                                       model  size_in_GB  \\\n",
+       "0              Xenova/ms-marco-MiniLM-L-6-v2        0.08   \n",
+       "1             Xenova/ms-marco-MiniLM-L-12-v2        0.12   \n",
+       "2            jinaai/jina-reranker-v1-tiny-en        0.13   \n",
+       "3           jinaai/jina-reranker-v1-turbo-en        0.15   \n",
+       "4                     BAAI/bge-reranker-base        1.04   \n",
+       "5  jinaai/jina-reranker-v2-base-multilingual        1.11   \n",
+       "\n",
+       "                                         description       license  \n",
+       "0  MiniLM-L-6-v2 model optimized for re-ranking t...    apache-2.0  \n",
+       "1  MiniLM-L-12-v2 model optimized for re-ranking ...    apache-2.0  \n",
+       "2  Designed for blazing-fast re-ranking with 8K c...    apache-2.0  \n",
+       "3  Designed for blazing-fast re-ranking with 8K c...    apache-2.0  \n",
+       "4  BGE reranker base model for cross-encoder re-r...           mit  \n",
+       "5  A multi-lingual reranker model for cross-encod...  cc-by-nc-4.0  "
+      ],
+      "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>model</th>\n",
+       "      <th>size_in_GB</th>\n",
+       "      <th>description</th>\n",
+       "      <th>license</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>Xenova/ms-marco-MiniLM-L-6-v2</td>\n",
+       "      <td>0.08</td>\n",
+       "      <td>MiniLM-L-6-v2 model optimized for re-ranking t...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>Xenova/ms-marco-MiniLM-L-12-v2</td>\n",
+       "      <td>0.12</td>\n",
+       "      <td>MiniLM-L-12-v2 model optimized for re-ranking ...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>jinaai/jina-reranker-v1-tiny-en</td>\n",
+       "      <td>0.13</td>\n",
+       "      <td>Designed for blazing-fast re-ranking with 8K c...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>jinaai/jina-reranker-v1-turbo-en</td>\n",
+       "      <td>0.15</td>\n",
+       "      <td>Designed for blazing-fast re-ranking with 8K c...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>BAAI/bge-reranker-base</td>\n",
+       "      <td>1.04</td>\n",
+       "      <td>BGE reranker base model for cross-encoder re-r...</td>\n",
+       "      <td>mit</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>5</th>\n",
+       "      <td>jinaai/jina-reranker-v2-base-multilingual</td>\n",
+       "      <td>1.11</td>\n",
+       "      <td>A multi-lingual reranker model for cross-encod...</td>\n",
+       "      <td>cc-by-nc-4.0</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "execution_count": 16
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "outputs": [],
+   "execution_count": null,
+   "source": ""
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": ".venv",
+   "display_name": "Python 3.8.18 ('base')",
    "language": "python",
    "name": "python3"
   },
@@ -729,9 +860,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/common/onnx_model.py

@@ -1,17 +1,7 @@
 import warnings
 from dataclasses import dataclass
 from pathlib import Path
-from typing import (
-    Any,
-    Dict,
-    Generic,
-    Iterable,
-    Optional,
-    Sequence,
-    Tuple,
-    Type,
-    TypeVar,
-)
+from typing import Any, Dict, Generic, Iterable, Optional, Sequence, Tuple, Type, TypeVar
 
 import numpy as np
 import onnxruntime as ort

fastembed/late_interaction/colbert.py

@@ -12,6 +12,7 @@ from fastembed.late_interaction.late_interaction_embedding_base import (
 )
 from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
 
+
 supported_colbert_models = [
     {
         "model": "colbert-ir/colbertv2.0",
@@ -41,7 +42,7 @@ supported_colbert_models = [
 class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
     QUERY_MARKER_TOKEN_ID = 1
     DOCUMENT_MARKER_TOKEN_ID = 2
-    MIN_QUERY_LENGTH = 32
+    MIN_QUERY_LENGTH = 31  # it's 32, we add one additional special token in the beginning
     MASK_TOKEN = "[MASK]"
 
     def _post_process_onnx_output(
@@ -67,15 +68,14 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
         return output.model_output.astype(np.float32)
 
     def _preprocess_onnx_input(
-        self, onnx_input: Dict[str, np.ndarray], is_doc: bool = True
+        self, onnx_input: Dict[str, np.ndarray], is_doc: bool = True, **kwargs: Any
     ) -> 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
+        marker_token = self.DOCUMENT_MARKER_TOKEN_ID if is_doc else self.QUERY_MARKER_TOKEN_ID
+        onnx_input["input_ids"] = np.insert(onnx_input["input_ids"], 1, marker_token, axis=1)
+        onnx_input["attention_mask"] = np.insert(onnx_input["attention_mask"], 1, 1, axis=1)
         return onnx_input
 
-    def tokenize(self, documents: List[str], is_doc: bool = True) -> List[Encoding]:
+    def tokenize(self, documents: List[str], is_doc: bool = True, **kwargs: Any) -> List[Encoding]:
         return (
             self._tokenize_documents(documents=documents)
             if is_doc
@@ -83,9 +83,6 @@ 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])
         # colbert authors recommend to pad queries with [MASK] tokens for query augmentation to improve performance
         if len(encoded[0].ids) < self.MIN_QUERY_LENGTH:
@@ -105,9 +102,6 @@ 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]
         encoded = self.tokenizer.encode_batch(documents)
         return encoded
 
@@ -232,7 +226,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
             **kwargs,
         )
 
-    def query_embed(self, query: Union[str, List[str]], **kwargs) -> Iterable[np.ndarray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
         if isinstance(query, str):
             query = [query]
 

fastembed/late_interaction/jina_colbert.py

@@ -5,6 +5,7 @@ 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",
@@ -24,7 +25,7 @@ supported_jina_colbert_models = [
 class JinaColbert(Colbert):
     QUERY_MARKER_TOKEN_ID = 250002
     DOCUMENT_MARKER_TOKEN_ID = 250003
-    MIN_QUERY_LENGTH = 32
+    MIN_QUERY_LENGTH = 31  # it's 32, we add one additional special token in the beginning
     MASK_TOKEN = "<mask>"
 
     @classmethod
@@ -41,13 +42,12 @@ class JinaColbert(Colbert):
         return supported_jina_colbert_models
 
     def _preprocess_onnx_input(
-        self, onnx_input: Dict[str, np.ndarray], is_doc: bool = True
+        self, onnx_input: Dict[str, np.ndarray], is_doc: bool = True, **kwargs: Any
     ) -> 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 = super()._preprocess_onnx_input(onnx_input, is_doc)
+
+        # the attention mask for jina-colbert-v2 is always 1 in queries
+        if not is_doc:
             onnx_input["attention_mask"][:] = 1
         return onnx_input
 

fastembed/parallel_processor.py

@@ -24,7 +24,7 @@ class QueueSignals(str, Enum):
 
 class Worker:
     @classmethod
-    def start(cls, **kwargs: Any) -> "Worker":
+    def start(cls, *args: Any, **kwargs: Any) -> "Worker":
         raise NotImplementedError()
 
     def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:

fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py

@@ -38,6 +38,36 @@ supported_onnx_models = [
         "description": "BGE reranker base model for cross-encoder re-ranking.",
         "license": "mit",
     },
+    {
+        "model": "jinaai/jina-reranker-v1-tiny-en",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "jinaai/jina-reranker-v1-tiny-en",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "Designed for blazing-fast re-ranking with 8K context length and fewer parameters than jina-reranker-v1-turbo-en.",
+        "license": "apache-2.0",
+    },
+    {
+        "model": "jinaai/jina-reranker-v1-turbo-en",
+        "size_in_GB": 0.15,
+        "sources": {
+            "hf": "jinaai/jina-reranker-v1-turbo-en",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "Designed for blazing-fast re-ranking with 8K context length.",
+        "license": "apache-2.0",
+    },
+    {
+        "model": "jinaai/jina-reranker-v2-base-multilingual",
+        "size_in_GB": 1.11,
+        "sources": {
+            "hf": "jinaai/jina-reranker-v2-base-multilingual",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "A multi-lingual reranker model for cross-encoder re-ranking with 1K context length and sliding window",
+        "license": "cc-by-nc-4.0",
+    },
 ]
 
 

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -4,7 +4,7 @@ from pathlib import Path
 import numpy as np
 from tokenizers import Encoding
 
-from fastembed.common.onnx_model import OnnxModel, OnnxProvider
+from fastembed.common.onnx_model import OnnxModel, OnnxProvider, OnnxOutputContext
 from fastembed.common.preprocessor_utils import load_tokenizer
 from fastembed.common.utils import iter_batch
 
@@ -34,7 +34,7 @@ class OnnxCrossEncoderModel(OnnxModel):
     def tokenize(self, query: str, documents: List[str], **kwargs) -> List[Encoding]:
         return self.tokenizer.encode_batch([(query, doc) for doc in documents])
 
-    def onnx_embed(self, query: str, documents: List[str], **kwargs) -> List[float]:
+    def onnx_embed(self, query: str, documents: List[str], **kwargs) -> OnnxOutputContext:
         tokenized_input = self.tokenize(query, documents, **kwargs)
 
         inputs = {
@@ -51,7 +51,7 @@ class OnnxCrossEncoderModel(OnnxModel):
 
         onnx_input = self._preprocess_onnx_input(inputs, **kwargs)
         outputs = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)
-        return outputs[0][:, 0].tolist()
+        return OnnxOutputContext(model_output=outputs[0][:, 0].tolist())
 
     def _rerank_documents(
         self, query: str, documents: Iterable[str], batch_size: int, **kwargs
@@ -59,7 +59,7 @@ class OnnxCrossEncoderModel(OnnxModel):
         if not hasattr(self, "model") or self.model is None:
             self.load_onnx_model()
         for batch in iter_batch(documents, batch_size):
-            yield from self.onnx_embed(query, batch, **kwargs)
+            yield from self.onnx_embed(query, batch, **kwargs).model_output
 
     def _preprocess_onnx_input(
         self, onnx_input: Dict[str, np.ndarray], **kwargs

fastembed/sparse/utils/tokenizer.py

@@ -5,6 +5,7 @@ from typing import List
 
 
 class SimpleTokenizer:
+    @staticmethod
     def tokenize(text: str) -> List[str]:
         text = re.sub(r"[^\w]", " ", text.lower())
         text = re.sub(r"\s+", " ", text)

pyproject.toml

@@ -1,5 +1,5 @@
 [tool.poetry]
-name = "fastembed"
+name = "fastembed-gpu"
 version = "0.4.2"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
@@ -13,15 +13,15 @@ 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,<1.20.0"
+onnxruntime-gpu = ">=1.17.0,<1.20.0"
 tqdm = "^4.66"
 requests = "^2.31"
 tokenizers = ">=0.15,<1.0"
 huggingface-hub = ">=0.20,<1.0"
 loguru = "^0.7.2"
 numpy = [
-    { version = ">=1.21, <2", python = "<3.12" },
-    { version = ">=1.26, <2", python = ">=3.12" }
+    { version = ">=1.21", python = "<3.12" },
+    { version = ">=1.26", python = ">=3.12" }
 ]
 pillow = "^10.3.0"
 mmh3 = "^4.1.0"

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:
-        shutil.rmtree(model.model._model_dir)
+        delete_model_cache(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:
-        shutil.rmtree(model.model._model_dir)
+        delete_model_cache(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:
-        shutil.rmtree(model.model._model_dir)
+        delete_model_cache(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:
-        shutil.rmtree(model.model._model_dir)
+        delete_model_cache(model.model._model_dir)
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions"])

tests/test_image_onnx_embeddings.py

@@ -1,5 +1,4 @@
 import os
-import shutil
 from io import BytesIO
 
 import numpy as np
@@ -9,6 +8,7 @@ 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:
-            shutil.rmtree(model.model._model_dir)
+            delete_model_cache(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:
-        shutil.rmtree(model.model._model_dir)
+        delete_model_cache(model.model._model_dir)
 
 
 @pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
@@ -102,11 +102,12 @@ 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:
-        shutil.rmtree(model.model._model_dir)
+        delete_model_cache(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 = [
@@ -115,3 +116,5 @@ 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,5 +1,4 @@
 import os
-import shutil
 
 import pytest
 import numpy as np
@@ -7,6 +6,7 @@ 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 = {
@@ -167,7 +167,7 @@ def test_batch_embedding():
             assert np.allclose(value[:, :abridged_dim], expected_result, atol=2e-3)
 
         if is_ci:
-            shutil.rmtree(model.model._model_dir)
+            delete_model_cache(model.model._model_dir)
 
 
 def test_single_embedding():
@@ -182,7 +182,7 @@ def test_single_embedding():
         assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
 
         if is_ci:
-            shutil.rmtree(model.model._model_dir)
+            delete_model_cache(model.model._model_dir)
 
 
 def test_single_embedding_query():
@@ -197,7 +197,7 @@ def test_single_embedding_query():
         assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
 
         if is_ci:
-            shutil.rmtree(model.model._model_dir)
+            delete_model_cache(model.model._model_dir)
 
 
 def test_parallel_processing():
@@ -219,7 +219,7 @@ def test_parallel_processing():
     assert np.allclose(embeddings, embeddings_3, atol=1e-3)
 
     if is_ci:
-        shutil.rmtree(model.model._model_dir)
+        delete_model_cache(model.model._model_dir)
 
 
 @pytest.mark.parametrize(
@@ -227,6 +227,8 @@ 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")
 
@@ -239,3 +241,6 @@ 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:
-            shutil.rmtree(model.model._model_dir)
+            delete_model_cache(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:
-            shutil.rmtree(model.model._model_dir)
+            delete_model_cache(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:
-        shutil.rmtree(model.model._model_dir)
+        delete_model_cache(model.model._model_dir)
 
 
 @pytest.fixture
@@ -116,7 +116,7 @@ def bm25_instance():
     model = Bm25("Qdrant/bm25", language="english")
     yield model
     if ci:
-        shutil.rmtree(model._model_dir)
+        delete_model_cache(model._model_dir)
 
 
 def test_stem_with_stopwords_and_punctuation(bm25_instance):
@@ -150,11 +150,13 @@ 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")
 
@@ -167,3 +169,6 @@ 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,14 +2,17 @@ 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]),
     "Xenova/ms-marco-MiniLM-L-12-v2": np.array([9.330912, -2.0380247]),
     "BAAI/bge-reranker-base": np.array([6.15733337, -3.65939403]),
+    "jinaai/jina-reranker-v1-tiny-en": np.array([2.5911, 0.1122]),
+    "jinaai/jina-reranker-v1-turbo-en": np.array([1.8295, -2.8908]),
+    "jinaai/jina-reranker-v2-base-multilingual": np.array([1.6533, -1.6455]),
 }
 
 
@@ -32,12 +35,16 @@ def test_rerank():
             scores, canonical_scores, atol=1e-3
         ), f"Model: {model_name}, Scores: {scores}, Expected: {canonical_scores}"
         if is_ci:
-            shutil.rmtree(model.model._model_dir)
+            delete_model_cache(model.model._model_dir)
 
 
 @pytest.mark.parametrize(
     "model_name",
-    ["Xenova/ms-marco-MiniLM-L-6-v2", "Xenova/ms-marco-MiniLM-L-12-v2", "BAAI/bge-reranker-base"],
+    [
+        model_desc["model"]
+        for model_desc in TextCrossEncoder.list_supported_models()
+        if model_desc["size_in_GB"] < 1 and model_desc["model"] in CANONICAL_SCORE_VALUES.keys()
+    ],
 )
 def test_batch_rerank(model_name):
     is_ci = os.getenv("CI")
@@ -55,7 +62,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:
-        shutil.rmtree(model.model._model_dir)
+        delete_model_cache(model.model._model_dir)
 
 
 @pytest.mark.parametrize(
@@ -63,9 +70,13 @@ 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:
-            shutil.rmtree(model.model._model_dir)
+            delete_model_cache(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:
-        shutil.rmtree(model.model._model_dir)
+        delete_model_cache(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:
-        shutil.rmtree(model.model._model_dir)
+        delete_model_cache(model.model._model_dir)
 
 
 @pytest.mark.parametrize(
@@ -136,6 +136,7 @@ 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"]
@@ -147,3 +148,6 @@ 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

@@ -0,0 +1,32 @@
+import shutil
+import traceback
+
+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.
+    """
+
+    def on_error(func, path, exc_info):
+        print("Failed to remove: ", path)
+        print("Exception: ", exc_info)
+        traceback.print_exception(*exc_info)
+
+    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():
+        # todo: PermissionDenied is raised on blobs removal in Windows, with blobs > 2GB
+        shutil.rmtree(model_dir, onerror=on_error)