wehub/flagopen--flagembedding
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:39:21 +08:00
commit bc34f6df14
1149 changed files with 328099 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/tutorial/7_Finetuning/7.1.1.ipynb
+468
View File
@@ -0,0 +1,468 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Data Preparation for Fine-tuning"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In this tutorial, we will show an example of the first step for fine-tuning: dataset preparation."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 0. Installation"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"% pip install -U datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Suppose we are willing to fine-tune our model for financial tasks. We found an open-source dataset that could be useful: [financial-qa-10k](https://huggingface.co/datasets/virattt/financial-qa-10K). Let's see how to properly prepare our dataset for fine-tuning."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The raw dataset has the following structure:\n",
"- 5 columns of: 'question', 'answer', 'context', 'ticker', and 'filing'.\n",
"- 7000 rows."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"Dataset({\n",
" features: ['question', 'answer', 'context', 'ticker', 'filing'],\n",
" num_rows: 7000\n",
"})"
]
},
"execution_count": 2,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from datasets import load_dataset\n",
"\n",
"ds = load_dataset(\"virattt/financial-qa-10K\", split=\"train\")\n",
"ds"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 1. Data for Fine-tuning"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Construct the dataset to the following format:\n",
"\n",
"``` python\n",
"{\"query\": str, \"pos\": List[str], \"neg\":List[str], \"pos_scores\": List[int], \"neg_scores\": List[int], \"prompt\": str, \"type\": str}\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"`query` is the query, and `pos` is a list of positive texts, `neg` is a list of negative texts. `pos_scores` is a list of scores corresponding to the query and pos, `neg_scores` is a list of scores corresponding to the `query` and `neg`, if you don't use knowledge distillation, it can be ignored. `prompt` is the prompt used for the query, it will cover query_instruction_for_retrieval. `type` is used for bge-en-icl, it includes `normal`, `symmetric_class`, `symmetric_clustering`, .etc. If you have no negative texts for a query, you can random sample some from the entire corpus as the negatives."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We select the columns 'question' and 'context' as our query and answer(pos), and rename the columns. Then add the 'id' column for later evaluation use."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'query': 'What area did NVIDIA initially focus on before expanding to other computationally intensive fields?',\n",
" 'pos': 'Since our original focus on PC graphics, we have expanded to several other large and important computationally intensive fields.',\n",
" 'id': '0'}"
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"ds = ds.select_columns(column_names=[\"question\", \"context\"])\n",
"ds = ds.rename_column(\"question\", \"query\")\n",
"ds = ds.rename_column(\"context\", \"pos\")\n",
"ds = ds.add_column(\"id\", [str(i) for i in range(len(ds))])\n",
"ds[0]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Negative examples are important during the training of embedding models. Our initial dataset does not come with negative texts. Thus we directly sample a few from the whole corpus."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"Map: 100%|██████████| 7000/7000 [00:00<00:00, 22336.83 examples/s]\n"
]
}
],
"source": [
"import numpy as np\n",
"\n",
"np.random.seed(520)\n",
"neg_num = 10\n",
"\n",
"def str_to_lst(data):\n",
" data[\"pos\"] = [data[\"pos\"]]\n",
" return data\n",
"\n",
"# sample negative texts\n",
"new_col = []\n",
"for i in range(len(ds)):\n",
" ids = np.random.randint(0, len(ds), size=neg_num)\n",
" while i in ids:\n",
" ids = np.random.randint(0, len(ds), size=neg_num)\n",
" neg = [ds[i.item()][\"pos\"] for i in ids]\n",
" new_col.append(neg)\n",
"ds = ds.add_column(\"neg\", new_col)\n",
"\n",
"# change the key of 'pos' to a list\n",
"ds = ds.map(str_to_lst)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Lastly, we add the prompt which is used for query. It will be the `query_instruction_for_retrieval` during inference."
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"instruction = \"Represent this sentence for searching relevant passages: \"\n",
"ds = ds.add_column(\"prompt\", [instruction]*len(ds))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now a single row of the dataset is:"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'query': 'What area did NVIDIA initially focus on before expanding to other computationally intensive fields?',\n",
" 'pos': ['Since our original focus on PC graphics, we have expanded to several other large and important computationally intensive fields.'],\n",
" 'id': '0',\n",
" 'neg': ['Kroger expects that its value creation model will deliver total shareholder return within a target range of 8% to 11% over time.',\n",
" 'CSB purchased First Mortgages of $2.9 billion during 2023.',\n",
" 'See Note 13 to our Consolidated Financial Statements for information on certain legal proceedings for which there are contingencies.',\n",
" 'Diluted earnings per share were $16.69 in fiscal 2022 compared to $15.53 in fiscal 2021.',\n",
" 'In the year ended December 31, 2023, Total net sales and revenue increased primarily due to: (1) increased net wholesale volumes primarily due to increased sales of crossover vehicles and full-size pickup trucks, partially offset by decreased sales of mid-size pickup trucks; (2) favorable Price as a result of low dealer inventory levels and strong demand for our products; (3) favorable Mix associated with increased sales of full-size pickup trucks and full-size SUVs and decreased sales of vans, passenger cars and mid-size pickup trucks, partially offset by increased sales of crossover vehicles; and (4) favorable Other due to increased sales of parts and accessories.',\n",
" 'As of December 31, 2023, we had 3,157 full-time employees.',\n",
" 'Item 3. Legal Proceedings. The information contained in Note 18 ‘‘Commitments and Contingencies’’ included in Item 8 of this 10-K is incorporated herein by reference.',\n",
" 'Under the amended 2019 Secured Facility, the maturity date is set to July 20, 2026.',\n",
" 'Accounts receivable for Las Vegas Sands Corp. on December 31, 2023, totaled $685 million, with a provision for credit losses of $201 million, resulting in a net balance of $484 million.',\n",
" 'Operating expenses as a percentage of segment net sales decreased 25 basis points for fiscal 2023 when compared to the previous fiscal year, primarily driven by strong sales growth and lower incremental COVID-19 related costs, partially offset by increased wage costs.'],\n",
" 'prompt': 'Represent this sentence for searching relevant passages: '}"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"ds[0]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Then we split the dataset into training set and testing set."
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"split = ds.train_test_split(test_size=0.1, shuffle=True, seed=520)\n",
"train = split[\"train\"]\n",
"test = split[\"test\"]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now we are ready to store the data for later fine-tuning:"
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"Creating json from Arrow format: 100%|██████████| 7/7 [00:00<00:00, 39.73ba/s]\n"
]
},
{
"data": {
"text/plain": [
"16583481"
]
},
"execution_count": 15,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"train.to_json(\"ft_data/training.json\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 2. Test Data for Evaluation"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The last step is to construct the testing dataset for evaluaton."
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"Dataset({\n",
" features: ['query', 'pos', 'id', 'neg', 'prompt'],\n",
" num_rows: 700\n",
"})"
]
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"test"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"First select the columns for queries:"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'id': '1289',\n",
" 'text': 'How does Starbucks recognize the interest and penalties related to income tax matters on their financial statements?'}"
]
},
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"queries = test.select_columns(column_names=[\"id\", \"query\"])\n",
"queries = queries.rename_column(\"query\", \"text\")\n",
"queries[0]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Then select the columns for corpus:"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [],
"source": [
"corpus = ds.select_columns(column_names=[\"id\", \"pos\"])\n",
"corpus = corpus.rename_column(\"pos\", \"text\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Finally, make the qrels that indicating the relations of queries and corresponding corpus\""
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"Flattening the indices: 100%|██████████| 700/700 [00:00<00:00, 180956.10 examples/s]\n"
]
},
{
"data": {
"text/plain": [
"{'qid': '1289', 'docid': '1289', 'relevance': 1}"
]
},
"execution_count": 13,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"qrels = test.select_columns([\"id\"])\n",
"qrels = qrels.rename_column(\"id\", \"qid\")\n",
"qrels = qrels.add_column(\"docid\", list(test[\"id\"]))\n",
"qrels = qrels.add_column(\"relevance\", [1]*len(test))\n",
"qrels[0]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Store the training set"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"Creating json from Arrow format: 100%|██████████| 1/1 [00:00<00:00, 210.42ba/s]\n",
"Creating json from Arrow format: 100%|██████████| 7/7 [00:00<00:00, 261.19ba/s]\n",
"Creating json from Arrow format: 100%|██████████| 1/1 [00:00<00:00, 591.08ba/s]\n"
]
},
{
"data": {
"text/plain": [
"30574"
]
},
"execution_count": 14,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"queries.to_json(\"ft_data/test_queries.jsonl\")\n",
"corpus.to_json(\"ft_data/corpus.jsonl\")\n",
"qrels.to_json(\"ft_data/test_qrels.jsonl\")"
]
}
],
"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
}
docs/source/tutorial/7_Finetuning/7.1.2.ipynb
+3734
View File
File diff suppressed because it is too large Load Diff
docs/source/tutorial/7_Finetuning/7.1.3.ipynb
+299
View File
@@ -0,0 +1,299 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Evaluate the Fine-tuned Model"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In the previous sections, we prepared the dataset and fine-tuned the model. In this tutorial, we will go through how to evaluate the model with the test dataset we constructed."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 0. Installation"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"% pip install -U datasets pytrec_eval FlagEmbedding"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 1. Load Data"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We first load data from the files we processed."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"from datasets import load_dataset\n",
"\n",
"queries = load_dataset(\"json\", data_files=\"ft_data/test_queries.jsonl\")[\"train\"]\n",
"corpus = load_dataset(\"json\", data_files=\"ft_data/corpus.jsonl\")[\"train\"]\n",
"qrels = load_dataset(\"json\", data_files=\"ft_data/test_qrels.jsonl\")[\"train\"]\n",
"\n",
"queries_text = queries[\"text\"]\n",
"corpus_text = [text for sub in corpus[\"text\"] for text in sub]"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"qrels_dict = {}\n",
"for line in qrels:\n",
" if line['qid'] not in qrels_dict:\n",
" qrels_dict[line['qid']] = {}\n",
" qrels_dict[line['qid']][line['docid']] = line['relevance']"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 2. Search"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Then we prepare a function to encode the text into embeddings and search the results:"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"import faiss\n",
"import numpy as np\n",
"from tqdm import tqdm\n",
"\n",
"\n",
"def search(model, queries_text, corpus_text):\n",
" \n",
" queries_embeddings = model.encode_queries(queries_text)\n",
" corpus_embeddings = model.encode_corpus(corpus_text)\n",
" \n",
" # create and store the embeddings in a Faiss index\n",
" dim = corpus_embeddings.shape[-1]\n",
" index = faiss.index_factory(dim, 'Flat', faiss.METRIC_INNER_PRODUCT)\n",
" corpus_embeddings = corpus_embeddings.astype(np.float32)\n",
" index.train(corpus_embeddings)\n",
" index.add(corpus_embeddings)\n",
" \n",
" query_size = len(queries_embeddings)\n",
"\n",
" all_scores = []\n",
" all_indices = []\n",
"\n",
" # search top 100 answers for all the queries\n",
" for i in tqdm(range(0, query_size, 32), desc=\"Searching\"):\n",
" j = min(i + 32, query_size)\n",
" query_embedding = queries_embeddings[i: j]\n",
" score, indice = index.search(query_embedding.astype(np.float32), k=100)\n",
" all_scores.append(score)\n",
" all_indices.append(indice)\n",
"\n",
" all_scores = np.concatenate(all_scores, axis=0)\n",
" all_indices = np.concatenate(all_indices, axis=0)\n",
" \n",
" # store the results into the format for evaluation\n",
" results = {}\n",
" for idx, (scores, indices) in enumerate(zip(all_scores, all_indices)):\n",
" results[queries[\"id\"][idx]] = {}\n",
" for score, index in zip(scores, indices):\n",
" if index != -1:\n",
" results[queries[\"id\"][idx]][corpus[\"id\"][index]] = float(score)\n",
" \n",
" return results"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 3. Evaluation"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"from FlagEmbedding.abc.evaluation.utils import evaluate_metrics, evaluate_mrr\n",
"from FlagEmbedding import FlagModel\n",
"\n",
"k_values = [10,100]\n",
"\n",
"raw_name = \"BAAI/bge-large-en-v1.5\"\n",
"finetuned_path = \"test_encoder_only_base_bge-large-en-v1.5\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The result for the original model:"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"pre tokenize: 100%|██████████| 3/3 [00:00<00:00, 129.75it/s]\n",
"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",
"Inference Embeddings: 100%|██████████| 3/3 [00:00<00:00, 11.08it/s]\n",
"pre tokenize: 100%|██████████| 28/28 [00:00<00:00, 164.29it/s]\n",
"Inference Embeddings: 100%|██████████| 28/28 [00:04<00:00, 6.09it/s]\n",
"Searching: 100%|██████████| 22/22 [00:08<00:00, 2.56it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"defaultdict(<class 'list'>, {'NDCG@10': 0.70405, 'NDCG@100': 0.73528})\n",
"defaultdict(<class 'list'>, {'MAP@10': 0.666, 'MAP@100': 0.67213})\n",
"defaultdict(<class 'list'>, {'Recall@10': 0.82286, 'Recall@100': 0.97286})\n",
"defaultdict(<class 'list'>, {'P@10': 0.08229, 'P@100': 0.00973})\n",
"defaultdict(<class 'list'>, {'MRR@10': 0.666, 'MRR@100': 0.67213})\n"
]
}
],
"source": [
"raw_model = FlagModel(\n",
" raw_name, \n",
" query_instruction_for_retrieval=\"Represent this sentence for searching relevant passages:\",\n",
" devices=[0],\n",
" use_fp16=False\n",
")\n",
"\n",
"results = search(raw_model, queries_text, corpus_text)\n",
"\n",
"eval_res = evaluate_metrics(qrels_dict, results, k_values)\n",
"mrr = evaluate_mrr(qrels_dict, results, k_values)\n",
"\n",
"for res in eval_res:\n",
" print(res)\n",
"print(mrr)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Then the result for the model after fine-tuning:"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"pre tokenize: 100%|██████████| 3/3 [00:00<00:00, 164.72it/s]\n",
"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",
"Inference Embeddings: 100%|██████████| 3/3 [00:00<00:00, 9.45it/s]\n",
"pre tokenize: 100%|██████████| 28/28 [00:00<00:00, 160.19it/s]\n",
"Inference Embeddings: 100%|██████████| 28/28 [00:04<00:00, 6.06it/s]\n",
"Searching: 100%|██████████| 22/22 [00:07<00:00, 2.80it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"defaultdict(<class 'list'>, {'NDCG@10': 0.84392, 'NDCG@100': 0.85792})\n",
"defaultdict(<class 'list'>, {'MAP@10': 0.81562, 'MAP@100': 0.81875})\n",
"defaultdict(<class 'list'>, {'Recall@10': 0.93143, 'Recall@100': 0.99429})\n",
"defaultdict(<class 'list'>, {'P@10': 0.09314, 'P@100': 0.00994})\n",
"defaultdict(<class 'list'>, {'MRR@10': 0.81562, 'MRR@100': 0.81875})\n"
]
}
],
"source": [
"ft_model = FlagModel(\n",
" finetuned_path, \n",
" query_instruction_for_retrieval=\"Represent this sentence for searching relevant passages:\",\n",
" devices=[0],\n",
" use_fp16=False\n",
")\n",
"\n",
"results = search(ft_model, queries_text, corpus_text)\n",
"\n",
"eval_res = evaluate_metrics(qrels_dict, results, k_values)\n",
"mrr = evaluate_mrr(qrels_dict, results, k_values)\n",
"\n",
"for res in eval_res:\n",
" print(res)\n",
"print(mrr)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can see an obvious improvement in all the metrics."
]
}
],
"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
}
docs/source/tutorial/7_Finetuning/7.2.1.ipynb
+393
View File
@@ -0,0 +1,393 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Hard Negatives"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Hard negatives are those negative samples that are particularly challenging for the model to distinguish from the positive ones. They are often close to the decision boundary or exhibit features that make them highly similar to the positive samples. Thus hard negative mining is widely used in machine learning tasks to make the model focus on subtle differences between similar instances, leading to better discrimination.\n",
"\n",
"In text retrieval system, a hard negative could be document that share some feature similarities with the query but does not truly satisfy the query's intent. During retrieval, those documents could rank higher than the real answers. Thus it's valuable to explicitly train the model on these hard negatives."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 1. Preparation"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"First, load an embedding model:"
]
},
{
"cell_type": "code",
"execution_count": 1,
"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"
]
}
],
"source": [
"from FlagEmbedding import FlagModel\n",
"\n",
"model = FlagModel('BAAI/bge-base-en-v1.5')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Then, load the queries and corpus from dataset:"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"from datasets import load_dataset\n",
"\n",
"corpus = load_dataset(\"BeIR/scifact\", \"corpus\")[\"corpus\"]\n",
"queries = load_dataset(\"BeIR/scifact\", \"queries\")[\"queries\"]\n",
"\n",
"corpus_ids = corpus.select_columns([\"_id\"])[\"_id\"]\n",
"corpus = corpus.select_columns([\"text\"])[\"text\"]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We create a dictionary maping auto generated ids (starting from 0) used by FAISS index, for later use."
]
},
{
"cell_type": "code",
"execution_count": 24,
"metadata": {},
"outputs": [],
"source": [
"corpus_ids_map = {}\n",
"for i in range(len(corpus)):\n",
" corpus_ids_map[i] = corpus_ids[i]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 2. Indexing"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Use the embedding model to encode the queries and corpus:"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"pre tokenize: 100%|██████████| 21/21 [00:00<00:00, 46.18it/s]\n",
"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",
"Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"/share/project/xzy/Envs/ft/lib/python3.11/site-packages/_distutils_hack/__init__.py:54: UserWarning: Reliance on distutils from stdlib is deprecated. Users must rely on setuptools to provide the distutils module. Avoid importing distutils or import setuptools first, and avoid setting SETUPTOOLS_USE_DISTUTILS=stdlib. Register concerns at https://github.com/pypa/setuptools/issues/new?template=distutils-deprecation.yml\n",
" warnings.warn(\n",
"Inference Embeddings: 0%| | 0/21 [00:00<?, ?it/s]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 5%|▍ | 1/21 [00:49<16:20, 49.00s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 10%|▉ | 2/21 [01:36<15:10, 47.91s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 14%|█▍ | 3/21 [02:16<13:23, 44.66s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 19%|█▉ | 4/21 [02:52<11:39, 41.13s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 24%|██▍ | 5/21 [03:23<09:58, 37.38s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 29%|██▊ | 6/21 [03:55<08:51, 35.44s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 33%|███▎ | 7/21 [04:24<07:47, 33.37s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 38%|███▊ | 8/21 [04:51<06:49, 31.51s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 43%|████▎ | 9/21 [05:16<05:52, 29.37s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 48%|████▊ | 10/21 [05:42<05:13, 28.51s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 52%|█████▏ | 11/21 [06:05<04:25, 26.59s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 57%|█████▋ | 12/21 [06:26<03:43, 24.85s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 62%|██████▏ | 13/21 [06:45<03:06, 23.35s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 67%|██████▋ | 14/21 [07:04<02:33, 21.89s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 71%|███████▏ | 15/21 [07:21<02:03, 20.54s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 76%|███████▌ | 16/21 [07:38<01:36, 19.30s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 81%|████████ | 17/21 [07:52<01:11, 17.87s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 86%|████████▌ | 18/21 [08:06<00:49, 16.58s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 90%|█████████ | 19/21 [08:18<00:30, 15.21s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 95%|█████████▌| 20/21 [08:28<00:13, 13.56s/it]Attempting to cast a BatchEncoding to type None. This is not supported.\n",
"Inference Embeddings: 100%|██████████| 21/21 [08:29<00:00, 24.26s/it]\n"
]
}
],
"source": [
"p_vecs = model.encode(corpus)"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"(5183, 768)"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"p_vecs.shape"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Then create a FAISS index"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"import torch, faiss\n",
"import numpy as np\n",
"\n",
"# create a basic flat index with dimension match our embedding\n",
"index = faiss.IndexFlatIP(len(p_vecs[0]))\n",
"# make sure the embeddings are float32\n",
"p_vecs = np.asarray(p_vecs, dtype=np.float32)\n",
"# use gpu to accelerate index searching\n",
"if torch.cuda.is_available():\n",
" co = faiss.GpuMultipleClonerOptions()\n",
" co.shard = True\n",
" co.useFloat16 = True\n",
" index = faiss.index_cpu_to_all_gpus(index, co=co)\n",
"# add all the embeddings to the index\n",
"index.add(p_vecs)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 3. Searching"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"For better demonstration, let's use a single query:"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'_id': '0',\n",
" 'title': '',\n",
" 'text': '0-dimensional biomaterials lack inductive properties.'}"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"query = queries[0]\n",
"query"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Get the id and content of that query, then use our embedding model to get its embedding vector."
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {},
"outputs": [],
"source": [
"q_id, q_text = query[\"_id\"], query[\"text\"]\n",
"# use the encode_queries() function to encode query\n",
"q_vec = model.encode_queries(queries=q_text)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Use the index to search for closest results:"
]
},
{
"cell_type": "code",
"execution_count": 31,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"['4346436',\n",
" '17388232',\n",
" '14103509',\n",
" '37437064',\n",
" '29638116',\n",
" '25435456',\n",
" '32532238',\n",
" '31715818',\n",
" '23763738',\n",
" '7583104',\n",
" '21456232',\n",
" '2121272',\n",
" '35621259',\n",
" '58050905',\n",
" '196664003']"
]
},
"execution_count": 31,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"_, ids = index.search(np.expand_dims(q_vec, axis=0), k=15)\n",
"# convert the auto ids back to ids in the original dataset\n",
"converted = [corpus_ids_map[id] for id in ids[0]]\n",
"converted"
]
},
{
"cell_type": "code",
"execution_count": 32,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'query-id': 0, 'corpus-id': 31715818, 'score': 1}"
]
},
"execution_count": 32,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"qrels = load_dataset(\"BeIR/scifact-qrels\")[\"train\"]\n",
"pos_id = qrels[0]\n",
"pos_id"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Lastly, we use the mothod of top-k shifted by N, which get the top 10 negatives after rank 5."
]
},
{
"cell_type": "code",
"execution_count": 44,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"['25435456',\n",
" '32532238',\n",
" '23763738',\n",
" '7583104',\n",
" '21456232',\n",
" '2121272',\n",
" '35621259',\n",
" '58050905',\n",
" '196664003']"
]
},
"execution_count": 44,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"negatives = [id for id in converted[5:] if int(id) != pos_id[\"corpus-id\"]]\n",
"negatives"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now we have select a group of hard negatives for the first query!"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"There are other methods to refine the process of choosing hard negatives. For example, the [implementation](https://github.com/FlagOpen/FlagEmbedding/blob/master/scripts/hn_mine.py) in our GitHub repo get the top 200 shifted by 10, which mean top 10-210. And then sample 15 from the 200 candidates. The reason is directly choosing the top K may introduce some false negatives, passages that somehow relative to the query but not exactly the answer to that query, into the negative set. This could influence model's performance."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"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
}