disable tmate

.github/ISSUE_TEMPLATE/bug-report.yml

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

.github/ISSUE_TEMPLATE/config.yaml

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

.github/ISSUE_TEMPLATE/feature-request.yml

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

.github/ISSUE_TEMPLATE/model-request.yml

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

.github/PULL_REQUEST_TEMPLATE.md

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

.github/workflows/python-tests.yml

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

NOTICE

@@ -5,8 +5,6 @@ This product includes software developed by Qdrant
 This distribution includes the following Jina AI models, each with its respective license:
 - jinaai/jina-colbert-v2
     - License: cc-by-nc-4.0
-- jinaai/jina-reranker-v2-base-multilingual
-    - License: cc-by-nc-4.0
 
 These models are developed by Jina (https://jina.ai/) and are subject to Jina AI's licensing terms.
 

README.md

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

docs/examples/Supported_Models.ipynb

@@ -2,36 +2,38 @@
  "cells": [
   {
    "cell_type": "code",
+   "execution_count": 1,
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-11-13T09:01:03.324551Z",
-     "start_time": "2024-11-13T09:01:03.234711Z"
+     "end_time": "2024-05-31T18:13:23.806907Z",
+     "start_time": "2024-05-31T18:13:23.797078Z"
     }
    },
+   "outputs": [],
    "source": [
     "%load_ext autoreload\n",
     "%autoreload 2"
-   ],
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "The autoreload extension is already loaded. To reload it, use:\n",
-      "  %reload_ext autoreload\n"
-     ]
-    }
-   ],
-   "execution_count": 10
+   ]
   },
   {
    "cell_type": "code",
+   "execution_count": 2,
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-11-13T09:01:04.505772Z",
-     "start_time": "2024-11-13T09:01:04.493296Z"
+     "end_time": "2024-05-31T18:14:31.147674Z",
+     "start_time": "2024-05-31T18:14:31.134015Z"
     }
    },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/hossam/.pyenv/versions/.venv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
    "source": [
     "import pandas as pd\n",
     "\n",
@@ -40,11 +42,8 @@
     "    TextEmbedding,\n",
     "    LateInteractionTextEmbedding,\n",
     "    ImageEmbedding,\n",
-    ")\n",
-    "from fastembed.rerank.cross_encoder import TextCrossEncoder"
-   ],
-   "outputs": [],
-   "execution_count": 11
+    ")"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -55,79 +54,16 @@
   },
   {
    "cell_type": "code",
+   "execution_count": 3,
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-11-13T09:01:05.812271Z",
-     "start_time": "2024-11-13T09:01:05.795846Z"
+     "end_time": "2024-05-31T18:13:25.863008Z",
+     "start_time": "2024-05-31T18:13:25.837795Z"
     }
    },
-   "source": [
-    "supported_models = (\n",
-    "    pd.DataFrame(TextEmbedding.list_supported_models())\n",
-    "    .sort_values(\"size_in_GB\")\n",
-    "    .drop(columns=[\"sources\", \"model_file\", \"additional_files\"])\n",
-    "    .reset_index(drop=True)\n",
-    ")\n",
-    "supported_models"
-   ],
    "outputs": [
     {
      "data": {
-      "text/plain": [
-       "                                                model   dim  \\\n",
-       "0                              BAAI/bge-small-en-v1.5   384   \n",
-       "1                              BAAI/bge-small-zh-v1.5   512   \n",
-       "2                 snowflake/snowflake-arctic-embed-xs   384   \n",
-       "3              sentence-transformers/all-MiniLM-L6-v2   384   \n",
-       "4                  jinaai/jina-embeddings-v2-small-en   512   \n",
-       "5                                   BAAI/bge-small-en   384   \n",
-       "6                  snowflake/snowflake-arctic-embed-s   384   \n",
-       "7                    nomic-ai/nomic-embed-text-v1.5-Q   768   \n",
-       "8                               BAAI/bge-base-en-v1.5   768   \n",
-       "9   sentence-transformers/paraphrase-multilingual-...   384   \n",
-       "10                          Qdrant/clip-ViT-B-32-text   512   \n",
-       "11                  jinaai/jina-embeddings-v2-base-de   768   \n",
-       "12                                   BAAI/bge-base-en   768   \n",
-       "13                 snowflake/snowflake-arctic-embed-m   768   \n",
-       "14                     nomic-ai/nomic-embed-text-v1.5   768   \n",
-       "15                  jinaai/jina-embeddings-v2-base-en   768   \n",
-       "16                       nomic-ai/nomic-embed-text-v1   768   \n",
-       "17            snowflake/snowflake-arctic-embed-m-long   768   \n",
-       "18                 mixedbread-ai/mxbai-embed-large-v1  1024   \n",
-       "19                jinaai/jina-embeddings-v2-base-code   768   \n",
-       "20  sentence-transformers/paraphrase-multilingual-...   768   \n",
-       "21                 snowflake/snowflake-arctic-embed-l  1024   \n",
-       "22                                 thenlper/gte-large  1024   \n",
-       "23                             BAAI/bge-large-en-v1.5  1024   \n",
-       "24                     intfloat/multilingual-e5-large  1024   \n",
-       "\n",
-       "                                          description     license  size_in_GB  \n",
-       "0   Text embeddings, Unimodal (text), English, 512...         mit       0.067  \n",
-       "1   Text embeddings, Unimodal (text), Chinese, 512...         mit       0.090  \n",
-       "2   Text embeddings, Unimodal (text), English, 512...  apache-2.0       0.090  \n",
-       "3   Text embeddings, Unimodal (text), English, 256...  apache-2.0       0.090  \n",
-       "4   Text embeddings, Unimodal (text), English, 819...  apache-2.0       0.120  \n",
-       "5   Text embeddings, Unimodal (text), English, 512...         mit       0.130  \n",
-       "6   Text embeddings, Unimodal (text), English, 512...  apache-2.0       0.130  \n",
-       "7   Text embeddings, Multimodal (text, image), Eng...  apache-2.0       0.130  \n",
-       "8   Text embeddings, Unimodal (text), English, 512...         mit       0.210  \n",
-       "9   Text embeddings, Unimodal (text), Multilingual...  apache-2.0       0.220  \n",
-       "10  Text embeddings, Multimodal (text&image), Engl...         mit       0.250  \n",
-       "11  Text embeddings, Unimodal (text), Multilingual...  apache-2.0       0.320  \n",
-       "12  Text embeddings, Unimodal (text), English, 512...         mit       0.420  \n",
-       "13  Text embeddings, Unimodal (text), English, 512...  apache-2.0       0.430  \n",
-       "14  Text embeddings, Multimodal (text, image), Eng...  apache-2.0       0.520  \n",
-       "15  Text embeddings, Unimodal (text), English, 819...  apache-2.0       0.520  \n",
-       "16  Text embeddings, Multimodal (text, image), Eng...  apache-2.0       0.520  \n",
-       "17  Text embeddings, Unimodal (text), English, 204...  apache-2.0       0.540  \n",
-       "18  Text embeddings, Unimodal (text), English, 512...  apache-2.0       0.640  \n",
-       "19  Text embeddings, Unimodal (text), Multilingual...  apache-2.0       0.640  \n",
-       "20  Text embeddings, Unimodal (text), Multilingual...  apache-2.0       1.000  \n",
-       "21  Text embeddings, Unimodal (text), English, 512...  apache-2.0       1.020  \n",
-       "22  Text embeddings, Unimodal (text), English, 512...         mit       1.200  \n",
-       "23  Text embeddings, Unimodal (text), English, 512...         mit       1.200  \n",
-       "24  Text embeddings, Unimodal (text), Multilingual...         mit       2.240  "
-      ],
       "text/html": [
        "<div>\n",
        "<style scoped>\n",
@@ -358,14 +294,77 @@
        "  </tbody>\n",
        "</table>\n",
        "</div>"
+      ],
+      "text/plain": [
+       "                                                model   dim  \\\n",
+       "0                              BAAI/bge-small-en-v1.5   384   \n",
+       "1                              BAAI/bge-small-zh-v1.5   512   \n",
+       "2                 snowflake/snowflake-arctic-embed-xs   384   \n",
+       "3              sentence-transformers/all-MiniLM-L6-v2   384   \n",
+       "4                  jinaai/jina-embeddings-v2-small-en   512   \n",
+       "5                                   BAAI/bge-small-en   384   \n",
+       "6                  snowflake/snowflake-arctic-embed-s   384   \n",
+       "7                    nomic-ai/nomic-embed-text-v1.5-Q   768   \n",
+       "8                               BAAI/bge-base-en-v1.5   768   \n",
+       "9   sentence-transformers/paraphrase-multilingual-...   384   \n",
+       "10                          Qdrant/clip-ViT-B-32-text   512   \n",
+       "11                  jinaai/jina-embeddings-v2-base-de   768   \n",
+       "12                                   BAAI/bge-base-en   768   \n",
+       "13                 snowflake/snowflake-arctic-embed-m   768   \n",
+       "14                     nomic-ai/nomic-embed-text-v1.5   768   \n",
+       "15                  jinaai/jina-embeddings-v2-base-en   768   \n",
+       "16                       nomic-ai/nomic-embed-text-v1   768   \n",
+       "17            snowflake/snowflake-arctic-embed-m-long   768   \n",
+       "18                 mixedbread-ai/mxbai-embed-large-v1  1024   \n",
+       "19                jinaai/jina-embeddings-v2-base-code   768   \n",
+       "20  sentence-transformers/paraphrase-multilingual-...   768   \n",
+       "21                 snowflake/snowflake-arctic-embed-l  1024   \n",
+       "22                                 thenlper/gte-large  1024   \n",
+       "23                             BAAI/bge-large-en-v1.5  1024   \n",
+       "24                     intfloat/multilingual-e5-large  1024   \n",
+       "\n",
+       "                                          description     license  size_in_GB  \n",
+       "0   Text embeddings, Unimodal (text), English, 512...         mit       0.067  \n",
+       "1   Text embeddings, Unimodal (text), Chinese, 512...         mit       0.090  \n",
+       "2   Text embeddings, Unimodal (text), English, 512...  apache-2.0       0.090  \n",
+       "3   Text embeddings, Unimodal (text), English, 256...  apache-2.0       0.090  \n",
+       "4   Text embeddings, Unimodal (text), English, 819...  apache-2.0       0.120  \n",
+       "5   Text embeddings, Unimodal (text), English, 512...         mit       0.130  \n",
+       "6   Text embeddings, Unimodal (text), English, 512...  apache-2.0       0.130  \n",
+       "7   Text embeddings, Multimodal (text, image), Eng...  apache-2.0       0.130  \n",
+       "8   Text embeddings, Unimodal (text), English, 512...         mit       0.210  \n",
+       "9   Text embeddings, Unimodal (text), Multilingual...  apache-2.0       0.220  \n",
+       "10  Text embeddings, Multimodal (text&image), Engl...         mit       0.250  \n",
+       "11  Text embeddings, Unimodal (text), Multilingual...  apache-2.0       0.320  \n",
+       "12  Text embeddings, Unimodal (text), English, 512...         mit       0.420  \n",
+       "13  Text embeddings, Unimodal (text), English, 512...  apache-2.0       0.430  \n",
+       "14  Text embeddings, Multimodal (text, image), Eng...  apache-2.0       0.520  \n",
+       "15  Text embeddings, Unimodal (text), English, 819...  apache-2.0       0.520  \n",
+       "16  Text embeddings, Multimodal (text, image), Eng...  apache-2.0       0.520  \n",
+       "17  Text embeddings, Unimodal (text), English, 204...  apache-2.0       0.540  \n",
+       "18  Text embeddings, Unimodal (text), English, 512...  apache-2.0       0.640  \n",
+       "19  Text embeddings, Unimodal (text), Multilingual...  apache-2.0       0.640  \n",
+       "20  Text embeddings, Unimodal (text), Multilingual...  apache-2.0       1.000  \n",
+       "21  Text embeddings, Unimodal (text), English, 512...  apache-2.0       1.020  \n",
+       "22  Text embeddings, Unimodal (text), English, 512...         mit       1.200  \n",
+       "23  Text embeddings, Unimodal (text), English, 512...         mit       1.200  \n",
+       "24  Text embeddings, Unimodal (text), Multilingual...         mit       2.240  "
       ]
      },
-     "execution_count": 12,
+     "execution_count": 3,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
-   "execution_count": 12
+   "source": [
+    "supported_models = (\n",
+    "    pd.DataFrame(TextEmbedding.list_supported_models())\n",
+    "    .sort_values(\"size_in_GB\")\n",
+    "    .drop(columns=[\"sources\", \"model_file\", \"additional_files\"])\n",
+    "    .reset_index(drop=True)\n",
+    ")\n",
+    "supported_models"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -376,42 +375,16 @@
   },
   {
    "cell_type": "code",
+   "execution_count": 4,
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-11-13T09:01:07.038954Z",
-     "start_time": "2024-11-13T09:01:07.019656Z"
+     "end_time": "2024-05-31T18:13:27.124747Z",
+     "start_time": "2024-05-31T18:13:27.096212Z"
     }
    },
-   "source": [
-    "(\n",
-    "    pd.DataFrame(SparseTextEmbedding.list_supported_models())\n",
-    "    .sort_values(\"size_in_GB\")\n",
-    "    .drop(columns=[\"sources\", \"model_file\", \"additional_files\"])\n",
-    "    .reset_index(drop=True)\n",
-    ")"
-   ],
    "outputs": [
     {
      "data": {
-      "text/plain": [
-       "                                     model  vocab_size  \\\n",
-       "0                              Qdrant/bm25         NaN   \n",
-       "1  Qdrant/bm42-all-minilm-l6-v2-attentions     30522.0   \n",
-       "2               prithivida/Splade_PP_en_v1     30522.0   \n",
-       "3                prithvida/Splade_PP_en_v1     30522.0   \n",
-       "\n",
-       "                                         description     license  size_in_GB  \\\n",
-       "0  BM25 as sparse embeddings meant to be used wit...  apache-2.0       0.010   \n",
-       "1  Light sparse embedding model, which assigns an...  apache-2.0       0.090   \n",
-       "2  Independent Implementation of SPLADE++ Model f...  apache-2.0       0.532   \n",
-       "3  Independent Implementation of SPLADE++ Model f...  apache-2.0       0.532   \n",
-       "\n",
-       "  requires_idf  \n",
-       "0         True  \n",
-       "1         True  \n",
-       "2          NaN  \n",
-       "3          NaN  "
-      ],
       "text/html": [
        "<div>\n",
        "<style scoped>\n",
@@ -479,14 +452,40 @@
        "  </tbody>\n",
        "</table>\n",
        "</div>"
+      ],
+      "text/plain": [
+       "                                     model  vocab_size  \\\n",
+       "0                              Qdrant/bm25         NaN   \n",
+       "1  Qdrant/bm42-all-minilm-l6-v2-attentions     30522.0   \n",
+       "2               prithivida/Splade_PP_en_v1     30522.0   \n",
+       "3                prithvida/Splade_PP_en_v1     30522.0   \n",
+       "\n",
+       "                                         description     license  size_in_GB  \\\n",
+       "0  BM25 as sparse embeddings meant to be used wit...  apache-2.0       0.010   \n",
+       "1  Light sparse embedding model, which assigns an...  apache-2.0       0.090   \n",
+       "2  Independent Implementation of SPLADE++ Model f...  apache-2.0       0.532   \n",
+       "3  Independent Implementation of SPLADE++ Model f...  apache-2.0       0.532   \n",
+       "\n",
+       "  requires_idf  \n",
+       "0         True  \n",
+       "1         True  \n",
+       "2          NaN  \n",
+       "3          NaN  "
       ]
      },
-     "execution_count": 13,
+     "execution_count": 4,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
-   "execution_count": 13
+   "source": [
+    "(\n",
+    "    pd.DataFrame(SparseTextEmbedding.list_supported_models())\n",
+    "    .sort_values(\"size_in_GB\")\n",
+    "    .drop(columns=[\"sources\", \"model_file\", \"additional_files\"])\n",
+    "    .reset_index(drop=True)\n",
+    ")"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -499,40 +498,17 @@
   },
   {
    "cell_type": "code",
+   "execution_count": 5,
    "metadata": {
-    "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-11-13T09:01:08.074442Z",
-     "start_time": "2024-11-13T09:01:08.056138Z"
-    }
+     "end_time": "2024-05-31T18:14:34.370252Z",
+     "start_time": "2024-05-31T18:14:34.354270Z"
+    },
+    "collapsed": false
    },
-   "source": [
-    "(\n",
-    "    pd.DataFrame(LateInteractionTextEmbedding.list_supported_models())\n",
-    "    .sort_values(\"size_in_GB\")\n",
-    "    .drop(columns=[\"sources\", \"model_file\"])\n",
-    "    .reset_index(drop=True)\n",
-    ")"
-   ],
    "outputs": [
     {
      "data": {
-      "text/plain": [
-       "                                   model  dim  \\\n",
-       "0  answerdotai/answerai-colbert-small-v1   96   \n",
-       "1                 colbert-ir/colbertv2.0  128   \n",
-       "2                 jinaai/jina-colbert-v2  128   \n",
-       "\n",
-       "                                         description       license  \\\n",
-       "0  Text embeddings, Unimodal (text), Multilingual...    apache-2.0   \n",
-       "1                             Late interaction model           mit   \n",
-       "2  New model that expands capabilities of colbert...  cc-by-nc-4.0   \n",
-       "\n",
-       "   size_in_GB        additional_files  \n",
-       "0        0.13                     NaN  \n",
-       "1        0.44                     NaN  \n",
-       "2        2.24  [onnx/model.onnx_data]  "
-      ],
       "text/html": [
        "<div>\n",
        "<style scoped>\n",
@@ -582,7 +558,7 @@
        "    <tr>\n",
        "      <th>2</th>\n",
        "      <td>jinaai/jina-colbert-v2</td>\n",
-       "      <td>128</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",
@@ -591,14 +567,37 @@
        "  </tbody>\n",
        "</table>\n",
        "</div>"
+      ],
+      "text/plain": [
+       "                                   model   dim  \\\n",
+       "0  answerdotai/answerai-colbert-small-v1    96   \n",
+       "1                 colbert-ir/colbertv2.0   128   \n",
+       "2                 jinaai/jina-colbert-v2  1024   \n",
+       "\n",
+       "                                         description       license  \\\n",
+       "0  Text embeddings, Unimodal (text), Multilingual...    apache-2.0   \n",
+       "1                             Late interaction model           mit   \n",
+       "2  New model that expands capabilities of colbert...  cc-by-nc-4.0   \n",
+       "\n",
+       "   size_in_GB        additional_files  \n",
+       "0        0.13                     NaN  \n",
+       "1        0.44                     NaN  \n",
+       "2        2.24  [onnx/model.onnx_data]  "
       ]
      },
-     "execution_count": 14,
+     "execution_count": 5,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
-   "execution_count": 14
+   "source": [
+    "(\n",
+    "    pd.DataFrame(LateInteractionTextEmbedding.list_supported_models())\n",
+    "    .sort_values(\"size_in_GB\")\n",
+    "    .drop(columns=[\"sources\", \"model_file\"])\n",
+    "    .reset_index(drop=True)\n",
+    ")"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -611,37 +610,17 @@
   },
   {
    "cell_type": "code",
+   "execution_count": 6,
    "metadata": {
-    "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-11-13T09:01:09.171647Z",
-     "start_time": "2024-11-13T09:01:09.150940Z"
-    }
+     "end_time": "2024-05-31T18:14:42.501881Z",
+     "start_time": "2024-05-31T18:14:42.484726Z"
+    },
+    "collapsed": false
    },
-   "source": [
-    "(\n",
-    "    pd.DataFrame(ImageEmbedding.list_supported_models())\n",
-    "    .sort_values(\"size_in_GB\")\n",
-    "    .drop(columns=[\"sources\", \"model_file\"])\n",
-    "    .reset_index(drop=True)\n",
-    ")"
-   ],
    "outputs": [
     {
      "data": {
-      "text/plain": [
-       "                         model   dim  \\\n",
-       "0         Qdrant/resnet50-onnx  2048   \n",
-       "1  Qdrant/clip-ViT-B-32-vision   512   \n",
-       "2       Qdrant/Unicom-ViT-B-32   512   \n",
-       "3       Qdrant/Unicom-ViT-B-16   768   \n",
-       "\n",
-       "                                         description     license  size_in_GB  \n",
-       "0      Image embeddings, Unimodal (image), 2016 year  apache-2.0        0.10  \n",
-       "1  Image embeddings, Multimodal (text&image), 202...         mit        0.34  \n",
-       "2  Image embeddings, Multimodal (text&image), 202...  apache-2.0        0.48  \n",
-       "3  Image embeddings (more detailed than Unicom-Vi...  apache-2.0        0.82  "
-      ],
       "text/html": [
        "<div>\n",
        "<style scoped>\n",
@@ -704,149 +683,39 @@
        "  </tbody>\n",
        "</table>\n",
        "</div>"
+      ],
+      "text/plain": [
+       "                         model   dim  \\\n",
+       "0         Qdrant/resnet50-onnx  2048   \n",
+       "1  Qdrant/clip-ViT-B-32-vision   512   \n",
+       "2       Qdrant/Unicom-ViT-B-32   512   \n",
+       "3       Qdrant/Unicom-ViT-B-16   768   \n",
+       "\n",
+       "                                         description     license  size_in_GB  \n",
+       "0      Image embeddings, Unimodal (image), 2016 year  apache-2.0        0.10  \n",
+       "1  Image embeddings, Multimodal (text&image), 202...         mit        0.34  \n",
+       "2  Image embeddings, Multimodal (text&image), 202...  apache-2.0        0.48  \n",
+       "3  Image embeddings (more detailed than Unicom-Vi...  apache-2.0        0.82  "
       ]
      },
-     "execution_count": 15,
+     "execution_count": 6,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
-   "execution_count": 15
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Supported Rerank Cross Encoder Models"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-11-13T09:01:10.313943Z",
-     "start_time": "2024-11-13T09:01:10.298428Z"
-    }
-   },
    "source": [
     "(\n",
-    "    pd.DataFrame(TextCrossEncoder.list_supported_models())\n",
+    "    pd.DataFrame(ImageEmbedding.list_supported_models())\n",
     "    .sort_values(\"size_in_GB\")\n",
     "    .drop(columns=[\"sources\", \"model_file\"])\n",
     "    .reset_index(drop=True)\n",
     ")"
-   ],
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "                                       model  size_in_GB  \\\n",
-       "0              Xenova/ms-marco-MiniLM-L-6-v2        0.08   \n",
-       "1             Xenova/ms-marco-MiniLM-L-12-v2        0.12   \n",
-       "2            jinaai/jina-reranker-v1-tiny-en        0.13   \n",
-       "3           jinaai/jina-reranker-v1-turbo-en        0.15   \n",
-       "4                     BAAI/bge-reranker-base        1.04   \n",
-       "5  jinaai/jina-reranker-v2-base-multilingual        1.11   \n",
-       "\n",
-       "                                         description       license  \n",
-       "0  MiniLM-L-6-v2 model optimized for re-ranking t...    apache-2.0  \n",
-       "1  MiniLM-L-12-v2 model optimized for re-ranking ...    apache-2.0  \n",
-       "2  Designed for blazing-fast re-ranking with 8K c...    apache-2.0  \n",
-       "3  Designed for blazing-fast re-ranking with 8K c...    apache-2.0  \n",
-       "4  BGE reranker base model for cross-encoder re-r...           mit  \n",
-       "5  A multi-lingual reranker model for cross-encod...  cc-by-nc-4.0  "
-      ],
-      "text/html": [
-       "<div>\n",
-       "<style scoped>\n",
-       "    .dataframe tbody tr th:only-of-type {\n",
-       "        vertical-align: middle;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe tbody tr th {\n",
-       "        vertical-align: top;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe thead th {\n",
-       "        text-align: right;\n",
-       "    }\n",
-       "</style>\n",
-       "<table border=\"1\" class=\"dataframe\">\n",
-       "  <thead>\n",
-       "    <tr style=\"text-align: right;\">\n",
-       "      <th></th>\n",
-       "      <th>model</th>\n",
-       "      <th>size_in_GB</th>\n",
-       "      <th>description</th>\n",
-       "      <th>license</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>Xenova/ms-marco-MiniLM-L-6-v2</td>\n",
-       "      <td>0.08</td>\n",
-       "      <td>MiniLM-L-6-v2 model optimized for re-ranking t...</td>\n",
-       "      <td>apache-2.0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1</th>\n",
-       "      <td>Xenova/ms-marco-MiniLM-L-12-v2</td>\n",
-       "      <td>0.12</td>\n",
-       "      <td>MiniLM-L-12-v2 model optimized for re-ranking ...</td>\n",
-       "      <td>apache-2.0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>2</th>\n",
-       "      <td>jinaai/jina-reranker-v1-tiny-en</td>\n",
-       "      <td>0.13</td>\n",
-       "      <td>Designed for blazing-fast re-ranking with 8K c...</td>\n",
-       "      <td>apache-2.0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>3</th>\n",
-       "      <td>jinaai/jina-reranker-v1-turbo-en</td>\n",
-       "      <td>0.15</td>\n",
-       "      <td>Designed for blazing-fast re-ranking with 8K c...</td>\n",
-       "      <td>apache-2.0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>4</th>\n",
-       "      <td>BAAI/bge-reranker-base</td>\n",
-       "      <td>1.04</td>\n",
-       "      <td>BGE reranker base model for cross-encoder re-r...</td>\n",
-       "      <td>mit</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>5</th>\n",
-       "      <td>jinaai/jina-reranker-v2-base-multilingual</td>\n",
-       "      <td>1.11</td>\n",
-       "      <td>A multi-lingual reranker model for cross-encod...</td>\n",
-       "      <td>cc-by-nc-4.0</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "</div>"
-      ]
-     },
-     "execution_count": 16,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "execution_count": 16
-  },
-  {
-   "metadata": {},
-   "cell_type": "code",
-   "outputs": [],
-   "execution_count": null,
-   "source": ""
+   ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3.8.18 ('base')",
+   "display_name": ".venv",
    "language": "python",
    "name": "python3"
   },
@@ -860,14 +729,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.8"
+   "version": "3.10.15"
   },
-  "orig_nbformat": 4,
-  "vscode": {
-   "interpreter": {
-    "hash": "c4a27af61e455bc18dcf16f5867a2ff0402fa12b01dd0f6ce3a79ae73ad15e91"
-   }
-  }
+  "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 2

fastembed/common/onnx_model.py

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

fastembed/late_interaction/colbert.py

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

fastembed/late_interaction/jina_colbert.py

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

fastembed/parallel_processor.py

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

fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py

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

fastembed/rerank/cross_encoder/onnx_text_model.py

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

fastembed/sparse/utils/tokenizer.py

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

pyproject.toml

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

tests/test_text_cross_encoder.py

@@ -10,9 +10,6 @@ CANONICAL_SCORE_VALUES = {
     "Xenova/ms-marco-MiniLM-L-6-v2": np.array([8.500708, -2.541011]),
     "Xenova/ms-marco-MiniLM-L-12-v2": np.array([9.330912, -2.0380247]),
     "BAAI/bge-reranker-base": np.array([6.15733337, -3.65939403]),
-    "jinaai/jina-reranker-v1-tiny-en": np.array([2.5911, 0.1122]),
-    "jinaai/jina-reranker-v1-turbo-en": np.array([1.8295, -2.8908]),
-    "jinaai/jina-reranker-v2-base-multilingual": np.array([1.6533, -1.6455]),
 }
 
 
@@ -40,11 +37,7 @@ def test_rerank():
 
 @pytest.mark.parametrize(
     "model_name",
-    [
-        model_desc["model"]
-        for model_desc in TextCrossEncoder.list_supported_models()
-        if model_desc["size_in_GB"] < 1 and model_desc["model"] in CANONICAL_SCORE_VALUES.keys()
-    ],
+    ["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")

tests/utils.py

@@ -1,5 +1,4 @@
 import shutil
-import traceback
 
 from pathlib import Path
 from typing import Union
@@ -15,12 +14,6 @@ def delete_model_cache(model_dir: Union[str, Path]) -> None:
     Args:
         model_dir (Union[str, Path]): The path to the model cache directory.
     """
-
-    def on_error(func, path, exc_info):
-        print("Failed to remove: ", path)
-        print("Exception: ", exc_info)
-        traceback.print_exception(*exc_info)
-
     if isinstance(model_dir, str):
         model_dir = Path(model_dir)
 
@@ -28,5 +21,4 @@ def delete_model_cache(model_dir: Union[str, Path]) -> None:
         model_dir = model_dir.parent.parent
 
     if model_dir.exists():
-        # todo: PermissionDenied is raised on blobs removal in Windows, with blobs > 2GB
-        shutil.rmtree(model_dir, onerror=on_error)
+        shutil.rmtree(model_dir)