chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,415 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Intro to Embedding"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"For text retrieval, pattern matching is the most intuitive way. People would use certain characters, words, phrases, or sentence patterns. However, not only for human, it is also extremely inefficient for computer to do pattern matching between a query and a collection of text files to find the possible results. \n",
|
||||
"\n",
|
||||
"For images and acoustic waves, there are rgb pixels and digital signals. Similarly, in order to accomplish more sophisticated tasks of natural language such as retrieval, classification, clustering, or semantic search, we need a way to represent text data. That's how text embedding comes in front of the stage."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 1. Background"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Traditional text embedding methods like one-hot encoding and bag-of-words (BoW) represent words and sentences as sparse vectors based on their statistical features, such as word appearance and frequency within a document. More advanced methods like TF-IDF and BM25 improve on these by considering a word's importance across an entire corpus, while n-gram techniques capture word order in small groups. However, these approaches suffer from the \"curse of dimensionality\" and fail to capture semantic similarity like \"cat\" and \"kitty\", difference like \"play the watch\" and \"watch the play\"."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# example of bag-of-words\n",
|
||||
"sentence1 = \"I love basketball\"\n",
|
||||
"sentence2 = \"I have a basketball match\"\n",
|
||||
"\n",
|
||||
"words = ['I', 'love', 'basketball', 'have', 'a', 'match']\n",
|
||||
"sen1_vec = [1, 1, 1, 0, 0, 0]\n",
|
||||
"sen2_vec = [1, 0, 1, 1, 1, 1]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"To overcome these limitations, dense word embeddings were developed, mapping words to vectors in a low-dimensional space that captures semantic and relational information. Early models like Word2Vec demonstrated the power of dense embeddings using neural networks. Subsequent advancements with neural network architectures like RNNs, LSTMs, and Transformers have enabled more sophisticated models such as BERT, RoBERTa, and GPT to excel in capturing complex word relationships and contexts. **BAAI General Embedding (BGE)** provide a series of open-source models that could satisfy all kinds of demands."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Get Embedding"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The first step of modern text retrieval is embedding the text. So let's take a look at how to use the embedding models."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Install the packages:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 10,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"%pip install -U FlagEmbedding sentence_transformers openai cohere"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import os \n",
|
||||
"os.environ['TRANSFORMERS_NO_ADVISORY_WARNINGS'] = 'true'\n",
|
||||
"# single GPU is better for small tasks\n",
|
||||
"os.environ['CUDA_VISIBLE_DEVICES'] = '0'"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"We'll use the following three sentences as the inputs:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"sentences = [\n",
|
||||
" \"That is a happy dog\",\n",
|
||||
" \"That is a very happy person\",\n",
|
||||
" \"Today is a sunny day\",\n",
|
||||
"]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Open-source Models"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"A huge portion of embedding models are in the open source community. The advantages of open-source models include:\n",
|
||||
"- Free, no extra cost. But make sure to check the License and your use case before using.\n",
|
||||
"- No frequency limit, can accelerate a lot if you have enough GPUs to parallelize.\n",
|
||||
"- Transparent and might be reproducible.\n",
|
||||
"\n",
|
||||
"Let's take a look at two representatives:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"#### BGE"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"BGE is a series of embedding models and rerankers published by BAAI. Several of them reached SOTA at the time they released."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"initial target device: 100%|██████████| 8/8 [00:31<00:00, 3.89s/it]\n",
|
||||
"Chunks: 100%|██████████| 3/3 [00:04<00:00, 1.61s/it]\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Embeddings:\n",
|
||||
"(3, 768)\n",
|
||||
"Similarity scores:\n",
|
||||
"[[1. 0.79 0.575 ]\n",
|
||||
" [0.79 0.9995 0.592 ]\n",
|
||||
" [0.575 0.592 0.999 ]]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from FlagEmbedding import FlagModel\n",
|
||||
"\n",
|
||||
"# Load BGE model\n",
|
||||
"model = FlagModel('BAAI/bge-base-en-v1.5')\n",
|
||||
"\n",
|
||||
"# encode the queries and corpus\n",
|
||||
"embeddings = model.encode(sentences)\n",
|
||||
"print(f\"Embeddings:\\n{embeddings.shape}\")\n",
|
||||
"\n",
|
||||
"scores = embeddings @ embeddings.T\n",
|
||||
"print(f\"Similarity scores:\\n{scores}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"#### Sentence Transformers"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Sentence Transformers is a library for sentence embeddings with a huge amount of embedding models and datasets for related tasks."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Embeddings:\n",
|
||||
"(3, 384)\n",
|
||||
"Similarity scores:\n",
|
||||
"[[0.99999976 0.6210502 0.24906276]\n",
|
||||
" [0.6210502 0.9999997 0.21061528]\n",
|
||||
" [0.24906276 0.21061528 0.9999999 ]]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from sentence_transformers import SentenceTransformer\n",
|
||||
"\n",
|
||||
"model = SentenceTransformer(\"all-MiniLM-L6-v2\")\n",
|
||||
"\n",
|
||||
"embeddings = model.encode(sentences, normalize_embeddings=True)\n",
|
||||
"print(f\"Embeddings:\\n{embeddings.shape}\")\n",
|
||||
"\n",
|
||||
"scores = embeddings @ embeddings.T\n",
|
||||
"print(f\"Similarity scores:\\n{scores}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Commercial Models"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"There are also plenty choices of commercial models. They have the advantages of:\n",
|
||||
"- Efficient memory usage, fast inference with no need of GPUs.\n",
|
||||
"- Systematic support, commercial models have closer connections with their other products.\n",
|
||||
"- Better training data, commercial models might be trained on larger, higher-quality datasets than some open-source models."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"#### OpenAI"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Along with GPT series, OpenAI has their own embedding models. Make sure to fill in your own API key in the field `\"YOUR_API_KEY\"`"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 19,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"import numpy as np\n",
|
||||
"\n",
|
||||
"os.environ[\"OPENAI_API_KEY\"] = \"YOUR_API_KEY\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Then run the following cells to get the embeddings. Check their official [documentation](https://platform.openai.com/docs/guides/embeddings) for more details."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 20,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from openai import OpenAI\n",
|
||||
"\n",
|
||||
"client = OpenAI()\n",
|
||||
"\n",
|
||||
"response = client.embeddings.create(input = sentences, model=\"text-embedding-3-small\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 21,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Embeddings:\n",
|
||||
"(3, 1536)\n",
|
||||
"Similarity scores:\n",
|
||||
"[[1.00000004 0.697673 0.34739798]\n",
|
||||
" [0.697673 1.00000005 0.31969923]\n",
|
||||
" [0.34739798 0.31969923 0.99999998]]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embeddings = np.asarray([response.data[i].embedding for i in range(len(sentences))])\n",
|
||||
"print(f\"Embeddings:\\n{embeddings.shape}\")\n",
|
||||
"\n",
|
||||
"scores = embeddings @ embeddings.T\n",
|
||||
"print(f\"Similarity scores:\\n{scores}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"#### Voyage AI"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Voyage AI provides embedding models and rerankers for different purpus and in various fields. Their API keys can be freely used in low frequency and token length."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 22,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"os.environ[\"VOYAGE_API_KEY\"] = \"YOUR_API_KEY\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Check their official [documentation](https://docs.voyageai.com/docs/api-key-and-installation) for more details."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 23,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import voyageai\n",
|
||||
"\n",
|
||||
"vo = voyageai.Client()\n",
|
||||
"\n",
|
||||
"result = vo.embed(sentences, model=\"voyage-large-2-instruct\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 24,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Embeddings:\n",
|
||||
"(3, 1024)\n",
|
||||
"Similarity scores:\n",
|
||||
"[[0.99999997 0.87282517 0.63276503]\n",
|
||||
" [0.87282517 0.99999998 0.64720015]\n",
|
||||
" [0.63276503 0.64720015 0.99999999]]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embeddings = np.asarray(result.embeddings)\n",
|
||||
"print(f\"Embeddings:\\n{embeddings.shape}\")\n",
|
||||
"\n",
|
||||
"scores = embeddings @ embeddings.T\n",
|
||||
"print(f\"Similarity scores:\\n{scores}\")"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "dev",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.12.7"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
@@ -0,0 +1,606 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "06cff9e4",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# BGE Series"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "880e229d",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"In this Part, we will walk through the BGE series and introduce how to use the BGE embedding models."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "2516fd49",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 1. BAAI General Embedding"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "2113ee71",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"BGE stands for BAAI General Embedding, it's a series of embeddings models developed and published by Beijing Academy of Artificial Intelligence (BAAI)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "16515b99",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"A full support of APIs and related usages of BGE is maintained in [FlagEmbedding](https://github.com/FlagOpen/FlagEmbedding) on GitHub.\n",
|
||||
"\n",
|
||||
"Run the following cell to install FlagEmbedding in your environment."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "88095fd0",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"%pip install -U FlagEmbedding"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"id": "a2376217",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import os \n",
|
||||
"os.environ['TRANSFORMERS_NO_ADVISORY_WARNINGS'] = 'true'\n",
|
||||
"# single GPU is better for small tasks\n",
|
||||
"os.environ['CUDA_VISIBLE_DEVICES'] = '0'"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "bc6e30a0",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The collection of BGE models can be found in [Huggingface collection](https://huggingface.co/collections/BAAI/bge-66797a74476eb1f085c7446d)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "67a16ccf",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 2. BGE Series Models"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "2e10034a",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### 2.1 BGE"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "0cdc6702",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The very first version of BGE has 6 models, with 'large', 'base', and 'small' for English and Chinese. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "04b75f72",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"| Model | Language | Parameters | Model Size | Description | Base Model |\n",
|
||||
"|:-------|:--------:|:--------------:|:--------------:|:-----------------:|:----------------:|\n",
|
||||
"| [BAAI/bge-large-en](https://huggingface.co/BAAI/bge-large-en) | English | 500M | 1.34 GB | Embedding Model which map text into vector | BERT |\n",
|
||||
"| [BAAI/bge-base-en](https://huggingface.co/BAAI/bge-base-en) | English | 109M | 438 MB | a base-scale model but with similar ability to `bge-large-en` | BERT |\n",
|
||||
"| [BAAI/bge-small-en](https://huggingface.co/BAAI/bge-small-en) | English | 33.4M | 133 MB | a small-scale model but with competitive performance | BERT |\n",
|
||||
"| [BAAI/bge-large-zh](https://huggingface.co/BAAI/bge-large-zh) | Chinese | 326M | 1.3 GB | Embedding Model which map text into vector | BERT |\n",
|
||||
"| [BAAI/bge-base-zh](https://huggingface.co/BAAI/bge-base-zh) | Chinese | 102M | 409 MB | a base-scale model but with similar ability to `bge-large-zh` | BERT |\n",
|
||||
"| [BAAI/bge-small-zh](https://huggingface.co/BAAI/bge-small-zh) | Chinese | 24M | 95.8 MB | a small-scale model but with competitive performance | BERT |"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "c9c45d17",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"For inference, simply import FlagModel from FlagEmbedding and initialize the model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "89e07751",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"[[0.84864 0.7946737 ]\n",
|
||||
" [0.760097 0.85449743]]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from FlagEmbedding import FlagModel\n",
|
||||
"\n",
|
||||
"# Load BGE model\n",
|
||||
"model = FlagModel(\n",
|
||||
" 'BAAI/bge-base-en',\n",
|
||||
" query_instruction_for_retrieval=\"Represent this sentence for searching relevant passages:\",\n",
|
||||
" query_instruction_format='{}{}',\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"queries = [\"query 1\", \"query 2\"]\n",
|
||||
"corpus = [\"passage 1\", \"passage 2\"]\n",
|
||||
"\n",
|
||||
"# encode the queries and corpus\n",
|
||||
"q_embeddings = model.encode_queries(queries)\n",
|
||||
"p_embeddings = model.encode_corpus(corpus)\n",
|
||||
"\n",
|
||||
"# compute the similarity scores\n",
|
||||
"scores = q_embeddings @ p_embeddings.T\n",
|
||||
"print(scores)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "6c8e69ed",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"For general encoding, use either `encode()`:\n",
|
||||
"```python\n",
|
||||
"FlagModel.encode(sentences, batch_size=256, max_length=512, convert_to_numpy=True)\n",
|
||||
"```\n",
|
||||
"or `encode_corpus()` that directly calls `encode()`:\n",
|
||||
"```python\n",
|
||||
"FlagModel.encode_corpus(corpus, batch_size=256, max_length=512, convert_to_numpy=True)\n",
|
||||
"```\n",
|
||||
"The *encode_queries()* function concatenate the `query_instruction_for_retrieval` with each of the input query to form the new sentences and then feed them to `encode()`.\n",
|
||||
"```python\n",
|
||||
"FlagModel.encode_queries(queries, batch_size=256, max_length=512, convert_to_numpy=True)\n",
|
||||
"```"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "2c86a5a3",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### 2.2 BGE v1.5"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "454ff7aa",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"BGE 1.5 alleviate the issue of the similarity distribution, and enhance retrieval ability without instruction."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "30b1f897",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"| Model | Language | Parameters | Model Size | Description | Base Model |\n",
|
||||
"|:-------|:--------:|:--------------:|:--------------:|:-----------------:|:----------------:|\n",
|
||||
"| [BAAI/bge-large-en-v1.5](https://huggingface.co/BAAI/bge-large-en-v1.5) | English | 335M | 1.34 GB | version 1.5 with more reasonable similarity distribution | BERT |\n",
|
||||
"| [BAAI/bge-base-en-v1.5](https://huggingface.co/BAAI/bge-base-en-v1.5) | English | 109M | 438 MB | version 1.5 with more reasonable similarity distribution | BERT |\n",
|
||||
"| [BAAI/bge-small-en-v1.5](https://huggingface.co/BAAI/bge-small-en-v1.5) | English | 33.4M | 133 MB | version 1.5 with more reasonable similarity distribution | BERT |\n",
|
||||
"| [BAAI/bge-large-zh-v1.5](https://huggingface.co/BAAI/bge-large-zh-v1.5) | Chinese | 326M | 1.3 GB | version 1.5 with more reasonable similarity distribution | BERT |\n",
|
||||
"| [BAAI/bge-base-zh-v1.5](https://huggingface.co/BAAI/bge-base-zh-v1.5) | Chinese | 102M | 409 MB | version 1.5 with more reasonable similarity distribution | BERT |\n",
|
||||
"| [BAAI/bge-small-zh-v1.5](https://huggingface.co/BAAI/bge-small-zh-v1.5) | Chinese | 24M | 95.8 MB | version 1.5 with more reasonable similarity distribution | BERT |"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ed00c504",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"You can use BGE 1.5 models exactly same to BGE v1 models."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"id": "9b17afcc",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"pre tokenize: 100%|██████████| 1/1 [00:00<00:00, 2252.58it/s]\n",
|
||||
"pre tokenize: 100%|██████████| 1/1 [00:00<00:00, 3575.71it/s]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"[[0.76 0.6714]\n",
|
||||
" [0.6177 0.7603]]\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"model = FlagModel(\n",
|
||||
" 'BAAI/bge-base-en-v1.5',\n",
|
||||
" query_instruction_for_retrieval=\"Represent this sentence for searching relevant passages:\",\n",
|
||||
" query_instruction_format='{}{}'\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"queries = [\"query 1\", \"query 2\"]\n",
|
||||
"corpus = [\"passage 1\", \"passage 2\"]\n",
|
||||
"\n",
|
||||
"# encode the queries and corpus\n",
|
||||
"q_embeddings = model.encode_queries(queries)\n",
|
||||
"p_embeddings = model.encode_corpus(corpus)\n",
|
||||
"\n",
|
||||
"# compute the similarity scores\n",
|
||||
"scores = q_embeddings @ p_embeddings.T\n",
|
||||
"print(scores)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "dcf2a82b",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### 2.3 BGE M3"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "cc5b5a5e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"BGE-M3 is the new version of BGE models that is distinguished for its versatility in:\n",
|
||||
"- Multi-Functionality: Simultaneously perform the three common retrieval functionalities of embedding model: dense retrieval, multi-vector retrieval, and sparse retrieval.\n",
|
||||
"- Multi-Linguality: Supports more than 100 working languages.\n",
|
||||
"- Multi-Granularity: Can proces inputs with different granularityies, spanning from short sentences to long documents of up to 8192 tokens.\n",
|
||||
"\n",
|
||||
"For more details, feel free to check out the [paper](https://arxiv.org/pdf/2402.03216)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "41348e03",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"| Model | Language | Parameters | Model Size | Description | Base Model |\n",
|
||||
"|:-------|:--------:|:--------------:|:--------------:|:-----------------:|:----------------:|\n",
|
||||
"| [BAAI/bge-m3](https://huggingface.co/BAAI/bge-m3) | Multilingual | 568M | 2.27 GB | Multi-Functionality(dense retrieval, sparse retrieval, multi-vector(colbert)), Multi-Linguality, and Multi-Granularity(8192 tokens) | XLM-RoBERTa |"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"id": "d4647625",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Fetching 30 files: 100%|██████████| 30/30 [00:00<00:00, 194180.74it/s]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from FlagEmbedding import BGEM3FlagModel\n",
|
||||
"\n",
|
||||
"model = BGEM3FlagModel('BAAI/bge-m3', use_fp16=True)\n",
|
||||
"\n",
|
||||
"sentences = [\"What is BGE M3?\", \"Defination of BM25\"]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "1f89f1a9",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"```python\n",
|
||||
"BGEM3FlagModel.encode(\n",
|
||||
" sentences, \n",
|
||||
" batch_size=12, \n",
|
||||
" max_length=8192, \n",
|
||||
" return_dense=True, \n",
|
||||
" return_sparse=False, \n",
|
||||
" return_colbert_vecs=False\n",
|
||||
")\n",
|
||||
"```\n",
|
||||
"It returns a dictionary like:\n",
|
||||
"```python\n",
|
||||
"{\n",
|
||||
" 'dense_vecs': # array of dense embeddings of inputs if return_dense=True, otherwise None,\n",
|
||||
" 'lexical_weights': # array of dictionaries with keys and values are ids of tokens and their corresponding weights if return_sparse=True, otherwise None,\n",
|
||||
" 'colbert_vecs': # array of multi-vector embeddings of inputs if return_cobert_vecs=True, otherwise None,'\n",
|
||||
"}\n",
|
||||
"```"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"id": "f0b11cf0",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"pre tokenize: 100%|██████████| 1/1 [00:00<00:00, 1148.18it/s]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# If you don't need such a long length of 8192 input tokens, you can set max_length to a smaller value to speed up encoding.\n",
|
||||
"embeddings = model.encode(\n",
|
||||
" sentences, \n",
|
||||
" max_length=10,\n",
|
||||
" return_dense=True, \n",
|
||||
" return_sparse=True, \n",
|
||||
" return_colbert_vecs=True\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"id": "72cba126",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"dense embedding:\n",
|
||||
"[[-0.03412 -0.04706 -0.00087 ... 0.04822 0.007614 -0.02957 ]\n",
|
||||
" [-0.01035 -0.04483 -0.02434 ... -0.008224 0.01497 0.011055]]\n",
|
||||
"sparse embedding:\n",
|
||||
"[defaultdict(<class 'int'>, {'4865': np.float16(0.0836), '83': np.float16(0.0814), '335': np.float16(0.1296), '11679': np.float16(0.2517), '276': np.float16(0.1699), '363': np.float16(0.2695), '32': np.float16(0.04077)}), defaultdict(<class 'int'>, {'262': np.float16(0.05014), '5983': np.float16(0.1367), '2320': np.float16(0.04517), '111': np.float16(0.0634), '90017': np.float16(0.2517), '2588': np.float16(0.3333)})]\n",
|
||||
"multi-vector:\n",
|
||||
"[array([[-8.68966337e-03, -4.89266850e-02, -3.03634931e-03, ...,\n",
|
||||
" -2.21243706e-02, 5.72856329e-02, 1.28355855e-02],\n",
|
||||
" [-8.92937183e-03, -4.67235669e-02, -9.52814799e-03, ...,\n",
|
||||
" -3.14785317e-02, 5.39088845e-02, 6.96671568e-03],\n",
|
||||
" [ 1.84195358e-02, -4.22310382e-02, 8.55499704e-04, ...,\n",
|
||||
" -1.97946690e-02, 3.84313315e-02, 7.71250250e-03],\n",
|
||||
" ...,\n",
|
||||
" [-2.55824160e-02, -1.65533274e-02, -4.21357416e-02, ...,\n",
|
||||
" -4.50234264e-02, 4.41286489e-02, -1.00052059e-02],\n",
|
||||
" [ 5.90990965e-07, -5.53734899e-02, 8.51499755e-03, ...,\n",
|
||||
" -2.29209941e-02, 6.04418293e-02, 9.39912070e-03],\n",
|
||||
" [ 2.57394509e-03, -2.92690992e-02, -1.89342294e-02, ...,\n",
|
||||
" -8.04431178e-03, 3.28964666e-02, 4.38723788e-02]], dtype=float32), array([[ 0.01724418, 0.03835401, -0.02309308, ..., 0.00141706,\n",
|
||||
" 0.02995041, -0.05990082],\n",
|
||||
" [ 0.00996325, 0.03922409, -0.03849588, ..., 0.00591671,\n",
|
||||
" 0.02722516, -0.06510868],\n",
|
||||
" [ 0.01781915, 0.03925728, -0.01710397, ..., 0.00801776,\n",
|
||||
" 0.03987768, -0.05070014],\n",
|
||||
" ...,\n",
|
||||
" [ 0.05478653, 0.00755799, 0.00328444, ..., -0.01648209,\n",
|
||||
" 0.02405782, 0.00363262],\n",
|
||||
" [ 0.00936953, 0.05028074, -0.02388872, ..., 0.02567679,\n",
|
||||
" 0.00791224, -0.03257877],\n",
|
||||
" [ 0.01803976, 0.0133922 , 0.00019365, ..., 0.0184015 ,\n",
|
||||
" 0.01373822, 0.00315539]], dtype=float32)]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"print(f\"dense embedding:\\n{embeddings['dense_vecs']}\")\n",
|
||||
"print(f\"sparse embedding:\\n{embeddings['lexical_weights']}\")\n",
|
||||
"print(f\"multi-vector:\\n{embeddings['colbert_vecs']}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "14d83caa",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### 2.4 BGE Multilingual Gemma2"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "fd4c67df",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"BGE Multilingual Gemma2 is a LLM-based Multi-Lingual embedding model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "abdca22e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"| Model | Language | Parameters | Model Size | Description | Base Model |\n",
|
||||
"|:-------|:--------:|:--------------:|:--------------:|:-----------------:|:----------------:|\n",
|
||||
"| [BAAI/bge-multilingual-gemma2](https://huggingface.co/BAAI/bge-multilingual-gemma2) | Multilingual | 9.24B | 37 GB | LLM-based multilingual embedding model with SOTA results on multilingual benchmarks | Gemma2-9B |"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"id": "8ec545bc",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Loading checkpoint shards: 100%|██████████| 4/4 [00:00<00:00, 6.34it/s]\n",
|
||||
"pre tokenize: 100%|██████████| 1/1 [00:00<00:00, 816.49it/s]\n",
|
||||
"pre tokenize: 100%|██████████| 1/1 [00:00<00:00, 718.33it/s]\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"[[0.559 0.01685 ]\n",
|
||||
" [0.0008683 0.5015 ]]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from FlagEmbedding import FlagLLMModel\n",
|
||||
"\n",
|
||||
"queries = [\"how much protein should a female eat\", \"summit define\"]\n",
|
||||
"documents = [\n",
|
||||
" \"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.\",\n",
|
||||
" \"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments.\"\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"model = FlagLLMModel('BAAI/bge-multilingual-gemma2', \n",
|
||||
" query_instruction_for_retrieval=\"Given a web search query, retrieve relevant passages that answer the query.\",\n",
|
||||
" use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation\n",
|
||||
"\n",
|
||||
"embeddings_1 = model.encode_queries(queries)\n",
|
||||
"embeddings_2 = model.encode_corpus(documents)\n",
|
||||
"similarity = embeddings_1 @ embeddings_2.T\n",
|
||||
"print(similarity)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "8b7b2aa4",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### 2.4 BGE ICL"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "7c9acb92",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"BGE ICL stands for in-context learning. By providing few-shot examples in the query, it can significantly enhance the model's ability to handle new tasks."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "cf6c9345",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"| Model | Language | Parameters | Model Size | Description | Base Model |\n",
|
||||
"|:-------|:--------:|:--------------:|:--------------:|:-----------------:|:----------------:|\n",
|
||||
"| [BAAI/bge-en-icl](https://huggingface.co/BAAI/bge-en-icl) | English | 7.11B | 28.5 GB | LLM-based English embedding model with excellent in-context learning ability. | Mistral-7B |"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"id": "4595bae7",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"documents = [\n",
|
||||
" \"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.\",\n",
|
||||
" \"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments.\"\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"examples = [\n",
|
||||
" {\n",
|
||||
" 'instruct': 'Given a web search query, retrieve relevant passages that answer the query.',\n",
|
||||
" 'query': 'what is a virtual interface',\n",
|
||||
" 'response': \"A virtual interface is a software-defined abstraction that mimics the behavior and characteristics of a physical network interface. It allows multiple logical network connections to share the same physical network interface, enabling efficient utilization of network resources. Virtual interfaces are commonly used in virtualization technologies such as virtual machines and containers to provide network connectivity without requiring dedicated hardware. They facilitate flexible network configurations and help in isolating network traffic for security and management purposes.\"\n",
|
||||
" },\n",
|
||||
" {\n",
|
||||
" 'instruct': 'Given a web search query, retrieve relevant passages that answer the query.',\n",
|
||||
" 'query': 'causes of back pain in female for a week',\n",
|
||||
" 'response': \"Back pain in females lasting a week can stem from various factors. Common causes include muscle strain due to lifting heavy objects or improper posture, spinal issues like herniated discs or osteoporosis, menstrual cramps causing referred pain, urinary tract infections, or pelvic inflammatory disease. Pregnancy-related changes can also contribute. Stress and lack of physical activity may exacerbate symptoms. Proper diagnosis by a healthcare professional is crucial for effective treatment and management.\"\n",
|
||||
" }\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"queries = [\"how much protein should a female eat\", \"summit define\"]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "ffb586c6",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Loading checkpoint shards: 100%|██████████| 3/3 [00:00<00:00, 6.55it/s]\n",
|
||||
"pre tokenize: 100%|██████████| 1/1 [00:00<00:00, 366.09it/s]\n",
|
||||
"pre tokenize: 100%|██████████| 1/1 [00:00<00:00, 623.69it/s]\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"[[0.6064 0.3018]\n",
|
||||
" [0.257 0.537 ]]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from FlagEmbedding import FlagICLModel\n",
|
||||
"import os\n",
|
||||
"\n",
|
||||
"model = FlagICLModel('BAAI/bge-en-icl', \n",
|
||||
" examples_for_task=examples, # set `examples_for_task=None` to use model without examples\n",
|
||||
" # examples_instruction_format=\"<instruct>{}\\n<query>{}\\n<response>{}\" # specify the format to use examples_for_task\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
"embeddings_1 = model.encode_queries(queries)\n",
|
||||
"embeddings_2 = model.encode_corpus(documents)\n",
|
||||
"similarity = embeddings_1 @ embeddings_2.T\n",
|
||||
"\n",
|
||||
"print(similarity)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "dev",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.12.7"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -0,0 +1,277 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# BGE Auto Embedder"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"FlagEmbedding provides a high level class `FlagAutoModel` that unify the inference of embedding models. Besides BGE series, it also supports other popular open-source embedding models such as E5, GTE, SFR, etc. In this tutorial, we will have an idea how to use it."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"% pip install FlagEmbedding"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 1. Usage"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"First, import `FlagAutoModel` from FlagEmbedding, and use the `from_finetuned()` function to initialize the model:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 10,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from FlagEmbedding import FlagAutoModel\n",
|
||||
"\n",
|
||||
"model = FlagAutoModel.from_finetuned(\n",
|
||||
" 'BAAI/bge-base-en-v1.5',\n",
|
||||
" query_instruction_for_retrieval=\"Represent this sentence for searching relevant passages: \",\n",
|
||||
" devices=\"cuda:0\", # if not specified, will use all available gpus or cpu when no gpu available\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Then use the model exactly same to `FlagModel` (`FlagM3Model` if using BGE M3, `FlagLLMModel` if using BGE Multilingual Gemma2, `FlagICLModel` if using BGE ICL)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 11,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"You're using a BertTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"[[0.76 0.6714]\n",
|
||||
" [0.6177 0.7603]]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"queries = [\"query 1\", \"query 2\"]\n",
|
||||
"corpus = [\"passage 1\", \"passage 2\"]\n",
|
||||
"\n",
|
||||
"# encode the queries and corpus\n",
|
||||
"q_embeddings = model.encode_queries(queries)\n",
|
||||
"p_embeddings = model.encode_corpus(corpus)\n",
|
||||
"\n",
|
||||
"# compute the similarity scores\n",
|
||||
"scores = q_embeddings @ p_embeddings.T\n",
|
||||
"print(scores)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 2. Explanation"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"`FlagAutoModel` use an OrderedDict `MODEL_MAPPING` to store all the supported models configuration:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 12,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"['bge-en-icl',\n",
|
||||
" 'bge-multilingual-gemma2',\n",
|
||||
" 'bge-m3',\n",
|
||||
" 'bge-large-en-v1.5',\n",
|
||||
" 'bge-base-en-v1.5',\n",
|
||||
" 'bge-small-en-v1.5',\n",
|
||||
" 'bge-large-zh-v1.5',\n",
|
||||
" 'bge-base-zh-v1.5',\n",
|
||||
" 'bge-small-zh-v1.5',\n",
|
||||
" 'bge-large-en',\n",
|
||||
" 'bge-base-en',\n",
|
||||
" 'bge-small-en',\n",
|
||||
" 'bge-large-zh',\n",
|
||||
" 'bge-base-zh',\n",
|
||||
" 'bge-small-zh',\n",
|
||||
" 'e5-mistral-7b-instruct',\n",
|
||||
" 'e5-large-v2',\n",
|
||||
" 'e5-base-v2',\n",
|
||||
" 'e5-small-v2',\n",
|
||||
" 'multilingual-e5-large-instruct',\n",
|
||||
" 'multilingual-e5-large',\n",
|
||||
" 'multilingual-e5-base',\n",
|
||||
" 'multilingual-e5-small',\n",
|
||||
" 'e5-large',\n",
|
||||
" 'e5-base',\n",
|
||||
" 'e5-small',\n",
|
||||
" 'gte-Qwen2-7B-instruct',\n",
|
||||
" 'gte-Qwen2-1.5B-instruct',\n",
|
||||
" 'gte-Qwen1.5-7B-instruct',\n",
|
||||
" 'gte-multilingual-base',\n",
|
||||
" 'gte-large-en-v1.5',\n",
|
||||
" 'gte-base-en-v1.5',\n",
|
||||
" 'gte-large',\n",
|
||||
" 'gte-base',\n",
|
||||
" 'gte-small',\n",
|
||||
" 'gte-large-zh',\n",
|
||||
" 'gte-base-zh',\n",
|
||||
" 'gte-small-zh',\n",
|
||||
" 'SFR-Embedding-2_R',\n",
|
||||
" 'SFR-Embedding-Mistral',\n",
|
||||
" 'Linq-Embed-Mistral']"
|
||||
]
|
||||
},
|
||||
"execution_count": 12,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from FlagEmbedding.inference.embedder.model_mapping import AUTO_EMBEDDER_MAPPING\n",
|
||||
"\n",
|
||||
"list(AUTO_EMBEDDER_MAPPING.keys())"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 13,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"EmbedderConfig(model_class=<class 'FlagEmbedding.inference.embedder.decoder_only.icl.ICLLLMEmbedder'>, pooling_method=<PoolingMethod.LAST_TOKEN: 'last_token'>, trust_remote_code=False, query_instruction_format='<instruct>{}\\n<query>{}')\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"print(AUTO_EMBEDDER_MAPPING['bge-en-icl'])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Taking a look at the value of each key, which is an object of `EmbedderConfig`. It consists four attributes:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"```python\n",
|
||||
"@dataclass\n",
|
||||
"class EmbedderConfig:\n",
|
||||
" model_class: Type[AbsEmbedder]\n",
|
||||
" pooling_method: PoolingMethod\n",
|
||||
" trust_remote_code: bool = False\n",
|
||||
" query_instruction_format: str = \"{}{}\"\n",
|
||||
"```"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Not only the BGE series, it supports other models such as E5 similarly:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 14,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"EmbedderConfig(model_class=<class 'FlagEmbedding.inference.embedder.decoder_only.icl.ICLLLMEmbedder'>, pooling_method=<PoolingMethod.LAST_TOKEN: 'last_token'>, trust_remote_code=False, query_instruction_format='<instruct>{}\\n<query>{}')\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"print(AUTO_EMBEDDER_MAPPING['bge-en-icl'])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 3. Customization"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"If you want to use your own models through `FlagAutoModel`, consider the following steps:\n",
|
||||
"\n",
|
||||
"1. Check the type of your embedding model and choose the appropriate model class, is it an encoder or a decoder?\n",
|
||||
"2. What kind of pooling method it uses? CLS token, mean pooling, or last token?\n",
|
||||
"3. Does your model needs `trust_remote_code=Ture` to ran?\n",
|
||||
"4. Is there a query instruction format for retrieval?\n",
|
||||
"\n",
|
||||
"After these four attributes are assured, add your model name as the key and corresponding EmbedderConfig as the value to `MODEL_MAPPING`. Now have a try!"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "dev",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.12.7"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
@@ -0,0 +1,419 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# BGE Explanation"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"In this section, we will go through BGE and BGE-v1.5's structure and how they generate embeddings."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 0. Installation"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Install the required packages in your environment."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 18,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"%pip install -U transformers FlagEmbedding"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 1. Encode sentences"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"To know how exactly a sentence is encoded, let's first load the tokenizer and model from HF transformers instead of FlagEmbedding"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 19,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from transformers import AutoTokenizer, AutoModel\n",
|
||||
"import torch\n",
|
||||
"\n",
|
||||
"tokenizer = AutoTokenizer.from_pretrained(\"BAAI/bge-base-en-v1.5\")\n",
|
||||
"model = AutoModel.from_pretrained(\"BAAI/bge-base-en-v1.5\")\n",
|
||||
"\n",
|
||||
"sentences = [\"embedding\", \"I love machine learning and nlp\"]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Run the following cell to check the model of bge-base-en-v1.5. It uses BERT-base as base model, with 12 encoder layers and hidden dimension of 768.\n",
|
||||
"\n",
|
||||
"Note that the corresponding models of BGE and BGE-v1.5 have same structures. For example, bge-base-en and bge-base-en-v1.5 have the same structure."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 20,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"BertModel(\n",
|
||||
" (embeddings): BertEmbeddings(\n",
|
||||
" (word_embeddings): Embedding(30522, 768, padding_idx=0)\n",
|
||||
" (position_embeddings): Embedding(512, 768)\n",
|
||||
" (token_type_embeddings): Embedding(2, 768)\n",
|
||||
" (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
|
||||
" (dropout): Dropout(p=0.1, inplace=False)\n",
|
||||
" )\n",
|
||||
" (encoder): BertEncoder(\n",
|
||||
" (layer): ModuleList(\n",
|
||||
" (0-11): 12 x BertLayer(\n",
|
||||
" (attention): BertAttention(\n",
|
||||
" (self): BertSelfAttention(\n",
|
||||
" (query): Linear(in_features=768, out_features=768, bias=True)\n",
|
||||
" (key): Linear(in_features=768, out_features=768, bias=True)\n",
|
||||
" (value): Linear(in_features=768, out_features=768, bias=True)\n",
|
||||
" (dropout): Dropout(p=0.1, inplace=False)\n",
|
||||
" )\n",
|
||||
" (output): BertSelfOutput(\n",
|
||||
" (dense): Linear(in_features=768, out_features=768, bias=True)\n",
|
||||
" (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
|
||||
" (dropout): Dropout(p=0.1, inplace=False)\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" (intermediate): BertIntermediate(\n",
|
||||
" (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
|
||||
" (intermediate_act_fn): GELUActivation()\n",
|
||||
" )\n",
|
||||
" (output): BertOutput(\n",
|
||||
" (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
|
||||
" (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
|
||||
" (dropout): Dropout(p=0.1, inplace=False)\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" (pooler): BertPooler(\n",
|
||||
" (dense): Linear(in_features=768, out_features=768, bias=True)\n",
|
||||
" (activation): Tanh()\n",
|
||||
" )\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
"execution_count": 20,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"model.eval()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"First, let's tokenize the sentences."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 21,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"{'input_ids': tensor([[ 101, 7861, 8270, 4667, 102, 0, 0, 0, 0],\n",
|
||||
" [ 101, 1045, 2293, 3698, 4083, 1998, 17953, 2361, 102]]), 'token_type_ids': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
|
||||
" [0, 0, 0, 0, 0, 0, 0, 0, 0]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 0, 0, 0, 0],\n",
|
||||
" [1, 1, 1, 1, 1, 1, 1, 1, 1]])}"
|
||||
]
|
||||
},
|
||||
"execution_count": 21,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"inputs = tokenizer(\n",
|
||||
" sentences, \n",
|
||||
" padding=True, \n",
|
||||
" truncation=True, \n",
|
||||
" return_tensors='pt', \n",
|
||||
" max_length=512\n",
|
||||
")\n",
|
||||
"inputs"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"From the results, we can see that each sentence begins with token 101 and ends with 102, which are the `[CLS]` and `[SEP]` special token used in BERT."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 22,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"torch.Size([2, 9, 768])"
|
||||
]
|
||||
},
|
||||
"execution_count": 22,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"last_hidden_state = model(**inputs, return_dict=True).last_hidden_state\n",
|
||||
"last_hidden_state.shape"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Here we implement the pooling function, with two choices of using `[CLS]`'s last hidden state, or the mean pooling of the whole last hidden state."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 23,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def pooling(last_hidden_state: torch.Tensor, pooling_method='cls', attention_mask: torch.Tensor = None):\n",
|
||||
" if pooling_method == 'cls':\n",
|
||||
" return last_hidden_state[:, 0]\n",
|
||||
" elif pooling_method == 'mean':\n",
|
||||
" s = torch.sum(last_hidden_state * attention_mask.unsqueeze(-1).float(), dim=1)\n",
|
||||
" d = attention_mask.sum(dim=1, keepdim=True).float()\n",
|
||||
" return s / d"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Different from more commonly used mean pooling, BGE is trained to use the last hidden state of `[CLS]` as the sentence embedding: \n",
|
||||
"\n",
|
||||
"`sentence_embeddings = model_output[0][:, 0]`\n",
|
||||
"\n",
|
||||
"If you use mean pooling, there will be a significant decrease in performance. Therefore, make sure to use the correct method to obtain sentence vectors."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 24,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"torch.Size([2, 768])"
|
||||
]
|
||||
},
|
||||
"execution_count": 24,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embeddings = pooling(\n",
|
||||
" last_hidden_state, \n",
|
||||
" pooling_method='cls', \n",
|
||||
" attention_mask=inputs['attention_mask']\n",
|
||||
")\n",
|
||||
"embeddings.shape"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Assembling them together, we get the whole encoding function:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 25,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def _encode(sentences, max_length=512, convert_to_numpy=True):\n",
|
||||
"\n",
|
||||
" # handle the case of single sentence and a list of sentences\n",
|
||||
" input_was_string = False\n",
|
||||
" if isinstance(sentences, str):\n",
|
||||
" sentences = [sentences]\n",
|
||||
" input_was_string = True\n",
|
||||
"\n",
|
||||
" inputs = tokenizer(\n",
|
||||
" sentences, \n",
|
||||
" padding=True, \n",
|
||||
" truncation=True, \n",
|
||||
" return_tensors='pt', \n",
|
||||
" max_length=max_length\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" last_hidden_state = model(**inputs, return_dict=True).last_hidden_state\n",
|
||||
" \n",
|
||||
" embeddings = pooling(\n",
|
||||
" last_hidden_state, \n",
|
||||
" pooling_method='cls', \n",
|
||||
" attention_mask=inputs['attention_mask']\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" # normalize the embedding vectors\n",
|
||||
" embeddings = torch.nn.functional.normalize(embeddings, dim=-1)\n",
|
||||
"\n",
|
||||
" # convert to numpy if needed\n",
|
||||
" if convert_to_numpy:\n",
|
||||
" embeddings = embeddings.detach().numpy()\n",
|
||||
"\n",
|
||||
" return embeddings[0] if input_was_string else embeddings"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 2. Comparison"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Now let's run the function we wrote to get the embeddings of the two sentences:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 26,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Embeddings:\n",
|
||||
"[[ 1.4549762e-02 -9.6840411e-03 3.7761475e-03 ... -8.5092714e-04\n",
|
||||
" 2.8417887e-02 6.3214332e-02]\n",
|
||||
" [ 3.3924331e-05 -3.2998275e-03 1.7206438e-02 ... 3.5703944e-03\n",
|
||||
" 1.8721525e-02 -2.0371782e-02]]\n",
|
||||
"Similarity scores:\n",
|
||||
"[[0.9999997 0.6077381]\n",
|
||||
" [0.6077381 0.9999999]]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embeddings = _encode(sentences)\n",
|
||||
"print(f\"Embeddings:\\n{embeddings}\")\n",
|
||||
"\n",
|
||||
"scores = embeddings @ embeddings.T\n",
|
||||
"print(f\"Similarity scores:\\n{scores}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Then, run the API provided in FlagEmbedding:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 27,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Embeddings:\n",
|
||||
"[[ 1.4549762e-02 -9.6840411e-03 3.7761475e-03 ... -8.5092714e-04\n",
|
||||
" 2.8417887e-02 6.3214332e-02]\n",
|
||||
" [ 3.3924331e-05 -3.2998275e-03 1.7206438e-02 ... 3.5703944e-03\n",
|
||||
" 1.8721525e-02 -2.0371782e-02]]\n",
|
||||
"Similarity scores:\n",
|
||||
"[[0.9999997 0.6077381]\n",
|
||||
" [0.6077381 0.9999999]]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from FlagEmbedding import FlagModel\n",
|
||||
"\n",
|
||||
"model = FlagModel('BAAI/bge-base-en-v1.5')\n",
|
||||
"\n",
|
||||
"embeddings = model.encode(sentences)\n",
|
||||
"print(f\"Embeddings:\\n{embeddings}\")\n",
|
||||
"\n",
|
||||
"scores = embeddings @ embeddings.T\n",
|
||||
"print(f\"Similarity scores:\\n{scores}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"As we expect, the two encoding functions return exactly the same results. The full implementation in FlagEmbedding handles large datasets by batching and contains GPU support and parallelization. Feel free to check the [source code](https://github.com/FlagOpen/FlagEmbedding/blob/master/FlagEmbedding/inference/embedder/encoder_only/base.py) for more details."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "dev",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.13.0"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
@@ -0,0 +1,414 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# BGE-M3"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 0. Installation"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Install the required packages in your environment."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"%pip install -U transformers FlagEmbedding accelerate"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 1. BGE-M3 structure"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from transformers import AutoTokenizer, AutoModel\n",
|
||||
"import torch, os\n",
|
||||
"\n",
|
||||
"tokenizer = AutoTokenizer.from_pretrained(\"BAAI/bge-m3\")\n",
|
||||
"raw_model = AutoModel.from_pretrained(\"BAAI/bge-m3\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The base model of BGE-M3 is [XLM-RoBERTa-large](https://huggingface.co/FacebookAI/xlm-roberta-large), which is a multilingual version of RoBERTa."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"XLMRobertaModel(\n",
|
||||
" (embeddings): XLMRobertaEmbeddings(\n",
|
||||
" (word_embeddings): Embedding(250002, 1024, padding_idx=1)\n",
|
||||
" (position_embeddings): Embedding(8194, 1024, padding_idx=1)\n",
|
||||
" (token_type_embeddings): Embedding(1, 1024)\n",
|
||||
" (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
|
||||
" (dropout): Dropout(p=0.1, inplace=False)\n",
|
||||
" )\n",
|
||||
" (encoder): XLMRobertaEncoder(\n",
|
||||
" (layer): ModuleList(\n",
|
||||
" (0-23): 24 x XLMRobertaLayer(\n",
|
||||
" (attention): XLMRobertaAttention(\n",
|
||||
" (self): XLMRobertaSelfAttention(\n",
|
||||
" (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
|
||||
" (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
|
||||
" (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
|
||||
" (dropout): Dropout(p=0.1, inplace=False)\n",
|
||||
" )\n",
|
||||
" (output): XLMRobertaSelfOutput(\n",
|
||||
" (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
|
||||
" (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
|
||||
" (dropout): Dropout(p=0.1, inplace=False)\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" (intermediate): XLMRobertaIntermediate(\n",
|
||||
" (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
|
||||
" (intermediate_act_fn): GELUActivation()\n",
|
||||
" )\n",
|
||||
" (output): XLMRobertaOutput(\n",
|
||||
" (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
|
||||
" (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
|
||||
" (dropout): Dropout(p=0.1, inplace=False)\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" (pooler): XLMRobertaPooler(\n",
|
||||
" (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
|
||||
" (activation): Tanh()\n",
|
||||
" )\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
"execution_count": 4,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"raw_model.eval()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 2. Multi-Functionality"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Fetching 30 files: 100%|██████████| 30/30 [00:00<00:00, 240131.91it/s]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from FlagEmbedding import BGEM3FlagModel\n",
|
||||
"\n",
|
||||
"model = BGEM3FlagModel('BAAI/bge-m3', use_fp16=True)\n",
|
||||
"\n",
|
||||
"sentences_1 = [\"What is BGE M3?\", \"Defination of BM25\"]\n",
|
||||
"sentences_2 = [\"BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.\", \n",
|
||||
" \"BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document\"]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### 2.1 Dense Retrieval"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Using BGE M3 for dense embedding has similar steps to BGE or BGE 1.5 models.\n",
|
||||
"\n",
|
||||
"Use the normalized hidden state of the special token [CLS] as the embedding:\n",
|
||||
"\n",
|
||||
"$$e_q = norm(H_q[0])$$\n",
|
||||
"\n",
|
||||
"Then compute the relevance score between the query and passage:\n",
|
||||
"\n",
|
||||
"$$s_{dense}=f_{sim}(e_p, e_q)$$\n",
|
||||
"\n",
|
||||
"where $e_p, e_q$ are the embedding vectors of passage and query, respectively.\n",
|
||||
"\n",
|
||||
"$f_{sim}$ is the score function (such as inner product and L2 distance) for comupting two embeddings' similarity."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"[[0.6259035 0.34749585]\n",
|
||||
" [0.349868 0.6782462 ]]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# If you don't need such a long length of 8192 input tokens, you can set max_length to a smaller value to speed up encoding.\n",
|
||||
"embeddings_1 = model.encode(sentences_1, max_length=10)['dense_vecs']\n",
|
||||
"embeddings_2 = model.encode(sentences_2, max_length=100)['dense_vecs']\n",
|
||||
"\n",
|
||||
"# compute the similarity scores\n",
|
||||
"s_dense = embeddings_1 @ embeddings_2.T\n",
|
||||
"print(s_dense)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### 2.2 Sparse Retrieval"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Set `return_sparse` to true to make the model return sparse vector. If a term token appears multiple times in the sentence, we only retain its max weight.\n",
|
||||
"\n",
|
||||
"BGE-M3 generates sparce embeddings by adding a linear layer and a ReLU activation function following the hidden states:\n",
|
||||
"\n",
|
||||
"$$w_{qt} = \\text{Relu}(W_{lex}^T H_q [i])$$\n",
|
||||
"\n",
|
||||
"where $W_{lex}$ representes the weights of linear layer and $H_q[i]$ is the encoder's output of the $i^{th}$ token."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"[{'What': 0.08362077, 'is': 0.081469566, 'B': 0.12964639, 'GE': 0.25186998, 'M': 0.17001738, '3': 0.26957875, '?': 0.040755156}, {'De': 0.050144322, 'fin': 0.13689369, 'ation': 0.045134712, 'of': 0.06342201, 'BM': 0.25167602, '25': 0.33353207}]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"output_1 = model.encode(sentences_1, return_sparse=True)\n",
|
||||
"output_2 = model.encode(sentences_2, return_sparse=True)\n",
|
||||
"\n",
|
||||
"# you can see the weight for each token:\n",
|
||||
"print(model.convert_id_to_token(output_1['lexical_weights']))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Based on the tokens' weights of query and passage, the relevance score between them is computed by the joint importance of the co-existed terms within the query and passage:\n",
|
||||
"\n",
|
||||
"$$s_{lex} = \\sum_{t\\in q\\cap p}(w_{qt} * w_{pt})$$\n",
|
||||
"\n",
|
||||
"where $w_{qt}, w_{pt}$ are the importance weights of each co-existed term $t$ in query and passage, respectively."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"0.19554448500275612\n",
|
||||
"0.00880391988903284\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# compute the scores via lexical mathcing\n",
|
||||
"s_lex_10_20 = model.compute_lexical_matching_score(output_1['lexical_weights'][0], output_2['lexical_weights'][0])\n",
|
||||
"s_lex_10_21 = model.compute_lexical_matching_score(output_1['lexical_weights'][0], output_2['lexical_weights'][1])\n",
|
||||
"\n",
|
||||
"print(s_lex_10_20)\n",
|
||||
"print(s_lex_10_21)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### 2.3 Multi-Vector"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The multi-vector method utilizes the entire output embeddings for the representation of query $E_q$ and passage $E_p$.\n",
|
||||
"\n",
|
||||
"$$E_q = norm(W_{mul}^T H_q)$$\n",
|
||||
"$$E_p = norm(W_{mul}^T H_p)$$\n",
|
||||
"\n",
|
||||
"where $W_{mul}$ is the learnable projection matrix."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 10,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"(8, 1024)\n",
|
||||
"(30, 1024)\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"output_1 = model.encode(sentences_1, return_dense=True, return_sparse=True, return_colbert_vecs=True)\n",
|
||||
"output_2 = model.encode(sentences_2, return_dense=True, return_sparse=True, return_colbert_vecs=True)\n",
|
||||
"\n",
|
||||
"print(f\"({len(output_1['colbert_vecs'][0])}, {len(output_1['colbert_vecs'][0][0])})\")\n",
|
||||
"print(f\"({len(output_2['colbert_vecs'][0])}, {len(output_2['colbert_vecs'][0][0])})\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Following ColBert, we use late-interaction to compute the fine-grained relevance score:\n",
|
||||
"\n",
|
||||
"$$s_{mul}=\\frac{1}{N}\\sum_{i=1}^N\\max_{j=1}^M E_q[i]\\cdot E_p^T[j]$$\n",
|
||||
"\n",
|
||||
"where $E_q, E_p$ are the entire output embeddings of query and passage, respectively.\n",
|
||||
"\n",
|
||||
"This is a summation of average of maximum similarity of each $v\\in E_q$ with vectors in $E_p$"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 11,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"0.7796662449836731\n",
|
||||
"0.4621177911758423\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"s_mul_10_20 = model.colbert_score(output_1['colbert_vecs'][0], output_2['colbert_vecs'][0]).item()\n",
|
||||
"s_mul_10_21 = model.colbert_score(output_1['colbert_vecs'][0], output_2['colbert_vecs'][1]).item()\n",
|
||||
"\n",
|
||||
"print(s_mul_10_20)\n",
|
||||
"print(s_mul_10_21)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### 2.4 Hybrid Ranking"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"BGE-M3's multi-functionality gives the possibility of hybrid ranking to improve retrieval. Firstly, due to the heavy cost of multi-vector method, we can retrieve the candidate results by either of the dense or sparse method. Then, to get the final result, we can rerank the candidates based on the integrated relevance score:\n",
|
||||
"\n",
|
||||
"$$s_{rank} = w_1\\cdot s_{dense}+w_2\\cdot s_{lex} + w_3\\cdot s_{mul}$$\n",
|
||||
"\n",
|
||||
"where the values chosen for $w_1, w_2$ and $w_3$ varies depending on the downstream scenario (here 1/3 is just for demonstration)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 12,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"0.5337047390639782\n",
|
||||
"0.27280585498859483\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"s_rank_10_20 = 1/3 * s_dense[0][0] + 1/3 * s_lex_10_20 + 1/3 * s_mul_10_20\n",
|
||||
"s_rank_10_21 = 1/3 * s_dense[0][1] + 1/3 * s_lex_10_21 + 1/3 * s_mul_10_21\n",
|
||||
"\n",
|
||||
"print(s_rank_10_20)\n",
|
||||
"print(s_rank_10_21)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "base",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.10.13"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
@@ -0,0 +1,346 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# BGE-EN-ICL"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"In this tutorial, we will go through BGE-EN-ICL, an LLM based embedding model with both strong zero-shot and few-shot embedding capability."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 0.Installation"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Install the required packages in your environment."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%pip install -U transformers FlagEmbedding"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 1. BGE-EN-ICL structure"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"/share/project/xzy/Envs/ft/lib/python3.11/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",
|
||||
"Loading checkpoint shards: 100%|██████████| 3/3 [00:00<00:00, 9.94it/s]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from transformers import AutoTokenizer, AutoModel\n",
|
||||
"import torch, os\n",
|
||||
"\n",
|
||||
"tokenizer = AutoTokenizer.from_pretrained(\"BAAI/bge-en-icl\")\n",
|
||||
"raw_model = AutoModel.from_pretrained(\"BAAI/bge-en-icl\")\n",
|
||||
"\n",
|
||||
"sentences = [\"embedding\", \"I love machine learning and nlp\"]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Different from the previous BGE embedding models which are encoder only models, BGE-EN-ICL use decoder only LLM, Mistral-7B, as the base model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"MistralModel(\n",
|
||||
" (embed_tokens): Embedding(32003, 4096)\n",
|
||||
" (layers): ModuleList(\n",
|
||||
" (0-31): 32 x MistralDecoderLayer(\n",
|
||||
" (self_attn): MistralSdpaAttention(\n",
|
||||
" (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
|
||||
" (k_proj): Linear(in_features=4096, out_features=1024, bias=False)\n",
|
||||
" (v_proj): Linear(in_features=4096, out_features=1024, bias=False)\n",
|
||||
" (o_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
|
||||
" (rotary_emb): MistralRotaryEmbedding()\n",
|
||||
" )\n",
|
||||
" (mlp): MistralMLP(\n",
|
||||
" (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)\n",
|
||||
" (up_proj): Linear(in_features=4096, out_features=14336, bias=False)\n",
|
||||
" (down_proj): Linear(in_features=14336, out_features=4096, bias=False)\n",
|
||||
" (act_fn): SiLU()\n",
|
||||
" )\n",
|
||||
" (input_layernorm): MistralRMSNorm((4096,), eps=1e-05)\n",
|
||||
" (post_attention_layernorm): MistralRMSNorm((4096,), eps=1e-05)\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" (norm): MistralRMSNorm((4096,), eps=1e-05)\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"raw_model.eval()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 2. New Pooling Method"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"BERT-like encoder only networks are considered with strong capacity for representation learning because of their bidirectional attention structure. Some previous work replace unidirectional attention with bidirectional attention during the embedding training phase. But this might creates a mismatch with the model's pre-training design, which could potentially undermine its in-context learning and generative properties.\n",
|
||||
"\n",
|
||||
"Thus BGE-EN-ICL introduces a [EOS] token's output embedding to address this issue."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"{'input_ids': tensor([[ 0, 0, 0, 0, 0, 0, 1, 28643, 2],\n",
|
||||
" [ 1, 315, 2016, 5599, 5168, 304, 307, 12312, 2]]), 'attention_mask': tensor([[0, 0, 0, 0, 0, 0, 1, 1, 1],\n",
|
||||
" [1, 1, 1, 1, 1, 1, 1, 1, 1]])}"
|
||||
]
|
||||
},
|
||||
"execution_count": 5,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"inputs = tokenizer(\n",
|
||||
" sentences,\n",
|
||||
" padding=True,\n",
|
||||
" return_tensors='pt',\n",
|
||||
")\n",
|
||||
"inputs"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"torch.Size([2, 9, 4096])"
|
||||
]
|
||||
},
|
||||
"execution_count": 7,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"last_hidden_state = raw_model(**inputs, return_dict=True).last_hidden_state\n",
|
||||
"last_hidden_state.shape"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The last token/[EOS] pooling method can be described as:\n",
|
||||
"\n",
|
||||
"Given the tokenized input sequence $T: [\\text{BOS}], t_1, ..., t_N$ is sent into the LLM:\n",
|
||||
"$$h_t = \\text{LLM}(T)[\\text{EOS}]$$\n",
|
||||
"where $h_t$ represents the text embedding taken from the output embedding of the special token [EOS]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def last_token_pool(last_hidden_states: torch.Tensor,\n",
|
||||
" attention_mask: torch.Tensor) -> torch.Tensor:\n",
|
||||
" \n",
|
||||
" left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])\n",
|
||||
" if left_padding:\n",
|
||||
" return last_hidden_states[:, -1]\n",
|
||||
" else:\n",
|
||||
" sequence_lengths = attention_mask.sum(dim=1) - 1\n",
|
||||
" batch_size = last_hidden_states.shape[0]\n",
|
||||
" return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 10,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"torch.Size([2, 4096])"
|
||||
]
|
||||
},
|
||||
"execution_count": 10,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embeddings = last_token_pool(\n",
|
||||
" last_hidden_state, \n",
|
||||
" attention_mask=inputs['attention_mask']\n",
|
||||
")\n",
|
||||
"embeddings.shape"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 3. In-Context Learning"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"BGE-EN-ICL integrate strong in-context learning of LLM into embedding model while still persisting strong zero-shot embedding capability."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"For zero-shot inference, it's exactly same to BGE v1&1.5. For few-shot inference, use the following way:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 11,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"examples = [\n",
|
||||
" {\n",
|
||||
" 'instruct': 'Given a web search query, retrieve relevant passages that answer the query.',\n",
|
||||
" 'query': 'what is a virtual interface',\n",
|
||||
" 'response': \"A virtual interface is a software-defined abstraction that mimics the behavior and characteristics of a physical network interface. It allows multiple logical network connections to share the same physical network interface, enabling efficient utilization of network resources. Virtual interfaces are commonly used in virtualization technologies such as virtual machines and containers to provide network connectivity without requiring dedicated hardware. They facilitate flexible network configurations and help in isolating network traffic for security and management purposes.\"\n",
|
||||
" },\n",
|
||||
" {\n",
|
||||
" 'instruct': 'Given a web search query, retrieve relevant passages that answer the query.',\n",
|
||||
" 'query': 'causes of back pain in female for a week',\n",
|
||||
" 'response': \"Back pain in females lasting a week can stem from various factors. Common causes include muscle strain due to lifting heavy objects or improper posture, spinal issues like herniated discs or osteoporosis, menstrual cramps causing referred pain, urinary tract infections, or pelvic inflammatory disease. Pregnancy-related changes can also contribute. Stress and lack of physical activity may exacerbate symptoms. Proper diagnosis by a healthcare professional is crucial for effective treatment and management.\"\n",
|
||||
" }\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"queries = [\"how much protein should a female eat\", \"summit define\"]\n",
|
||||
"documents = [\n",
|
||||
" \"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.\",\n",
|
||||
" \"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments.\"\n",
|
||||
"]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 16,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Loading checkpoint shards: 100%|██████████| 3/3 [00:00<00:00, 4.59it/s]\n",
|
||||
"pre tokenize: 100%|██████████| 1/1 [00:00<00:00, 501.41it/s]\n",
|
||||
"You're using a LlamaTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"[[0.6064 0.302 ]\n",
|
||||
" [0.257 0.5366]]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from FlagEmbedding import FlagICLModel\n",
|
||||
"\n",
|
||||
"model = FlagICLModel('BAAI/bge-en-icl', \n",
|
||||
" examples_for_task=examples, # set `examples_for_task=None` to use model without examples\n",
|
||||
" examples_instruction_format=\"<instruct>{}\\n<query>{}\\n<response>{}\", # specify the format to use examples_for_task\n",
|
||||
" devices=[0],\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
"embeddings_1 = model.encode_queries(queries)\n",
|
||||
"embeddings_2 = model.encode_corpus(documents)\n",
|
||||
"similarity = embeddings_1 @ embeddings_2.T\n",
|
||||
"\n",
|
||||
"print(similarity)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "ft",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.10"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
Reference in New Issue
Block a user