Tsk 375 add jina rerankers (#379)

* feat: Added jina reranker models * chore: Added jina reranker canonical score values * chore: added rounding of the output for easier reproducability * chore: Added jina reranker models in batch test * chore: remove redundant np.round * chore: test only <1gb files in local * chore: Updated docs to add rerankers * fix: recompute canonical values with fp16 * new: extend NOTICE with jina reranker v2 --------- Co-authored-by: George Panchuk <[email protected]>
qdrant · Nov 13, 2024 · 860e2ad · 860e2ad
1 parent d141e29
commit 860e2ad
Show file tree

Hide file tree

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

diff --git a/docs/examples/Supported_Models.ipynb b/docs/examples/Supported_Models.ipynb
@@ -4,8 +4,8 @@
    "cell_type": "code",
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-11-13T08:52:46.651285Z",
-     "start_time": "2024-11-13T08:52:46.633741Z"
+     "end_time": "2024-11-13T09:01:03.324551Z",
+     "start_time": "2024-11-13T09:01:03.234711Z"
     }
    },
    "source": [
@@ -22,14 +22,14 @@
      ]
     }
    ],
-   "execution_count": 3
+   "execution_count": 10
   },
   {
    "cell_type": "code",
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-11-13T08:52:46.973001Z",
-     "start_time": "2024-11-13T08:52:46.962113Z"
+     "end_time": "2024-11-13T09:01:04.505772Z",
+     "start_time": "2024-11-13T09:01:04.493296Z"
     }
    },
    "source": [
@@ -44,7 +44,7 @@
     "from fastembed.rerank.cross_encoder import TextCrossEncoder"
    ],
    "outputs": [],
-   "execution_count": 4
+   "execution_count": 11
   },
   {
    "cell_type": "markdown",
@@ -57,8 +57,8 @@
    "cell_type": "code",
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-11-13T08:52:49.705020Z",
-     "start_time": "2024-11-13T08:52:49.665915Z"
+     "end_time": "2024-11-13T09:01:05.812271Z",
+     "start_time": "2024-11-13T09:01:05.795846Z"
     }
    },
    "source": [
@@ -360,12 +360,12 @@
        "</div>"
       ]
      },
-     "execution_count": 5,
+     "execution_count": 12,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
-   "execution_count": 5
+   "execution_count": 12
   },
   {
    "cell_type": "markdown",
@@ -378,8 +378,8 @@
    "cell_type": "code",
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-11-13T08:52:51.690851Z",
-     "start_time": "2024-11-13T08:52:51.677404Z"
+     "end_time": "2024-11-13T09:01:07.038954Z",
+     "start_time": "2024-11-13T09:01:07.019656Z"
     }
    },
    "source": [
@@ -481,12 +481,12 @@
        "</div>"
       ]
      },
-     "execution_count": 6,
+     "execution_count": 13,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
-   "execution_count": 6
+   "execution_count": 13
   },
   {
    "cell_type": "markdown",
@@ -502,8 +502,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-11-13T08:52:52.866135Z",
-     "start_time": "2024-11-13T08:52:52.852411Z"
+     "end_time": "2024-11-13T09:01:08.074442Z",
+     "start_time": "2024-11-13T09:01:08.056138Z"
     }
    },
    "source": [
@@ -593,12 +593,12 @@
        "</div>"
       ]
      },
-     "execution_count": 7,
+     "execution_count": 14,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
-   "execution_count": 7
+   "execution_count": 14
   },
   {
    "cell_type": "markdown",
@@ -614,8 +614,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-11-13T08:52:54.060642Z",
-     "start_time": "2024-11-13T08:52:54.043437Z"
+     "end_time": "2024-11-13T09:01:09.171647Z",
+     "start_time": "2024-11-13T09:01:09.150940Z"
     }
    },
    "source": [
@@ -706,26 +706,26 @@
        "</div>"
       ]
      },
-     "execution_count": 8,
+     "execution_count": 15,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
-   "execution_count": 8
+   "execution_count": 15
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Supported Rerankers Models"
+    "## Supported Rerank Cross Encoder Models"
    ]
   },
   {
    "cell_type": "code",
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-11-13T08:52:55.933033Z",
-     "start_time": "2024-11-13T08:52:55.915572Z"
+     "end_time": "2024-11-13T09:01:10.313943Z",
+     "start_time": "2024-11-13T09:01:10.298428Z"
     }
    },
    "source": [
@@ -740,15 +740,21 @@
     {
      "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          BAAI/bge-reranker-base        1.04   \n",
+       "                                       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  BGE reranker base model for cross-encoder re-r...         mit  "
+       "                                         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",
@@ -792,22 +798,43 @@
        "    </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": 9,
+     "execution_count": 16,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
-   "execution_count": 9
+   "execution_count": 16
   },
   {
    "metadata": {},
@@ -819,7 +846,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": ".venv",
+   "display_name": "Python 3.8.18 ('base')",
    "language": "python",
    "name": "python3"
   },
@@ -833,9 +860,14 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.8"
   },
-  "orig_nbformat": 4
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "c4a27af61e455bc18dcf16f5867a2ff0402fa12b01dd0f6ce3a79ae73ad15e91"
+   }
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2

diff --git a/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py b/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py
@@ -38,6 +38,36 @@
         "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",
+    },
 ]
 
 

diff --git a/tests/test_text_cross_encoder.py b/tests/test_text_cross_encoder.py
@@ -10,6 +10,9 @@
     "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]),
 }
 
 
@@ -37,7 +40,11 @@ def test_rerank():
 
 @pytest.mark.parametrize(
     "model_name",
-    ["Xenova/ms-marco-MiniLM-L-6-v2", "Xenova/ms-marco-MiniLM-L-12-v2", "BAAI/bge-reranker-base"],
+    [
+        model_desc["model"]
+        for model_desc in TextCrossEncoder.list_supported_models()
+        if model_desc["size_in_GB"] < 1 and model_desc["model"] in CANONICAL_SCORE_VALUES.keys()
+    ],
 )
 def test_batch_rerank(model_name):
     is_ci = os.getenv("CI")