From 1deb83032876a133cd4dcf1ad5098d01cfb98e3a Mon Sep 17 00:00:00 2001 From: Hossam Hagag <90828745+hh-space-invader@users.noreply.github.com> Date: Sun, 3 Nov 2024 11:02:41 +0200 Subject: [PATCH] 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 --- docs/examples/Supported_Models.ipynb | 268 +++++++++++------- fastembed/late_interaction/colbert.py | 16 +- fastembed/late_interaction/jina_colbert.py | 62 ++++ .../late_interaction_text_embedding.py | 5 +- tests/test_late_interaction_embeddings.py | 51 +++- 5 files changed, 291 insertions(+), 111 deletions(-) create mode 100644 fastembed/late_interaction/jina_colbert.py diff --git a/docs/examples/Supported_Models.ipynb b/docs/examples/Supported_Models.ipynb index dc40582..2cdb8cc 100644 --- a/docs/examples/Supported_Models.ipynb +++ b/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": 5, + "execution_count": 3, "metadata": { "ExecuteTime": { "end_time": "2024-05-31T18:13:25.863008Z", @@ -86,6 +86,7 @@ " model\n", " dim\n", " description\n", + " license\n", " size_in_GB\n", " \n", " \n", @@ -94,175 +95,200 @@ " 0\n", " BAAI/bge-small-en-v1.5\n", " 384\n", - " Fast and Default English model\n", + " Text embeddings, Unimodal (text), English, 512...\n", + " mit\n", " 0.067\n", " \n", " \n", " 1\n", " BAAI/bge-small-zh-v1.5\n", " 512\n", - " Fast and recommended Chinese model\n", + " Text embeddings, Unimodal (text), Chinese, 512...\n", + " mit\n", " 0.090\n", " \n", " \n", " 2\n", " snowflake/snowflake-arctic-embed-xs\n", " 384\n", - " Based on all-MiniLM-L6-v2 model with only 22m ...\n", + " Text embeddings, Unimodal (text), English, 512...\n", + " apache-2.0\n", " 0.090\n", " \n", " \n", " 3\n", " sentence-transformers/all-MiniLM-L6-v2\n", " 384\n", - " Sentence Transformer model, MiniLM-L6-v2\n", + " Text embeddings, Unimodal (text), English, 256...\n", + " apache-2.0\n", " 0.090\n", " \n", " \n", " 4\n", " jinaai/jina-embeddings-v2-small-en\n", " 512\n", - " English embedding model supporting 8192 sequen...\n", + " Text embeddings, Unimodal (text), English, 819...\n", + " apache-2.0\n", " 0.120\n", " \n", " \n", " 5\n", " BAAI/bge-small-en\n", " 384\n", - " Fast English model\n", + " Text embeddings, Unimodal (text), English, 512...\n", + " mit\n", " 0.130\n", " \n", " \n", " 6\n", " snowflake/snowflake-arctic-embed-s\n", " 384\n", - " Based on infloat/e5-small-unsupervised, does n...\n", + " Text embeddings, Unimodal (text), English, 512...\n", + " apache-2.0\n", " 0.130\n", " \n", " \n", " 7\n", " nomic-ai/nomic-embed-text-v1.5-Q\n", " 768\n", - " Quantized 8192 context length english model\n", + " Text embeddings, Multimodal (text, image), Eng...\n", + " apache-2.0\n", " 0.130\n", " \n", " \n", " 8\n", " BAAI/bge-base-en-v1.5\n", " 768\n", - " Base English model, v1.5\n", + " Text embeddings, Unimodal (text), English, 512...\n", + " mit\n", " 0.210\n", " \n", " \n", " 9\n", " sentence-transformers/paraphrase-multilingual-...\n", " 384\n", - " Sentence Transformer model, paraphrase-multili...\n", + " Text embeddings, Unimodal (text), Multilingual...\n", + " apache-2.0\n", " 0.220\n", " \n", " \n", " 10\n", " Qdrant/clip-ViT-B-32-text\n", " 512\n", - " CLIP text encoder\n", + " Text embeddings, Multimodal (text&image), Engl...\n", + " mit\n", " 0.250\n", " \n", " \n", " 11\n", " jinaai/jina-embeddings-v2-base-de\n", " 768\n", - " German embedding model supporting 8192 sequenc...\n", + " Text embeddings, Unimodal (text), Multilingual...\n", + " apache-2.0\n", " 0.320\n", " \n", " \n", " 12\n", " BAAI/bge-base-en\n", " 768\n", - " Base English model\n", + " Text embeddings, Unimodal (text), English, 512...\n", + " mit\n", " 0.420\n", " \n", " \n", " 13\n", " snowflake/snowflake-arctic-embed-m\n", " 768\n", - " Based on intfloat/e5-base-unsupervised model, ...\n", + " Text embeddings, Unimodal (text), English, 512...\n", + " apache-2.0\n", " 0.430\n", " \n", " \n", " 14\n", " nomic-ai/nomic-embed-text-v1.5\n", " 768\n", - " 8192 context length english model\n", + " Text embeddings, Multimodal (text, image), Eng...\n", + " apache-2.0\n", " 0.520\n", " \n", " \n", " 15\n", " jinaai/jina-embeddings-v2-base-en\n", " 768\n", - " English embedding model supporting 8192 sequen...\n", + " Text embeddings, Unimodal (text), English, 819...\n", + " apache-2.0\n", " 0.520\n", " \n", " \n", " 16\n", " nomic-ai/nomic-embed-text-v1\n", " 768\n", - " 8192 context length english model\n", + " Text embeddings, Multimodal (text, image), Eng...\n", + " apache-2.0\n", " 0.520\n", " \n", " \n", " 17\n", " snowflake/snowflake-arctic-embed-m-long\n", " 768\n", - " Based on nomic-ai/nomic-embed-text-v1-unsuperv...\n", + " Text embeddings, Unimodal (text), English, 204...\n", + " apache-2.0\n", " 0.540\n", " \n", " \n", " 18\n", " mixedbread-ai/mxbai-embed-large-v1\n", " 1024\n", - " MixedBread Base sentence embedding model, does...\n", + " Text embeddings, Unimodal (text), English, 512...\n", + " apache-2.0\n", " 0.640\n", " \n", " \n", " 19\n", " jinaai/jina-embeddings-v2-base-code\n", " 768\n", - " Source code embedding model supporting 8192 se...\n", + " Text embeddings, Unimodal (text), Multilingual...\n", + " apache-2.0\n", " 0.640\n", " \n", " \n", " 20\n", " sentence-transformers/paraphrase-multilingual-...\n", " 768\n", - " Sentence-transformers model for tasks like clu...\n", + " Text embeddings, Unimodal (text), Multilingual...\n", + " apache-2.0\n", " 1.000\n", " \n", " \n", " 21\n", " snowflake/snowflake-arctic-embed-l\n", " 1024\n", - " Based on intfloat/e5-large-unsupervised, large...\n", + " Text embeddings, Unimodal (text), English, 512...\n", + " apache-2.0\n", " 1.020\n", " \n", " \n", " 22\n", " thenlper/gte-large\n", " 1024\n", - " Large general text embeddings model\n", + " Text embeddings, Unimodal (text), English, 512...\n", + " mit\n", " 1.200\n", " \n", " \n", " 23\n", " BAAI/bge-large-en-v1.5\n", " 1024\n", - " Large English model, v1.5\n", + " Text embeddings, Unimodal (text), English, 512...\n", + " mit\n", " 1.200\n", " \n", " \n", " 24\n", " intfloat/multilingual-e5-large\n", " 1024\n", - " Multilingual model, e5-large. Recommend using ...\n", + " Text embeddings, Unimodal (text), Multilingual...\n", + " mit\n", " 2.240\n", " \n", " \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 @@ " model\n", " vocab_size\n", " description\n", + " license\n", " size_in_GB\n", " requires_idf\n", " \n", @@ -391,6 +418,7 @@ " Qdrant/bm25\n", " NaN\n", " BM25 as sparse embeddings meant to be used wit...\n", + " apache-2.0\n", " 0.010\n", " True\n", " \n", @@ -399,22 +427,25 @@ " Qdrant/bm42-all-minilm-l6-v2-attentions\n", " 30522.0\n", " Light sparse embedding model, which assigns an...\n", + " apache-2.0\n", " 0.090\n", " True\n", " \n", " \n", " 2\n", - " prithvida/Splade_PP_en_v1\n", + " prithivida/Splade_PP_en_v1\n", " 30522.0\n", - " Misspelled version of the model. Retained for ...\n", + " Independent Implementation of SPLADE++ Model f...\n", + " apache-2.0\n", " 0.532\n", " NaN\n", " \n", " \n", " 3\n", - " prithivida/Splade_PP_en_v1\n", + " prithvida/Splade_PP_en_v1\n", " 30522.0\n", " Independent Implementation of SPLADE++ Model f...\n", + " apache-2.0\n", " 0.532\n", " NaN\n", " \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 @@ " model\n", " dim\n", " description\n", + " license\n", " size_in_GB\n", + " additional_files\n", " \n", " \n", " \n", " \n", " 0\n", + " answerdotai/answerai-colbert-small-v1\n", + " 96\n", + " Text embeddings, Unimodal (text), Multilingual...\n", + " apache-2.0\n", + " 0.13\n", + " NaN\n", + " \n", + " \n", + " 1\n", " colbert-ir/colbertv2.0\n", " 128\n", " Late interaction model\n", + " mit\n", " 0.44\n", + " NaN\n", + " \n", + " \n", + " 2\n", + " jinaai/jina-colbert-v2\n", + " 1024\n", + " New model that expands capabilities of colbert...\n", + " cc-by-nc-4.0\n", + " 2.24\n", + " [onnx/model.onnx_data]\n", " \n", " \n", "\n", "" ], "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 @@ " model\n", " dim\n", " description\n", + " license\n", " size_in_GB\n", " \n", " \n", @@ -580,28 +652,32 @@ " 0\n", " Qdrant/resnet50-onnx\n", " 2048\n", - " ResNet-50 from `Deep Residual Learning for Ima...\n", + " Image embeddings, Unimodal (image), 2016 year\n", + " apache-2.0\n", " 0.10\n", " \n", " \n", " 1\n", " Qdrant/clip-ViT-B-32-vision\n", " 512\n", - " CLIP vision encoder based on ViT-B/32\n", + " Image embeddings, Multimodal (text&image), 202...\n", + " mit\n", " 0.34\n", " \n", " \n", " 2\n", " Qdrant/Unicom-ViT-B-32\n", " 512\n", - " Unicom Unicom-ViT-B-32 from open-metric-learning\n", + " Image embeddings, Multimodal (text&image), 202...\n", + " apache-2.0\n", " 0.48\n", " \n", " \n", " 3\n", " Qdrant/Unicom-ViT-B-16\n", " 768\n", - " Unicom Unicom-ViT-B-16 from open-metric-learning\n", + " Image embeddings (more detailed than Unicom-Vi...\n", + " apache-2.0\n", " 0.82\n", " \n", " \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 diff --git a/fastembed/late_interaction/colbert.py b/fastembed/late_interaction/colbert.py index 6e7f34d..686595d 100644 --- a/fastembed/late_interaction/colbert.py +++ b/fastembed/late_interaction/colbert.py @@ -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] diff --git a/fastembed/late_interaction/jina_colbert.py b/fastembed/late_interaction/jina_colbert.py new file mode 100644 index 0000000..42232ae --- /dev/null +++ b/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 = "" + + @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, + ) diff --git a/fastembed/late_interaction/late_interaction_text_embedding.py b/fastembed/late_interaction/late_interaction_text_embedding.py index 459783e..c2f05ba 100644 --- a/fastembed/late_interaction/late_interaction_text_embedding.py +++ b/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]]: diff --git a/tests/test_late_interaction_embeddings.py b/tests/test_late_interaction_embeddings.py index c8d62ca..191bc22 100644 --- a/tests/test_late_interaction_embeddings.py +++ b/tests/test_late_interaction_embeddings.py @@ -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)