2024/07/FederatedLearning/Part02/L2/L2_Centralized_LLM_Fine-tuning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "64c1a2de",
   "metadata": {},
   "source": [
    "# Lesson 2: Centralized LLM Fine-tuning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c16707d",
   "metadata": {},
   "source": [
    "Welcome to Lesson 2!\n",
    "\n",
    "To access the `requirements.txt` and `utils` files for this course, go to `File` and click `Open`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af1d0e65",
   "metadata": {},
   "source": [
    "#### 1. Import packages and utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "000642f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
    "warnings.filterwarnings(\"ignore\", category=FutureWarning)\n",
    "import torch\n",
    "from datasets import load_dataset\n",
    "from trl import SFTTrainer\n",
    "from transformers import TrainingArguments\n",
    "from utils.utils import *\n",
    "from utils.LLM import LLM_pretrained, LLM_cen_partial\n",
    "from utils.LLM import get_fireworks_api_key,load_env"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e69d454",
   "metadata": {},
   "source": [
    "> Note: Throughout this course, we use Hydra, a framework for managing and configuring files. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6db1fbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load config\n",
    "cfg = get_config(\"centralized\")\n",
    "\n",
    "# Inspect the config\n",
    "print_config(cfg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a917f369",
   "metadata": {},
   "source": [
    "#### 2. Load the dataset\n",
    "  \n",
    "**Note:** For more information about the dataset, check here [medAlpaca](https://arxiv.org/abs/2304.08247)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f9fd8a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainset_full = load_dataset(cfg.dataset.name, split='train')\n",
    "train_test = trainset_full.train_test_split(test_size=0.9, seed=1234)\n",
    "train_dataset = train_test[\"train\"]\n",
    "\n",
    "train_dataset = format_dataset(train_dataset)\n",
    "\n",
    "print(train_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7679e745",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print an example from the dataset\n",
    "example_index = 9\n",
    "\n",
    "data_point = train_dataset[example_index]\n",
    "\n",
    "data_point"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2315d3c4",
   "metadata": {},
   "source": [
    "#### 3. Asking the LLM \n",
    "\n",
    "\n",
    "* First attempt \n",
    "\n",
    "Ask a pre-trained LLM a question in a specific domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c20902da",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm_pretrained = LLM_pretrained()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e6d9be1",
   "metadata": {},
   "source": [
    "* Write a prompt to ask the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf56f92d",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"How to predict the weather\"\n",
    "llm_pretrained.eval(prompt, verbose = False)\n",
    "llm_pretrained.print_response(verbose = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38b13c69",
   "metadata": {},
   "source": [
    "* Evaluate pre-trained model on Medical Q&A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b44d1846",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm_pretrained.eval(data_point['instruction'], verbose=True)\n",
    "\n",
    "llm_pretrained.print_response()\n",
    "\n",
    "ex_response = format_string(data_point['response'])\n",
    "print(f\"Expected output:\\n\\t{ex_response}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7d2f9f1",
   "metadata": {},
   "source": [
    "* Second attempt:\n",
    "\n",
    "Ask a fine tuned LLM a question in a specific domain."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ceeeb1e",
   "metadata": {},
   "source": [
    "* Set the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "265fcae9",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = get_model(cfg.model)\n",
    "\n",
    "\n",
    "trainable, all_parameters = model.get_nb_trainable_parameters()\n",
    "print(f\"Trainable parameters: {trainable}\")\n",
    "print(f\"All parameters: {all_parameters}\")\n",
    "print(f\"Trainable (%): {100*trainable / all_parameters:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06e06af5",
   "metadata": {},
   "source": [
    "* Define the tokenizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c22c284",
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "tokenizer, \n",
    "data_collator, \n",
    "format_prompts_fn\n",
    ") = get_tokenizer_and_data_collator_and_propt_formatting(\n",
    "        cfg.model.name, cfg.model.use_fast_tokenizer, cfg.train.padding_side\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78c5c437",
   "metadata": {},
   "source": [
    "* Define the finetune_centralised function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14ab4a64",
   "metadata": {},
   "outputs": [],
   "source": [
    "save_centralized = \"./my_centralized_model\"\n",
    "def finetune_centralised():\n",
    "\n",
    "    # The notebooks you are running in this course does not\n",
    "    # come with a GPU. However, we don't harcode \"cpu\" as the\n",
    "    # device to use in case you wish to download this notebook\n",
    "    # and run it on your own GPU.\n",
    "    use_cuda = torch.cuda.is_available()\n",
    "    training_arguments = TrainingArguments(\n",
    "                         **cfg.train.training_arguments,\n",
    "                         use_cpu=not(use_cuda),\n",
    "                         output_dir=save_centralized,\n",
    "    )\n",
    "\n",
    "    # Construct trainer\n",
    "    trainer = SFTTrainer(\n",
    "        tokenizer=tokenizer,\n",
    "        data_collator=data_collator,\n",
    "        formatting_func=format_prompts_fn,\n",
    "        max_seq_length=cfg.train.seq_length,\n",
    "        model=model,\n",
    "        args=training_arguments,\n",
    "        train_dataset=train_dataset,\n",
    "    )\n",
    "\n",
    "    # Do local training\n",
    "    trainer.train()\n",
    "\n",
    "    # Save the checkpoint\n",
    "    model.save_pretrained(save_centralized)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "218de263",
   "metadata": {},
   "source": [
    "* Let's do some fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd8a0f40",
   "metadata": {},
   "outputs": [],
   "source": [
    "finetune_centralised()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3ca3dec",
   "metadata": {},
   "source": [
    "> Note: The previous fine-tuning was performed on a smaller model. To fine-tune a larger model (e.g., a 7 billion parameter model), you can use the same code with a different configuration. Refer to the code snippet below:\n",
    "\n",
    "\n",
    "```\n",
    " cfg = get_config(\"centralized_full\")\n",
    "```\n",
    "\n",
    "It is recommended to execute all prior cells with this configuration. However, please ensure that your machine has at least one GPU."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef22e176",
   "metadata": {},
   "source": [
    "* Evaluate the Centrally fine-tuned LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6582cc5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm_cen = LLM_cen_partial()\n",
    "example_index = 9\n",
    "\n",
    "data_point = train_dataset[example_index]\n",
    "\n",
    "llm_cen.eval(data_point['instruction'], verbose=True)\n",
    "\n",
    "llm_cen.print_response()\n",
    "\n",
    "ex_response = format_string(data_point['response'])\n",
    "print(f\"Expected output:\\n\\t{ex_response}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e32f9bf",
   "metadata": {},
   "source": [
    "#### 4. Visualize results of prompting with pretrained LLM and fine-tuned LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec0ae775",
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_results(results=['7b/pretrained', '7b/cen_10'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2568a9f",
   "metadata": {},
   "source": [
    "#### Extra! \n",
    "\n",
    "#### Generate the data structure for systematic evaluation\n",
    "\n",
    "We used this code to evaluate the performance of the finetuned 7B LLM you tested in the previous code cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8128e02b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lauch evalution code\n",
    "# from utils.utils import inference, evaluate\n",
    "\n",
    "# Step 1: generate answers\n",
    "# To run inference on the pre-trained model:\n",
    "# inference(base_model_name_path=cfg.model.name, run_name=\"pretrained\")\n",
    "\n",
    "# To run inference on the centralised finetuned model:\n",
    "# inference(\n",
    "#           base_model_name_path=cfg.model.name, \n",
    "#           peft_path=path/to/your/checkpoint/directory, \n",
    "#           run_name=\"centralised_finetuned\",\n",
    "# )\n",
    "\n",
    "\n",
    "# Step 2: evaluation --- accuracy value will be printed\n",
    "# evaluate(run_name=\"pretrained\")\n",
    "# evaluate(run_name=\"centralised_finetuned\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}