finetune-bge-embeddings.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "545c9100-5380-40c7-a75d-fd107665a84b",
   "metadata": {},
   "source": [
    "# Fine-Tune BGE Embedding Model Using Synthetic Data from Amazon Bedrock\n",
    "\n",
    "## Introduction\n",
    "\n",
    "Have you ever faced the challenge of obtaining high-quality data for fine-tuning your machine learning models? Generating synthetic data can provide a robust solution, especially when real-world data is scarce or sensitive. For instance, when developing a medical search engine, obtaining a large dataset of real user queries and relevant documents is often infeasible due to privacy concerns surrounding personal health information. However, synthetic data generation techniques can be employed to create realistic query-document pairs that resemble authentic user searches and relevant medical content, enabling the training of accurate retrieval models while preserving user privacy.\n",
    "\n",
    "In this blog post, we'll demonstrate how to leverage [Amazon Bedrock](https://aws.amazon.com/bedrock/) to create synthetic data, fine-tune a [BAAI General Embeddings (BGE) model](https://github.com/FlagOpen/FlagEmbedding?tab=readme-ov-file#bge-embedding), and deploy it using [Amazon SageMaker](https://aws.amazon.com/sagemaker/).\n",
    "\n",
    "Amazon Bedrock is a fully managed service that offers a choice of high-performing foundation models (FMs) from leading AI companies like AI21 Labs, Anthropic, Cohere, Meta, Stability AI, and Amazon via a single API, along with a broad set of capabilities you need to build generative AI applications with security, privacy, and responsible AI.\n",
    "\n",
    "### What is the BGE Model?\n",
    "\n",
    "BGE stands for Beijing Academy of Artificial Intelligence (BAAI) General Embeddings. It is a family of embedding models with a BERT-like architecture, designed to produce high-quality embeddings from text data. The BGE models come in three sizes:\n",
    "\n",
    "- `bge-small-en-v1.5`: 0.13GB, 384 embedding dimension\n",
    "- `bge-base-en-v1.5`: 0.44GB, 768 embedding dimension\n",
    "- `bge-large-en-v1.5`: 1.34GB, 1024 embedding dimension\n",
    "\n",
    "For comparing two pieces of text, the BGE model functions as a bi-encoder architecture, processing each piece of text through the same model in parallel to obtain their embeddings.\n",
    "\n",
    "### Why Use Synthetic Data?\n",
    "\n",
    "Generating synthetic data can significantly enhance the performance of your models by providing ample, high-quality training data without the constraints of traditional data collection methods. This blog post will guide you through generating synthetic data using Amazon Bedrock, fine-tuning a BGE model, evaluating its performance, and deploying it with Amazon SageMaker.\n",
    "\n",
    "### Solution Overview\n",
    "\n",
    "In this blog post, we will:\n",
    "\n",
    "1. Set up a SageMaker Studio Environment with the necessary IAM policies.\n",
    "2. Create a Conda environment for dependencies.\n",
    "3. Generate synthetic data using Meta Llama3 on Amazon Bedrock.\n",
    "4. Fine-tune the BGE embedding model with the generated data.\n",
    "5. Evaluate and compare the fine-tuned model.\n",
    "6. Deploy the model using Amazon SageMaker and [Hugging Face Text Embeddings Inference (TEI)](https://huggingface.co/docs/text-embeddings-inference/en/index)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "356bec95-8dd6-4cde-b917-4cf9925f9686",
   "metadata": {},
   "source": [
    "## Step 1: Generate Synthetic Data Using Amazon Bedrock\n",
    "\n",
    "We’ll start by adapting [LlamaIndex’s embedding model fine-tuning guide](https://docs.llamaindex.ai/en/stable/examples/finetuning/embeddings/finetune_embedding/) to use Amazon Bedrock to generate synthetic data for fine-tuning.\n",
    "\n",
    "### Download the data\n",
    "\n",
    "For our training and validation data, we’ll use the same Uber and Lyft 10K documents that are provided by the referenced LlamaIndex guide. \n",
    "\n",
    "Note that we're using \"!\" in Jupyter notebooks to execute shell commands directly from the notebook cells, which is convenient for performing system operations that might be more cumbersome or longer to implement directly in Python code.\n",
    "\n",
    "First we’ll download the data to be used for training and evaluation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5600e0e-a67c-4394-854c-0cf9b4d51822",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p 'data/10k/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/uber_2021.pdf' -O 'data/10k/uber_2021.pdf'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/lyft_2021.pdf' -O 'data/10k/lyft_2021.pdf'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce2bd360-1ffc-4795-90b2-fcb41388c00a",
   "metadata": {},
   "source": [
    "Next, we'll select PDFs from the corpus to train and evaluate on. You can replace these with your own data files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d928f51d-f791-4cc1-815a-753afe418db3",
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAIN_FILES = [\"./data/10k/lyft_2021.pdf\"]\n",
    "VAL_FILES = [\"./data/10k/uber_2021.pdf\"]\n",
    "\n",
    "TRAIN_CORPUS_FPATH = \"./data/train_corpus.json\"\n",
    "VAL_CORPUS_FPATH = \"./data/val_corpus.json\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "baa512bb-5a6c-4fae-a8b4-a4c455857ba1",
   "metadata": {},
   "source": [
    "Next, we create the corpus of text chunks by using LlamaIndex functionality to load the PDFs, and then parsing them into plain text chunks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "443d4dfc-b050-4d9f-bf66-5ee43e4915ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "from llama_index.core.node_parser import SentenceSplitter\n",
    "\n",
    "def load_corpus(files, verbose=False, chunk_size=1024, chunk_overlap=20):\n",
    "    if verbose:\n",
    "        print(f\"Loading files {files}\")\n",
    "\n",
    "    reader = SimpleDirectoryReader(input_files=files)\n",
    "    docs = reader.load_data()\n",
    "    if verbose:\n",
    "        print(f\"Loaded {len(docs)} docs\")\n",
    "\n",
    "    parser = SentenceSplitter(\n",
    "        chunk_size=chunk_size,\n",
    "        chunk_overlap=chunk_overlap,\n",
    "    )\n",
    "    nodes = parser.get_nodes_from_documents(docs, show_progress=verbose)\n",
    "\n",
    "    if verbose:\n",
    "        print(f\"Parsed {len(nodes)} nodes\")\n",
    "\n",
    "    return nodes\n",
    "\n",
    "chunk_size = 1024\n",
    "chunk_overlap = 20\n",
    "\n",
    "train_nodes = load_corpus(TRAIN_FILES, verbose=True, chunk_size=chunk_size, chunk_overlap=chunk_overlap)\n",
    "val_nodes = load_corpus(VAL_FILES, verbose=True, chunk_size=chunk_size, chunk_overlap=chunk_overlap)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07ac3646-f5dc-4475-b60e-8f6472a892b8",
   "metadata": {},
   "source": [
    "After parsing the data, we’ll run a couple cleaning operations to discard samples that aren’t useful for training and validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4a73a03-18a8-400e-abea-788045df2a9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove nodes where len(node.text) less than min_len\n",
    "min_len = 20\n",
    "train_nodes = [node for node in train_nodes if len(node.text) >= min_len]\n",
    "val_nodes = [node for node in val_nodes if len(node.text) >= min_len]\n",
    "\n",
    "print(len(train_nodes), len(val_nodes))\n",
    "\n",
    "# Remove nodes where text is not valid UTF-8\n",
    "def is_valid_utf8(string):\n",
    "    \"\"\"\n",
    "    Checks if a string is valid UTF-8.\n",
    "\n",
    "    :param string: String to be checked.\n",
    "    :return: True if the string is valid UTF-8, otherwise False.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        # Try encoding and decoding to check for UTF-8 validity\n",
    "        string.encode('utf-8').decode('utf-8')\n",
    "        return True\n",
    "    except UnicodeEncodeError:\n",
    "        return False\n",
    "\n",
    "train_nodes = [node for node in train_nodes if is_valid_utf8(node.text)]\n",
    "val_nodes = [node for node in val_nodes if is_valid_utf8(node.text)]\n",
    "\n",
    "print(len(train_nodes), len(val_nodes))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "132fbae8-fb8c-4234-891e-ce5eaee262c5",
   "metadata": {},
   "source": [
    "### Set Up Your LLM\n",
    "\n",
    "To generate synthetic data, we will use a Large Language Model (LLM) from Amazon Bedrock. In this case, we will use the Meta Llama3-70B-Instruct model, which offers great performance for low cost.\n",
    "\n",
    "First, set up the LLM and define the prompt template for generating questions based on the context:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73a497a5-5ae4-46dc-8fda-9fe1ae112119",
   "metadata": {},
   "outputs": [],
   "source": [
    "from botocore.config import Config\n",
    "from langchain_aws import ChatBedrock\n",
    "from llama_index.llms.langchain import LangChainLLM\n",
    "\n",
    "config = Config(retries={\"max_attempts\": 10, \"mode\": \"adaptive\"})\n",
    "\n",
    "model = ChatBedrock(model_id=\"meta.llama3-70b-instruct-v1:0\", config=config)\n",
    "llm = LangChainLLM(llm=model)\n",
    "\n",
    "qa_generate_prompt_tmpl = \"\"\"\n",
    "You are an AI assistant helping a teacher generate questions for an upcoming quiz. I will provide you with some context material in <context> tags. Your task is to generate {{NUM_QUESTIONS}} question(s) based only on the given context. The questions should be diverse and cover different parts of the context. Do not rely on any outside knowledge.\n",
    "\n",
    "Here are some important guidelines to follow:\n",
    "- Number each question, starting from 1.\n",
    "- Put each question on its own line, with no blank lines in between.\n",
    "- Do not generate any multiple choice questions.\n",
    "- Restrict the questions to information provided in the context only.\n",
    "- Do not reference the text, just directly ask a question that can be answered with it.\n",
    "    - For example, do not add phrases like, \"as referenced in the context\"\n",
    "- Do not output anything other than the questions.\n",
    "\n",
    "Here is the context to use for generating questions:\n",
    "<context>\n",
    "{context_str}\n",
    "</context>\n",
    "\n",
    "Please generate your numbered questions in the following format:\n",
    "\n",
    "1. Question 1\n",
    "2. Question 2\n",
    "3. Question 3\n",
    "\n",
    "The number of questions should match the {num_questions_per_chunk} value exactly.\n",
    "\n",
    "Assistant: Here is a list of questions to ask from the context above, where each question is a single line:\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bd79685-0910-4852-bd38-0d7b8e911a42",
   "metadata": {},
   "source": [
    "### Generate Synthetic QA Pairs\n",
    "\n",
    "Now, use the LLM to generate synthetic question-answer pairs. Here is a function to generate these pairs in parallel for a given set of nodes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1933119-8d63-4f9a-86f2-24b25a687ae9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import uuid\n",
    "from typing import List\n",
    "from tqdm import tqdm\n",
    "\n",
    "from llama_index.core.llms import LLM\n",
    "from llama_index.core.schema import MetadataMode, TextNode\n",
    "from llama_index.finetuning.embeddings.common import DEFAULT_QA_GENERATE_PROMPT_TMPL\n",
    "from llama_index.finetuning import EmbeddingQAFinetuneDataset\n",
    "\n",
    "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
    "\n",
    "def generate_qa_embedding_pairs_parallel(\n",
    "    nodes: List[TextNode],\n",
    "    llm: LLM,\n",
    "    qa_generate_prompt_tmpl: str = DEFAULT_QA_GENERATE_PROMPT_TMPL,\n",
    "    num_questions_per_chunk: int = 2,\n",
    "    max_workers: int = 8  # Number of parallel threads\n",
    ") -> EmbeddingQAFinetuneDataset:\n",
    "    \"\"\"Generate examples given a set of nodes in parallel with a progress bar.\"\"\"\n",
    "    node_dict = {\n",
    "        node.node_id: node.get_content(metadata_mode=MetadataMode.NONE)\n",
    "        for node in nodes\n",
    "    }\n",
    "\n",
    "    queries = {}\n",
    "    relevant_docs = {}\n",
    "\n",
    "    def process_node(node_id, text):\n",
    "        query = qa_generate_prompt_tmpl.format(\n",
    "            context_str=text, num_questions_per_chunk=num_questions_per_chunk\n",
    "        )\n",
    "        response = llm.complete(query)\n",
    "\n",
    "        result = str(response).strip().split(\"\\n\")\n",
    "        questions = [\n",
    "            re.sub(r\"^\\d+[\\).\\s]\", \"\", question).strip() for question in result\n",
    "        ]\n",
    "        questions = [question for question in questions if len(question) > 0]\n",
    "\n",
    "        local_queries = {}\n",
    "        local_relevant_docs = {}\n",
    "        for question in questions:\n",
    "            question_id = str(uuid.uuid4())\n",
    "            local_queries[question_id] = question\n",
    "            local_relevant_docs[question_id] = [node_id]\n",
    "\n",
    "        return local_queries, local_relevant_docs\n",
    "\n",
    "    with ThreadPoolExecutor(max_workers=max_workers) as executor:\n",
    "        # Submit all node processing jobs\n",
    "        future_to_node = {executor.submit(process_node, node_id, text): node_id for node_id, text in node_dict.items()}\n",
    "\n",
    "        # Process futures as they complete\n",
    "        for future in tqdm(as_completed(future_to_node), total=len(future_to_node), desc=\"Processing nodes\"):\n",
    "            node_id = future_to_node[future]\n",
    "            try:\n",
    "                local_queries, local_relevant_docs = future.result()\n",
    "                queries.update(local_queries)\n",
    "                relevant_docs.update(local_relevant_docs)\n",
    "            except Exception as exc:\n",
    "                print(f'Node {node_id} generated an exception: {exc}')\n",
    "\n",
    "    # Construct dataset\n",
    "    return EmbeddingQAFinetuneDataset(\n",
    "        queries=queries, corpus=node_dict, relevant_docs=relevant_docs\n",
    "    )\n",
    "\n",
    "# Usage\n",
    "train_dataset = generate_qa_embedding_pairs_parallel(train_nodes, llm=llm, qa_generate_prompt_tmpl=qa_generate_prompt_tmpl)\n",
    "val_dataset = generate_qa_embedding_pairs_parallel(val_nodes, llm=llm, qa_generate_prompt_tmpl=qa_generate_prompt_tmpl)\n",
    "\n",
    "# Create directories if they don't exist\n",
    "if not os.path.exists(\"./data\"):\n",
    "    os.makedirs(\"./data\")\n",
    "\n",
    "TRAIN_CORPUS_FPATH = \"./data/spgi_train_corpus.json\"\n",
    "VAL_CORPUS_FPATH = \"./data/spgi_val_corpus.json\"\n",
    "\n",
    "train_dataset.save_json(TRAIN_CORPUS_FPATH)\n",
    "val_dataset.save_json(VAL_CORPUS_FPATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "743e954e-3297-4d9e-923f-f9073476bdd4",
   "metadata": {},
   "source": [
    "### Load the Data\n",
    "\n",
    "Optionally, you can load the data using the EmbeddingQAFinetuneDataset class. This is useful if you want to verify the contents of your datasets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afaf61b5-d829-4ef7-b386-fadf2c408f53",
   "metadata": {},
   "outputs": [],
   "source": [
    "# [Optional] Load\n",
    "train_dataset = EmbeddingQAFinetuneDataset.from_json(TRAIN_CORPUS_FPATH)\n",
    "val_dataset = EmbeddingQAFinetuneDataset.from_json(VAL_CORPUS_FPATH)\n",
    "\n",
    "# Check the first 20 query-docid pairs to see if the questions generated look reasonable\n",
    "for p in train_dataset.query_docid_pairs[:20]:\n",
    "    print(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "296ae849-22b1-484a-a22d-c1c6acfb208c",
   "metadata": {},
   "source": [
    "### Save the Generated Data\n",
    "\n",
    "After generating the synthetic data, save it in JSONL format for later use in fine-tuning the BGE model.\n",
    "\n",
    "Note that instead of using LlamaIndex for fine-tuning, we'll use the [FlagEmbedding](https://github.com/FlagOpen/FlagEmbedding) library, so that we can take advantage of additional features this library provides such as Hard Negative Mining and Model Merging. For FlagEmbedding, data should be a json file, where each line is a dict like this:\n",
    "\n",
    "```json\n",
    "{\"query\": str, \"pos\": List[str], \"neg\": List[str]}\n",
    "```\n",
    "\n",
    "We'll take the chunks we've parsed and convert to this format. Negatives can be hard-negative mined in the next step if not available.\n",
    "\n",
    "Here’s how you can write the positive samples to a JSONL file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54d11c48-0786-4e26-87fe-17d33b521dd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import random\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "def write_jsonl(dataset: EmbeddingQAFinetuneDataset, filepath: str):\n",
    "    \"\"\"Write dataset to jsonl file.\"\"\"\n",
    "    # Collect the data\n",
    "    data = []\n",
    "    for query, doc_id in tqdm(dataset.query_docid_pairs):\n",
    "        assert len(doc_id) == 1\n",
    "        text = dataset.corpus[doc_id[0]]\n",
    "        data.append({\"query\": query, \"text\": text})\n",
    "    # Shuffle data\n",
    "    random.shuffle(data)\n",
    "    # Write to jsonl file\n",
    "    with open(filepath, \"w\") as f:\n",
    "        for sample in data:\n",
    "            query = sample[\"query\"]\n",
    "            text = sample[\"text\"]\n",
    "            f.write(json.dumps({\"query\": query, \"pos\": [text]}) + \"\\n\")\n",
    "\n",
    "# Define file paths\n",
    "TRAIN_DATA_FILEPATH = \"./data/train_data.jsonl\"\n",
    "VAL_DATA_FILEPATH = \"./data/val_data.jsonl\"\n",
    "\n",
    "# Usage\n",
    "write_jsonl(train_dataset, TRAIN_DATA_FILEPATH)\n",
    "write_jsonl(val_dataset, VAL_DATA_FILEPATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0263d49-f44f-46e7-aca5-7be54f597179",
   "metadata": {},
   "source": [
    "This process ensures that you have a well-structured synthetic dataset ready for the next step of fine-tuning the BGE model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28bad356-2093-4caf-bcf0-0d85b2ddac48",
   "metadata": {},
   "source": [
    "## Step 2: Fine-Tune the BGE Embedding Model\n",
    "\n",
    "### Model Selection\n",
    "\n",
    "To fine-tune the BGE model, we need to choose the appropriate model size based on our requirements and available resources. The BGE models come in three sizes:\n",
    "\n",
    "- `bge-small-en-v1.5`: 0.13GB, 384 embedding dimension\n",
    "- - `bge-base-en-v1.5`: 0.44GB, 768 embedding dimension\n",
    "- `bge-large-en-v1.5`: 1.34GB, 1024 embedding dimension\n",
    "\n",
    "For this demonstration, we will use the `bge-base-en-v1.5` model as it offers a good balance between performance and resource requirements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7c12e59-4e09-4dfe-b246-aea371c207f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import snapshot_download\n",
    "\n",
    "model_name = \"BAAI/bge-base-en-v1.5\"\n",
    "base_model_local_dir = model_name.split(\"/\")[-1]\n",
    "\n",
    "# We'll opt to download the model locally\n",
    "downloaded_model_dir = snapshot_download(\n",
    "    model_name,\n",
    "    local_dir=base_model_local_dir,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "207b956e-91f8-4bdf-8018-597f95d3ce3a",
   "metadata": {},
   "source": [
    "### Defining Instructions\n",
    "\n",
    "The BGE model can make use of instructions during the fine-tuning process. We can do this by prepending instructions to the passages and/or queries used for fine-tuning. In this example, we'll define a retrieval instruction for the query, and skip using a passage instruction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "534fa60f-2cba-4f79-87fb-101e73d8a4c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_instruction_for_retrieval = \"Represent the Financial question for retrieving supporting documents: \"\n",
    "\n",
    "# We could also use an instruction for retrieval, but it's not necessary\n",
    "passage_instruction_for_retrieval = \"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "793617c4-8f52-472c-b1bb-169dc461596e",
   "metadata": {},
   "source": [
    "### Generate Hard Negatives\n",
    "\n",
    "Hard negative mining is an essential step in fine-tuning embedding models. It helps improve the model's ability to distinguish between similar but not identical text pairs, thereby enhancing the overall quality of the embeddings. In this step, we will generate hard negatives to fine-tune our BGE model.\n",
    "\n",
    "We will use a pre-defined script to perform hard negative mining. The script utilizes the `FlagEmbedding` module to mine hard negatives from our dataset. Combining Python code with running scripts directly from a Jupyter notebook cell provides a flexible and efficient workflow. In the below code, `sys.executable` points to the Python executable for this Conda environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73b88659-3d1d-43e7-acd6-9401f3443ac1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "# Define math for saving hard negatives\n",
    "TRAIN_DATA_HARD_NEGATIVE_FILEPATH = \"./data/train_data_minedHN.jsonl\"\n",
    "\n",
    "# Hard Negative Mining\n",
    "!{sys.executable} -m FlagEmbedding.baai_general_embedding.finetune.hn_mine \\\n",
    "--model_name_or_path \"{base_model_local_dir}\" \\\n",
    "--input_file \"{TRAIN_DATA_FILEPATH}\" \\\n",
    "--output_file \"{TRAIN_DATA_HARD_NEGATIVE_FILEPATH}\" \\\n",
    "--range_for_sampling 3-200 \\\n",
    "--use_gpu_for_searching \\\n",
    "--query_instruction_for_retrieval \"{query_instruction_for_retrieval}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c7dcc3a-e4f3-4cbe-bd86-aa2982500faa",
   "metadata": {},
   "source": [
    "Running these scripts directly from a Jupyter notebook cell allows for seamless integration and automation within your workflow."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b4a86a3-bb73-4003-8654-1700c5d34fb9",
   "metadata": {},
   "source": [
    "### Fine-Tuning the Model\n",
    "\n",
    "Now we'll kick off the fine-tuning process. The model is trained with CrossEntropyLoss using in-batch negative sampling. This process involves adjusting the model's weights based on our custom dataset, which includes both positive and hard negative pairs.\n",
    "\n",
    "#### Important Training Arguments\n",
    "\n",
    "Here are some critical arguments to consider when fine-tuning the model:\n",
    "\n",
    "- `per_device_train_batch_size`: Batch size in training. Generally, a larger batch size will yield better performance. You can increase it by enabling options such as --fp16, --deepspeed ./df_config.json (you can refer to ds_config.json), and --gradient_checkpointing.\n",
    "- `train_group_size`: The number of positive and negative samples for a query during training. This argument controls the number of negatives (#negatives = train_group_size - 1). Ensure that the number of negatives is not larger than the number of negatives in the data (\"neg\": List[str]). In-batch negatives are also used in fine-tuning.\n",
    "- `negatives_cross_device`: Shares the negatives across all GPUs, extending the number of negatives.\n",
    "- `learning_rate`: Select an appropriate learning rate for your model. Recommended values are 1e-5, 2e-5, or 3e-5 for large, base, or small-scale models, respectively.\n",
    "- `temperature`: Influences the distribution of similarity scores.\n",
    "- `query_max_len`: Maximum length for queries. Set this according to the average length of queries in your data.\n",
    "- `passage_max_len`: Maximum length for passages. Set this according to the average length of passages in your data.\n",
    "- `query_instruction_for_retrieval`: Instruction for queries, which will be added to each query. You can also set it to \"\" to add nothing to the query.\n",
    "- `use_inbatch_neg`: Uses passages in the same batch as negatives. The default value is True.\n",
    "\n",
    "For more training arguments please refer to [transformers.TrainingArguments](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6addf19b-4ced-4f49-a6b2-8bc60024d464",
   "metadata": {},
   "source": [
    "#### Training Command\n",
    "\n",
    "We'll try using `gradient_accumulation_steps` in order to have a larger effective batch size. Here is the command used to initiate the training process:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb2459f5-e018-4da0-956f-311bb0da32f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Set the torchrun path\n",
    "torchrun_path = os.path.join(os.path.dirname(sys.executable), \"torchrun\")\n",
    "\n",
    "# Disable parallelism to avoid deadlocks\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "\n",
    "# Define the output path for fine-tuning\n",
    "finetuned_model_dir = \"./finetuned_model\"\n",
    "\n",
    "# Kick off the FlagEmbedding fine-tuning script using torchrun, passing the necessary arguments\n",
    "!{torchrun_path} --nproc_per_node 1 \\\n",
    "-m FlagEmbedding.baai_general_embedding.finetune.run \\\n",
    "--output_dir \"{finetuned_model_dir}\" \\\n",
    "--model_name_or_path \"{base_model_local_dir}\" \\\n",
    "--train_data \"{TRAIN_DATA_HARD_NEGATIVE_FILEPATH}\" \\\n",
    "--learning_rate 1e-5 \\\n",
    "--fp16 \\\n",
    "--num_train_epochs 20 \\\n",
    "--per_device_train_batch_size 32 \\\n",
    "--gradient_accumulation_steps 10 \\\n",
    "--dataloader_drop_last True \\\n",
    "--normlized True \\\n",
    "--temperature 0.02 \\\n",
    "--query_max_len 256 \\\n",
    "--passage_max_len 256 \\\n",
    "--train_group_size 2 \\\n",
    "--negatives_cross_device \\\n",
    "--logging_steps 10 \\\n",
    "--query_instruction_for_retrieval \"{query_instruction_for_retrieval}\" \\\n",
    "--passage_instruction_for_retrieval \"{passage_instruction_for_retrieval}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51af1e9e-3348-4482-a6e2-ed5bcf5e1be6",
   "metadata": {},
   "source": [
    "### Inspect the trained model\n",
    "\n",
    "After fine-tuning, we can load the model and print it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d316bc3-b3ed-4172-b947-2b43a92f767d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModel\n",
    "\n",
    "# Load model from HuggingFace Hub\n",
    "model = AutoModel.from_pretrained(finetuned_model_dir)\n",
    "\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "205f5972-c024-402d-9d8c-4d0f92467850",
   "metadata": {},
   "source": [
    "## Step 3: Model Merging\n",
    "\n",
    "Next we'll use `LM-Cocktail` to merge the fine-tuned weights with the original weights. This creates new parameters by calculating a weighted average of two or more models' parameters. Despite being simple, this tends to give better performance by mitigating the problem of catastrophic forgetting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f12798bc-ea25-4b01-aea9-404570016f35",
   "metadata": {},
   "outputs": [],
   "source": [
    "from LM_Cocktail import mix_models\n",
    "\n",
    "mixed_model_output_dir = \"./mixed_model\"\n",
    "\n",
    "# Mix fine-tuned model and base model; then save it to output_path\n",
    "model = mix_models(\n",
    "    model_names_or_paths=[base_model_local_dir, finetuned_model_dir],\n",
    "    model_type='encoder',\n",
    "    weights=[0.25, 0.75],  # you can change the weights to get a better trade-off.\n",
    "    output_path=mixed_model_output_dir\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe37816a-c194-4f56-b1e6-272f14b4852e",
   "metadata": {},
   "source": [
    "## Step 4: Test the Model Locally\n",
    "\n",
    "We can optionally test the model locally now before deploying. we will test the model out with pairs of queries and documents that we'd expect to score high in similarity, and then with pairs we'd expect to score low in similarity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa041658-ba5a-4533-ab45-e008c51c289b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "from transformers import AutoModel, AutoTokenizer\n",
    "import torch\n",
    "\n",
    "def model_fn(model_dir):\n",
    "    # load tokenizer and model from model_dir\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_dir, use_safetensors=True)\n",
    "    model = AutoModel.from_pretrained(model_dir, use_safetensors=True)\n",
    "    model.eval()\n",
    "\n",
    "    return tokenizer, model\n",
    "\n",
    "def predict_fn(data, tokenizer_and_model):\n",
    "    # unpack tokenizer and model\n",
    "    tokenizer, model = tokenizer_and_model\n",
    "\n",
    "    # process input\n",
    "    inputs = data.pop(\"inputs\", data)\n",
    "    instruction = data.pop(\"instruction\", None)\n",
    "\n",
    "    # Add instruction to query if instruction is provided\n",
    "    # for s2p(short query to long passage) retrieval task, add an instruction to query (DON'T add instruction for passages, unless fine-tuned)\n",
    "    if instruction is not None:\n",
    "        inputs = [instruction + q for q in inputs]\n",
    "\n",
    "    # Tokenize input sentences\n",
    "    encoded_input = tokenizer(inputs, padding=True, truncation=True, return_tensors='pt')\n",
    "\n",
    "    # Compute token embeddings\n",
    "    with torch.no_grad():\n",
    "        model_output = model(**encoded_input)\n",
    "        # Perform pooling. In this case, cls pooling.\n",
    "        embeddings = model_output[0][:, 0]\n",
    "    # normalize embeddings\n",
    "    embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)\n",
    "\n",
    "    return {\"embeddings\": embeddings}\n",
    "\n",
    "# Load the fine-tuned mixed model\n",
    "model, tokenizer = model_fn(mixed_model_output_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2bb2d85-22e1-41e4-b9e5-babd4fcee14f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Similar documents and queries\n",
    "similar_documents = [\n",
    "    \"The company's revenue for the fiscal year 2022 was $500 million, a 10% increase from the previous year. This growth was primarily driven by strong sales in the company's core product lines and successful expansion into new markets.\",\n",
    "    \"The company's net income for the fiscal year 2022 was $75 million, a 15% increase from the previous year. This improvement was mainly attributed to higher revenue, cost optimization initiatives, and operational efficiencies.\",\n",
    "    \"The company's total assets as of December 31, 2022, were $1.2 billion, a 5% increase from the previous year. This increase was primarily due to the acquisition of XYZ Company and investments in property, plant, and equipment.\"\n",
    "]\n",
    "\n",
    "similar_queries = [\n",
    "    \"What was the company's revenue for the fiscal year 2022?\",\n",
    "    \"How much did the company's net income increase in fiscal year 2022 compared to the previous year?\",\n",
    "    \"What was the value of the company's total assets as of December 31, 2022?\"\n",
    "]\n",
    "\n",
    "# Different documents and queries\n",
    "different_documents = [\n",
    "    \"The company's research and development expenses for the fiscal year 2022 were $50 million, representing 10% of the total revenue. The company continues to invest in innovative technologies and product development to maintain its competitive edge in the market.\",\n",
    "    \"The company's board of directors declared a quarterly dividend of $0.50 per share, payable on September 15, 2023, to shareholders of record as of August 31, 2023. This represents a 5% increase from the previous quarter's dividend.\",\n",
    "    \"The company's cash and cash equivalents balance as of December 31, 2022, was $200 million, a decrease of 20% from the previous year. This decrease was primarily due to the repayment of long-term debt and share repurchases during the year.\"\n",
    "]\n",
    "\n",
    "different_queries = [\n",
    "    \"What was the average age of the company's employees as of December 31, 2022?\",\n",
    "    \"How many patents did the company file in the fiscal year 2022?\",\n",
    "    \"What was the company's market share in the European market for the fiscal year 2022?\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf53a8af-ba4c-4af7-8843-3f24fd0e03c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get embeddings for similar documents and queries\n",
    "similar_document_results = predict_fn({\"inputs\": similar_documents, \"instruction\": passage_instruction_for_retrieval}, (model, tokenizer))\n",
    "similar_query_results = predict_fn({\"inputs\": similar_queries, \"instruction\": query_instruction_for_retrieval}, (model, tokenizer))\n",
    "\n",
    "# Get embeddings for different documents and queries\n",
    "different_document_results = predict_fn({\"inputs\": different_documents, \"instruction\": passage_instruction_for_retrieval}, (model, tokenizer))\n",
    "different_query_results = predict_fn({\"inputs\": different_queries, \"instruction\": query_instruction_for_retrieval}, (model, tokenizer))\n",
    "\n",
    "# Extract the embeddings\n",
    "similar_document_embeddings = similar_document_results['embeddings'].numpy()\n",
    "similar_query_embeddings = similar_query_results['embeddings'].numpy()\n",
    "different_document_embeddings = different_document_results['embeddings'].numpy()\n",
    "different_query_embeddings = different_query_results['embeddings'].numpy()\n",
    "\n",
    "# Calculate cosine similarity for similar document-query pairs\n",
    "cosine_sim_pairs_similar = []\n",
    "for doc_emb, query_emb in zip(similar_document_embeddings, similar_query_embeddings):\n",
    "    cosine_sim = cosine_similarity([doc_emb], [query_emb])\n",
    "    cosine_sim_pairs_similar.append(cosine_sim[0][0])\n",
    "\n",
    "# Calculate cosine similarity for different document-query pairs\n",
    "cosine_sim_pairs_different = []\n",
    "for doc_emb, query_emb in zip(different_document_embeddings, different_query_embeddings):\n",
    "    cosine_sim = cosine_similarity([doc_emb], [query_emb])\n",
    "    cosine_sim_pairs_different.append(cosine_sim[0][0])\n",
    "\n",
    "# Print the cosine similarity for each similar document-query pair\n",
    "print(\"Cosine Similarity for Similar Document-Query Pairs:\")\n",
    "for i, sim in enumerate(cosine_sim_pairs_similar):\n",
    "    print(f\"Pair {i+1}: Cosine Similarity = {sim}\")\n",
    "\n",
    "# Print the cosine similarity for each different document-query pair\n",
    "print(\"\\nCosine Similarity for Different Document-Query Pairs:\")\n",
    "for i, sim in enumerate(cosine_sim_pairs_different):\n",
    "    print(f\"Pair {i+1}: Cosine Similarity = {sim}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9c5f593-a158-4635-a737-9835386c2b42",
   "metadata": {},
   "source": [
    "## Step 5: Model Evaluation\n",
    "\n",
    "Next, we can evaluate the model and compare it to the base model (or to Titan, Cohere, etc. embedding models).\n",
    "\n",
    "Note: For black-box API embedding models, we can evaluate Hit Rate with the below evaluate function, but we can't apply the `evaluate_st` and the `InformationRetrievalEvaluator`, since it only works for `sentence-transformers` compatible models.\n",
    "\n",
    "### Evaluation Procedure 1: Hit Rate\n",
    "\n",
    "The first evaluation procedure we employ is a straightforward metric known as the Hit Rate. The process is as follows:\n",
    "\n",
    "1. We take each pair consisting of a query and its corresponding relevant document (query, relevant_doc).\n",
    "2. Using the query, we retrieve the top-K documents from the system being evaluated.\n",
    "3. If the retrieved results include the relevant_doc, we consider it a hit.\n",
    "\n",
    "This simple and intuitive approach serves as the initial step in our evaluation process, allowing us to assess the performance of various embedding models, including our own open-source and fine-tuned embedding models as well as proprietary API-based embedding models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcecbaee-65ae-4975-8228-ef855d2dc14d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy\n",
    "\n",
    "import pandas as pd\n",
    "from langchain_aws import ChatBedrock\n",
    "from llama_index.core import VectorStoreIndex\n",
    "from llama_index.core.schema import TextNode\n",
    "from llama_index.core import Settings\n",
    "from llama_index.embeddings.bedrock import BedrockEmbedding\n",
    "from llama_index.embeddings.huggingface import HuggingFaceEmbedding\n",
    "from llama_index.llms.langchain import LangChainLLM\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "def evaluate(\n",
    "    dataset,\n",
    "    embed_model,\n",
    "    top_k=1,\n",
    "    query_instruction_for_retrieval = \"\",\n",
    "    passage_instruction_for_retrieval = \"\",\n",
    "    verbose=False,\n",
    "):\n",
    "\n",
    "    model = ChatBedrock(model_id=\"meta.llama3-70b-instruct-v1:0\")\n",
    "    Settings.llm = LangChainLLM(llm=model) # Not used, but required\n",
    "    Settings.embed_model = embed_model\n",
    "\n",
    "    ds = deepcopy(dataset)\n",
    "\n",
    "    corpus = ds.corpus\n",
    "    queries = ds.queries\n",
    "    relevant_docs = ds.relevant_docs\n",
    "\n",
    "    # Add instruction to query if instruction is provided\n",
    "    for k,v in queries.items():\n",
    "        queries[k] = query_instruction_for_retrieval + v\n",
    "    # Add instruction to passage if instruction is provided\n",
    "    for k,v in corpus.items():\n",
    "        corpus[k] = passage_instruction_for_retrieval + v\n",
    "\n",
    "    nodes = [TextNode(id_=id_, text=text) for id_, text in corpus.items()]\n",
    "    index = VectorStoreIndex(nodes, show_progress=True)\n",
    "    retriever = index.as_retriever(similarity_top_k=top_k)\n",
    "\n",
    "    eval_results = []\n",
    "    for query_id, query in tqdm(queries.items()):\n",
    "        retrieved_nodes = retriever.retrieve(query)\n",
    "        retrieved_ids = [node.node.node_id for node in retrieved_nodes]\n",
    "        expected_id = relevant_docs[query_id][0]\n",
    "        is_hit = expected_id in retrieved_ids  # assume 1 relevant doc\n",
    "\n",
    "        eval_result = {\n",
    "            \"is_hit\": is_hit,\n",
    "            \"retrieved\": retrieved_ids,\n",
    "            \"expected\": expected_id,\n",
    "            \"query\": query_id,\n",
    "        }\n",
    "        eval_results.append(eval_result)\n",
    "    return eval_results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f98d3a6-0661-4636-949e-69e792bf28b2",
   "metadata": {},
   "source": [
    "### Evaluation Procedure 2: InformationRetrievalEvaluator\n",
    "\n",
    "For the second evaluation procedure, we utilize the `InformationRetrievalEvaluator` from the `sentence-transformers` library. This evaluator offers a more comprehensive suite of metrics, enabling us to perform a detailed analysis of the embedding models' performance.\n",
    "\n",
    "However, it's important to note that the `InformationRetrievalEvaluator` is only compatible with `sentence-transformers` models. This means that we can evaluate our open-source and fine-tuned embedding models using this procedure, but we cannot apply it to the proprietary API-based embedding models, such as the Bedrock Titan embedding models.\n",
    "\n",
    "By employing the `InformationRetrievalEvaluator`, we gain access to a wide range of evaluation metrics, providing a more in-depth understanding of how our `sentence-transformers` compatible models perform in various aspects of information retrieval tasks.\n",
    "\n",
    "This second evaluation procedure complements the hit rate metric used in the first procedure, allowing us to form a well-rounded assessment of the embedding models' capabilities and limitations.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22d4bdac-c04e-41e5-bfb6-ba3ceb1aa040",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "from sentence_transformers.evaluation import InformationRetrievalEvaluator\n",
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "def evaluate_st(\n",
    "    dataset,\n",
    "    model_id,\n",
    "    name,\n",
    "    query_instruction_for_retrieval = \"\",\n",
    "    passage_instruction_for_retrieval = \"\",\n",
    "):\n",
    "    \"\"\" Evaluate sentence-transformer models with InformationRetrievalEvaluator \"\"\"\n",
    "\n",
    "    ds = deepcopy(dataset)\n",
    "\n",
    "    corpus = ds.corpus\n",
    "    queries = ds.queries\n",
    "    relevant_docs = ds.relevant_docs\n",
    "\n",
    "    # Add instruction to query if instruction is provided\n",
    "    for k,v in queries.items():\n",
    "        queries[k] = query_instruction_for_retrieval + v\n",
    "    # Add instruction to passage if instruction is provided\n",
    "    for k,v in corpus.items():\n",
    "        corpus[k] = passage_instruction_for_retrieval + v\n",
    "\n",
    "    evaluator = InformationRetrievalEvaluator(\n",
    "        queries, corpus, relevant_docs, name=name\n",
    "    )\n",
    "    model = SentenceTransformer(model_id)\n",
    "    output_path = \"results/\"\n",
    "    Path(output_path).mkdir(exist_ok=True, parents=True)\n",
    "    return evaluator(model, output_path=output_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cce7972-8f33-459d-8f01-f724aa026509",
   "metadata": {},
   "source": [
    "### Evaluate the model(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26618d47-0343-47ed-984b-ab3cb77b9889",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "def evaluate_models(models, val_dataset):\n",
    "    results = []\n",
    "    for name, model_info in models.items():\n",
    "        print(f\"Evaluating model: {name}\")\n",
    "        model = model_info['model']\n",
    "        query_instruction = model_info.get('query_instruction', '')\n",
    "        passage_instruction = model_info.get('passage_instruction', '')\n",
    "\n",
    "        # For all models, evaluate using the basic hit rate evaluation function\n",
    "        val_results = evaluate(val_dataset, model, query_instruction_for_retrieval=query_instruction, passage_instruction_for_retrieval=passage_instruction)\n",
    "        df = pd.DataFrame(val_results)\n",
    "        hit_rate = df[\"is_hit\"].mean()\n",
    "\n",
    "        result = {\"model\": name, \"hit_rate\": hit_rate}\n",
    "\n",
    "        # For BGE models, additionally evaluate using sentence-transformers\n",
    "        if \"bge\" in name.lower():\n",
    "            csv_file = f\"results/Information-Retrieval_evaluation_{name.lower()}_results.csv\"\n",
    "            if os.path.exists(csv_file):\n",
    "                os.remove(csv_file)\n",
    "            _ = evaluate_st(val_dataset, model.model_name, name=name, query_instruction_for_retrieval=query_instruction, passage_instruction_for_retrieval=passage_instruction)\n",
    "            ir_results = pd.read_csv(csv_file, index_col=0, header=0)\n",
    "            row_dict = ir_results.to_dict('records')[0]\n",
    "            result.update(row_dict)\n",
    "\n",
    "        results.append(result)\n",
    "\n",
    "    df = pd.DataFrame(results)\n",
    "    return df.sort_values('hit_rate', ascending=False)\n",
    "\n",
    "\n",
    "models = {\n",
    "    \"titan-v1\": {  # Bedrock Titan Text V1 model\n",
    "        \"model\": BedrockEmbedding(model=\"amazon.titan-embed-text-v1\")\n",
    "    },\n",
    "    \"titan-v2\": {  # Bedrock Titan Text V2 model\n",
    "        \"model\": BedrockEmbedding(model=\"amazon.titan-embed-text-v2:0\")\n",
    "    },\n",
    "    \"bge\": { # Base BGE model with no instructions\n",
    "        \"model\": HuggingFaceEmbedding(base_model_local_dir),\n",
    "    },\n",
    "    \"ft-bge\": {  # Fine-tuned BGE model\n",
    "        \"model\": HuggingFaceEmbedding(finetuned_model_dir),\n",
    "        \"query_instruction\": query_instruction_for_retrieval,\n",
    "        \"passage_instruction\": passage_instruction_for_retrieval,\n",
    "    },\n",
    "    \"ft-mixed-bge\": {  # Fine-tuned mixed BGE model\n",
    "        \"model\": HuggingFaceEmbedding(mixed_model_output_dir),\n",
    "        \"query_instruction\": query_instruction_for_retrieval,\n",
    "        \"passage_instruction\": passage_instruction_for_retrieval,\n",
    "    },\n",
    "}\n",
    "\n",
    "results_df = evaluate_models(models, val_dataset)\n",
    "\n",
    "results_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2971e7b9-7987-42db-bea1-3043209da1fd",
   "metadata": {},
   "source": [
    "Note that better performance can be achieved by using the large version of the BGE model.\n",
    "\n",
    "We can see from the results that the fine-tuned model has the best Hit Rate, though we may want to use\n",
    "the mixed model, since it's close in performance, and should have greater generalization capability."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6927845-b6e7-4a80-b9cf-6f2f1b0ea427",
   "metadata": {},
   "source": [
    "## Upload the Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f57f5b1-17bf-44c4-9ab0-875a2a3042aa",
   "metadata": {},
   "source": [
    "#### Option 1. Upload to S3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "276afd5f-2fbe-4b12-9f80-dba510cbbb3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "from sagemaker.s3 import S3Uploader\n",
    "\n",
    "sagemaker_session = sagemaker.Session()\n",
    "s3_bucket = sagemaker_session.default_bucket()\n",
    "s3_prefix = f\"models/{base_model_local_dir}\"\n",
    "\n",
    "desired_s3_uri = f\"s3://{s3_bucket}/{s3_prefix}\"\n",
    "\n",
    "s3_model_uri = S3Uploader.upload(local_path=mixed_model_output_dir, desired_s3_uri=desired_s3_uri)\n",
    "print(s3_model_uri)\n",
    "\n",
    "upload_method = \"s3\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a20f9492-449b-4f00-99c9-b4046c813375",
   "metadata": {},
   "source": [
    "#### Option 2. Upload to Hugging Face Hub"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "def4ffdd-29c6-472e-8db8-0c76bc60ba86",
   "metadata": {},
   "source": [
    "Run the following, and paste in a valid HuggingFace Access Token with Write permission:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cee9644c-4145-4134-85b1-ebc40634adfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from huggingface_hub import notebook_login\n",
    "\n",
    "# notebook_login()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42ffef16-9705-4282-a653-82e20756386f",
   "metadata": {},
   "source": [
    "Now that we’re authenticated, we can upload our fine-tuned model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30bad3a3-8b7d-4e6b-8b03-a54b1db6cb00",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from huggingface_hub import HfApi\n",
    "# from transformers import AutoModel, AutoTokenizer\n",
    "\n",
    "# # Initialize the API\n",
    "# api = HfApi()\n",
    "\n",
    "# # Retrieve username\n",
    "# hf_username = api.whoami()[\"name\"]\n",
    "\n",
    "# model_id = f\"{hf_username}/{base_model_local_dir}\"\n",
    "# model = AutoModel.from_pretrained(mixed_model_output_dir)\n",
    "# tokenizer = AutoTokenizer.from_pretrained(mixed_model_output_dir)\n",
    "\n",
    "# model.push_to_hub(model_id)\n",
    "# tokenizer.push_to_hub(model_id)\n",
    "\n",
    "# upload_method = \"hf_hub\"\n",
    "\n",
    "# print(model_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60059a6f-8ffa-487f-81a6-0d11c3c414e4",
   "metadata": {},
   "source": [
    "## Part 6: Model Deployment\n",
    "\n",
    "In this section, we'll demonstrate how to deploy the model to Amazon SageMaker using a Hugging Face Text Embedding Inference Container. Text Embeddings Inference (TEI) is a high-performance toolkit for deploying and serving popular text embeddings and sequence classification models, including support for FlagEmbedding models. It will provide us with the fastest serving framework to deploy our model on SageMaker."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a43c7c72-0e84-4b8b-83c9-8190117173b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import sagemaker\n",
    "from datetime import datetime\n",
    "\n",
    "sagemaker_session = sagemaker.Session()\n",
    "# sagemaker session bucket -> used for uploading data, models and logs\n",
    "# sagemaker will automatically create this bucket if it does not exist\n",
    "sagemaker_session_bucket = None\n",
    "if sagemaker_session_bucket is None and sagemaker_session is not None:\n",
    "    # set to default bucket if a bucket name is not given\n",
    "    sagemaker_session_bucket = sagemaker_session.default_bucket()\n",
    "\n",
    "try:\n",
    "    role = sagemaker.get_execution_role()\n",
    "except ValueError:\n",
    "    iam = boto3.client('iam')\n",
    "    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']\n",
    "\n",
    "sagemaker_session = sagemaker.Session(default_bucket=sagemaker_session_bucket)\n",
    "\n",
    "print(f\"sagemaker role arn: {role}\")\n",
    "print(f\"sagemaker session region: {sagemaker_session.boto_region_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dde54315-26d5-46b5-81e1-3ee6e38e0183",
   "metadata": {},
   "source": [
    "### Retrieve the new Hugging Face Embedding Container"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0dee5de9-9838-4724-b1a3-30b34dacfac9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.huggingface import get_huggingface_llm_image_uri\n",
    "\n",
    "instance_type = \"ml.g5.xlarge\"\n",
    "\n",
    "# retrieve the image uri based on instance type\n",
    "def get_image_uri(instance_type):\n",
    "  key = \"huggingface-tei\" if instance_type.startswith(\"ml.g\") or instance_type.startswith(\"ml.p\") else \"huggingface-tei-cpu\"\n",
    "  return get_huggingface_llm_image_uri(key, version=\"1.2.3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12f594db-2b9c-40b5-bfc6-22e70ea8aa84",
   "metadata": {},
   "source": [
    "### Deploy model to Amazon SageMaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3930b235-81bc-4604-872f-da088fc42fb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from sagemaker.huggingface import HuggingFaceModel\n",
    "\n",
    "# sagemaker config\n",
    "instance_type = \"ml.g5.xlarge\"\n",
    "\n",
    "# Define Model and Endpoint configuration parameter\n",
    "config = {\n",
    "  \"POOLING\": \"cls\",\n",
    "  \"MAX_CONCURRENT_REQUESTS\": json.dumps(512),  # The maximum amount of concurrent requests for this particular deployment.\n",
    "  \"MAX_BATCH_TOKENS\": json.dumps(16384),  # Limits the number of tokens that can be processed in parallel during the generation\n",
    "}\n",
    "\n",
    "# Decide whether to deploy the model from S3 or Hugging Face Hub\n",
    "if upload_method == \"s3\":\n",
    "    config[\"HF_MODEL_ID\"] = \"/opt/ml/model\"\n",
    "elif upload_method == \"hf_hub\":\n",
    "    model_data = None\n",
    "    config[\"HF_MODEL_ID\"] = model_id\n",
    "\n",
    "print(config)\n",
    "\n",
    "# create HuggingFaceModel with the image uri\n",
    "emb_model = HuggingFaceModel(\n",
    "  role=role,\n",
    "  image_uri=get_image_uri(instance_type),\n",
    "  env=config,\n",
    "  model_data={'S3DataSource':{'S3Uri': s3_model_uri + \"/\",'S3DataType': 'S3Prefix','CompressionType': 'None'}},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db40aeef-a544-48c0-830c-3192ccbe079b",
   "metadata": {},
   "outputs": [],
   "source": [
    "resource_name = \"Demo-{}-{}\"\n",
    "\n",
    "tei_endpoint = emb_model.deploy(\n",
    "  initial_instance_count=1,\n",
    "  instance_type=instance_type,\n",
    "  endpoint_name=resource_name.format(\"TEI-Endpoint\", datetime.now().strftime(\"%Y-%m-%d-%H-%M-%S\")),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "488ead6d-909f-46a2-b3fc-7157eaea262d",
   "metadata": {},
   "source": [
    "### 5. Set autoscaling (optional)\n",
    "\n",
    "We can also add autoscaling:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b215c083-c024-4603-b211-dd836f6e95f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "asg = boto3.client('application-autoscaling')\n",
    "\n",
    "# Resource type is variant and the unique identifier is the resource ID.\n",
    "resource_id=f\"endpoint/{tei_endpoint.endpoint_name}/variant/AllTraffic\"\n",
    "\n",
    "# scaling configuration\n",
    "response = asg.register_scalable_target(\n",
    "    ServiceNamespace='sagemaker',\n",
    "    ResourceId=resource_id,\n",
    "    ScalableDimension='sagemaker:variant:DesiredInstanceCount', \n",
    "    MinCapacity=1,\n",
    "    MaxCapacity=4\n",
    ")\n",
    "\n",
    "response = asg.put_scaling_policy(\n",
    "    PolicyName=f'Request-ScalingPolicy-{tei_endpoint.endpoint_name}',\n",
    "    ServiceNamespace='sagemaker',\n",
    "    ResourceId=resource_id,\n",
    "    ScalableDimension='sagemaker:variant:DesiredInstanceCount',\n",
    "    PolicyType='TargetTrackingScaling',\n",
    "    TargetTrackingScalingPolicyConfiguration={\n",
    "        'TargetValue': 10.0, # Threshold\n",
    "        'PredefinedMetricSpecification': {\n",
    "            'PredefinedMetricType': 'SageMakerVariantInvocationsPerInstance',\n",
    "        },\n",
    "        'ScaleInCooldown': 300, # duration until scale in\n",
    "        'ScaleOutCooldown': 60 # duration between scale out\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b42dd8d7-2acb-41b0-8e88-346ba1e3c8bd",
   "metadata": {},
   "source": [
    "## Part 7. Model testing\n",
    "\n",
    "Finally, we can run the model, and optionally add instructions. Note that if the model was fine-tuned with instructions for queries and/or passages, you should match what was done during fine-tuning. In our case we used instructions for queries, but not for passages, so we can do the same while performing inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eab01599-5a6e-4685-b1d1-7aebbd0a0061",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_instruction_for_retrieval = \"Represent the Financial question for retrieving supporting documents: \"\n",
    "\n",
    "queries = [query_instruction_for_retrieval + query for query in similar_queries]\n",
    "\n",
    "example = {\n",
    "    \"inputs\": queries,\n",
    "    \"truncate\": True,\n",
    "}\n",
    "\n",
    "results = tei_endpoint.predict(example)\n",
    "\n",
    "# print some results\n",
    "print(f\"length of embeddings: {len(results[0])}\")\n",
    "print(f\"first 5 elements of embeddings: {results[0][:5]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a770ea77-f286-4fb0-a67a-e68de1fba861",
   "metadata": {},
   "source": [
    "And we can get the average response time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3e531ff-08af-4163-ad7b-6445b83d5a73",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "emb = tei_endpoint.predict(example)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9bac5ce-64ed-400c-afeb-03ef882daa75",
   "metadata": {},
   "source": [
    "Also we can test the throughput:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "120f1d5d-4121-4b88-b7cc-1121938153ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "import threading\n",
    "import time\n",
    "\n",
    "number_of_threads = 10\n",
    "number_of_requests = int(3900 // number_of_threads)\n",
    "print(f\"number of threads: {number_of_threads}\")\n",
    "print(f\"number of requests per thread: {number_of_requests}\")\n",
    "\n",
    "def send_requests():\n",
    "    for _ in range(number_of_requests):\n",
    "        # input counted at https://huggingface.co/spaces/Xenova/the-tokenizer-playground for 100 tokens\n",
    "        tei_endpoint.predict(data={\"inputs\": \"Hugging Face is a company and a popular platform in the field of natural language processing (NLP) and machine learning. They are known for their contributions to the development of state-of-the-art models for various NLP tasks and for providing a platform that facilitates the sharing and usage of pre-trained models. One of the key offerings from Hugging Face is the Transformers library, which is an open-source library for working with a variety of pre-trained transformer models, including those for text generation, translation, summarization, question answering, and more. The library is widely used in the research and development of NLP applications and is supported by a large and active community. Hugging Face also provides a model hub where users can discover, share, and download pre-trained models. Additionally, they offer tools and frameworks to make it easier for developers to integrate and use these models in their own projects. The company has played a significant role in advancing the field of NLP and making cutting-edge models more accessible to the broader community. Hugging Face also provides a model hub where users can discover, share, and download pre-trained models. Additionally, they offer tools and frameworks to make it easier for developers and ma\"})\n",
    "\n",
    "# Create multiple threads\n",
    "threads = [threading.Thread(target=send_requests) for _ in range(number_of_threads)]\n",
    "# start all threads\n",
    "start = time.time()\n",
    "[t.start() for t in threads]\n",
    "# wait for all threads to finish\n",
    "[t.join() for t in threads]\n",
    "print(f\"total time: {round(time.time() - start)} seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60e12bed-e266-45f4-a9ab-d90675d0eb09",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"https://console.aws.amazon.com/cloudwatch/home?region={sagemaker_session.boto_region_name}#metricsV2:graph=~(metrics~(~(~'AWS*2fSageMaker~'ModelLatency~'EndpointName~'{tei_endpoint.endpoint_name}~'VariantName~'AllTraffic))~view~'timeSeries~stacked~false~region~'{sagemaker_session.boto_region_name}~start~'-PT5M~end~'P0D~stat~'Average~period~30);query=~'*7bAWS*2fSageMaker*2cEndpointName*2cVariantName*7d*20{tei_endpoint.endpoint_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "438e03fb-b234-4733-b9e8-ffaa22f546c3",
   "metadata": {},
   "source": [
    "### LangChain integration\n",
    "\n",
    "As a bonus, we can create a custom LangChain content handler and instantiate a SagemakerEndpointEmbeddings model for use with LangChain applications:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b26c96ce-7b5f-48a9-abec-3f05a8fbd7d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from typing import Dict, List\n",
    "\n",
    "from botocore.config import Config\n",
    "from langchain.embeddings import SagemakerEndpointEmbeddings\n",
    "from langchain.embeddings.sagemaker_endpoint import EmbeddingsContentHandler\n",
    "\n",
    "class BGEContentHandler(EmbeddingsContentHandler):\n",
    "    \"\"\"\n",
    "    We'll create a custom LangChain content handler for the BGE model which is deployed to a SageMaker endpoint.\n",
    "    \"\"\"\n",
    "\n",
    "    content_type = \"application/json\"\n",
    "    accepts = \"application/json\"\n",
    "\n",
    "    def transform_input(self, inputs: list[str], model_kwargs: Dict) -> bytes:\n",
    "        \"\"\"\n",
    "        Transforms the input into bytes that can be consumed by SageMaker endpoint.\n",
    "        Args:\n",
    "            inputs: List of input strings.\n",
    "            model_kwargs: Additional keyword arguments to be passed to the endpoint.\n",
    "        Returns:\n",
    "            The transformed bytes input.\n",
    "        \"\"\"\n",
    "        input_str = json.dumps({\"inputs\": inputs, **model_kwargs})\n",
    "        return input_str.encode(\"utf-8\")\n",
    "\n",
    "    def transform_output(self, output: bytes) -> List[List[float]]:\n",
    "        \"\"\"\n",
    "        Transforms the bytes output from the endpoint into a list of embeddings.\n",
    "        Args:\n",
    "            output: The bytes output from SageMaker endpoint.\n",
    "        Returns:\n",
    "            The transformed output - list of embeddings\n",
    "        Note:\n",
    "            The length of the outer list is the number of input strings.\n",
    "            The length of the inner lists is the embedding dimension.\n",
    "        \"\"\"\n",
    "        return json.loads(output.read().decode(\"utf-8\"))\n",
    "\n",
    "content_handler = BGEContentHandler()\n",
    "\n",
    "sagemaker_config = Config(connect_timeout=180, read_timeout=180, retries={\"max_attempts\": 30})\n",
    "\n",
    "sagemaker_client = boto3.client(\n",
    "    service_name=\"sagemaker-runtime\",\n",
    "    region_name=sagemaker_session.boto_region_name,\n",
    "    config=sagemaker_config,\n",
    ")\n",
    "\n",
    "embeddings = SagemakerEndpointEmbeddings(\n",
    "    endpoint_name=tei_endpoint.endpoint_name,\n",
    "    client=sagemaker_client,\n",
    "    content_handler=content_handler,\n",
    ")\n",
    "\n",
    "embeddings.model_kwargs = {\"truncate\": True}\n",
    "\n",
    "results = embeddings.embed_documents(queries)\n",
    "\n",
    "assert len(results) == len(queries)\n",
    "assert len(results[0]) == 768 or len(results[0]) == 1024\n",
    "\n",
    "print(results[0][:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de3958bb-9b3a-4a8c-8119-cdabd34912c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get embeddings for similar documents and queries\n",
    "similar_document_embeddings = embeddings.embed_documents(similar_documents)\n",
    "similar_query_embeddings = embeddings.embed_documents(similar_queries)\n",
    "\n",
    "# Get embeddings for different documents and queries\n",
    "different_document_embeddings = embeddings.embed_documents(different_documents)\n",
    "different_query_embeddings = embeddings.embed_documents(different_queries)\n",
    "\n",
    "# Calculate cosine similarity for similar document-query pairs\n",
    "cosine_sim_pairs_similar = []\n",
    "for doc_emb, query_emb in zip(similar_document_embeddings, similar_query_embeddings):\n",
    "    cosine_sim = cosine_similarity([doc_emb], [query_emb])\n",
    "    cosine_sim_pairs_similar.append(cosine_sim[0][0])\n",
    "\n",
    "# Calculate cosine similarity for different document-query pairs\n",
    "cosine_sim_pairs_different = []\n",
    "for doc_emb, query_emb in zip(different_document_embeddings, different_query_embeddings):\n",
    "    cosine_sim = cosine_similarity([doc_emb], [query_emb])\n",
    "    cosine_sim_pairs_different.append(cosine_sim[0][0])\n",
    "\n",
    "# Print the cosine similarity for each similar document-query pair\n",
    "print(\"Cosine Similarity for Similar Document-Query Pairs:\")\n",
    "for i, sim in enumerate(cosine_sim_pairs_similar):\n",
    "    print(f\"Pair {i+1}: Cosine Similarity = {sim}\")\n",
    "\n",
    "# Print the cosine similarity for each different document-query pair\n",
    "print(\"\\nCosine Similarity for Different Document-Query Pairs:\")\n",
    "for i, sim in enumerate(cosine_sim_pairs_different):\n",
    "    print(f\"Pair {i+1}: Cosine Similarity = {sim}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bbe42d1-c8f0-47f1-a575-fbd8da215a1a",
   "metadata": {},
   "source": [
    "## Clean up\n",
    "\n",
    "Finally, when we're done with the model, we can delete it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40a09145-b36b-4505-b0d9-8e66cf2e518d",
   "metadata": {},
   "outputs": [],
   "source": [
    "tei_endpoint.delete_model()\n",
    "tei_endpoint.delete_endpoint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c40b7f99-d812-4b17-98b8-c91613d56078",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "In this post, we walked through the process of fine-tuning a BGE embedding model using synthetic data generated from Amazon Bedrock. We covered key steps including generating high-quality synthetic data, fine-tuning the model, evaluating performance, and deploying the optimized model using Amazon SageMaker.\n",
    "\n",
    "By leveraging synthetic data and advanced fine-tuning techniques like hard negative mining and model merging, you can significantly enhance the performance of embedding models for your specific use cases. This approach is especially valuable when real-world data is limited or difficult to obtain.\n",
    "\n",
    "To get started, we encourage you to experiment with the code and techniques demonstrated in this post. Adapt them to your own datasets and models to unlock performance improvements in your applications. You can find all the code used in this post in our GitHub repository.\n",
    "\n",
    "Related Links:\n",
    "\n",
    "* [LlamaIndex Finetune Embeddings Example](https://docs.llamaindex.ai/en/stable/examples/finetuning/embeddings/finetune_embedding/)\n",
    "* [BAAI/bge-base-en-v1.5](https://huggingface.co/BAAI/bge-base-en-v1.5)\n",
    "* [FlagEmbedding](https://github.com/FlagOpen/FlagEmbedding)\n",
    "* [LM-Cocktail](https://github.com/FlagOpen/FlagEmbedding/tree/master/LM_Cocktail)\n",
    "* [AWS Machine Learning Blog](https://aws.amazon.com/blogs/machine-learning/)\n",
    "* [GitHub Repository](URL_TO_GITHUB_REPOSITORY)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ft-embedding-blog",
   "language": "python",
   "name": "ft-embedding-blog"
  },
  "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}