disable tmate

* chore: Add notice file for jina ai models

* chore: Update notice

* Update NOTICE

Co-authored-by: George <george.panchuk@qdrant.tech>

* chore: added jina embeddings v3

* chore: removed unsupported models

---------

Co-authored-by: George <george.panchuk@qdrant.tech>

.github/workflows/python-tests.yml

@@ -1,10 +1,9 @@
 name: Tests
+run-name: Tests (gpu)
 
 on:
   push:
     branches: [ master, main, gpu ]
-  schedule:
-    - cron: 0 0 * * *
   pull_request:
 
 env:
@@ -23,8 +22,6 @@ jobs:
           - '3.12.x'
         os:
           - ubuntu-latest
-          - macos-latest
-          - windows-latest
 
     runs-on: ${{ matrix.os }}
 
@@ -36,19 +33,14 @@ 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
           poetry config virtualenvs.create false
           poetry install --no-interaction --no-ansi --without docs
 
-      - name: Install Test Dependencies
-        run: pip install pytest pytest-md pytest-emoji
-
       - name: Run pytest
-        uses: pavelzw/pytest-action@v2
-        with:
-            verbose: true
-            emoji: true
-            job-summary: true
-            report-title: 'FastEmbed Test Report'
+        run: |
+          poetry run pytest

NOTICE

@@ -0,0 +1,12 @@
+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.

README.md

@@ -8,9 +8,9 @@ The default text embedding (`TextEmbedding`) model is Flag Embedding, presented
 
 1. Light: FastEmbed is a lightweight library with few external dependencies. We don't require a GPU and don't download GBs of PyTorch dependencies, and instead use the ONNX Runtime. This makes it a great candidate for serverless runtimes like AWS Lambda. 
 
-2. Fast: FastEmbed is designed for speed. We use the ONNX Runtime, which is faster than PyTorch. We also use data-parallelism for encoding large datasets.
+2. Fast: FastEmbed is designed for speed. We use the ONNX Runtime, which is faster than PyTorch. We also use data parallelism for encoding large datasets.
 
-3. Accurate: FastEmbed is better than OpenAI Ada-002. We also [supported](https://qdrant.github.io/fastembed/examples/Supported_Models/) an ever expanding set of models, including a few multilingual models.
+3. Accurate: FastEmbed is better than OpenAI Ada-002. We also [support](https://qdrant.github.io/fastembed/examples/Supported_Models/) an ever-expanding set of models, including a few multilingual models.
 
 ## 🚀 Installation
 
@@ -54,7 +54,7 @@ The list of all the available models can be found [here](https://qdrant.github.i
 from fastembed import TextEmbedding
 
 model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
-embeddings = list(embedding_model.embed(documents))
+embeddings = list(model.embed(documents))
 
 # [
 #   array([-0.1115,  0.0097,  0.0052,  0.0195, ...], dtype=float32),
@@ -73,7 +73,7 @@ embeddings = list(embedding_model.embed(documents))
 from fastembed import SparseTextEmbedding
 
 model = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1")
-embeddings = list(embedding_model.embed(documents))
+embeddings = list(model.embed(documents))
 
 # [
 #   SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
@@ -88,7 +88,7 @@ embeddings = list(embedding_model.embed(documents))
 from fastembed import SparseTextEmbedding
 
 model = SparseTextEmbedding(model_name="Qdrant/bm42-all-minilm-l6-v2-attentions")
-embeddings = list(embedding_model.embed(documents))
+embeddings = list(model.embed(documents))
 
 # [
 #   SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
@@ -104,7 +104,7 @@ embeddings = list(embedding_model.embed(documents))
 from fastembed import LateInteractionTextEmbedding
 
 model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0")
-embeddings = list(embedding_model.embed(documents))
+embeddings = list(model.embed(documents))
 
 # [
 #   array([
@@ -129,7 +129,7 @@ images = [
 ]
 
 model = ImageEmbedding(model_name="Qdrant/clip-ViT-B-32-vision")
-embeddings = list(embedding_model.embed(images))
+embeddings = list(model.embed(images))
 
 # [
 #   array([-0.1115,  0.0097,  0.0052,  0.0195, ...], dtype=float32),
@@ -147,7 +147,7 @@ It requires installation of the `fastembed-gpu` package.
 pip install fastembed-gpu
 ```
 
-Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for the detailed instructions and CUDA 12.x support.
+Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for detailed instructions, CUDA 12.x support and troubleshooting of the common issues.
 
 ```python
 from fastembed import TextEmbedding
@@ -209,4 +209,4 @@ search_result = client.query(
     query_text="This is a query document"
 )
 print(search_result)
-```
+```

RELEASE.md

@@ -12,7 +12,7 @@ This is a guide how to release `fastembed` and `fastembed-gpu` packages.
 ```bash
 git checkout gpu
 git rebase main
-git push origin gpu
+git push -f origin gpu
 ```
 
 4. Draft release notes

docs/examples/FastEmbed_GPU.ipynb

@@ -37,28 +37,7 @@
     "id": "3aiGPqjCeGzo"
    },
    "source": [
-    "**NOTE**: `onnxruntime-gpu` and `onnxruntime` can't be installed in the same environment. If you have `onnxruntime` installed, you would need to uninstall it before installing `onnxruntime-gpu`. Same is true for `fastembed` and `fastembed-gpu`.\n",
-    "\n",
-    "### CUDA 12.x support\n",
-    "\n",
-    "By default `onnxruntime-gpu` is shipped with CUDA 11.8 support.\n",
-    "CUDA 12.x support requires installation of `onnxruntime-gpu` with providing of a direct url:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "OoSfWFFZeJ5t",
-    "outputId": "417b9332-6a7b-4000-c74b-4ed2b5b76590"
-   },
-   "outputs": [],
-   "source": [
-    "!pip install onnxruntime-gpu -i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ -qq\n",
-    "!pip install fastembed-gpu -qqq"
+    "**NOTE**: `onnxruntime-gpu` and `onnxruntime` can't be installed in the same environment. If you have `onnxruntime` installed, you would need to uninstall it before installing `onnxruntime-gpu`. Same is true for `fastembed` and `fastembed-gpu`."
    ]
   },
   {
@@ -67,9 +46,65 @@
     "id": "3xx3r-9jgAMi"
    },
    "source": [
+    "### CUDA 12.x support\n",
     "You can check your CUDA version using such commands as `nvidia-smi` or `nvcc --version`\n",
     "\n",
-    "Google Colab notebooks have CUDA 12.x."
+    "Starting from version 1.19.0, onnxruntime-gpu ships with support for CUDA 12.x by default.\n",
+    "\n",
+    "Google Colab notebooks have by default CUDA 12.x and CuDNN 8.x.\n",
+    "\n",
+    "Latest version of `onnxruntime-gpu` requires CuDNN 9.x, in order to install it you can run the following command: "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!sudo apt install cudnn9\n",
+    "!pip install fastembed-gpu -qqq"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "If it necessary to work with CuDNN 8, you can consider locking `onnxruntime-gpu` to 1.18.0 with CUDA 12.x by this command:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install onnxruntime-gpu==1.18.0 -i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ -qq\n",
+    "!pip install fastembed-gpu -qqq"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### CUDA 11.x support\n",
+    "To use latest version of `onnxruntime-gpu` with CUDA 11.x, you can run the following command:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install onnxruntime-gpu -i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-11/pypi/simple/ -qq"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "**NOTE**: Ensure that CuDNN 9.x is installed when working with the latest `onnxruntime-gpu`, whether using CUDA 11.x or 12.x."
    ]
   },
   {
@@ -82,7 +117,80 @@
     "\n",
     "FastEmbed does not include CUDA drivers and CuDNN libraries.\n",
     "You would need to take care of the environment setup on your own.\n",
-    "Dependencies required for the chosen onnxruntime version can be found [here](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#requirements)"
+    "The dependencies required for the chosen onnxruntime version are listed in the [CUDA Execution Provider requirements](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#requirements)."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Setting up fastembed-gpu on GCP\n",
+    "\n",
+    "#### CUDA drivers\n",
+    "[CUDA 11.8 toolkit](https://developer.nvidia.com/cuda-11-8-0-download-archive) or [CUDA 12.x toolkit](https://developer.nvidia.com/cuda-downloads) has to be installed if they haven't yet been set up.\n",
+    "\n",
+    "#### Example of setting up CUDA 12.x on Ubuntu 22.04\n",
+    "Make sure to download an archive which has been created for your particular platform, CPU architecture and OS distribution.\n",
+    "\n",
+    "For Ubuntu 22.04 with x86_64 CPU architecture the following [archive](https://developer.nvidia.com/cuda-downloads?target_os=Linux&target_arch=x86_64&Distribution=Ubuntu&target_version=22.04&target_type=deb_network) has to be downloaded.\n",
+    "\n",
+    "```bash\n",
+    "wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb\n",
+    "sudo dpkg -i cuda-keyring_1.1-1_all.deb\n",
+    "sudo apt-get update\n",
+    "sudo apt-get -y install cuda\n",
+    "```\n",
+    "**NOTE**: Specific CUDA libraries can be found in the [meta packages section](https://docs.nvidia.com/cuda/cuda-installation-guide-linux/#meta-packages) in the CUDA installation guide.\n",
+    "\n",
+    "**NOTE**: When installing CUDA, the environment variable might not be set by default. Make sure to add the following line to your environment variables:\n",
+    "```bash\n",
+    "LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH\n",
+    "```\n",
+    "This will ensure that the CUDA libraries are properly linked.\n",
+    "\n",
+    "#### CuDNN 9.x\n",
+    "CuDNN 9.x library can be installed via the following [archive](https://developer.nvidia.com/rdp/cudnn-archive).\n",
+    "\n",
+    "#### Example of setting up CuDNN 9.x on Ubuntu 22.04\n",
+    "CuDNN 9.x for Ubuntu 22.04 x86_64 [archive](https://developer.nvidia.com/cudnn-downloads?target_os=Linux&target_arch=x86_64&Distribution=Ubuntu&target_version=22.04&target_type=deb_network) can be downloaded and installed in the following way:\n",
+    "```bash\n",
+    "wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb\n",
+    "sudo dpkg -i cuda-keyring_1.1-1_all.deb\n",
+    "sudo apt-get update\n",
+    "sudo apt-get -y install cudnn\n",
+    "```\n",
+    "**NOTE**: When installing CuDNN, you can choose specific version, cudnn-cuda-11 or cudnn-cuda-12"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Common issues\n",
+    "\n",
+    "The following are some common issues that may arise while using `fastembed-gpu` if not installed properly:\n",
+    "\n",
+    "CUDA library is not installed:\n",
+    "```bash\n",
+    "FAIL : Failed to load library libonnxruntime_providers_cuda.so with error: libcublasLt.so.x: cannot open shared object file: No such file or directory\n",
+    "```\n",
+    "\n",
+    "\n",
+    "CuDNN library is not installed:\n",
+    "```bash\n",
+    "FAIL : Failed to load library libonnxruntime_providers_cuda.so with error: libcudnn.so.x: cannot open shared object file: No such file or directory\n",
+    "```\n",
+    "\n",
+    "\n",
+    "CUDA library path is not set:\n",
+    "```bash\n",
+    "FAIL : Failed to load library libonnxruntime_providers_cuda.so with error: libcufft.so.x: failed to map segment from shared object\n",
+    "```\n",
+    "\n",
+    "Make sure to add the following line to your environment variables:\n",
+    "```bash\n",
+    "LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH\n",
+    "```"
    ]
   },
   {

docs/examples/Supported_Models.ipynb

@@ -2,23 +2,14 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 1,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2024-05-31T18:13:23.806907Z",
      "start_time": "2024-05-31T18:13:23.797078Z"
     }
    },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "The autoreload extension is already loaded. To reload it, use:\n",
-      "  %reload_ext autoreload\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "%load_ext autoreload\n",
     "%autoreload 2"
@@ -26,14 +17,23 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2024-05-31T18:14:31.147674Z",
      "start_time": "2024-05-31T18:14:31.134015Z"
     }
    },
-   "outputs": [],
+   "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"
+     ]
+    }
+   ],
    "source": [
     "import pandas as pd\n",
     "\n",
@@ -54,7 +54,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 3,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2024-05-31T18:13:25.863008Z",
@@ -86,6 +86,7 @@
        "      <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",
@@ -94,161 +95,200 @@
        "      <th>0</th>\n",
        "      <td>BAAI/bge-small-en-v1.5</td>\n",
        "      <td>384</td>\n",
-       "      <td>Fast and Default English model</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
+       "      <td>mit</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>Fast and recommended Chinese model</td>\n",
+       "      <td>Text embeddings, Unimodal (text), Chinese, 512...</td>\n",
+       "      <td>mit</td>\n",
        "      <td>0.090</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>2</th>\n",
-       "      <td>sentence-transformers/all-MiniLM-L6-v2</td>\n",
+       "      <td>snowflake/snowflake-arctic-embed-xs</td>\n",
        "      <td>384</td>\n",
-       "      <td>Sentence Transformer model, MiniLM-L6-v2</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
+       "      <td>apache-2.0</td>\n",
        "      <td>0.090</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>3</th>\n",
-       "      <td>snowflake/snowflake-arctic-embed-xs</td>\n",
+       "      <td>sentence-transformers/all-MiniLM-L6-v2</td>\n",
        "      <td>384</td>\n",
-       "      <td>Based on all-MiniLM-L6-v2 model with only 22m ...</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 256...</td>\n",
+       "      <td>apache-2.0</td>\n",
        "      <td>0.090</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>4</th>\n",
        "      <td>jinaai/jina-embeddings-v2-small-en</td>\n",
        "      <td>512</td>\n",
-       "      <td>English embedding model supporting 8192 sequen...</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 819...</td>\n",
+       "      <td>apache-2.0</td>\n",
        "      <td>0.120</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>5</th>\n",
-       "      <td>snowflake/snowflake-arctic-embed-s</td>\n",
+       "      <td>BAAI/bge-small-en</td>\n",
        "      <td>384</td>\n",
-       "      <td>Based on infloat/e5-small-unsupervised, does n...</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
+       "      <td>mit</td>\n",
        "      <td>0.130</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>6</th>\n",
-       "      <td>BAAI/bge-small-en</td>\n",
+       "      <td>snowflake/snowflake-arctic-embed-s</td>\n",
        "      <td>384</td>\n",
-       "      <td>Fast English model</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
+       "      <td>apache-2.0</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>Quantized 8192 context length english model</td>\n",
+       "      <td>Text embeddings, Multimodal (text, image), Eng...</td>\n",
+       "      <td>apache-2.0</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>Base English model, v1.5</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
+       "      <td>mit</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>Sentence Transformer model, paraphrase-multili...</td>\n",
+       "      <td>Text embeddings, Unimodal (text), Multilingual...</td>\n",
+       "      <td>apache-2.0</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>CLIP text encoder</td>\n",
+       "      <td>Text embeddings, Multimodal (text&amp;image), Engl...</td>\n",
+       "      <td>mit</td>\n",
        "      <td>0.250</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>11</th>\n",
-       "      <td>BAAI/bge-base-en</td>\n",
+       "      <td>jinaai/jina-embeddings-v2-base-de</td>\n",
        "      <td>768</td>\n",
-       "      <td>Base English model</td>\n",
-       "      <td>0.420</td>\n",
+       "      <td>Text embeddings, Unimodal (text), Multilingual...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.320</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>12</th>\n",
-       "      <td>snowflake/snowflake-arctic-embed-m</td>\n",
+       "      <td>BAAI/bge-base-en</td>\n",
        "      <td>768</td>\n",
-       "      <td>Based on intfloat/e5-base-unsupervised model, ...</td>\n",
-       "      <td>0.430</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
+       "      <td>mit</td>\n",
+       "      <td>0.420</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>13</th>\n",
-       "      <td>nomic-ai/nomic-embed-text-v1</td>\n",
+       "      <td>snowflake/snowflake-arctic-embed-m</td>\n",
        "      <td>768</td>\n",
-       "      <td>8192 context length english model</td>\n",
-       "      <td>0.520</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.430</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>14</th>\n",
-       "      <td>jinaai/jina-embeddings-v2-base-en</td>\n",
+       "      <td>nomic-ai/nomic-embed-text-v1.5</td>\n",
        "      <td>768</td>\n",
-       "      <td>English embedding model supporting 8192 sequen...</td>\n",
+       "      <td>Text embeddings, Multimodal (text, image), Eng...</td>\n",
+       "      <td>apache-2.0</td>\n",
        "      <td>0.520</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>15</th>\n",
-       "      <td>nomic-ai/nomic-embed-text-v1.5</td>\n",
+       "      <td>jinaai/jina-embeddings-v2-base-en</td>\n",
        "      <td>768</td>\n",
-       "      <td>8192 context length english model</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 819...</td>\n",
+       "      <td>apache-2.0</td>\n",
        "      <td>0.520</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>16</th>\n",
-       "      <td>snowflake/snowflake-arctic-embed-m-long</td>\n",
+       "      <td>nomic-ai/nomic-embed-text-v1</td>\n",
        "      <td>768</td>\n",
-       "      <td>Based on nomic-ai/nomic-embed-text-v1-unsuperv...</td>\n",
-       "      <td>0.540</td>\n",
+       "      <td>Text embeddings, Multimodal (text, image), Eng...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.520</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>17</th>\n",
-       "      <td>mixedbread-ai/mxbai-embed-large-v1</td>\n",
-       "      <td>1024</td>\n",
-       "      <td>MixedBread Base sentence embedding model, does...</td>\n",
-       "      <td>0.640</td>\n",
+       "      <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>0.540</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>18</th>\n",
-       "      <td>sentence-transformers/paraphrase-multilingual-...</td>\n",
-       "      <td>768</td>\n",
-       "      <td>Sentence-transformers model for tasks like clu...</td>\n",
-       "      <td>1.000</td>\n",
+       "      <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>0.640</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>19</th>\n",
-       "      <td>snowflake/snowflake-arctic-embed-l</td>\n",
-       "      <td>1024</td>\n",
-       "      <td>Based on intfloat/e5-large-unsupervised, large...</td>\n",
-       "      <td>1.020</td>\n",
+       "      <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>0.640</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>20</th>\n",
-       "      <td>BAAI/bge-large-en-v1.5</td>\n",
-       "      <td>1024</td>\n",
-       "      <td>Large English model, v1.5</td>\n",
-       "      <td>1.200</td>\n",
+       "      <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>1.000</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>21</th>\n",
-       "      <td>thenlper/gte-large</td>\n",
+       "      <td>snowflake/snowflake-arctic-embed-l</td>\n",
        "      <td>1024</td>\n",
-       "      <td>Large general text embeddings model</td>\n",
-       "      <td>1.200</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
+       "      <td>apache-2.0</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>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>1.200</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>24</th>\n",
        "      <td>intfloat/multilingual-e5-large</td>\n",
        "      <td>1024</td>\n",
-       "      <td>Multilingual model, e5-large. Recommend using ...</td>\n",
+       "      <td>Text embeddings, Unimodal (text), Multilingual...</td>\n",
+       "      <td>mit</td>\n",
        "      <td>2.240</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
@@ -259,55 +299,59 @@
        "                                                model   dim  \\\n",
        "0                              BAAI/bge-small-en-v1.5   384   \n",
        "1                              BAAI/bge-small-zh-v1.5   512   \n",
-       "2              sentence-transformers/all-MiniLM-L6-v2   384   \n",
-       "3                 snowflake/snowflake-arctic-embed-xs   384   \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                  snowflake/snowflake-arctic-embed-s   384   \n",
-       "6                                   BAAI/bge-small-en   384   \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                                   BAAI/bge-base-en   768   \n",
-       "12                 snowflake/snowflake-arctic-embed-m   768   \n",
-       "13                       nomic-ai/nomic-embed-text-v1   768   \n",
-       "14                  jinaai/jina-embeddings-v2-base-en   768   \n",
-       "15                     nomic-ai/nomic-embed-text-v1.5   768   \n",
-       "16            snowflake/snowflake-arctic-embed-m-long   768   \n",
-       "17                 mixedbread-ai/mxbai-embed-large-v1  1024   \n",
-       "18  sentence-transformers/paraphrase-multilingual-...   768   \n",
-       "19                 snowflake/snowflake-arctic-embed-l  1024   \n",
-       "20                             BAAI/bge-large-en-v1.5  1024   \n",
-       "21                                 thenlper/gte-large  1024   \n",
-       "22                     intfloat/multilingual-e5-large  1024   \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  size_in_GB  \n",
-       "0                      Fast and Default English model       0.067  \n",
-       "1                  Fast and recommended Chinese model       0.090  \n",
-       "2            Sentence Transformer model, MiniLM-L6-v2       0.090  \n",
-       "3   Based on all-MiniLM-L6-v2 model with only 22m ...       0.090  \n",
-       "4   English embedding model supporting 8192 sequen...       0.120  \n",
-       "5   Based on infloat/e5-small-unsupervised, does n...       0.130  \n",
-       "6                                  Fast English model       0.130  \n",
-       "7         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                                 Base English model       0.420  \n",
-       "12  Based on intfloat/e5-base-unsupervised model, ...       0.430  \n",
-       "13                  8192 context length english model       0.520  \n",
-       "14  English embedding model supporting 8192 sequen...       0.520  \n",
-       "15                  8192 context length english model       0.520  \n",
-       "16  Based on nomic-ai/nomic-embed-text-v1-unsuperv...       0.540  \n",
-       "17  MixedBread Base sentence embedding model, does...       0.640  \n",
-       "18  Sentence-transformers model for tasks like clu...       1.000  \n",
-       "19  Based on intfloat/e5-large-unsupervised, large...       1.020  \n",
-       "20                          Large English model, v1.5       1.200  \n",
-       "21                Large general text embeddings model       1.200  \n",
-       "22  Multilingual model, e5-large. Recommend using ...       2.240  "
+       "                                          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": 6,
+     "execution_count": 3,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -331,7 +375,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 4,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2024-05-31T18:13:27.124747Z",
@@ -363,30 +407,47 @@
        "      <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",
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>0</th>\n",
-       "      <td>Qdrant/bm42-all-minilm-l6-v2-attentions</td>\n",
-       "      <td>30522</td>\n",
-       "      <td>Light sparse embedding model, which assigns an...</td>\n",
-       "      <td>0.090</td>\n",
+       "      <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",
        "    <tr>\n",
        "      <th>1</th>\n",
-       "      <td>prithvida/Splade_PP_en_v1</td>\n",
-       "      <td>30522</td>\n",
-       "      <td>Misspelled version of the model. Retained for ...</td>\n",
-       "      <td>0.532</td>\n",
+       "      <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>30522</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",
+       "    <tr>\n",
+       "      <th>3</th>\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>0.532</td>\n",
+       "      <td>NaN</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
@@ -394,24 +455,32 @@
       ],
       "text/plain": [
        "                                     model  vocab_size  \\\n",
-       "0  Qdrant/bm42-all-minilm-l6-v2-attentions       30522   \n",
-       "1                prithvida/Splade_PP_en_v1       30522   \n",
-       "2               prithivida/Splade_PP_en_v1       30522   \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  size_in_GB  \n",
-       "0  Light sparse embedding model, which assigns an...       0.090  \n",
-       "1  Misspelled version of the model. Retained for ...       0.532  \n",
-       "2  Independent Implementation of SPLADE++ Model f...       0.532  "
+       "                                         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": 8,
+     "execution_count": 4,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "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",
@@ -429,7 +498,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 5,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2024-05-31T18:14:34.370252Z",
@@ -462,34 +531,68 @@
        "      <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             description  size_in_GB\n",
-       "0  colbert-ir/colbertv2.0  128  Late interaction model        0.44"
+       "                                   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": 10,
+     "execution_count": 5,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "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",
@@ -507,7 +610,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 6,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2024-05-31T18:14:42.501881Z",
@@ -540,6 +643,7 @@
        "      <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",
@@ -548,16 +652,34 @@
        "      <th>0</th>\n",
        "      <td>Qdrant/resnet50-onnx</td>\n",
        "      <td>2048</td>\n",
-       "      <td>ResNet-50 from `Deep Residual Learning for Ima...</td>\n",
+       "      <td>Image embeddings, Unimodal (image), 2016 year</td>\n",
+       "      <td>apache-2.0</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>CLIP vision encoder based on ViT-B/32</td>\n",
+       "      <td>Image embeddings, Multimodal (text&amp;image), 202...</td>\n",
+       "      <td>mit</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>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>0.82</td>\n",
+       "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
@@ -566,20 +688,25 @@
        "                         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  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  "
+       "                                         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": 12,
+     "execution_count": 6,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "(\n",
-    "  pd.DataFrame(ImageEmbedding.list_supported_models()).sort_values(\"size_in_GB\")\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",
     ")"
@@ -588,7 +715,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3.8.18 ('base')",
+   "display_name": ".venv",
    "language": "python",
    "name": "python3"
   },
@@ -602,14 +729,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.4"
+   "version": "3.10.15"
   },
-  "orig_nbformat": 4,
-  "vscode": {
-   "interpreter": {
-    "hash": "c4a27af61e455bc18dcf16f5867a2ff0402fa12b01dd0f6ce3a79ae73ad15e91"
-   }
-  }
+  "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 2

fastembed/common/__init__.py

@@ -1,3 +1,3 @@
-from fastembed.common.types import ImageInput, OnnxProvider, PathInput
+from fastembed.common.types import ImageInput, OnnxProvider, PathInput, PilInput
 
-__all__ = ["OnnxProvider", "ImageInput", "PathInput"]
+__all__ = ["OnnxProvider", "ImageInput", "PathInput", "PilInput"]

fastembed/common/model_management.py

@@ -1,4 +1,5 @@
 import os
+import time
 import shutil
 import tarfile
 from pathlib import Path
@@ -42,9 +43,7 @@ class ModelManagement:
         raise ValueError(f"Model {model_name} is not supported in {cls.__name__}.")
 
     @classmethod
-    def download_file_from_gcs(
-        cls, url: str, output_path: str, show_progress: bool = True
-    ) -> str:
+    def download_file_from_gcs(cls, url: str, output_path: str, show_progress: bool = True) -> str:
         """
         Downloads a file from Google Cloud Storage.
 
@@ -73,11 +72,9 @@ class ModelManagement:
 
         # Warn if the total size is zero
         if total_size_in_bytes == 0:
-            print(
-                f"Warning: Content-length header is missing or zero in the response from {url}."
-            )
+            print(f"Warning: Content-length header is missing or zero in the response from {url}.")
 
-        show_progress = total_size_in_bytes and show_progress
+        show_progress = bool(total_size_in_bytes and show_progress)
 
         with tqdm(
             total=total_size_in_bytes,
@@ -98,6 +95,7 @@ class ModelManagement:
         hf_source_repo: str,
         cache_dir: Optional[str] = None,
         extra_patterns: Optional[List[str]] = None,
+        local_files_only: bool = False,
         **kwargs,
     ) -> str:
         """
@@ -107,6 +105,7 @@ class ModelManagement:
             cache_dir (Optional[str]): The path to the cache directory.
             extra_patterns (Optional[List[str]]): extra patterns to allow in the snapshot download, typically
                 includes the required model files.
+            local_files_only (bool, optional): Whether to only use local files. Defaults to False.
         Returns:
             Path: The path to the model directory.
         """
@@ -124,7 +123,8 @@ class ModelManagement:
             repo_id=hf_source_repo,
             allow_patterns=allow_patterns,
             cache_dir=cache_dir,
-            local_files_only=kwargs.get("local_files_only", False),
+            local_files_only=local_files_only,
+            **kwargs,
         )
 
     @classmethod
@@ -151,7 +151,9 @@ class ModelManagement:
             # Open the tar.gz file
             with tarfile.open(targz_path, "r:gz") as tar:
                 # Extract all files into the cache directory
-                tar.extractall(path=cache_dir)
+                tar.extractall(
+                    path=cache_dir,
+                )
         except tarfile.TarError as e:
             # If any error occurs while opening or extracting the tar.gz file,
             # delete the cache directory (if it was created in this function)
@@ -164,10 +166,9 @@ class ModelManagement:
 
     @classmethod
     def retrieve_model_gcs(
-        cls, model_name: str, source_url: str, cache_dir: str
+        cls, model_name: str, source_url: str, cache_dir: str, local_files_only: bool = False
     ) -> Path:
         fast_model_name = f"fast-{model_name.split('/')[-1]}"
-
         cache_tmp_dir = Path(cache_dir) / "tmp"
         model_tmp_dir = cache_tmp_dir / fast_model_name
         model_dir = Path(cache_dir) / fast_model_name
@@ -186,26 +187,32 @@ class ModelManagement:
         if model_tar_gz.exists():
             model_tar_gz.unlink()
 
-        cls.download_file_from_gcs(
-            source_url,
-            output_path=str(model_tar_gz),
-        )
+        if not local_files_only:
+            cls.download_file_from_gcs(
+                source_url,
+                output_path=str(model_tar_gz),
+            )
 
-        cls.decompress_to_cache(
-            targz_path=str(model_tar_gz), cache_dir=str(cache_tmp_dir)
-        )
-        assert (
-            model_tmp_dir.exists()
-        ), f"Could not find {model_tmp_dir} in {cache_tmp_dir}"
+            cls.decompress_to_cache(targz_path=str(model_tar_gz), cache_dir=str(cache_tmp_dir))
+            assert model_tmp_dir.exists(), f"Could not find {model_tmp_dir} in {cache_tmp_dir}"
 
-        model_tar_gz.unlink()
-        # Rename from tmp to final name is atomic
-        model_tmp_dir.rename(model_dir)
+            model_tar_gz.unlink()
+            # Rename from tmp to final name is atomic
+            model_tmp_dir.rename(model_dir)
+        else:
+            logger.error(
+                f"Could not find the model tar.gz file at {model_dir} and local_files_only=True."
+            )
+            raise ValueError(
+                f"Could not find the model tar.gz file at {model_dir} and local_files_only=True."
+            )
 
         return model_dir
 
     @classmethod
-    def download_model(cls, model: Dict[str, Any], cache_dir: Path, **kwargs) -> Path:
+    def download_model(
+        cls, model: Dict[str, Any], cache_dir: Path, retries: int = 3, **kwargs
+    ) -> Path:
         """
         Downloads a model from HuggingFace Hub or Google Cloud Storage.
 
@@ -225,34 +232,58 @@ class ModelManagement:
                 }
                 ```
             cache_dir (str): The path to the cache directory.
+            retries: (int): The number of times to retry (including the first attempt)
 
         Returns:
             Path: The path to the downloaded model directory.
         """
-
+        local_files_only = kwargs.get("local_files_only", False)
+        retries = 1 if local_files_only else retries
         hf_source = model.get("sources", {}).get("hf")
         url_source = model.get("sources", {}).get("url")
 
-        if hf_source:
-            extra_patterns = [model["model_file"]]
-            extra_patterns.extend(model.get("additional_files", []))
+        sleep = 3.0
+        while retries > 0:
+            retries -= 1
 
-            try:
-                return Path(
-                    cls.download_files_from_huggingface(
-                        hf_source,
-                        cache_dir=str(cache_dir),
-                        extra_patterns=extra_patterns,
-                        local_files_only=kwargs.get("local_files_only", False),
+            if hf_source:
+                extra_patterns = [model["model_file"]]
+                extra_patterns.extend(model.get("additional_files", []))
+
+                try:
+                    return Path(
+                        cls.download_files_from_huggingface(
+                            hf_source,
+                            cache_dir=str(cache_dir),
+                            extra_patterns=extra_patterns,
+                            **kwargs,
+                        )
                     )
-                )
-            except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
+                except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
+                    if not local_files_only:
+                        logger.error(
+                            f"Could not download model from HuggingFace: {e} "
+                            "Falling back to other sources."
+                        )
+            if url_source or local_files_only:
+                try:
+                    return cls.retrieve_model_gcs(
+                        model["model"],
+                        url_source,
+                        str(cache_dir),
+                        local_files_only=local_files_only,
+                    )
+                except Exception:
+                    if not local_files_only:
+                        logger.error(f"Could not download model from url: {url_source}")
+
+            if local_files_only:
+                logger.error("Could not find model in cache_dir")
+            else:
                 logger.error(
-                    f"Could not download model from HuggingFace: {e}"
-                    "Falling back to other sources."
+                    f"Could not download model from either source, sleeping for {sleep} seconds, {retries} retries left."
                 )
+            time.sleep(sleep)
+            sleep *= 3
 
-        if url_source:
-            return cls.retrieve_model_gcs(model["model"], url_source, str(cache_dir))
-
-        raise ValueError(f"Could not download model {model['model']} from any source.")
+        raise ValueError(f"Could not load model {model['model']} from any source.")

fastembed/common/onnx_model.py

@@ -50,19 +50,28 @@ class OnnxModel(Generic[T]):
         """
         return onnx_input
 
-    def load_onnx_model(
+    def _load_onnx_model(
         self,
         model_dir: Path,
         model_file: str,
         threads: Optional[int],
         providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_id: Optional[int] = None,
     ) -> None:
         model_path = model_dir / model_file
         # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
 
-        onnx_providers = (
-            ["CPUExecutionProvider"] if providers is None else list(providers)
-        )
+        if providers is not None:
+            onnx_providers = list(providers)
+        elif cuda:
+            if device_id is None:
+                onnx_providers = ["CUDAExecutionProvider"]
+            else:
+                onnx_providers = [("CUDAExecutionProvider", {"device_id": device_id})]
+        else:
+            onnx_providers = ["CPUExecutionProvider"]
+
         available_providers = ort.get_available_providers()
         requested_provider_names = []
         for provider in onnx_providers:
@@ -94,6 +103,9 @@ class OnnxModel(Generic[T]):
                     RuntimeWarning,
                 )
 
+    def load_onnx_model(self) -> None:
+        raise NotImplementedError("Subclasses must implement this method")
+
     def onnx_embed(self, *args, **kwargs) -> OnnxOutputContext:
         raise NotImplementedError("Subclasses must implement this method")
 

fastembed/common/preprocessor_utils.py

@@ -1,7 +1,6 @@
 import json
 from pathlib import Path
 from typing import Tuple
-
 from tokenizers import AddedToken, Tokenizer
 
 from fastembed.image.transform.operators import Compose
@@ -18,7 +17,7 @@ def load_special_tokens(model_dir: Path) -> dict:
     return tokens_map
 
 
-def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tuple[Tokenizer, dict]:
+def load_tokenizer(model_dir: Path) -> Tuple[Tokenizer, dict]:
     config_path = model_dir / "config.json"
     if not config_path.exists():
         raise ValueError(f"Could not find config.json in {model_dir}")
@@ -36,13 +35,20 @@ def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tuple[Tokenizer, d
 
     with open(str(tokenizer_config_path)) as tokenizer_config_file:
         tokenizer_config = json.load(tokenizer_config_file)
+        assert (
+            "model_max_length" in tokenizer_config or "max_length" in tokenizer_config
+        ), "Models without model_max_length or max_length are not supported."
+        if "model_max_length" not in tokenizer_config:
+            max_context = tokenizer_config["max_length"]
+        elif "max_length" not in tokenizer_config:
+            max_context = tokenizer_config["model_max_length"]
+        else:
+            max_context = min(tokenizer_config["model_max_length"], tokenizer_config["max_length"])
 
     tokens_map = load_special_tokens(model_dir)
 
     tokenizer = Tokenizer.from_file(str(tokenizer_path))
-    tokenizer.enable_truncation(
-        max_length=min(tokenizer_config["model_max_length"], max_length)
-    )
+    tokenizer.enable_truncation(max_length=max_context)
     tokenizer.enable_padding(
         pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"]
     )

fastembed/common/types.py

@@ -1,5 +1,6 @@
 import os
 import sys
+from PIL import Image
 from typing import Any, Dict, Iterable, Tuple, Union
 
 if sys.version_info >= (3, 10):
@@ -9,6 +10,7 @@ else:
 
 
 PathInput: TypeAlias = Union[str, os.PathLike]
-ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput]]
+PilInput: TypeAlias = Union[Image.Image, Iterable[Image.Image]]
+ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput], PilInput]
 
 OnnxProvider: TypeAlias = Union[str, Tuple[str, Dict[Any, Any]]]

fastembed/common/utils.py

@@ -3,8 +3,11 @@ import tempfile
 from itertools import islice
 from pathlib import Path
 from typing import Generator, Iterable, Optional, Union
-
+import unicodedata
+import sys
 import numpy as np
+import re
+from typing import Set
 
 
 def normalize(input_array, p=2, dim=1, eps=1e-12) -> np.ndarray:
@@ -37,7 +40,16 @@ def define_cache_dir(cache_dir: Optional[str] = None) -> Path:
         cache_path = Path(os.getenv("FASTEMBED_CACHE_PATH", default_cache_dir))
     else:
         cache_path = Path(cache_dir)
-
     cache_path.mkdir(parents=True, exist_ok=True)
 
     return cache_path
+
+
+def get_all_punctuation() -> Set[str]:
+    return set(
+        chr(i) for i in range(sys.maxunicode) if unicodedata.category(chr(i)).startswith("P")
+    )
+
+
+def remove_non_alphanumeric(text: str) -> str:
+    return re.sub(r"[^\w\s]", " ", text, flags=re.UNICODE)

fastembed/image/image_embedding.py

@@ -25,6 +25,7 @@ class ImageEmbedding(ImageEmbeddingBase):
                         "model": "Qdrant/clip-ViT-B-32-vision",
                         "dim": 512,
                         "description": "CLIP vision encoder based on ViT-B/32",
+                        "license": "mit",
                         "size_in_GB": 0.33,
                         "sources": {
                             "hf": "Qdrant/clip-ViT-B-32-vision",
@@ -45,28 +46,30 @@ class ImageEmbedding(ImageEmbeddingBase):
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
         providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[List[int]] = None,
+        lazy_load: bool = False,
         **kwargs,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
-
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
             supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
-            if any(
-                model_name.lower() == model["model"].lower()
-                for model in supported_models
-            ):
+            if any(model_name.lower() == model["model"].lower() for model in supported_models):
                 self.model = EMBEDDING_MODEL_TYPE(
                     model_name,
                     cache_dir,
                     threads=threads,
                     providers=providers,
+                    cuda=cuda,
+                    device_ids=device_ids,
+                    lazy_load=lazy_load,
                     **kwargs,
                 )
                 return
 
         raise ValueError(
-            f"Model {model_name} is not supported in TextEmbedding."
-            "Please check the supported models using `TextEmbedding.list_supported_models()`"
+            f"Model {model_name} is not supported in ImageEmbedding."
+            "Please check the supported models using `ImageEmbedding.list_supported_models()`"
         )
 
     def embed(

fastembed/image/image_embedding_base.py

@@ -30,7 +30,7 @@ class ImageEmbeddingBase(ModelManagement):
         Embeds a list of images into a list of embeddings.
 
         Args:
-            images - The list of image paths to preprocess and embed.
+            images: The list of image paths to preprocess and embed.
             batch_size: Batch size for encoding
             parallel:
                 If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.

fastembed/image/onnx_embedding.py

@@ -12,7 +12,8 @@ supported_onnx_models = [
     {
         "model": "Qdrant/clip-ViT-B-32-vision",
         "dim": 512,
-        "description": "CLIP vision encoder based on ViT-B/32",
+        "description": "Image embeddings, Multimodal (text&image), 2021 year",
+        "license": "mit",
         "size_in_GB": 0.34,
         "sources": {
             "hf": "Qdrant/clip-ViT-B-32-vision",
@@ -22,13 +23,36 @@ supported_onnx_models = [
     {
         "model": "Qdrant/resnet50-onnx",
         "dim": 2048,
-        "description": "ResNet-50 from `Deep Residual Learning for Image Recognition <https://arxiv.org/abs/1512.03385>`__.",
+        "description": "Image embeddings, Unimodal (image), 2016 year",
+        "license": "apache-2.0",
         "size_in_GB": 0.1,
         "sources": {
             "hf": "Qdrant/resnet50-onnx",
         },
         "model_file": "model.onnx",
     },
+    {
+        "model": "Qdrant/Unicom-ViT-B-16",
+        "dim": 768,
+        "description": "Image embeddings (more detailed than Unicom-ViT-B-32), Multimodal (text&image), 2023 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.82,
+        "sources": {
+            "hf": "Qdrant/Unicom-ViT-B-16",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "Qdrant/Unicom-ViT-B-32",
+        "dim": 512,
+        "description": "Image embeddings, Multimodal (text&image), 2023 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.48,
+        "sources": {
+            "hf": "Qdrant/Unicom-ViT-B-32",
+        },
+        "model_file": "model.onnx",
+    },
 ]
 
 
@@ -39,6 +63,10 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
         providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[List[int]] = None,
+        lazy_load: bool = False,
+        device_id: Optional[int] = None,
         **kwargs,
     ):
         """
@@ -48,24 +76,56 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
                                        Can be set using the `FASTEMBED_CACHE_PATH` env variable.
                                        Defaults to `fastembed_cache` in the system's temp directory.
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
+            providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
+                Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
+            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to False.
+            device_ids (Optional[List[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
+                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
 
         Raises:
             ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
         """
 
         super().__init__(model_name, cache_dir, threads, **kwargs)
+        self.providers = providers
+        self.lazy_load = lazy_load
 
-        model_description = self._get_model_description(model_name)
+        # List of device ids, that can be used for data parallel processing in workers
+        self.device_ids = device_ids
+        self.cuda = cuda
+
+        # This device_id will be used if we need to load model in current process
+        if device_id is not None:
+            self.device_id = device_id
+        elif self.device_ids is not None:
+            self.device_id = self.device_ids[0]
+        else:
+            self.device_id = None
+
+        self.model_description = self._get_model_description(model_name)
         self.cache_dir = define_cache_dir(cache_dir)
-        model_dir = self.download_model(
-            model_description, self.cache_dir, local_files_only=self._local_files_only
+        self._model_dir = self.download_model(
+            self.model_description, self.cache_dir, local_files_only=self._local_files_only
         )
 
-        self.load_onnx_model(
-            model_dir=model_dir,
-            model_file=model_description["model_file"],
-            threads=threads,
-            providers=providers,
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    def load_onnx_model(self) -> None:
+        """
+        Load the onnx model.
+        """
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description["model_file"],
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
         )
 
     @classmethod
@@ -100,12 +160,16 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
         Returns:
             List of embeddings, one per document
         """
+
         yield from self._embed_images(
             model_name=self.model_name,
             cache_dir=str(self.cache_dir),
             images=images,
             batch_size=batch_size,
             parallel=parallel,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
             **kwargs,
         )
 
@@ -128,4 +192,9 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
 
 class OnnxImageEmbeddingWorker(ImageEmbeddingWorker):
     def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxImageEmbedding:
-        return OnnxImageEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)
+        return OnnxImageEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/image/onnx_image_model.py

@@ -7,7 +7,7 @@ from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type
 import numpy as np
 from PIL import Image
 
-from fastembed.common import ImageInput, OnnxProvider, PathInput
+from fastembed.common import ImageInput, OnnxProvider
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_preprocessor
 from fastembed.common.utils import iter_batch
@@ -36,27 +36,37 @@ class OnnxImageModel(OnnxModel[T]):
         """
         return onnx_input
 
-    def load_onnx_model(
+    def _load_onnx_model(
         self,
         model_dir: Path,
         model_file: str,
         threads: Optional[int],
         providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_id: Optional[int] = None,
     ) -> None:
-        super().load_onnx_model(
+        super()._load_onnx_model(
             model_dir=model_dir,
             model_file=model_file,
             threads=threads,
             providers=providers,
+            cuda=cuda,
+            device_id=device_id,
         )
         self.processor = load_preprocessor(model_dir=model_dir)
 
+    def load_onnx_model(self) -> None:
+        raise NotImplementedError("Subclasses must implement this method")
+
     def _build_onnx_input(self, encoded: np.ndarray) -> Dict[str, np.ndarray]:
         return {node.name: encoded for node in self.model.get_inputs()}
 
-    def onnx_embed(self, images: List[PathInput], **kwargs) -> OnnxOutputContext:
+    def onnx_embed(self, images: List[ImageInput], **kwargs) -> OnnxOutputContext:
         with contextlib.ExitStack():
-            image_files = [Image.open(image) for image in images]
+            image_files = [
+                Image.open(image) if not isinstance(image, Image.Image) else image
+                for image in images
+            ]
             encoded = self.processor(image_files)
         onnx_input = self._build_onnx_input(encoded)
         onnx_input = self._preprocess_onnx_input(onnx_input)
@@ -71,31 +81,44 @@ class OnnxImageModel(OnnxModel[T]):
         images: ImageInput,
         batch_size: int = 256,
         parallel: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[List[int]] = None,
         **kwargs,
     ) -> Iterable[T]:
         is_small = False
 
-        if isinstance(images, str) or isinstance(images, Path):
+        if isinstance(images, (str, Path, Image.Image)):
             images = [images]
             is_small = True
 
-        if isinstance(images, list):
-            if len(images) < batch_size:
-                is_small = True
-
-        if parallel == 0:
-            parallel = os.cpu_count()
+        if isinstance(images, list) and len(images) < batch_size:
+            is_small = True
 
         if parallel is None or is_small:
+            if not hasattr(self, "model") or self.model is None:
+                self.load_onnx_model()
+
             for batch in iter_batch(images, batch_size):
                 yield from self._post_process_onnx_output(self.onnx_embed(batch))
         else:
-            start_method = (
-                "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
-            )
-            params = {"model_name": model_name, "cache_dir": cache_dir, **kwargs}
+            if parallel == 0:
+                parallel = os.cpu_count()
+
+            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+            params = {
+                "model_name": model_name,
+                "cache_dir": cache_dir,
+                "providers": providers,
+                **kwargs,
+            }
+
             pool = ParallelWorkerPool(
-                parallel, self._get_worker_class(), start_method=start_method
+                num_workers=parallel or 1,
+                worker=self._get_worker_class(),
+                cuda=cuda,
+                device_ids=device_ids,
+                start_method=start_method,
             )
             for batch in pool.ordered_map(iter_batch(images, batch_size), **params):
                 yield from self._post_process_onnx_output(batch)

fastembed/late_interaction/colbert.py

@@ -17,12 +17,24 @@ supported_colbert_models = [
         "model": "colbert-ir/colbertv2.0",
         "dim": 128,
         "description": "Late interaction model",
+        "license": "mit",
         "size_in_GB": 0.44,
         "sources": {
             "hf": "colbert-ir/colbertv2.0",
         },
         "model_file": "model.onnx",
-    }
+    },
+    {
+        "model": "answerdotai/answerai-colbert-small-v1",
+        "dim": 96,
+        "description": "Text embeddings, Unimodal (text), Multilingual (~100 languages), 512 input tokens truncation, 2024 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "answerdotai/answerai-colbert-small-v1",
+        },
+        "model_file": "vespa_colbert.onnx",
+    },
 ]
 
 
@@ -38,6 +50,11 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
         if not is_doc:
             return output.model_output.astype(np.float32)
 
+        if output.input_ids is None or output.attention_mask is None:
+            raise ValueError(
+                "input_ids and attention_mask must be provided for document post-processing"
+            )
+
         for i, token_sequence in enumerate(output.input_ids):
             for j, token_id in enumerate(token_sequence):
                 if token_id in self.skip_list or token_id == self.pad_token_id:
@@ -66,9 +83,10 @@ 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
-        query = [f". {query}"]
-        encoded = self.tokenizer.encode_batch(query)
+        # "@ " 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:
             prev_padding = None
@@ -79,7 +97,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:
@@ -87,8 +105,9 @@ 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
-        documents = [". " + doc for doc in documents]
+        # "@ " 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
 
@@ -107,6 +126,10 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
         providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[List[int]] = None,
+        lazy_load: bool = False,
+        device_id: Optional[int] = None,
         **kwargs,
     ):
         """
@@ -116,29 +139,60 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
                                        Can be set using the `FASTEMBED_CACHE_PATH` env variable.
                                        Defaults to `fastembed_cache` in the system's temp directory.
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
+            providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
+                Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
+            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to False.
+            device_ids (Optional[List[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
+                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
 
         Raises:
             ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
         """
 
         super().__init__(model_name, cache_dir, threads, **kwargs)
+        self.providers = providers
+        self.lazy_load = lazy_load
 
-        model_description = self._get_model_description(model_name)
+        # List of device ids, that can be used for data parallel processing in workers
+        self.device_ids = device_ids
+        self.cuda = cuda
+
+        # This device_id will be used if we need to load model in current process
+        if device_id is not None:
+            self.device_id = device_id
+        elif self.device_ids is not None:
+            self.device_id = self.device_ids[0]
+        else:
+            self.device_id = None
+
+        self.model_description = self._get_model_description(model_name)
         self.cache_dir = define_cache_dir(cache_dir)
 
-        model_dir = self.download_model(
-            model_description, self.cache_dir, local_files_only=self._local_files_only
+        self._model_dir = self.download_model(
+            self.model_description, self.cache_dir, local_files_only=self._local_files_only
         )
+        self.mask_token_id = None
+        self.pad_token_id = None
+        self.skip_list = set()
 
-        self.load_onnx_model(
-            model_dir=model_dir,
-            model_file=model_description["model_file"],
-            threads=threads,
-            providers=providers,
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    def load_onnx_model(self) -> None:
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description["model_file"],
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
         )
-        self.mask_token_id = self.special_token_to_id["[MASK]"]
+        self.mask_token_id = self.special_token_to_id[self.MASK_TOKEN]
         self.pad_token_id = self.tokenizer.padding["pad_id"]
-
         self.skip_list = {
             self.tokenizer.encode(symbol, add_special_tokens=False).ids[0]
             for symbol in string.punctuation
@@ -172,13 +226,19 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
             documents=documents,
             batch_size=batch_size,
             parallel=parallel,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
             **kwargs,
         )
 
-    def query_embed(self, query: Union[str, List[str]], **kwargs) -> np.ndarray:
+    def query_embed(self, query: Union[str, List[str]], **kwargs) -> Iterable[np.ndarray]:
         if isinstance(query, str):
             query = [query]
 
+        if not hasattr(self, "model") or self.model is None:
+            self.load_onnx_model()
+
         for text in query:
             yield from self._post_process_onnx_output(
                 self.onnx_embed([text], is_doc=False), is_doc=False
@@ -191,4 +251,9 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
 
 class ColbertEmbeddingWorker(TextEmbeddingWorker):
     def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> Colbert:
-        return Colbert(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)
+        return Colbert(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/late_interaction/jina_colbert.py

@@ -0,0 +1,62 @@
+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,15 +4,14 @@ 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,
-    ]
+    EMBEDDINGS_REGISTRY: List[Type[LateInteractionTextEmbeddingBase]] = [Colbert, JinaColbert]
 
     @classmethod
     def list_supported_models(cls) -> List[Dict[str, Any]]:
@@ -26,14 +25,16 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
                 ```
                 [
                     {
-                        "model": "prithvida/SPLADE_PP_en_v1",
-                        "vocab_size": 30522,
-                        "description": "Independent Implementation of SPLADE++ Model for English",
-                        "size_in_GB": 0.532,
+                        "model": "colbert-ir/colbertv2.0",
+                        "dim": 128,
+                        "description": "Late interaction model",
+                        "license": "mit",
+                        "size_in_GB": 0.44,
                         "sources": {
-                            "hf": "qdrant/SPLADE_PP_en_v1",
+                            "hf": "colbert-ir/colbertv2.0",
                         },
-                    }
+                        "model_file": "model.onnx",
+                    },
                 ]
                 ```
         """
@@ -48,24 +49,30 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
         providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[List[int]] = None,
+        lazy_load: bool = False,
         **kwargs,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
-
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
             supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
-            if any(
-                model_name.lower() == model["model"].lower()
-                for model in supported_models
-            ):
+            if any(model_name.lower() == model["model"].lower() for model in supported_models):
                 self.model = EMBEDDING_MODEL_TYPE(
-                    model_name, cache_dir, threads, providers=providers, **kwargs
+                    model_name,
+                    cache_dir,
+                    threads=threads,
+                    providers=providers,
+                    cuda=cuda,
+                    device_ids=device_ids,
+                    lazy_load=lazy_load,
+                    **kwargs,
                 )
                 return
 
         raise ValueError(
-            f"Model {model_name} is not supported in SparseTextEmbedding."
-            "Please check the supported models using `SparseTextEmbedding.list_supported_models()`"
+            f"Model {model_name} is not supported in LateInteractionTextEmbedding."
+            "Please check the supported models using `LateInteractionTextEmbedding.list_supported_models()`"
         )
 
     def embed(
@@ -92,9 +99,7 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
         """
         yield from self.model.embed(documents, batch_size, parallel, **kwargs)
 
-    def query_embed(
-        self, query: Union[str, Iterable[str]], **kwargs
-    ) -> Iterable[np.ndarray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
         """
         Embeds queries
 

fastembed/parallel_processor.py

@@ -8,6 +8,7 @@ from multiprocessing.process import BaseProcess
 from multiprocessing.sharedctypes import Synchronized as BaseValue
 from queue import Empty
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Type
+from copy import deepcopy
 
 # Single item should be processed in less than:
 processing_timeout = 10 * 60  # seconds
@@ -47,7 +48,9 @@ def _worker(
     if kwargs is None:
         kwargs = {}
 
-    logging.info(f"Reader worker: {worker_id} PID: {os.getpid()}")
+    logging.info(
+        f"Reader worker: {worker_id} PID: {os.getpid()} Device: {kwargs.get('device_id', 'CPU')}"
+    )
     try:
         worker = worker_class.start(**kwargs)
 
@@ -73,7 +76,9 @@ def _worker(
         # See:
         # https://docs.python.org/3.6/library/multiprocessing.html?highlight=process#pipes-and-queues
         # https://docs.python.org/3.6/library/multiprocessing.html?highlight=process#programming-guidelines
+        input_queue.close()
         output_queue.close()
+        input_queue.join_thread()
         output_queue.join_thread()
 
         with num_active_workers.get_lock():
@@ -84,7 +89,12 @@ def _worker(
 
 class ParallelWorkerPool:
     def __init__(
-        self, num_workers: int, worker: Type[Worker], start_method: Optional[str] = None
+        self,
+        num_workers: int,
+        worker: Type[Worker],
+        start_method: Optional[str] = None,
+        device_ids: Optional[List[int]] = None,
+        cuda: bool = False,
     ):
         self.worker_class = worker
         self.num_workers = num_workers
@@ -93,7 +103,9 @@ class ParallelWorkerPool:
         self.ctx: BaseContext = get_context(start_method)
         self.processes: List[BaseProcess] = []
         self.queue_size = self.num_workers * max_internal_batch_size
-
+        self.emergency_shutdown = False
+        self.device_ids = device_ids
+        self.cuda = cuda
         self.num_active_workers: Optional[BaseValue] = None
 
     def start(self, **kwargs: Any) -> None:
@@ -105,6 +117,12 @@ class ParallelWorkerPool:
         self.num_active_workers = ctx_value
 
         for worker_id in range(0, self.num_workers):
+            worker_kwargs = deepcopy(kwargs)
+            if self.device_ids:
+                device_id = self.device_ids[worker_id % len(self.device_ids)]
+                worker_kwargs["device_id"] = device_id
+                worker_kwargs["cuda"] = self.cuda
+
             assert hasattr(self.ctx, "Process")
             process = self.ctx.Process(
                 target=_worker,
@@ -114,15 +132,13 @@ class ParallelWorkerPool:
                     self.output_queue,
                     self.num_active_workers,
                     worker_id,
-                    kwargs.copy(),
+                    worker_kwargs,
                 ),
             )
             process.start()
             self.processes.append(process)
 
-    def ordered_map(
-        self, stream: Iterable[Any], *args: Any, **kwargs: Any
-    ) -> Iterable[Any]:
+    def ordered_map(self, stream: Iterable[Any], *args: Any, **kwargs: Any) -> Iterable[Any]:
         buffer = defaultdict(Any)
         next_expected = 0
 
@@ -144,6 +160,7 @@ class ParallelWorkerPool:
             pushed = 0
             read = 0
             for idx, item in enumerate(stream):
+                self.check_worker_health()
                 if pushed - read < self.queue_size:
                     try:
                         out_item = self.output_queue.get_nowait()
@@ -170,6 +187,7 @@ class ParallelWorkerPool:
                 self.input_queue.put(QueueSignals.stop)
 
             while read < pushed:
+                self.check_worker_health()
                 out_item = self.output_queue.get(timeout=processing_timeout)
                 if out_item == QueueSignals.error:
                     self.join_or_terminate()
@@ -179,8 +197,27 @@ class ParallelWorkerPool:
         finally:
             assert self.input_queue is not None, "Input queue is None"
             assert self.output_queue is not None, "Output queue is None"
+            self.join()
             self.input_queue.close()
             self.output_queue.close()
+            if self.emergency_shutdown:
+                self.input_queue.cancel_join_thread()
+                self.output_queue.cancel_join_thread()
+            else:
+                self.input_queue.join_thread()
+                self.output_queue.join_thread()
+
+    def check_worker_health(self) -> None:
+        """
+        Checks if any worker process has terminated unexpectedly
+        """
+        for process in self.processes:
+            if not process.is_alive() and process.exitcode != 0:
+                self.emergency_shutdown = True
+                self.join_or_terminate()
+                raise RuntimeError(
+                    f"Worker PID: {process.pid} terminated unexpectedly with code {process.exitcode}"
+                )
 
     def join_or_terminate(self, timeout: Optional[int] = 1) -> None:
         """
@@ -210,4 +247,5 @@ class ParallelWorkerPool:
         https://eli.thegreenplace.net/2009/06/12/safely-using-destructors-in-python/.
         """
         for process in self.processes:
-            process.terminate()
+            if process.is_alive():
+                process.terminate()

fastembed/rerank/cross_encoder/__init__.py

@@ -0,0 +1,3 @@
+from fastembed.rerank.cross_encoder.text_cross_encoder import TextCrossEncoder
+
+__all__ = ["TextCrossEncoder"]

fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py

@@ -0,0 +1,147 @@
+from typing import List, Iterable, Dict, Any, Sequence, Optional
+
+from loguru import logger
+
+from fastembed.common import OnnxProvider
+from fastembed.rerank.cross_encoder.onnx_text_model import OnnxCrossEncoderModel
+from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase
+from fastembed.common.utils import define_cache_dir
+
+supported_onnx_models = [
+    {
+        "model": "Xenova/ms-marco-MiniLM-L-6-v2",
+        "size_in_GB": 0.08,
+        "sources": {
+            "hf": "Xenova/ms-marco-MiniLM-L-6-v2",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "MiniLM-L-6-v2 model optimized for re-ranking tasks.",
+        "license": "apache-2.0",
+    },
+    {
+        "model": "Xenova/ms-marco-MiniLM-L-12-v2",
+        "size_in_GB": 0.12,
+        "sources": {
+            "hf": "Xenova/ms-marco-MiniLM-L-12-v2",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "MiniLM-L-12-v2 model optimized for re-ranking tasks.",
+        "license": "apache-2.0",
+    },
+    {
+        "model": "BAAI/bge-reranker-base",
+        "size_in_GB": 1.04,
+        "sources": {
+            "hf": "BAAI/bge-reranker-base",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "BGE reranker base model for cross-encoder re-ranking.",
+        "license": "mit",
+    },
+]
+
+
+class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
+    @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_onnx_models
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[List[int]] = None,
+        lazy_load: bool = False,
+        device_id: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            model_name (str): The name of the model to use.
+            cache_dir (str, optional): The path to the cache directory.
+                                       Can be set using the `FASTEMBED_CACHE_PATH` env variable.
+                                       Defaults to `fastembed_cache` in the system's temp directory.
+            threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
+            providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
+                Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
+            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to False.
+            device_ids (Optional[List[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
+                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+
+        Raises:
+            ValueError: If the model_name is not in the format <org>/<model> e.g. Xenova/ms-marco-MiniLM-L-6-v2.
+        """
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+        self.providers = providers
+        self.lazy_load = lazy_load
+
+        # List of device ids, that can be used for data parallel processing in workers
+        self.device_ids = device_ids
+        self.cuda = cuda
+
+        if self.device_ids is not None and len(self.device_ids) > 1:
+            logger.warning(
+                "Parallel execution is currently not supported for cross encoders, "
+                f"only the first device will be used for inference: {self.device_ids[0]}."
+            )
+
+        # This device_id will be used if we need to load model in current process
+        if device_id is not None:
+            self.device_id = device_id
+        elif self.device_ids is not None:
+            self.device_id = self.device_ids[0]
+        else:
+            self.device_id = None
+
+        self.model_description = self._get_model_description(model_name)
+        self.cache_dir = define_cache_dir(cache_dir)
+        self._model_dir = self.download_model(
+            self.model_description, self.cache_dir, local_files_only=self._local_files_only
+        )
+
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    def load_onnx_model(self) -> None:
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description["model_file"],
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
+        )
+
+    def rerank(
+        self,
+        query: str,
+        documents: Iterable[str],
+        batch_size: int = 64,
+        **kwargs,
+    ) -> Iterable[float]:
+        """Reranks documents based on their relevance to a given query.
+
+        Args:
+            query (str): The query string to which document relevance is calculated.
+            documents (Iterable[str]): Iterable of documents to be reranked.
+            batch_size (int, optional): The number of documents processed in each batch. Higher batch sizes improve speed
+                                        but require more memory. Default is 64.
+        Returns:
+            Iterable[float]: An iterable of relevance scores for each document.
+        """
+
+        yield from self._rerank_documents(
+            query=query, documents=documents, batch_size=batch_size, **kwargs
+        )

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -0,0 +1,70 @@
+from typing import Sequence, Optional, List, Dict, Iterable
+from pathlib import Path
+
+import numpy as np
+from tokenizers import Encoding
+
+from fastembed.common.onnx_model import OnnxModel, OnnxProvider
+from fastembed.common.preprocessor_utils import load_tokenizer
+from fastembed.common.utils import iter_batch
+
+
+class OnnxCrossEncoderModel(OnnxModel):
+    ONNX_OUTPUT_NAMES: Optional[List[str]] = None
+
+    def _load_onnx_model(
+        self,
+        model_dir: Path,
+        model_file: str,
+        threads: Optional[int],
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_id: Optional[int] = None,
+    ) -> None:
+        super()._load_onnx_model(
+            model_dir=model_dir,
+            model_file=model_file,
+            threads=threads,
+            providers=providers,
+            cuda=cuda,
+            device_id=device_id,
+        )
+        self.tokenizer, _ = load_tokenizer(model_dir=model_dir)
+
+    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]:
+        tokenized_input = self.tokenize(query, documents, **kwargs)
+
+        inputs = {
+            "input_ids": np.array([enc.ids for enc in tokenized_input], dtype=np.int64),
+            "attention_mask": np.array(
+                [enc.attention_mask for enc in tokenized_input], dtype=np.int64
+            ),
+        }
+        input_names = {node.name for node in self.model.get_inputs()}
+        if "token_type_ids" in input_names:
+            inputs["token_type_ids"] = np.array(
+                [enc.type_ids for enc in tokenized_input], dtype=np.int64
+            )
+
+        onnx_input = self._preprocess_onnx_input(inputs, **kwargs)
+        outputs = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)
+        return outputs[0][:, 0].tolist()
+
+    def _rerank_documents(
+        self, query: str, documents: Iterable[str], batch_size: int, **kwargs
+    ) -> Iterable[float]:
+        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)
+
+    def _preprocess_onnx_input(
+        self, onnx_input: Dict[str, np.ndarray], **kwargs
+    ) -> Dict[str, np.ndarray]:
+        """
+        Preprocess the onnx input.
+        """
+        return onnx_input

fastembed/rerank/cross_encoder/text_cross_encoder.py

@@ -0,0 +1,87 @@
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Type
+
+from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase
+from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
+from fastembed.common import OnnxProvider
+
+
+class TextCrossEncoder(TextCrossEncoderBase):
+    CROSS_ENCODER_REGISTRY: List[Type[TextCrossEncoderBase]] = [
+        OnnxTextCrossEncoder,
+    ]
+
+    @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.
+
+            Example:
+                ```
+                [
+                    {
+                        "model": "Xenova/ms-marco-MiniLM-L-6-v2",
+                        "size_in_GB": 0.08,
+                        "sources": {
+                            "hf": "Xenova/ms-marco-MiniLM-L-6-v2",
+                        },
+                        "model_file": "onnx/model.onnx",
+                        "description": "MiniLM-L-6-v2 model optimized for re-ranking tasks.",
+                        "license": "apache-2.0",
+                    }
+                ]
+                ```
+        """
+        result = []
+        for encoder in cls.CROSS_ENCODER_REGISTRY:
+            result.extend(encoder.list_supported_models())
+        return result
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[List[int]] = None,
+        lazy_load: bool = False,
+        **kwargs,
+    ):
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+
+        for CROSS_ENCODER_TYPE in self.CROSS_ENCODER_REGISTRY:
+            supported_models = CROSS_ENCODER_TYPE.list_supported_models()
+            if any(model_name.lower() == model["model"].lower() for model in supported_models):
+                self.model = CROSS_ENCODER_TYPE(
+                    model_name=model_name,
+                    cache_dir=cache_dir,
+                    threads=threads,
+                    providers=providers,
+                    cuda=cuda,
+                    device_ids=device_ids,
+                    lazy_load=lazy_load,
+                    **kwargs,
+                )
+                return
+
+        raise ValueError(
+            f"Model {model_name} is not supported in TextCrossEncoder."
+            "Please check the supported models using `TextCrossEncoder.list_supported_models()`"
+        )
+
+    def rerank(
+        self, query: str, documents: Iterable[str], batch_size: int = 64, **kwargs
+    ) -> Iterable[float]:
+        """Rerank a list of documents based on a query.
+
+        Args:
+            query: Query to rerank the documents against
+            documents: Iterator of documents to rerank
+            batch_size: Batch size for reranking
+
+        Returns:
+            Iterable of scores for each document
+        """
+        yield from self.model.rerank(query, documents, batch_size=batch_size, **kwargs)

fastembed/rerank/cross_encoder/text_cross_encoder_base.py

@@ -0,0 +1,37 @@
+from typing import Iterable, Optional
+
+from fastembed.common.model_management import ModelManagement
+
+
+class TextCrossEncoderBase(ModelManagement):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        **kwargs,
+    ):
+        self.model_name = model_name
+        self.cache_dir = cache_dir
+        self.threads = threads
+        self._local_files_only = kwargs.pop("local_files_only", False)
+
+    def rerank(
+        self,
+        query: str,
+        documents: Iterable[str],
+        batch_size: int = 64,
+        **kwargs,
+    ) -> Iterable[float]:
+        """Reranks a list of documents given a query.
+
+        Args:
+            query (str): The query to rerank the documents.
+            documents (Iterable[str]): The list of texts to rerank.
+            batch_size (int): The batch size to use for reranking.
+            **kwargs: Additional keyword argument to pass to the rerank method.
+
+        Yields:
+            Iterable[float]: The scores of reranked the documents.
+        """
+        raise NotImplementedError("This method should be overridden by subclasses")

fastembed/sparse/bm25.py

@@ -1,5 +1,4 @@
 import os
-import string
 from collections import defaultdict
 from multiprocessing import get_all_start_methods
 from pathlib import Path
@@ -7,33 +6,67 @@ from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Union
 
 import mmh3
 import numpy as np
-from snowballstemmer import stemmer as get_stemmer
-
-from fastembed.common.utils import define_cache_dir, iter_batch
+from py_rust_stemmers import SnowballStemmer
+from fastembed.common.utils import (
+    define_cache_dir,
+    iter_batch,
+    get_all_punctuation,
+    remove_non_alphanumeric,
+)
 from fastembed.parallel_processor import ParallelWorkerPool, Worker
 from fastembed.sparse.sparse_embedding_base import (
     SparseEmbedding,
     SparseTextEmbeddingBase,
 )
-from fastembed.sparse.utils.tokenizer import WordTokenizer
+from fastembed.sparse.utils.tokenizer import SimpleTokenizer
+
+supported_languages = [
+    "arabic",
+    "azerbaijani",
+    "basque",
+    "bengali",
+    "catalan",
+    "chinese",
+    "danish",
+    "dutch",
+    "english",
+    "finnish",
+    "french",
+    "german",
+    "greek",
+    "hebrew",
+    "hinglish",
+    "hungarian",
+    "indonesian",
+    "italian",
+    "kazakh",
+    "nepali",
+    "norwegian",
+    "portuguese",
+    "romanian",
+    "russian",
+    "slovene",
+    "spanish",
+    "swedish",
+    "tajik",
+    "turkish",
+]
 
 supported_bm25_models = [
     {
         "model": "Qdrant/bm25",
         "description": "BM25 as sparse embeddings meant to be used with Qdrant",
+        "license": "apache-2.0",
         "size_in_GB": 0.01,
         "sources": {
             "hf": "Qdrant/bm25",
         },
         "model_file": "mock.file",  # bm25 does not require a model, so we just use a mock
-        "additional_files": ["stopwords.txt"],
+        "additional_files": [f"{lang}.txt" for lang in supported_languages],
+        "requires_idf": True,
     },
 ]
 
-MODEL_TO_LANGUAGE = {
-    "Qdrant/bm25": "english",
-}
-
 
 class Bm25(SparseTextEmbeddingBase):
     """Implements traditional BM25 in a form of sparse embeddings.
@@ -70,10 +103,17 @@ class Bm25(SparseTextEmbeddingBase):
         k: float = 1.2,
         b: float = 0.75,
         avg_len: float = 256.0,
+        language: str = "english",
+        token_max_length: int = 40,
         **kwargs,
     ):
         super().__init__(model_name, cache_dir, **kwargs)
 
+        if language not in supported_languages:
+            raise ValueError(f"{language} language is not supported")
+        else:
+            self.language = language
+
         self.k = k
         self.b = b
         self.avg_len = avg_len
@@ -81,14 +121,16 @@ class Bm25(SparseTextEmbeddingBase):
         model_description = self._get_model_description(model_name)
         self.cache_dir = define_cache_dir(cache_dir)
 
-        model_dir = self.download_model(
+        self._model_dir = self.download_model(
             model_description, self.cache_dir, local_files_only=self._local_files_only
         )
 
-        self.punctuation = set(string.punctuation)
-        self.stopwords = set(self._load_stopwords(model_dir))
-        self.stemmer = get_stemmer(MODEL_TO_LANGUAGE[model_name])
-        self.tokenizer = WordTokenizer
+        self.token_max_length = token_max_length
+        self.punctuation = set(get_all_punctuation())
+        self.stopwords = set(self._load_stopwords(self._model_dir, self.language))
+
+        self.stemmer = SnowballStemmer(language)
+        self.tokenizer = SimpleTokenizer
 
     @classmethod
     def list_supported_models(cls) -> List[Dict[str, Any]]:
@@ -100,8 +142,8 @@ class Bm25(SparseTextEmbeddingBase):
         return supported_bm25_models
 
     @classmethod
-    def _load_stopwords(cls, model_dir: Path) -> List[str]:
-        stopwords_path = model_dir / "stopwords.txt"
+    def _load_stopwords(cls, model_dir: Path, language: str) -> List[str]:
+        stopwords_path = model_dir / f"{language}.txt"
         if not stopwords_path.exists():
             return []
 
@@ -126,13 +168,13 @@ class Bm25(SparseTextEmbeddingBase):
             if len(documents) < batch_size:
                 is_small = True
 
-        if parallel == 0:
-            parallel = os.cpu_count()
-
         if parallel is None or is_small:
             for batch in iter_batch(documents, batch_size):
                 yield from self.raw_embed(batch)
         else:
+            if parallel == 0:
+                parallel = os.cpu_count()
+
             start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
             params = {
                 "model_name": model_name,
@@ -142,7 +184,9 @@ class Bm25(SparseTextEmbeddingBase):
                 "avg_len": self.avg_len,
             }
             pool = ParallelWorkerPool(
-                parallel, self._get_worker_class(), start_method=start_method
+                num_workers=parallel or 1,
+                worker=self._get_worker_class(),
+                start_method=start_method,
             )
             for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
                 for record in batch:
@@ -184,10 +228,13 @@ class Bm25(SparseTextEmbeddingBase):
             if token in self.punctuation:
                 continue
 
-            if token in self.stopwords:
+            if token.lower() in self.stopwords:
                 continue
 
-            stemmed_token = self.stemmer.stemWord(token)
+            if len(token) > self.token_max_length:
+                continue
+
+            stemmed_token = self.stemmer.stem_word(token.lower())
 
             if stemmed_token:
                 stemmed_tokens.append(stemmed_token)
@@ -199,6 +246,7 @@ class Bm25(SparseTextEmbeddingBase):
     ) -> List[SparseEmbedding]:
         embeddings = []
         for document in documents:
+            document = remove_non_alphanumeric(document)
             tokens = self.tokenizer.tokenize(document)
             stemmed_tokens = self._stem(tokens)
             token_id2value = self._term_frequency(stemmed_tokens)
@@ -247,6 +295,7 @@ class Bm25(SparseTextEmbeddingBase):
             query = [query]
 
         for text in query:
+            text = remove_non_alphanumeric(text)
             tokens = self.tokenizer.tokenize(text)
             stemmed_tokens = self._stem(tokens)
             token_ids = np.array(

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 snowballstemmer import stemmer as get_stemmer
+from py_rust_stemmers import SnowballStemmer
 
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
@@ -21,12 +21,14 @@ supported_bm42_models = [
         "model": "Qdrant/bm42-all-minilm-l6-v2-attentions",
         "vocab_size": 30522,
         "description": "Light sparse embedding model, which assigns an importance score to each token in the text",
+        "license": "apache-2.0",
         "size_in_GB": 0.09,
         "sources": {
             "hf": "Qdrant/all_miniLM_L6_v2_with_attentions",
         },
         "model_file": "model.onnx",
         "additional_files": ["stopwords.txt"],
+        "requires_idf": True,
     },
 ]
 
@@ -60,6 +62,10 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         threads: Optional[int] = None,
         providers: Optional[Sequence[OnnxProvider]] = None,
         alpha: float = 0.5,
+        cuda: bool = False,
+        device_ids: Optional[List[int]] = None,
+        lazy_load: bool = False,
+        device_id: Optional[int] = None,
         **kwargs,
     ):
         """
@@ -73,38 +79,67 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             alpha (float, optional): Parameter, that defines the importance of the token weight in the document
                 versus the importance of the token frequency in the corpus. Defaults to 0.5, based on empirical testing.
                 It is recommended to only change this parameter based on training data for a specific dataset.
+            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to False.
+            device_ids (Optional[List[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
+                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
 
         Raises:
             ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
         """
 
         super().__init__(model_name, cache_dir, threads, **kwargs)
+        self.providers = providers
+        self.lazy_load = lazy_load
 
-        model_description = self._get_model_description(model_name)
+        # List of device ids, that can be used for data parallel processing in workers
+        self.device_ids = device_ids
+        self.cuda = cuda
+
+        # This device_id will be used if we need to load model in current process
+        if device_id is not None:
+            self.device_id = device_id
+        elif self.device_ids is not None:
+            self.device_id = self.device_ids[0]
+        else:
+            self.device_id = None
+
+        self.model_description = self._get_model_description(model_name)
         self.cache_dir = define_cache_dir(cache_dir)
 
-        model_dir = self.download_model(
-            model_description, self.cache_dir, local_files_only=self._local_files_only
-        )
-
-        self.load_onnx_model(
-            model_dir=model_dir,
-            model_file=model_description["model_file"],
-            threads=threads,
-            providers=providers,
+        self._model_dir = self.download_model(
+            self.model_description, self.cache_dir, local_files_only=self._local_files_only
         )
 
         self.invert_vocab = {}
 
+        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.alpha = alpha
+
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    def load_onnx_model(self) -> None:
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description["model_file"],
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
+        )
         for token, idx in self.tokenizer.get_vocab().items():
             self.invert_vocab[idx] = token
-
         self.special_tokens = set(self.special_token_to_id.keys())
         self.special_tokens_ids = set(self.special_token_to_id.values())
-        self.punctuation = set(string.punctuation)
-        self.stopwords = set(self._load_stopwords(model_dir))
-        self.stemmer = get_stemmer(MODEL_TO_LANGUAGE[model_name])
-        self.alpha = alpha
+        self.stopwords = set(self._load_stopwords(self._model_dir))
 
     def _filter_pair_tokens(self, tokens: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:
         result = []
@@ -117,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.stemWord(token)
+            processed_token = self.stemmer.stem_word(token)
             result.append((processed_token, value))
         return result
 
@@ -179,6 +214,9 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         return new_vector
 
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[SparseEmbedding]:
+        if output.input_ids is None:
+            raise ValueError("input_ids must be provided for document post-processing")
+
         token_ids_batch = output.input_ids
 
         # attention_value shape: (batch_size, num_heads, num_tokens, num_tokens)
@@ -253,6 +291,9 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             documents=documents,
             batch_size=batch_size,
             parallel=parallel,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
             alpha=self.alpha,
         )
 
@@ -273,6 +314,9 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         if isinstance(query, str):
             query = [query]
 
+        if not hasattr(self, "model") or self.model is None:
+            self.load_onnx_model()
+
         for text in query:
             encoded = self.tokenizer.encode(text)
             document_tokens_with_ids = enumerate(encoded.tokens)
@@ -289,4 +333,8 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
 
 class Bm42TextEmbeddingWorker(TextEmbeddingWorker):
     def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> Bm42:
-        return Bm42(model_name=model_name, cache_dir=cache_dir, **kwargs)
+        return Bm42(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            **kwargs,
+        )

fastembed/sparse/sparse_text_embedding.py

@@ -8,6 +8,7 @@ from fastembed.sparse.sparse_embedding_base import (
     SparseTextEmbeddingBase,
 )
 from fastembed.sparse.splade_pp import SpladePP
+import warnings
 
 
 class SparseTextEmbedding(SparseTextEmbeddingBase):
@@ -28,6 +29,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
                         "model": "prithvida/SPLADE_PP_en_v1",
                         "vocab_size": 30522,
                         "description": "Independent Implementation of SPLADE++ Model for English",
+                        "license": "apache-2.0",
                         "size_in_GB": 0.532,
                         "sources": {
                             "hf": "qdrant/SPLADE_PP_en_v1",
@@ -47,21 +49,32 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
         providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[List[int]] = None,
+        lazy_load: bool = False,
         **kwargs,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
+        if model_name == "prithvida/Splade_PP_en_v1":
+            warnings.warn(
+                "The right spelling is prithivida/Splade_PP_en_v1. "
+                "Support of this name will be removed soon, please fix the model_name",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            model_name = "prithivida/Splade_PP_en_v1"
 
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
             supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
-            if any(
-                model_name.lower() == model["model"].lower()
-                for model in supported_models
-            ):
+            if any(model_name.lower() == model["model"].lower() for model in supported_models):
                 self.model = EMBEDDING_MODEL_TYPE(
                     model_name,
                     cache_dir,
                     threads=threads,
                     providers=providers,
+                    cuda=cuda,
+                    device_ids=device_ids,
+                    lazy_load=lazy_load,
                     **kwargs,
                 )
                 return
@@ -95,9 +108,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         """
         yield from self.model.embed(documents, batch_size, parallel, **kwargs)
 
-    def query_embed(
-        self, query: Union[str, Iterable[str]], **kwargs
-    ) -> Iterable[SparseEmbedding]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]:
         """
         Embeds queries
 

fastembed/sparse/splade_pp.py

@@ -1,7 +1,6 @@
 from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union
 
 import numpy as np
-
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import define_cache_dir
@@ -13,9 +12,10 @@ from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
 
 supported_splade_models = [
     {
-        "model": "prithvida/Splade_PP_en_v1",
+        "model": "prithivida/Splade_PP_en_v1",
         "vocab_size": 30522,
-        "description": "Misspelled version of the model. Retained for backward compatibility. Independent Implementation of SPLADE++ Model for English",
+        "description": "Independent Implementation of SPLADE++ Model for English.",
+        "license": "apache-2.0",
         "size_in_GB": 0.532,
         "sources": {
             "hf": "Qdrant/SPLADE_PP_en_v1",
@@ -23,9 +23,10 @@ supported_splade_models = [
         "model_file": "model.onnx",
     },
     {
-        "model": "prithivida/Splade_PP_en_v1",
+        "model": "prithvida/Splade_PP_en_v1",
         "vocab_size": 30522,
-        "description": "Independent Implementation of SPLADE++ Model for English",
+        "description": "Independent Implementation of SPLADE++ Model for English.",
+        "license": "apache-2.0",
         "size_in_GB": 0.532,
         "sources": {
             "hf": "Qdrant/SPLADE_PP_en_v1",
@@ -37,6 +38,9 @@ supported_splade_models = [
 
 class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[SparseEmbedding]:
+        if output.attention_mask is None:
+            raise ValueError("attention_mask must be provided for document post-processing")
+
         relu_log = np.log(1 + np.maximum(output.model_output, 0))
 
         weighted_log = relu_log * np.expand_dims(output.attention_mask, axis=-1)
@@ -65,6 +69,10 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
         providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[List[int]] = None,
+        lazy_load: bool = False,
+        device_id: Optional[int] = None,
         **kwargs,
     ):
         """
@@ -74,25 +82,53 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
                                        Can be set using the `FASTEMBED_CACHE_PATH` env variable.
                                        Defaults to `fastembed_cache` in the system's temp directory.
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
+            providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
+                Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
+            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to False.
+            device_ids (Optional[List[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
+                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
 
         Raises:
             ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
         """
-
         super().__init__(model_name, cache_dir, threads, **kwargs)
+        self.providers = providers
+        self.lazy_load = lazy_load
 
-        model_description = self._get_model_description(model_name)
+        # List of device ids, that can be used for data parallel processing in workers
+        self.device_ids = device_ids
+        self.cuda = cuda
+
+        # This device_id will be used if we need to load model in current process
+        if device_id is not None:
+            self.device_id = device_id
+        elif self.device_ids is not None:
+            self.device_id = self.device_ids[0]
+        else:
+            self.device_id = None
+
+        self.model_description = self._get_model_description(model_name)
         self.cache_dir = define_cache_dir(cache_dir)
 
-        model_dir = self.download_model(
-            model_description, self.cache_dir, local_files_only=self._local_files_only
+        self._model_dir = self.download_model(
+            self.model_description, self.cache_dir, local_files_only=self._local_files_only
         )
 
-        self.load_onnx_model(
-            model_dir=model_dir,
-            model_file=model_description["model_file"],
-            threads=threads,
-            providers=providers,
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    def load_onnx_model(self) -> None:
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description["model_file"],
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
         )
 
     def embed(
@@ -123,6 +159,10 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             documents=documents,
             batch_size=batch_size,
             parallel=parallel,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            **kwargs,
         )
 
     @classmethod
@@ -132,4 +172,9 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
 
 class SpladePPEmbeddingWorker(TextEmbeddingWorker):
     def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> SpladePP:
-        return SpladePP(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)
+        return SpladePP(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/sparse/utils/tokenizer.py

@@ -4,6 +4,14 @@ import re
 from typing import List
 
 
+class SimpleTokenizer:
+    def tokenize(text: str) -> List[str]:
+        text = re.sub(r"[^\w]", " ", text.lower())
+        text = re.sub(r"\s+", " ", text)
+
+        return text.strip().split()
+
+
 class WordTokenizer:
     """The tokenizer is "destructive" such that the regexes applied will munge the
     input string to a state beyond re-construction.
@@ -68,8 +76,7 @@ class WordTokenizer:
         )
     ]
     CONTRACTIONS3 = [
-        re.compile(pattern)
-        for pattern in (r"(?i) ('t)(?#X)(is)\b", r"(?i) ('t)(?#X)(was)\b")
+        re.compile(pattern) for pattern in (r"(?i) ('t)(?#X)(is)\b", r"(?i) ('t)(?#X)(was)\b")
     ]
 
     @classmethod

fastembed/text/clip_embedding.py

@@ -10,7 +10,8 @@ supported_clip_models = [
     {
         "model": "Qdrant/clip-ViT-B-32-text",
         "dim": 512,
-        "description": "CLIP text encoder",
+        "description": "Text embeddings, Multimodal (text&image), English, 77 input tokens truncation, Prefixes for queries/documents: not necessary, 2021 year",
+        "license": "mit",
         "size_in_GB": 0.25,
         "sources": {
             "hf": "Qdrant/clip-ViT-B-32-text",
@@ -34,16 +35,20 @@ class CLIPOnnxEmbedding(OnnxTextEmbedding):
         """
         return supported_clip_models
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext
-    ) -> Iterable[np.ndarray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
         return output.model_output
 
 
 class CLIPEmbeddingWorker(OnnxTextEmbeddingWorker):
     def init_embedding(
-        self, model_name: str, cache_dir: str, **kwargs
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs,
     ) -> OnnxTextEmbedding:
         return CLIPOnnxEmbedding(
-            model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
         )

fastembed/text/e5_onnx_embedding.py

@@ -9,7 +9,8 @@ supported_multilingual_e5_models = [
     {
         "model": "intfloat/multilingual-e5-large",
         "dim": 1024,
-        "description": "Multilingual model, e5-large. Recommend using this model for non-English languages",
+        "description": "Text embeddings, Unimodal (text), Multilingual (~100 languages), 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "mit",
         "size_in_GB": 2.24,
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
@@ -21,7 +22,8 @@ supported_multilingual_e5_models = [
     {
         "model": "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
         "dim": 768,
-        "description": "Sentence-transformers model for tasks like clustering or semantic search",
+        "description": "Text embeddings, Unimodal (text), Multilingual (~50 languages), 384 input tokens truncation, Prefixes for queries/documents: not necessary, 2021 year.",
+        "license": "apache-2.0",
         "size_in_GB": 1.00,
         "sources": {
             "hf": "xenova/paraphrase-multilingual-mpnet-base-v2",
@@ -57,8 +59,14 @@ class E5OnnxEmbedding(OnnxTextEmbedding):
 
 class E5OnnxEmbeddingWorker(OnnxTextEmbeddingWorker):
     def init_embedding(
-        self, model_name: str, cache_dir: str, **kwargs
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs,
     ) -> E5OnnxEmbedding:
         return E5OnnxEmbedding(
-            model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
         )

fastembed/text/jina_onnx_embedding.py

@@ -1,76 +0,0 @@
-from typing import Any, Dict, Iterable, List, Type
-
-import numpy as np
-
-from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.utils import normalize
-from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
-from fastembed.text.onnx_text_model import TextEmbeddingWorker
-
-supported_jina_models = [
-    {
-        "model": "jinaai/jina-embeddings-v2-base-en",
-        "dim": 768,
-        "description": "English embedding model supporting 8192 sequence length",
-        "size_in_GB": 0.52,
-        "sources": {"hf": "xenova/jina-embeddings-v2-base-en"},
-        "model_file": "onnx/model.onnx",
-    },
-    {
-        "model": "jinaai/jina-embeddings-v2-small-en",
-        "dim": 512,
-        "description": "English embedding model supporting 8192 sequence length",
-        "size_in_GB": 0.12,
-        "sources": {"hf": "xenova/jina-embeddings-v2-small-en"},
-        "model_file": "onnx/model.onnx",
-    },
-    {
-        "model": "jinaai/jina-embeddings-v2-base-de",
-        "dim": 768,
-        "description": "German embedding model supporting 8192 sequence length",
-        "size_in_GB": 0.32,
-        "sources": {"hf": "jinaai/jina-embeddings-v2-base-de"},
-        "model_file": "onnx/model_fp16.onnx",
-    },
-]
-
-
-class JinaOnnxEmbedding(OnnxTextEmbedding):
-    @classmethod
-    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
-        return JinaEmbeddingWorker
-
-    @classmethod
-    def mean_pooling(cls, model_output, attention_mask) -> np.ndarray:
-        token_embeddings = model_output
-        input_mask_expanded = (np.expand_dims(attention_mask, axis=-1)).astype(float)
-
-        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
-        mask_sum = np.clip(np.sum(input_mask_expanded, axis=1), a_min=1e-9, a_max=None)
-
-        return sum_embeddings / mask_sum
-
-    @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_models
-
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext
-    ) -> Iterable[np.ndarray]:
-        embeddings = output.model_output
-        attn_mask = output.attention_mask
-        return normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
-
-
-class JinaEmbeddingWorker(OnnxTextEmbeddingWorker):
-    def init_embedding(
-        self, model_name: str, cache_dir: str, **kwargs
-    ) -> OnnxTextEmbedding:
-        return JinaOnnxEmbedding(
-            model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
-        )

fastembed/text/mini_lm_embedding.py

@@ -1,58 +0,0 @@
-from typing import Any, Dict, Iterable, List, Type
-
-import numpy as np
-
-from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.utils import normalize
-from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
-from fastembed.text.onnx_text_model import TextEmbeddingWorker
-
-supported_mini_lm_models = [
-    {
-        "model": "sentence-transformers/all-MiniLM-L6-v2",
-        "dim": 384,
-        "description": "Sentence Transformer model, MiniLM-L6-v2",
-        "size_in_GB": 0.09,
-        "sources": {
-            "url": "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz",
-            "hf": "qdrant/all-MiniLM-L6-v2-onnx",
-        },
-        "model_file": "model.onnx",
-    }
-]
-
-
-class MiniLMOnnxEmbedding(OnnxTextEmbedding):
-    @classmethod
-    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
-        return MiniLMEmbeddingWorker
-
-    @classmethod
-    def mean_pooling(cls, model_output: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
-        token_embeddings = model_output
-        input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
-        input_mask_expanded = np.tile(input_mask_expanded, (1, 1, token_embeddings.shape[-1]))
-        input_mask_expanded = input_mask_expanded.astype(float)
-        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
-        sum_mask = np.sum(input_mask_expanded, axis=1)
-        pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)
-        return pooled_embeddings
-
-    @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_mini_lm_models
-
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
-        embeddings = output.model_output
-        attn_mask = output.attention_mask
-        return normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
-
-
-class MiniLMEmbeddingWorker(OnnxTextEmbeddingWorker):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxTextEmbedding:
-        return MiniLMOnnxEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)

fastembed/text/onnx_embedding.py

@@ -12,7 +12,8 @@ supported_onnx_models = [
     {
         "model": "BAAI/bge-base-en",
         "dim": 768,
-        "description": "Base English model",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2023 year.",
+        "license": "mit",
         "size_in_GB": 0.42,
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz",
@@ -22,7 +23,8 @@ supported_onnx_models = [
     {
         "model": "BAAI/bge-base-en-v1.5",
         "dim": 768,
-        "description": "Base English model, v1.5",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: not so necessary, 2023 year.",
+        "license": "mit",
         "size_in_GB": 0.21,
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
@@ -33,7 +35,8 @@ supported_onnx_models = [
     {
         "model": "BAAI/bge-large-en-v1.5",
         "dim": 1024,
-        "description": "Large English model, v1.5",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: not so necessary, 2023 year.",
+        "license": "mit",
         "size_in_GB": 1.20,
         "sources": {
             "hf": "qdrant/bge-large-en-v1.5-onnx",
@@ -43,7 +46,8 @@ supported_onnx_models = [
     {
         "model": "BAAI/bge-small-en",
         "dim": 384,
-        "description": "Fast English model",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2023 year.",
+        "license": "mit",
         "size_in_GB": 0.13,
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
@@ -53,7 +57,8 @@ supported_onnx_models = [
     {
         "model": "BAAI/bge-small-en-v1.5",
         "dim": 384,
-        "description": "Fast and Default English model",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: not so necessary, 2023 year.",
+        "license": "mit",
         "size_in_GB": 0.067,
         "sources": {
             "hf": "qdrant/bge-small-en-v1.5-onnx-q",
@@ -63,7 +68,8 @@ supported_onnx_models = [
     {
         "model": "BAAI/bge-small-zh-v1.5",
         "dim": 512,
-        "description": "Fast and recommended Chinese model",
+        "description": "Text embeddings, Unimodal (text), Chinese, 512 input tokens truncation, Prefixes for queries/documents: not so necessary, 2023 year.",
+        "license": "mit",
         "size_in_GB": 0.09,
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
@@ -73,47 +79,19 @@ supported_onnx_models = [
     {
         "model": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
         "dim": 384,
-        "description": "Sentence Transformer model, paraphrase-multilingual-MiniLM-L12-v2",
+        "description": "Text embeddings, Unimodal (text), Multilingual (~50 languages), 512 input tokens truncation, Prefixes for queries/documents: not necessary, 2019 year.",
+        "license": "apache-2.0",
         "size_in_GB": 0.22,
         "sources": {
             "hf": "qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q",
         },
         "model_file": "model_optimized.onnx",
     },
-    {
-        "model": "nomic-ai/nomic-embed-text-v1",
-        "dim": 768,
-        "description": "8192 context length english model",
-        "size_in_GB": 0.52,
-        "sources": {
-            "hf": "nomic-ai/nomic-embed-text-v1",
-        },
-        "model_file": "onnx/model.onnx",
-    },
-    {
-        "model": "nomic-ai/nomic-embed-text-v1.5",
-        "dim": 768,
-        "description": "8192 context length english model",
-        "size_in_GB": 0.52,
-        "sources": {
-            "hf": "nomic-ai/nomic-embed-text-v1.5",
-        },
-        "model_file": "onnx/model.onnx",
-    },
-    {
-        "model": "nomic-ai/nomic-embed-text-v1.5-Q",
-        "dim": 768,
-        "description": "Quantized 8192 context length english model",
-        "size_in_GB": 0.13,
-        "sources": {
-            "hf": "nomic-ai/nomic-embed-text-v1.5",
-        },
-        "model_file": "onnx/model_quantized.onnx",
-    },
     {
         "model": "thenlper/gte-large",
         "dim": 1024,
-        "description": "Large general text embeddings model",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: not necessary, 2023 year.",
+        "license": "mit",
         "size_in_GB": 1.20,
         "sources": {
             "hf": "qdrant/gte-large-onnx",
@@ -123,7 +101,8 @@ supported_onnx_models = [
     {
         "model": "mixedbread-ai/mxbai-embed-large-v1",
         "dim": 1024,
-        "description": "MixedBread Base sentence embedding model, does well on MTEB",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
         "size_in_GB": 0.64,
         "sources": {
             "hf": "mixedbread-ai/mxbai-embed-large-v1",
@@ -133,7 +112,8 @@ supported_onnx_models = [
     {
         "model": "snowflake/snowflake-arctic-embed-xs",
         "dim": 384,
-        "description": "Based on all-MiniLM-L6-v2 model with only 22m parameters, ideal for latency/TCO budgets.",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
         "size_in_GB": 0.09,
         "sources": {
             "hf": "snowflake/snowflake-arctic-embed-xs",
@@ -143,7 +123,8 @@ supported_onnx_models = [
     {
         "model": "snowflake/snowflake-arctic-embed-s",
         "dim": 384,
-        "description": "Based on infloat/e5-small-unsupervised, does not trade off retrieval accuracy for its small size.",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
         "size_in_GB": 0.13,
         "sources": {
             "hf": "snowflake/snowflake-arctic-embed-s",
@@ -153,7 +134,8 @@ supported_onnx_models = [
     {
         "model": "snowflake/snowflake-arctic-embed-m",
         "dim": 768,
-        "description": "Based on intfloat/e5-base-unsupervised model, provides the best retrieval without slowing down inference.",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
         "size_in_GB": 0.43,
         "sources": {
             "hf": "Snowflake/snowflake-arctic-embed-m",
@@ -163,7 +145,8 @@ supported_onnx_models = [
     {
         "model": "snowflake/snowflake-arctic-embed-m-long",
         "dim": 768,
-        "description": "Based on nomic-ai/nomic-embed-text-v1-unsupervised model, 8192 context-length model",
+        "description": "Text embeddings, Unimodal (text), English, 2048 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
         "size_in_GB": 0.54,
         "sources": {
             "hf": "snowflake/snowflake-arctic-embed-m-long",
@@ -173,7 +156,8 @@ supported_onnx_models = [
     {
         "model": "snowflake/snowflake-arctic-embed-l",
         "dim": 1024,
-        "description": "Based on intfloat/e5-large-unsupervised, large model for most accurate retrieval.",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
         "size_in_GB": 1.02,
         "sources": {
             "hf": "snowflake/snowflake-arctic-embed-l",
@@ -202,6 +186,10 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
         providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[List[int]] = None,
+        lazy_load: bool = False,
+        device_id: Optional[int] = None,
         **kwargs,
     ):
         """
@@ -211,25 +199,43 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
                                        Can be set using the `FASTEMBED_CACHE_PATH` env variable.
                                        Defaults to `fastembed_cache` in the system's temp directory.
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
+            providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
+                Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
+            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to False.
+            device_ids (Optional[List[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
+                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
 
         Raises:
             ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
         """
-
         super().__init__(model_name, cache_dir, threads, **kwargs)
+        self.providers = providers
+        self.lazy_load = lazy_load
 
-        model_description = self._get_model_description(model_name)
+        # List of device ids, that can be used for data parallel processing in workers
+        self.device_ids = device_ids
+        self.cuda = cuda
+
+        # This device_id will be used if we need to load model in current process
+        if device_id is not None:
+            self.device_id = device_id
+        elif self.device_ids is not None:
+            self.device_id = self.device_ids[0]
+        else:
+            self.device_id = None
+
+        self.model_description = self._get_model_description(model_name)
         self.cache_dir = define_cache_dir(cache_dir)
-        model_dir = self.download_model(
-            model_description, self.cache_dir, local_files_only=self._local_files_only
+        self._model_dir = self.download_model(
+            self.model_description, self.cache_dir, local_files_only=self._local_files_only
         )
 
-        self.load_onnx_model(
-            model_dir=model_dir,
-            model_file=model_description["model_file"],
-            threads=threads,
-            providers=providers,
-        )
+        if not self.lazy_load:
+            self.load_onnx_model()
 
     def embed(
         self,
@@ -259,6 +265,9 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
             documents=documents,
             batch_size=batch_size,
             parallel=parallel,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
             **kwargs,
         )
 
@@ -278,6 +287,16 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
         embeddings = output.model_output
         return normalize(embeddings[:, 0]).astype(np.float32)
 
+    def load_onnx_model(self) -> None:
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description["model_file"],
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
+        )
+
 
 class OnnxTextEmbeddingWorker(TextEmbeddingWorker):
     def init_embedding(
@@ -286,4 +305,9 @@ class OnnxTextEmbeddingWorker(TextEmbeddingWorker):
         cache_dir: str,
         **kwargs,
     ) -> OnnxTextEmbedding:
-        return OnnxTextEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)
+        return OnnxTextEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/text/onnx_text_model.py

@@ -36,21 +36,28 @@ class OnnxTextModel(OnnxModel[T]):
         """
         return onnx_input
 
-    def load_onnx_model(
+    def _load_onnx_model(
         self,
         model_dir: Path,
         model_file: str,
         threads: Optional[int],
         providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_id: Optional[int] = None,
     ) -> None:
-        super().load_onnx_model(
+        super()._load_onnx_model(
             model_dir=model_dir,
             model_file=model_file,
             threads=threads,
             providers=providers,
+            cuda=cuda,
+            device_id=device_id,
         )
         self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
 
+    def load_onnx_model(self) -> None:
+        raise NotImplementedError("Subclasses must implement this method")
+
     def tokenize(self, documents: List[str], **kwargs) -> List[Encoding]:
         return self.tokenizer.encode_batch(documents)
 
@@ -89,6 +96,9 @@ class OnnxTextModel(OnnxModel[T]):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[List[int]] = None,
         **kwargs,
     ) -> Iterable[T]:
         is_small = False
@@ -101,19 +111,29 @@ class OnnxTextModel(OnnxModel[T]):
             if len(documents) < batch_size:
                 is_small = True
 
-        if parallel == 0:
-            parallel = os.cpu_count()
-
         if parallel is None or is_small:
+            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._post_process_onnx_output(self.onnx_embed(batch))
         else:
-            start_method = (
-                "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
-            )
-            params = {"model_name": model_name, "cache_dir": cache_dir, **kwargs}
+            if parallel == 0:
+                parallel = os.cpu_count()
+
+            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+            params = {
+                "model_name": model_name,
+                "cache_dir": cache_dir,
+                "providers": providers,
+                **kwargs,
+            }
+
             pool = ParallelWorkerPool(
-                parallel, self._get_worker_class(), start_method=start_method
+                num_workers=parallel or 1,
+                worker=self._get_worker_class(),
+                cuda=cuda,
+                device_ids=device_ids,
+                start_method=start_method,
             )
             for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
                 yield from self._post_process_onnx_output(batch)

fastembed/text/pooled_embedding.py

@@ -0,0 +1,92 @@
+from typing import Any, Dict, Iterable, List, Type
+
+import numpy as np
+
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
+
+supported_pooled_models = [
+    {
+        "model": "nomic-ai/nomic-embed-text-v1.5",
+        "dim": 768,
+        "description": "Text embeddings, Multimodal (text, image), English, 8192 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.52,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1.5",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "nomic-ai/nomic-embed-text-v1.5-Q",
+        "dim": 768,
+        "description": "Text embeddings, Multimodal (text, image), English, 8192 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1.5",
+        },
+        "model_file": "onnx/model_quantized.onnx",
+    },
+    {
+        "model": "nomic-ai/nomic-embed-text-v1",
+        "dim": 768,
+        "description": "Text embeddings, Multimodal (text, image), English, 8192 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.52,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+]
+
+
+class PooledEmbedding(OnnxTextEmbedding):
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
+        return PooledEmbeddingWorker
+
+    @classmethod
+    def mean_pooling(cls, model_output: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
+        token_embeddings = model_output
+        input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
+        input_mask_expanded = np.tile(input_mask_expanded, (1, 1, token_embeddings.shape[-1]))
+        input_mask_expanded = input_mask_expanded.astype(float)
+        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
+        sum_mask = np.sum(input_mask_expanded, axis=1)
+        pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)
+        return pooled_embeddings
+
+    @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_pooled_models
+
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
+        if output.attention_mask is None:
+            raise ValueError("attention_mask must be provided for document post-processing")
+
+        embeddings = output.model_output
+        attn_mask = output.attention_mask
+        return self.mean_pooling(embeddings, attn_mask).astype(np.float32)
+
+
+class PooledEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs,
+    ) -> OnnxTextEmbedding:
+        return PooledEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/text/pooled_normalized_embedding.py

@@ -0,0 +1,98 @@
+from typing import Any, Dict, Iterable, List, Type
+
+import numpy as np
+
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.utils import normalize
+from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
+from fastembed.text.pooled_embedding import PooledEmbedding
+
+supported_pooled_normalized_models = [
+    {
+        "model": "sentence-transformers/all-MiniLM-L6-v2",
+        "dim": 384,
+        "description": "Text embeddings, Unimodal (text), English, 256 input tokens truncation, Prefixes for queries/documents: not necessary, 2021 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.09,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz",
+            "hf": "qdrant/all-MiniLM-L6-v2-onnx",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "jinaai/jina-embeddings-v2-base-en",
+        "dim": 768,
+        "description": "Text embeddings, Unimodal (text), English, 8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2023 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.52,
+        "sources": {"hf": "xenova/jina-embeddings-v2-base-en"},
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "jinaai/jina-embeddings-v2-small-en",
+        "dim": 512,
+        "description": "Text embeddings, Unimodal (text), English, 8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2023 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.12,
+        "sources": {"hf": "xenova/jina-embeddings-v2-small-en"},
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "jinaai/jina-embeddings-v2-base-de",
+        "dim": 768,
+        "description": "Text embeddings, Unimodal (text), Multilingual (German, English), 8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.32,
+        "sources": {"hf": "jinaai/jina-embeddings-v2-base-de"},
+        "model_file": "onnx/model_fp16.onnx",
+    },
+    {
+        "model": "jinaai/jina-embeddings-v2-base-code",
+        "dim": 768,
+        "description": "Text embeddings, Unimodal (text), Multilingual (English, 30 programming languages), 8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.64,
+        "sources": {"hf": "jinaai/jina-embeddings-v2-base-code"},
+        "model_file": "onnx/model.onnx",
+    },
+]
+
+
+class PooledNormalizedEmbedding(PooledEmbedding):
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
+        return PooledNormalizedEmbeddingWorker
+
+    @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_pooled_normalized_models
+
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
+        if output.attention_mask is None:
+            raise ValueError("attention_mask must be provided for document post-processing")
+
+        embeddings = output.model_output
+        attn_mask = output.attention_mask
+        return normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
+
+
+class PooledNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs,
+    ) -> OnnxTextEmbedding:
+        return PooledNormalizedEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/text/text_embedding.py

@@ -5,8 +5,8 @@ import numpy as np
 from fastembed.common import OnnxProvider
 from fastembed.text.clip_embedding import CLIPOnnxEmbedding
 from fastembed.text.e5_onnx_embedding import E5OnnxEmbedding
-from fastembed.text.jina_onnx_embedding import JinaOnnxEmbedding
-from fastembed.text.mini_lm_embedding import MiniLMOnnxEmbedding
+from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
+from fastembed.text.pooled_embedding import PooledEmbedding
 from fastembed.text.onnx_embedding import OnnxTextEmbedding
 from fastembed.text.text_embedding_base import TextEmbeddingBase
 
@@ -15,9 +15,9 @@ class TextEmbedding(TextEmbeddingBase):
     EMBEDDINGS_REGISTRY: List[Type[TextEmbeddingBase]] = [
         OnnxTextEmbedding,
         E5OnnxEmbedding,
-        JinaOnnxEmbedding,
         CLIPOnnxEmbedding,
-        MiniLMOnnxEmbedding,
+        PooledNormalizedEmbedding,
+        PooledEmbedding,
     ]
 
     @classmethod
@@ -35,6 +35,7 @@ class TextEmbedding(TextEmbeddingBase):
                         "model": "intfloat/multilingual-e5-large",
                         "dim": 1024,
                         "description": "Multilingual model, e5-large. Recommend using this model for non-English languages",
+                        "license": "mit",
                         "size_in_GB": 2.24,
                         "sources": {
                             "gcp": "https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
@@ -55,21 +56,23 @@ class TextEmbedding(TextEmbeddingBase):
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
         providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[List[int]] = None,
+        lazy_load: bool = False,
         **kwargs,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
-
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
             supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
-            if any(
-                model_name.lower() == model["model"].lower()
-                for model in supported_models
-            ):
+            if any(model_name.lower() == model["model"].lower() for model in supported_models):
                 self.model = EMBEDDING_MODEL_TYPE(
-                    model_name,
-                    cache_dir,
+                    model_name=model_name,
+                    cache_dir=cache_dir,
                     threads=threads,
                     providers=providers,
+                    cuda=cuda,
+                    device_ids=device_ids,
+                    lazy_load=lazy_load,
                     **kwargs,
                 )
                 return

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
-name = "fastembed"
-version = "0.3.3"
+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>"]
 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 = "^1.17.0"
+onnxruntime-gpu = ">=1.17.0,<1.20.0"
 tqdm = "^4.66"
 requests = "^2.31"
 tokenizers = ">=0.15,<1.0"
@@ -24,9 +24,8 @@ numpy = [
     { version = ">=1.26, <2", python = ">=3.12" }
 ]
 pillow = "^10.3.0"
-snowballstemmer = "^2.2.0"
-PyStemmer = "^2.2.0"
-mmh3 = "^4.0"
+mmh3 = "^4.1.0"
+py-rust-stemmers = "^0.1.0"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.4.2"

tests/__init__.py

@@ -0,0 +1,4 @@
+import os
+
+# disable DeprecationWarning https://github.com/jupyter/jupyter_core/issues/398
+os.environ["JUPYTER_PLATFORM_DIRS"] = "1"

tests/profiling.py

@@ -105,9 +105,7 @@ def calculate_time_stats(
 # %%
 hf_stats = calculate_time_stats(hf.embed, documents, k=2)
 print(f"Huggingface Transformers (Average, Max, Min): {hf_stats}")
-fst_stats = calculate_time_stats(
-    lambda x: list(embedding_model.embed(x)), documents, k=2
-)
+fst_stats = calculate_time_stats(lambda x: list(embedding_model.embed(x)), documents, k=2)
 print(f"FastEmbed (Average, Max, Min): {fst_stats}")
 
 

tests/test_attention_embeddings.py

@@ -1,13 +1,15 @@
+import os
+
 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"]
-)
+@pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"])
 def test_attention_embeddings(model_name):
+    is_ci = os.getenv("CI")
     model = SparseTextEmbedding(model_name=model_name)
 
     output = list(
@@ -64,14 +66,17 @@ def test_attention_embeddings(model_name):
         assert len(result.indices) == len(result.values)
         assert len(result.indices) == 2
 
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
 
-@pytest.mark.parametrize(
-    "model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"]
-)
+
+@pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"])
 def test_parallel_processing(model_name):
+    is_ci = os.getenv("CI")
+
     model = SparseTextEmbedding(model_name=model_name)
 
-    docs = ["hello world", "attention embedding"] * 100
+    docs = ["hello world", "attention embedding", "Mangez-vous vraiment des grenouilles?"] * 100
     embeddings = list(model.embed(docs, batch_size=10, parallel=2))
 
     embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
@@ -85,3 +90,70 @@ def test_parallel_processing(model_name):
         assert np.allclose(emb_1.indices, emb_3.indices)
         assert np.allclose(emb_1.values, emb_2.values)
         assert np.allclose(emb_1.values, emb_3.values)
+
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
+
+
+@pytest.mark.parametrize("model_name", ["Qdrant/bm25"])
+def test_multilanguage(model_name):
+    is_ci = os.getenv("CI")
+
+    docs = ["Mangez-vous vraiment des grenouilles?", "Je suis au lit"]
+
+    model = SparseTextEmbedding(model_name=model_name, language="french")
+    embeddings = list(model.embed(docs))[:2]
+    assert embeddings[0].values.shape == (3,)
+    assert embeddings[0].indices.shape == (3,)
+
+    assert embeddings[1].values.shape == (1,)
+    assert embeddings[1].indices.shape == (1,)
+
+    model = SparseTextEmbedding(model_name=model_name, language="english")
+    embeddings = list(model.embed(docs))[:2]
+    assert embeddings[0].values.shape == (5,)
+    assert embeddings[0].indices.shape == (5,)
+
+    assert embeddings[1].values.shape == (4,)
+    assert embeddings[1].indices.shape == (4,)
+
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
+
+
+@pytest.mark.parametrize("model_name", ["Qdrant/bm25"])
+def test_special_characters(model_name):
+    is_ci = os.getenv("CI")
+
+    docs = [
+        "Über den größten Flüssen Österreichs äußern sich Experten häufig: Öko-Systeme müssen geschützt werden!",
+        "L'élève français s'écrie : « Où est mon crayon ? J'ai besoin de finir cet exercice avant la récréation!",
+        "Într-o zi însorită, Ștefan și Ioana au mâncat mămăligă cu brânză și au băut țuică la cabană.",
+        "Üzgün öğretmen öğrencilere seslendi: Lütfen gürültü yapmayın, sınavınızı bitirmeye çalışıyorum!",
+        "Ο Ξενοφών είπε: «Ψάχνω για ένα ωραίο δώρο για τη γιαγιά μου. Ίσως ένα φυτό ή ένα βιβλίο;»",
+        "Hola! ¿Cómo estás? Estoy muy emocionado por el cumpleaños de mi hermano, ¡va a ser increíble! También quiero comprar un pastel de chocolate con fresas y un regalo especial: un libro titulado «Cien años de soledad",
+    ]
+
+    model = SparseTextEmbedding(model_name=model_name, language="english")
+    embeddings = list(model.embed(docs))
+    for idx, shape in enumerate([14, 18, 15, 10, 15]):
+        assert embeddings[idx].values.shape == (shape,)
+        assert embeddings[idx].indices.shape == (shape,)
+
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
+
+
+@pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions"])
+def test_lazy_load(model_name):
+    model = SparseTextEmbedding(model_name=model_name, lazy_load=True)
+    assert not hasattr(model.model, "model")
+    docs = ["hello world", "flag embedding"]
+    list(model.embed(docs))
+    assert hasattr(model.model, "model")
+
+    model = SparseTextEmbedding(model_name=model_name, lazy_load=True)
+    list(model.query_embed(docs))
+
+    model = SparseTextEmbedding(model_name=model_name, lazy_load=True)
+    list(model.passage_embed(docs))

tests/test_image_onnx_embeddings.py

@@ -1,16 +1,26 @@
 import os
+from io import BytesIO
 
 import numpy as np
 import pytest
+import requests
+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]),
     "Qdrant/resnet50-onnx": np.array(
         [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.01046245, 0.01171397, 0.00705971, 0.0]
     ),
+    "Qdrant/Unicom-ViT-B-16": np.array(
+        [0.0170, -0.0361, 0.0125, -0.0428, -0.0232, 0.0232, -0.0602, -0.0333, 0.0155, 0.0497]
+    ),
+    "Qdrant/Unicom-ViT-B-32": np.array(
+        [0.0418, 0.0550, 0.0003, 0.0253, -0.0185, 0.0016, -0.0368, -0.0402, -0.0891, -0.0186]
+    ),
 }
 
 
@@ -25,10 +35,15 @@ def test_embedding():
 
         model = ImageEmbedding(model_name=model_desc["model"])
 
-        images = [TEST_MISC_DIR / "image.jpeg", str(TEST_MISC_DIR / "small_image.jpeg")]
+        images = [
+            TEST_MISC_DIR / "image.jpeg",
+            str(TEST_MISC_DIR / "small_image.jpeg"),
+            Image.open((TEST_MISC_DIR / "small_image.jpeg")),
+            Image.open(BytesIO(requests.get("https://qdrant.tech/img/logo.png").content)),
+        ]
         embeddings = list(model.embed(images))
         embeddings = np.stack(embeddings, axis=0)
-        assert embeddings.shape == (2, dim)
+        assert embeddings.shape == (len(images), dim)
 
         canonical_vector = CANONICAL_VECTOR_VALUES[model_desc["model"]]
 
@@ -36,29 +51,44 @@ def test_embedding():
             embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
         ), model_desc["model"]
 
+        assert np.allclose(embeddings[1], embeddings[2]), model_desc["model"]
+
+        if is_ci:
+            delete_model_cache(model.model._model_dir)
+
 
 @pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
 def test_batch_embedding(n_dims, model_name):
+    is_ci = os.getenv("CI")
     model = ImageEmbedding(model_name=model_name)
     n_images = 32
-    images = [TEST_MISC_DIR / "image.jpeg", str(TEST_MISC_DIR / "small_image.jpeg")] * (
-        n_images // 2
-    )
+    test_images = [
+        TEST_MISC_DIR / "image.jpeg",
+        str(TEST_MISC_DIR / "small_image.jpeg"),
+        Image.open(TEST_MISC_DIR / "small_image.jpeg"),
+    ]
+    images = test_images * n_images
 
     embeddings = list(model.embed(images, batch_size=10))
     embeddings = np.stack(embeddings, axis=0)
 
-    assert embeddings.shape == (n_images, n_dims)
+    assert embeddings.shape == (len(test_images) * n_images, n_dims)
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
 
 
 @pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
 def test_parallel_processing(n_dims, model_name):
+    is_ci = os.getenv("CI")
     model = ImageEmbedding(model_name=model_name)
 
     n_images = 32
-    images = [TEST_MISC_DIR / "image.jpeg", str(TEST_MISC_DIR / "small_image.jpeg")] * (
-        n_images // 2
-    )
+    test_images = [
+        TEST_MISC_DIR / "image.jpeg",
+        str(TEST_MISC_DIR / "small_image.jpeg"),
+        Image.open(TEST_MISC_DIR / "small_image.jpeg"),
+    ]
+    images = test_images * n_images
     embeddings = list(model.embed(images, batch_size=10, parallel=2))
     embeddings = np.stack(embeddings, axis=0)
 
@@ -68,6 +98,23 @@ def test_parallel_processing(n_dims, model_name):
     embeddings_3 = list(model.embed(images, batch_size=10, parallel=0))
     embeddings_3 = np.stack(embeddings_3, axis=0)
 
-    assert embeddings.shape == (n_images, n_dims)
+    assert embeddings.shape == (n_images * len(test_images), n_dims)
     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)
+
+
+@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 = [
+        TEST_MISC_DIR / "image.jpeg",
+        str(TEST_MISC_DIR / "small_image.jpeg"),
+    ]
+    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,8 +1,12 @@
+import os
+
+import pytest
 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 = {
@@ -14,7 +18,25 @@ CANONICAL_COLUMN_VALUES = {
             [0.0846, 0.0122, 0.0032, -0.0109, -0.1041],
             [0.0477, 0.1078, -0.0314, 0.016, 0.0156],
         ]
-    )
+    ),
+    "answerdotai/answerai-colbert-small-v1": np.array(
+        [
+            [-0.07281, 0.04632, -0.04711, 0.00762, -0.07374],
+            [-0.04464, 0.04426, -0.074, 0.01801, -0.05233],
+            [0.09936, -0.05123, -0.04925, -0.05276, -0.08944],
+            [0.01644, 0.0203, -0.03789, 0.03165, -0.06501],
+            [-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 = {
@@ -53,13 +75,86 @@ CANONICAL_QUERY_VALUES = {
             [0.1022, 0.0228, -0.0174, -0.0102, -0.065],
             [0.1043, 0.0231, -0.0144, -0.0246, -0.067],
         ]
-    )
+    ),
+    "answerdotai/answerai-colbert-small-v1": np.array(
+        [
+            [-0.07284, 0.04657, -0.04746, 0.00786, -0.07342],
+            [-0.0473, 0.04615, -0.07551, 0.01591, -0.0517],
+            [0.09658, -0.0506, -0.04593, -0.05225, -0.09086],
+            [0.01815, 0.0165, -0.03366, 0.03214, -0.07019],
+            [-0.07284, 0.04657, -0.04746, 0.00787, -0.07342],
+            [-0.07748, 0.04493, -0.055, 0.00481, -0.0486],
+            [-0.0803, 0.04229, -0.0589, 0.00379, -0.04506],
+            [-0.08477, 0.03724, -0.06162, 0.00578, -0.04554],
+            [-0.08392, 0.03805, -0.06202, 0.00899, -0.0409],
+            [-0.07945, 0.04163, -0.06151, 0.00569, -0.04432],
+            [-0.08469, 0.03985, -0.05765, 0.00485, -0.04485],
+            [-0.08306, 0.04111, -0.05774, 0.00583, -0.04325],
+            [-0.08244, 0.04597, -0.05842, 0.00433, -0.04025],
+            [-0.08385, 0.04745, -0.05845, 0.00469, -0.04002],
+            [-0.08402, 0.05014, -0.05941, 0.00692, -0.03452],
+            [-0.08303, 0.05693, -0.05701, 0.00504, -0.03565],
+            [-0.08216, 0.05516, -0.05687, 0.0057, -0.03748],
+            [-0.08051, 0.05751, -0.05647, 0.00283, -0.03645],
+            [-0.08172, 0.05608, -0.06064, 0.00252, -0.03533],
+            [-0.08073, 0.06144, -0.06373, 0.00935, -0.03154],
+            [-0.06651, 0.06697, -0.06769, 0.01717, -0.03369],
+            [-0.06526, 0.06931, -0.06935, 0.0139, -0.03702],
+            [-0.05435, 0.05829, -0.06593, 0.01708, -0.04559],
+            [-0.03648, 0.05234, -0.06759, 0.02057, -0.05053],
+            [-0.03461, 0.05032, -0.06747, 0.02216, -0.05209],
+            [-0.03444, 0.04835, -0.06812, 0.02296, -0.05276],
+            [-0.03292, 0.04853, -0.06811, 0.02348, -0.05303],
+            [-0.03349, 0.04783, -0.06846, 0.02393, -0.05334],
+            [-0.03485, 0.04677, -0.06826, 0.02362, -0.05326],
+            [-0.03408, 0.04744, -0.06931, 0.02302, -0.05288],
+            [-0.03444, 0.04838, -0.06945, 0.02133, -0.05277],
+            [-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"]
 
 
 def test_batch_embedding():
+    is_ci = os.getenv("CI")
     docs_to_embed = docs * 10
 
     for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
@@ -69,10 +164,14 @@ def test_batch_embedding():
 
         for value in result:
             token_num, abridged_dim = expected_result.shape
-            assert np.allclose(value[:, :abridged_dim], expected_result, atol=10e-4)
+            assert np.allclose(value[:, :abridged_dim], expected_result, atol=2e-3)
+
+        if is_ci:
+            delete_model_cache(model.model._model_dir)
 
 
 def test_single_embedding():
+    is_ci = os.getenv("CI")
     docs_to_embed = docs
 
     for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
@@ -80,10 +179,14 @@ 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=10e-4)
+        assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
+
+        if is_ci:
+            delete_model_cache(model.model._model_dir)
 
 
 def test_single_embedding_query():
+    is_ci = os.getenv("CI")
     queries_to_embed = docs
 
     for model_name, expected_result in CANONICAL_QUERY_VALUES.items():
@@ -91,10 +194,14 @@ 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=10e-4)
+        assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
+
+        if is_ci:
+            delete_model_cache(model.model._model_dir)
 
 
 def test_parallel_processing():
+    is_ci = os.getenv("CI")
     model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0")
     token_dim = 128
     docs = ["hello world", "flag embedding"] * 100
@@ -110,3 +217,30 @@ def test_parallel_processing():
     assert embeddings.shape[0] == len(docs) and embeddings.shape[-1] == token_dim
     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)
+
+
+@pytest.mark.parametrize(
+    "model_name",
+    ["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")
+
+    docs = ["hello world", "flag embedding"]
+    list(model.embed(docs))
+    assert hasattr(model.model, "model")
+
+    model = LateInteractionTextEmbedding(model_name=model_name, lazy_load=True)
+    list(model.query_embed(docs))
+
+    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_multi_gpu.py

@@ -0,0 +1,220 @@
+import pytest
+from fastembed import (
+    TextEmbedding,
+    SparseTextEmbedding,
+    LateInteractionTextEmbedding,
+    ImageEmbedding,
+)
+from fastembed.rerank.cross_encoder import TextCrossEncoder
+from tests.config import TEST_MISC_DIR
+
+CACHE_DIR = "../model_cache"
+
+
+@pytest.mark.skip(reason="Requires a multi-gpu server")
+@pytest.mark.parametrize("device_id", [None, 0, 1])
+def test_gpu_via_providers(device_id):
+    docs = ["hello world", "flag embedding"]
+
+    device_id = device_id if device_id is not None else 0
+    providers = (
+        ["CUDAExecutionProvider"]
+        if device_id is None
+        else [("CUDAExecutionProvider", {"device_id": device_id})]
+    )
+    embedding_model = TextEmbedding(
+        "sentence-transformers/all-MiniLM-L6-v2",
+        providers=providers,
+        cache_dir=CACHE_DIR,
+    )
+    list(embedding_model.embed(docs))
+    options = embedding_model.model.model.get_provider_options()
+    assert options["CUDAExecutionProvider"]["device_id"] == str(device_id)
+
+    embedding_model = SparseTextEmbedding(
+        "prithvida/Splade_PP_en_v1",
+        providers=providers,
+        cache_dir=CACHE_DIR,
+    )
+    list(embedding_model.embed(docs))
+    options = embedding_model.model.model.get_provider_options()
+    assert options["CUDAExecutionProvider"]["device_id"] == str(device_id)
+
+    embedding_model = SparseTextEmbedding(
+        "Qdrant/bm42-all-minilm-l6-v2-attentions",
+        providers=providers,
+        cache_dir=CACHE_DIR,
+    )
+    list(embedding_model.embed(docs))
+    options = embedding_model.model.model.get_provider_options()
+    assert options["CUDAExecutionProvider"]["device_id"] == str(device_id)
+
+    embedding_model = LateInteractionTextEmbedding(
+        "colbert-ir/colbertv2.0",
+        providers=providers,
+        cache_dir=CACHE_DIR,
+    )
+    list(embedding_model.embed(docs))
+    options = embedding_model.model.model.get_provider_options()
+    assert options["CUDAExecutionProvider"]["device_id"] == str(device_id)
+
+    embedding_model = ImageEmbedding(
+        model_name="Qdrant/clip-ViT-B-32-vision",
+        providers=providers,
+        cache_dir=CACHE_DIR,
+    )
+    images = [
+        TEST_MISC_DIR / "image.jpeg",
+        str(TEST_MISC_DIR / "small_image.jpeg"),
+    ]
+    list(embedding_model.embed(images))
+    options = embedding_model.model.model.get_provider_options()
+    assert options["CUDAExecutionProvider"]["device_id"] == str(device_id)
+
+    model = TextCrossEncoder(
+        model_name="Xenova/ms-marco-MiniLM-L-6-v2",
+        providers=providers,
+        cache_dir=CACHE_DIR,
+    )
+    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))
+    options = embedding_model.model.model.get_provider_options()
+    assert options["CUDAExecutionProvider"]["device_id"] == str(device_id)
+
+
+@pytest.mark.skip(reason="Requires a multi-gpu server")
+@pytest.mark.parametrize("device_ids", [None, [0], [1], [0, 1]])
+def test_gpu_cuda_device_ids(device_ids):
+    docs = ["hello world", "flag embedding"]
+    device_id = device_ids[0] if device_ids else 0
+    embedding_model = TextEmbedding(
+        "sentence-transformers/all-MiniLM-L6-v2",
+        cuda=True,
+        device_ids=device_ids,
+        cache_dir=CACHE_DIR,
+    )
+    list(embedding_model.embed(docs))
+    options = embedding_model.model.model.get_provider_options()
+    assert options["CUDAExecutionProvider"]["device_id"] == str(
+        device_id
+    ), f"Text embedding: {options}"
+
+    embedding_model = SparseTextEmbedding(
+        "prithvida/Splade_PP_en_v1",
+        cuda=True,
+        device_ids=device_ids,
+        cache_dir=CACHE_DIR,
+    )
+    list(embedding_model.embed(docs))
+    options = embedding_model.model.model.get_provider_options()
+    assert options["CUDAExecutionProvider"]["device_id"] == str(
+        device_id
+    ), f"Sparse text embedding: {options}"
+
+    embedding_model = SparseTextEmbedding(
+        "Qdrant/bm42-all-minilm-l6-v2-attentions",
+        cuda=True,
+        device_ids=device_ids,
+        cache_dir=CACHE_DIR,
+    )
+    list(embedding_model.embed(docs))
+    options = embedding_model.model.model.get_provider_options()
+    assert options["CUDAExecutionProvider"]["device_id"] == str(device_id), f"Bm42: {options}"
+
+    embedding_model = LateInteractionTextEmbedding(
+        "colbert-ir/colbertv2.0",
+        cuda=True,
+        device_ids=device_ids,
+        cache_dir=CACHE_DIR,
+    )
+    list(embedding_model.embed(docs))
+    options = embedding_model.model.model.get_provider_options()
+    assert options["CUDAExecutionProvider"]["device_id"] == str(
+        device_id
+    ), f"Late interaction text embedding: {options}"
+
+    embedding_model = ImageEmbedding(
+        model_name="Qdrant/clip-ViT-B-32-vision",
+        cuda=True,
+        device_ids=device_ids,
+        cache_dir=CACHE_DIR,
+    )
+    images = [
+        TEST_MISC_DIR / "image.jpeg",
+        str(TEST_MISC_DIR / "small_image.jpeg"),
+    ]
+    list(embedding_model.embed(images))
+    options = embedding_model.model.model.get_provider_options()
+    assert options["CUDAExecutionProvider"]["device_id"] == str(
+        device_id
+    ), f"Image embedding: {options}"
+
+    if device_ids is None or len(device_ids) == 1:
+        model = TextCrossEncoder(
+            model_name="Xenova/ms-marco-MiniLM-L-6-v2",
+            cuda=True,
+            device_ids=device_ids,
+            cache_dir=CACHE_DIR,
+        )
+        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))
+        options = embedding_model.model.model.get_provider_options()
+        assert options["CUDAExecutionProvider"]["device_id"] == str(
+            device_id
+        ), f"Text cross encoder: {options}"
+
+
+@pytest.mark.skip(reason="Requires a multi-gpu server")
+@pytest.mark.parametrize(
+    "device_ids,parallel", [(None, None), (None, 2), ([1], None), ([1], 1), ([1], 2), ([0, 1], 2)]
+)
+def test_multi_gpu_parallel_inference(device_ids, parallel):
+    docs = ["hello world", "flag embedding"] * 100
+    batch_size = 5
+
+    embedding_model = TextEmbedding(
+        "sentence-transformers/all-MiniLM-L6-v2",
+        cuda=True,
+        device_ids=device_ids,
+        cache_dir=CACHE_DIR,
+        lazy_load=True,
+    )
+    list(embedding_model.embed(docs, batch_size=batch_size, parallel=parallel))
+
+    embedding_model = SparseTextEmbedding(
+        "prithvida/Splade_PP_en_v1",
+        cuda=True,
+        device_ids=device_ids,
+        cache_dir=CACHE_DIR,
+    )
+    list(embedding_model.embed(docs, batch_size=batch_size, parallel=parallel))
+
+    embedding_model = SparseTextEmbedding(
+        "Qdrant/bm42-all-minilm-l6-v2-attentions",
+        cuda=True,
+        device_ids=device_ids,
+        cache_dir=CACHE_DIR,
+    )
+    list(embedding_model.embed(docs, batch_size=batch_size, parallel=parallel))
+
+    embedding_model = LateInteractionTextEmbedding(
+        "colbert-ir/colbertv2.0",
+        cuda=True,
+        device_ids=device_ids,
+        cache_dir=CACHE_DIR,
+    )
+    list(embedding_model.embed(docs, batch_size=batch_size, parallel=parallel))
+
+    embedding_model = ImageEmbedding(
+        model_name="Qdrant/clip-ViT-B-32-vision",
+        cuda=True,
+        device_ids=device_ids,
+        cache_dir=CACHE_DIR,
+    )
+    images = [
+        TEST_MISC_DIR / "image.jpeg",
+        str(TEST_MISC_DIR / "small_image.jpeg"),
+    ] * 100
+    list(embedding_model.embed(images, batch_size=batch_size, parallel=parallel))

tests/test_sparse_embeddings.py

@@ -1,6 +1,11 @@
-import pytest
+import os
 
+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": {
@@ -45,6 +50,7 @@ docs = ["Hello World"]
 
 
 def test_batch_embedding():
+    is_ci = os.getenv("CI")
     docs_to_embed = docs * 10
 
     for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
@@ -54,9 +60,12 @@ 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)
 
 
 def test_single_embedding():
+    is_ci = os.getenv("CI")
     for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
         model = SparseTextEmbedding(model_name=model_name)
 
@@ -67,11 +76,12 @@ 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)
 
 
 def test_parallel_processing():
-    import numpy as np
-
+    is_ci = os.getenv("CI")
     model = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1")
     docs = ["hello world", "flag embedding"] * 30
     sparse_embeddings_duo = list(model.embed(docs, batch_size=10, parallel=2))
@@ -93,9 +103,72 @@ def test_parallel_processing():
             == sparse_embedding_duo.indices.tolist()
             == sparse_embedding_all.indices.tolist()
         )
-        assert np.allclose(
-            sparse_embedding.values, sparse_embedding_duo.values, atol=1e-3
-        )
-        assert np.allclose(
-            sparse_embedding.values, sparse_embedding_all.values, atol=1e-3
-        )
+        assert np.allclose(sparse_embedding.values, sparse_embedding_duo.values, atol=1e-3)
+        assert np.allclose(sparse_embedding.values, sparse_embedding_all.values, atol=1e-3)
+
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
+
+
+@pytest.fixture
+def bm25_instance():
+    ci = os.getenv("CI", True)
+    model = Bm25("Qdrant/bm25", language="english")
+    yield model
+    if ci:
+        delete_model_cache(model._model_dir)
+
+
+def test_stem_with_stopwords_and_punctuation(bm25_instance):
+    # Setup
+    bm25_instance.stopwords = {"the", "is", "a"}
+    bm25_instance.punctuation = {".", ",", "!"}
+
+    # Test data
+    tokens = ["The", "quick", "brown", "fox", "is", "a", "test", "sentence", ".", "!"]
+
+    # Execute
+    result = bm25_instance._stem(tokens)
+
+    # Assert
+    expected = ["quick", "brown", "fox", "test", "sentenc"]
+    assert result == expected, f"Expected {expected}, but got {result}"
+
+
+def test_stem_case_insensitive_stopwords(bm25_instance):
+    # Setup
+    bm25_instance.stopwords = {"the", "is", "a"}
+    bm25_instance.punctuation = {".", ",", "!"}
+
+    # Test data
+    tokens = ["THE", "Quick", "Brown", "Fox", "IS", "A", "Test", "Sentence", ".", "!"]
+
+    # Execute
+    result = bm25_instance._stem(tokens)
+
+    # Assert
+    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")
+
+    docs = ["hello world", "flag embedding"]
+    list(model.embed(docs))
+    assert hasattr(model.model, "model")
+
+    model = SparseTextEmbedding(model_name=model_name, lazy_load=True)
+    list(model.query_embed(docs))
+
+    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

@@ -0,0 +1,75 @@
+import os
+
+import numpy as np
+import pytest
+
+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]),
+}
+
+
+def test_rerank():
+    is_ci = os.getenv("CI")
+
+    for model_desc in TextCrossEncoder.list_supported_models():
+        if not is_ci and model_desc["size_in_GB"] > 1:
+            continue
+
+        model_name = model_desc["model"]
+        model = TextCrossEncoder(model_name=model_name)
+
+        query = "What is the capital of France?"
+        documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
+        scores = np.array(list(model.rerank(query, documents)))
+
+        canonical_scores = CANONICAL_SCORE_VALUES[model_name]
+        assert np.allclose(
+            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)
+
+
+@pytest.mark.parametrize(
+    "model_name",
+    ["Xenova/ms-marco-MiniLM-L-6-v2", "Xenova/ms-marco-MiniLM-L-12-v2", "BAAI/bge-reranker-base"],
+)
+def test_batch_rerank(model_name):
+    is_ci = os.getenv("CI")
+
+    model = TextCrossEncoder(model_name=model_name)
+
+    query = "What is the capital of France?"
+    documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."] * 50
+    scores = np.array(list(model.rerank(query, documents, batch_size=10)))
+
+    canonical_scores = np.tile(CANONICAL_SCORE_VALUES[model_name], 50)
+
+    assert scores.shape == canonical_scores.shape, f"Unexpected shape for model {model_name}"
+    assert np.allclose(
+        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)
+
+
+@pytest.mark.parametrize(
+    "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

@@ -4,6 +4,7 @@ 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]),
@@ -32,29 +33,20 @@ CANONICAL_VECTOR_VALUES = {
     "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2": np.array(
         [0.0094, 0.0184, 0.0328, 0.0072, -0.0351]
     ),
-    "intfloat/multilingual-e5-large": np.array(
-        [0.0098, 0.0045, 0.0066, -0.0354, 0.0070]
-    ),
+    "intfloat/multilingual-e5-large": np.array([0.0098, 0.0045, 0.0066, -0.0354, 0.0070]),
     "sentence-transformers/paraphrase-multilingual-mpnet-base-v2": np.array(
         [-0.01341097, 0.0416553, -0.00480805, 0.02844842, 0.0505299]
     ),
-    "jinaai/jina-embeddings-v2-small-en": np.array(
-        [-0.0455, -0.0428, -0.0122, 0.0613, 0.0015]
-    ),
-    "jinaai/jina-embeddings-v2-base-en": np.array(
-        [-0.0332, -0.0509, 0.0287, -0.0043, -0.0077]
-    ),
-    "jinaai/jina-embeddings-v2-base-de": np.array(
-      [-0.0085, 0.0417, 0.0342, 0.0309, -0.0149]
-    ),
-    "nomic-ai/nomic-embed-text-v1": np.array(
-        [0.0061, 0.0103, -0.0296, -0.0242, -0.0170]
-    ),
+    "jinaai/jina-embeddings-v2-small-en": np.array([-0.0455, -0.0428, -0.0122, 0.0613, 0.0015]),
+    "jinaai/jina-embeddings-v2-base-en": np.array([-0.0332, -0.0509, 0.0287, -0.0043, -0.0077]),
+    "jinaai/jina-embeddings-v2-base-de": np.array([-0.0085, 0.0417, 0.0342, 0.0309, -0.0149]),
+    "jinaai/jina-embeddings-v2-base-code": np.array([0.0145, -0.0164, 0.0136, -0.0170, 0.0734]),
+    "nomic-ai/nomic-embed-text-v1": np.array([0.3708, 0.2031, -0.3406, -0.2114, -0.3230]),
     "nomic-ai/nomic-embed-text-v1.5": np.array(
-        [-1.6531514e-02, 8.5380634e-05, -1.8171231e-01, -3.9333291e-03, 1.2763254e-02]
+        [-0.15407836, -0.03053198, -3.9138033, 0.1910364, 0.13224715]
     ),
     "nomic-ai/nomic-embed-text-v1.5-Q": np.array(
-        [-0.01554983, 0.0129992, -0.17909265, -0.01062993, 0.00512859]
+        [-0.12525563, 0.38030425, -3.961622, 0.04176439, -0.0758301]
     ),
     "thenlper/gte-large": np.array(
         [-0.01920587, 0.00113156, -0.00708992, -0.00632304, -0.04025577]
@@ -62,21 +54,13 @@ CANONICAL_VECTOR_VALUES = {
     "mixedbread-ai/mxbai-embed-large-v1": np.array(
         [0.02295546, 0.03196154, 0.016512, -0.04031524, -0.0219634]
     ),
-    "snowflake/snowflake-arctic-embed-xs": np.array(
-        [0.0092, 0.0619, 0.0196, 0.009, -0.0114]
-    ),
-    "snowflake/snowflake-arctic-embed-s": np.array(
-        [-0.0416, -0.0867, 0.0209, 0.0554, -0.0272]
-    ),
-    "snowflake/snowflake-arctic-embed-m": np.array(
-        [-0.0329, 0.0364, 0.0481, 0.0016, 0.0328]
-    ),
+    "snowflake/snowflake-arctic-embed-xs": np.array([0.0092, 0.0619, 0.0196, 0.009, -0.0114]),
+    "snowflake/snowflake-arctic-embed-s": np.array([-0.0416, -0.0867, 0.0209, 0.0554, -0.0272]),
+    "snowflake/snowflake-arctic-embed-m": np.array([-0.0329, 0.0364, 0.0481, 0.0016, 0.0328]),
     "snowflake/snowflake-arctic-embed-m-long": np.array(
         [0.0080, -0.0266, -0.0335, 0.0282, 0.0143]
     ),
-    "snowflake/snowflake-arctic-embed-l": np.array(
-        [0.0189, -0.0673, 0.0183, 0.0124, 0.0146]
-    ),
+    "snowflake/snowflake-arctic-embed-l": np.array([0.0189, -0.0673, 0.0183, 0.0124, 0.0146]),
     "Qdrant/clip-ViT-B-32-text": np.array([0.0083, 0.0103, -0.0138, 0.0199, -0.0069]),
 }
 
@@ -91,7 +75,6 @@ def test_embedding():
         dim = model_desc["dim"]
 
         model = TextEmbedding(model_name=model_desc["model"])
-
         docs = ["hello world", "flag embedding"]
         embeddings = list(model.embed(docs))
         embeddings = np.stack(embeddings, axis=0)
@@ -101,6 +84,8 @@ def test_embedding():
         assert np.allclose(
             embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
         ), model_desc["model"]
+        if is_ci:
+            delete_model_cache(model.model._model_dir)
 
 
 @pytest.mark.parametrize(
@@ -108,6 +93,7 @@ def test_embedding():
     [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")],
 )
 def test_batch_embedding(n_dims, model_name):
+    is_ci = os.getenv("CI")
     model = TextEmbedding(model_name=model_name)
 
     docs = ["hello world", "flag embedding"] * 100
@@ -115,6 +101,8 @@ def test_batch_embedding(n_dims, model_name):
     embeddings = np.stack(embeddings, axis=0)
 
     assert embeddings.shape == (200, n_dims)
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
 
 
 @pytest.mark.parametrize(
@@ -122,6 +110,7 @@ def test_batch_embedding(n_dims, model_name):
     [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")],
 )
 def test_parallel_processing(n_dims, model_name):
+    is_ci = os.getenv("CI")
     model = TextEmbedding(model_name=model_name)
 
     docs = ["hello world", "flag embedding"] * 100
@@ -137,3 +126,28 @@ def test_parallel_processing(n_dims, model_name):
     assert embeddings.shape == (200, n_dims)
     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)
+
+
+@pytest.mark.parametrize(
+    "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"]
+    list(model.embed(docs))
+    assert hasattr(model.model, "model")
+
+    model = TextEmbedding(model_name=model_name, lazy_load=True)
+    list(model.query_embed(docs))
+
+    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,24 @@
+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)