mirror of
https://github.com/qdrant/fastembed.git
synced 2026-07-23 11:20:51 -05:00
Tsk 374 add jina colbert v2 (#378)
* feat: Added support for jina-colbert-v2 * chore: Generalized query marker and document marker * nit: remove github action on dispatch * chore: updated license * fix: Fix attention mask to be all 1 in xlmrobertatokenizer * feat: Added class for JinaColbertV2 * feat: Added jina colbert * chore: Change tolerance of the test * chore: Changed encoding of attention mask to 1 to be only in queries * chore: Changed the replacable token to be ' @' as its considered as one token * chore: Removed redundant functions * chore: Updated supported models docs * nit: Remove print statement * nit: visual stuff * fix: Fix dimention of jina colbert in description * fix: canonical query and document values for jina colbert
This commit is contained in:
Hossam Hagag
@@ -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": 5,
|
||||
"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,175 +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>snowflake/snowflake-arctic-embed-xs</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, 512...</td>\n",
|
||||
" <td>apache-2.0</td>\n",
|
||||
" <td>0.090</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>3</th>\n",
|
||||
" <td>sentence-transformers/all-MiniLM-L6-v2</td>\n",
|
||||
" <td>384</td>\n",
|
||||
" <td>Sentence Transformer model, MiniLM-L6-v2</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>BAAI/bge-small-en</td>\n",
|
||||
" <td>384</td>\n",
|
||||
" <td>Fast English model</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>snowflake/snowflake-arctic-embed-s</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>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&image), Engl...</td>\n",
|
||||
" <td>mit</td>\n",
|
||||
" <td>0.250</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>11</th>\n",
|
||||
" <td>jinaai/jina-embeddings-v2-base-de</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>German embedding model supporting 8192 sequenc...</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>BAAI/bge-base-en</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>Base English model</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>snowflake/snowflake-arctic-embed-m</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>Based on intfloat/e5-base-unsupervised model, ...</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>nomic-ai/nomic-embed-text-v1.5</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>8192 context length english model</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>jinaai/jina-embeddings-v2-base-en</td>\n",
|
||||
" <td>768</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.520</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>16</th>\n",
|
||||
" <td>nomic-ai/nomic-embed-text-v1</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>8192 context length english model</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>snowflake/snowflake-arctic-embed-m-long</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>Based on nomic-ai/nomic-embed-text-v1-unsuperv...</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>mixedbread-ai/mxbai-embed-large-v1</td>\n",
|
||||
" <td>1024</td>\n",
|
||||
" <td>MixedBread Base sentence embedding model, does...</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>jinaai/jina-embeddings-v2-base-code</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>Source code embedding model supporting 8192 se...</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>sentence-transformers/paraphrase-multilingual-...</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>Sentence-transformers model for tasks like clu...</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>snowflake/snowflake-arctic-embed-l</td>\n",
|
||||
" <td>1024</td>\n",
|
||||
" <td>Based on intfloat/e5-large-unsupervised, large...</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>Large general text embeddings model</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>Large English model, v1.5</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",
|
||||
@@ -297,35 +323,35 @@
|
||||
"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 Based on all-MiniLM-L6-v2 model with only 22m ... 0.090 \n",
|
||||
"3 Sentence Transformer model, MiniLM-L6-v2 0.090 \n",
|
||||
"4 English embedding model supporting 8192 sequen... 0.120 \n",
|
||||
"5 Fast English model 0.130 \n",
|
||||
"6 Based on infloat/e5-small-unsupervised, does n... 0.130 \n",
|
||||
"7 Quantized 8192 context length english model 0.130 \n",
|
||||
"8 Base English model, v1.5 0.210 \n",
|
||||
"9 Sentence Transformer model, paraphrase-multili... 0.220 \n",
|
||||
"10 CLIP text encoder 0.250 \n",
|
||||
"11 German embedding model supporting 8192 sequenc... 0.320 \n",
|
||||
"12 Base English model 0.420 \n",
|
||||
"13 Based on intfloat/e5-base-unsupervised model, ... 0.430 \n",
|
||||
"14 8192 context length english model 0.520 \n",
|
||||
"15 English embedding model supporting 8192 sequen... 0.520 \n",
|
||||
"16 8192 context length english model 0.520 \n",
|
||||
"17 Based on nomic-ai/nomic-embed-text-v1-unsuperv... 0.540 \n",
|
||||
"18 MixedBread Base sentence embedding model, does... 0.640 \n",
|
||||
"19 Source code embedding model supporting 8192 se... 0.640 \n",
|
||||
"20 Sentence-transformers model for tasks like clu... 1.000 \n",
|
||||
"21 Based on intfloat/e5-large-unsupervised, large... 1.020 \n",
|
||||
"22 Large general text embeddings model 1.200 \n",
|
||||
"23 Large English model, v1.5 1.200 \n",
|
||||
"24 Multilingual model, e5-large. Recommend using ... 2.240 "
|
||||
" 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": 5,
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
@@ -381,6 +407,7 @@
|
||||
" <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",
|
||||
@@ -391,6 +418,7 @@
|
||||
" <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",
|
||||
@@ -399,22 +427,25 @@
|
||||
" <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>prithvida/Splade_PP_en_v1</td>\n",
|
||||
" <td>prithivida/Splade_PP_en_v1</td>\n",
|
||||
" <td>30522.0</td>\n",
|
||||
" <td>Misspelled version of the model. Retained for ...</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>prithivida/Splade_PP_en_v1</td>\n",
|
||||
" <td>prithvida/Splade_PP_en_v1</td>\n",
|
||||
" <td>30522.0</td>\n",
|
||||
" <td>Independent Implementation of SPLADE++ Model f...</td>\n",
|
||||
" <td>apache-2.0</td>\n",
|
||||
" <td>0.532</td>\n",
|
||||
" <td>NaN</td>\n",
|
||||
" </tr>\n",
|
||||
@@ -426,14 +457,20 @@
|
||||
" model vocab_size \\\n",
|
||||
"0 Qdrant/bm25 NaN \n",
|
||||
"1 Qdrant/bm42-all-minilm-l6-v2-attentions 30522.0 \n",
|
||||
"2 prithvida/Splade_PP_en_v1 30522.0 \n",
|
||||
"3 prithivida/Splade_PP_en_v1 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 requires_idf \n",
|
||||
"0 BM25 as sparse embeddings meant to be used wit... 0.010 True \n",
|
||||
"1 Light sparse embedding model, which assigns an... 0.090 True \n",
|
||||
"2 Misspelled version of the model. Retained for ... 0.532 NaN \n",
|
||||
"3 Independent Implementation of SPLADE++ Model f... 0.532 NaN "
|
||||
" description license size_in_GB \\\n",
|
||||
"0 BM25 as sparse embeddings meant to be used wit... apache-2.0 0.010 \n",
|
||||
"1 Light sparse embedding model, which assigns an... apache-2.0 0.090 \n",
|
||||
"2 Independent Implementation of SPLADE++ Model f... apache-2.0 0.532 \n",
|
||||
"3 Independent Implementation of SPLADE++ Model f... apache-2.0 0.532 \n",
|
||||
"\n",
|
||||
" requires_idf \n",
|
||||
"0 True \n",
|
||||
"1 True \n",
|
||||
"2 NaN \n",
|
||||
"3 NaN "
|
||||
]
|
||||
},
|
||||
"execution_count": 4,
|
||||
@@ -443,7 +480,7 @@
|
||||
],
|
||||
"source": [
|
||||
"(\n",
|
||||
" pd.DataFrame(SparseTextEmbedding.list_supported_models())\n",
|
||||
" pd.DataFrame(SparseTextEmbedding.list_supported_models())\n",
|
||||
" .sort_values(\"size_in_GB\")\n",
|
||||
" .drop(columns=[\"sources\", \"model_file\", \"additional_files\"])\n",
|
||||
" .reset_index(drop=True)\n",
|
||||
@@ -494,24 +531,58 @@
|
||||
" <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": 5,
|
||||
@@ -521,7 +592,7 @@
|
||||
],
|
||||
"source": [
|
||||
"(\n",
|
||||
" pd.DataFrame(LateInteractionTextEmbedding.list_supported_models())\n",
|
||||
" pd.DataFrame(LateInteractionTextEmbedding.list_supported_models())\n",
|
||||
" .sort_values(\"size_in_GB\")\n",
|
||||
" .drop(columns=[\"sources\", \"model_file\"])\n",
|
||||
" .reset_index(drop=True)\n",
|
||||
@@ -572,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",
|
||||
@@ -580,28 +652,32 @@
|
||||
" <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&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>Unicom Unicom-ViT-B-32 from open-metric-learning</td>\n",
|
||||
" <td>Image embeddings, Multimodal (text&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>Unicom Unicom-ViT-B-16 from open-metric-learning</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",
|
||||
@@ -615,11 +691,11 @@
|
||||
"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 \n",
|
||||
"2 Unicom Unicom-ViT-B-32 from open-metric-learning 0.48 \n",
|
||||
"3 Unicom Unicom-ViT-B-16 from open-metric-learning 0.82 "
|
||||
" description license size_in_GB \n",
|
||||
"0 Image embeddings, Unimodal (image), 2016 year apache-2.0 0.10 \n",
|
||||
"1 Image embeddings, Multimodal (text&image), 202... mit 0.34 \n",
|
||||
"2 Image embeddings, Multimodal (text&image), 202... apache-2.0 0.48 \n",
|
||||
"3 Image embeddings (more detailed than Unicom-Vi... apache-2.0 0.82 "
|
||||
]
|
||||
},
|
||||
"execution_count": 6,
|
||||
@@ -629,7 +705,8 @@
|
||||
],
|
||||
"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",
|
||||
")"
|
||||
@@ -638,7 +715,7 @@
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3.8.18 ('base')",
|
||||
"display_name": ".venv",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
@@ -652,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
|
||||
|
||||
@@ -83,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
|
||||
@@ -96,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:
|
||||
@@ -104,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
|
||||
|
||||
@@ -189,7 +191,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
|
||||
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]
|
||||
|
||||
62
fastembed/late_interaction/jina_colbert.py
Normal file
62
fastembed/late_interaction/jina_colbert.py
Normal file
@@ -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,
|
||||
)
|
||||
@@ -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]]:
|
||||
|
||||
@@ -28,6 +28,15 @@ CANONICAL_COLUMN_VALUES = {
|
||||
[-0.07281, 0.04633, -0.04711, 0.00762, -0.07374],
|
||||
]
|
||||
),
|
||||
"jinaai/jina-colbert-v2": np.array(
|
||||
[
|
||||
[0.0742, 0.0591, -0.2403, -0.1774, 0.02],
|
||||
[0.1318, 0.0882, -0.1138, -0.2066, 0.146],
|
||||
[-0.0183, -0.1354, -0.0139, -0.1079, -0.051],
|
||||
[0.0003, -0.1184, -0.07, -0.0479, -0.0649],
|
||||
[0.0766, 0.0452, -0.2343, -0.183, 0.0058],
|
||||
]
|
||||
),
|
||||
}
|
||||
|
||||
CANONICAL_QUERY_VALUES = {
|
||||
@@ -103,6 +112,42 @@ CANONICAL_QUERY_VALUES = {
|
||||
[-0.03473, 0.04792, -0.07033, 0.02196, -0.05314],
|
||||
]
|
||||
),
|
||||
"jinaai/jina-colbert-v2": np.array(
|
||||
[
|
||||
[0.0477, 0.0255, -0.2224, -0.1085, -0.03],
|
||||
[0.0206, -0.0845, -0.0075, -0.1712, 0.0156],
|
||||
[-0.0056, -0.0957, -0.0147, -0.1277, -0.0225],
|
||||
[0.0486, -0.0499, -0.1609, 0.0194, 0.0274],
|
||||
[0.0481, 0.0253, -0.2278, -0.1126, -0.0294],
|
||||
[0.0599, -0.0678, -0.0956, -0.0757, 0.0236],
|
||||
[0.0592, -0.0862, -0.0621, -0.1084, 0.0155],
|
||||
[0.0874, -0.0714, -0.0772, -0.1414, 0.037],
|
||||
[0.1009, -0.0552, -0.0669, -0.163, 0.0493],
|
||||
[0.1135, -0.047, -0.0576, -0.1699, 0.0538],
|
||||
[0.1228, -0.0428, -0.0507, -0.1725, 0.0562],
|
||||
[0.1291, -0.0388, -0.042, -0.1753, 0.0569],
|
||||
[0.1365, -0.0337, -0.0326, -0.1786, 0.0574],
|
||||
[0.1439, -0.026, -0.024, -0.1831, 0.0574],
|
||||
[0.1527, -0.0099, -0.0179, -0.1874, 0.057],
|
||||
[0.1555, 0.0186, -0.023, -0.1801, 0.0539],
|
||||
[0.1389, 0.054, -0.0345, -0.1636, 0.0429],
|
||||
[0.1058, 0.0862, -0.0418, -0.1455, 0.0222],
|
||||
[0.0713, 0.1061, -0.0438, -0.1288, 0.0002],
|
||||
[0.0453, 0.1143, -0.0457, -0.1119, -0.019],
|
||||
[0.0346, 0.1131, -0.0487, -0.0952, -0.0338],
|
||||
[0.0355, 0.1073, -0.0493, -0.0823, -0.0438],
|
||||
[0.0424, 0.1041, -0.0459, -0.0761, -0.048],
|
||||
[0.048, 0.102, -0.0421, -0.0718, -0.0477],
|
||||
[0.0474, 0.0989, -0.0413, -0.0654, -0.0431],
|
||||
[0.0434, 0.095, -0.0415, -0.0589, -0.0345],
|
||||
[0.0408, 0.0897, -0.0405, -0.0554, -0.0197],
|
||||
[0.0433, 0.0811, -0.0407, -0.0545, 0.0055],
|
||||
[0.0514, 0.0629, -0.0446, -0.0549, 0.0368],
|
||||
[0.058, 0.048, -0.0527, -0.0607, 0.0568],
|
||||
[0.0561, 0.0447, -0.0661, -0.0702, 0.0764],
|
||||
[0.0204, -0.0856, -0.0386, -0.1232, -0.0332],
|
||||
]
|
||||
),
|
||||
}
|
||||
|
||||
docs = ["Hello World"]
|
||||
@@ -119,7 +164,7 @@ 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:
|
||||
shutil.rmtree(model.model._model_dir)
|
||||
@@ -134,7 +179,7 @@ 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:
|
||||
shutil.rmtree(model.model._model_dir)
|
||||
@@ -149,7 +194,7 @@ 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:
|
||||
shutil.rmtree(model.model._model_dir)
|
||||
|
||||
Reference in New Issue
Block a user