discharge_notes_summarization.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# HealTac 2024 Tutorial\n",
        "## Instruction Tuning for Discharge Notes Summarization\n",
        "\n",
        "- Yunsoo Kim (yunsoo.kim.23@ucl.ac.uk), Jinge Wu (jinge.wu.20@ucl.ac.uk), Honghan Wu (honghan.wu@ucl.ac.uk)\n",
        "\n",
        "<a target=\"_blank\" href=\"https://colab.research.google.com/github/knowlab/healtac2024_tutorial/blob/main/discharge_notes_summarization.ipynb\">\n",
        "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
        "</a>"
      ],
      "metadata": {
        "id": "PgS8fQu0dju7"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Set the runtime to be T4 GPU.  \n",
        "\n",
        "We will get started with installing packages and downloading the model because they take some time."
      ],
      "metadata": {
        "id": "vcXEvrKFeVG5"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Instruction Fine-tuning\n",
        "\n",
        "![](https://substackcdn.com/image/fetch/f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fceb4e28d-424d-450d-a29e-598d883d1fb2_1229x500.png)"
      ],
      "metadata": {
        "id": "tKHt5hZ9RS0g"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## PEFT: Parameter Efficient Fine-tuning\n",
        "\n",
        "### LoRA: Low Rank Adaptation For Fine-tuning Large Models\n",
        "\n",
        "LoRA presented an effective solution to this problem by decomposing the update matrix during finetuing.\n",
        "\n",
        "LoRA proposes representing ( Δ W ) as the product of two smaller matrices, ( A ) and ( B ), with a lower rank. The updated weight matrix ( W’ ) thus becomes:\n",
        "\n",
        "$[ W’ = W + BA ]$\n",
        "\n",
        "In this equation, ( W ) remains frozen (i.e., it is not updated during training). The matrices ( B ) and ( A ) are of lower dimensionality, with their product ( BA ) representing a low-rank approximation of ( Δ W ).\n",
        "\n",
        "![](https://miro.medium.com/v2/resize:fit:720/format:webp/1*Vp90_0_4B3eOGuzU8lF-dg.png)"
      ],
      "metadata": {
        "id": "ukxbSVXURVQg"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### QLoRA: Efficient Fine-tuning With Quantization\n",
        "\n",
        "Quantization is the process of representing weights, bias and activations in neural networks using lower-precision data types, such as 8-bit integers (int8), instead of the conventional 32-bit floating point (float32) representation. By doing so, it significantly reduces the memory footprint and computational demands during inference, enabling deployment on resource-constrained devices.\n",
        "\n",
        "QLoRA develops quantization of the parameters down to 4-bit with Double Quantization of the scaling factors down to 8-bit. This brings down the memory requirements to fine-tune a 65B parameter model from ~780GB to <48GB, which makes it achievable on a single GPU machine.\n",
        "\n",
        "![](https://cdn-images-1.medium.com/max/800/0*F1QQgKRBxpmgWqq6)\n",
        "\n",
        "![](https://cdn-images-1.medium.com/max/800/0*clxn__Qp8ekzrr_t)"
      ],
      "metadata": {
        "id": "zW8TtsFORVOJ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Gradient Accumulation\n",
        "\n",
        "We can’t fit a large batch size but how about we keep track and collect gradients calculated for each batch and act on them together? Gradient Accumulation is a technique used to improve memory efficiency and stabilize training in neural networks by accumulating gradients over multiple batches before updating the model parameters.\n",
        "\n",
        "\n",
        "![](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r1.8/tutorials/experts/source_zh_cn/others/images/GradientAccumulation1.png)"
      ],
      "metadata": {
        "id": "ql0-y6LSRVLs"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Hardware Requirement (Estimated)\n",
        "\n",
        "| Method            | Bits | 7B    | 13B   | 30B   | 70B    | 110B   | 8x7B  | 8x22B  |\n",
        "| ----------------- | ---- | ----- | ----- | ----- | ------ | ------ | ----- | ------ |\n",
        "| Full              | AMP  | 120GB | 240GB | 600GB | 1200GB | 2000GB | 900GB | 2400GB |\n",
        "| Full              | 16   | 60GB  | 120GB | 300GB | 600GB  | 900GB  | 400GB | 1200GB |\n",
        "| Freeze            | 16   | 20GB  | 40GB  | 80GB  | 200GB  | 360GB  | 160GB | 400GB  |\n",
        "| LoRA/GaLore/BAdam | 16   | 16GB  | 32GB  | 64GB  | 160GB  | 240GB  | 120GB | 320GB  |\n",
        "| QLoRA             | 8    | 10GB  | 20GB  | 40GB  | 80GB   | 140GB  | 60GB  | 160GB  |\n",
        "| QLoRA             | 4    | 6GB   | 12GB  | 24GB  | 48GB   | 72GB   | 30GB  | 96GB   |\n",
        "| QLoRA             | 2    | 4GB   | 8GB   | 16GB  | 24GB   | 48GB   | 18GB  | 48GB   |"
      ],
      "metadata": {
        "id": "8vOnSZP8RVJF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Run nvidia-smi to check the gpu resource\n",
        "!nvidia-smi"
      ],
      "metadata": {
        "id": "Aoe8KWfRfKQC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JLdwofkkdfEt"
      },
      "outputs": [],
      "source": [
        "# Install the required packages\n",
        "!pip install -q accelerate==0.25.0 peft==0.6.2 bitsandbytes==0.41.1 transformers==4.36.2 trl==0.7.4 einops gradio"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Import packages\n",
        "import torch\n",
        "from datasets import load_dataset # loading the dataset\n",
        "from transformers import (\n",
        "    AutoModelForCausalLM,\n",
        "    AutoTokenizer,\n",
        "    BitsAndBytesConfig,\n",
        "    TrainingArguments,\n",
        ")\n",
        "from peft import LoraConfig # for LoRA\n",
        "from trl import SFTTrainer, DataCollatorForCompletionOnlyLM # Trainer and DataCollator\n",
        "import gradio as gr # for deployment"
      ],
      "metadata": {
        "id": "7D953H5Xedjr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Define Quantization Config\n",
        "quantization_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_compute_dtype=torch.float16,\n",
        "    bnb_4bit_quant_type=\"nf4\",\n",
        ")\n",
        "\n",
        "# Load Model\n",
        "# We will use Phi-2\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    \"microsoft/phi-2\",\n",
        "    trust_remote_code=True,\n",
        "    quantization_config=quantization_config,\n",
        "    device_map=\"auto\",\n",
        "    revision=\"refs/pr/23\"\n",
        ")\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained('microsoft/phi-2')\n",
        "tokenizer.pad_token = tokenizer.eos_token\n",
        "tokenizer.padding_sight = \"right\"\n",
        "\n",
        "# Load Dataset\n",
        "dataset = load_dataset(\"bluesky333/synthetic_discharge_summ\")"
      ],
      "metadata": {
        "id": "g4dbN3Bzefu8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "PEFT Markup Description\n",
        "\n",
        "LoRA\n",
        "\n",
        "Quantization"
      ],
      "metadata": {
        "id": "VYRWZT__j7RP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Let's have a look at the train dataset\n",
        "print(dataset['train'])\n",
        "dataset['train'][0]"
      ],
      "metadata": {
        "id": "OMoZKdyTe4Hd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# the test dataset\n",
        "print(dataset['test'])\n",
        "dataset['test'][0]"
      ],
      "metadata": {
        "id": "8EI5xgRN2LF_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# We make this dataset to phi-2 compatible\n",
        "# Phi-2 instruction-answer format: \"Instruct: <prompt>\\nOutput:\"\n",
        "\n",
        "# Make your own prompt!\n",
        "prompt_template=\"\"\"Instruct: Please write down your own prompt.\n",
        "For instance, you can insert the note as {{note}}\n",
        "{note}\n",
        "Model should answer to {{question}} based on the note.\n",
        "{question}\n",
        "You should maintain the phi-2 format\n",
        "Accordingly, the last line must be like the below.\n",
        "Do not forget to insert a new line between your prompt and 'Output'!\n",
        "Output: {answer}\n",
        "\"\"\"\n",
        "\n",
        "prompt_template=\"\"\"Instruct: Answer the question about the following clinical note. \\n{note}. Question: \\n{question}.\n",
        "Output: {answer}\n",
        "\"\"\"\n",
        "\n",
        "\n",
        "# Should get Dict[List] as input, return list of prompts\n",
        "def format_dataset(samples):\n",
        "    outputs = []\n",
        "    for _, note, question, answer in zip(*samples.values()):\n",
        "        out = prompt_template.format(note=note, question=question, answer=answer)\n",
        "        outputs.append(out)\n",
        "    return outputs\n",
        "\n",
        "sample_input = format_dataset({k: [v] for k, v in dataset['train'][0].items()})[0]\n",
        "print(sample_input)\n",
        "print(\"*\"*20)\n",
        "\n",
        "# Sanity Check\n",
        "prompt_len = len(tokenizer.encode(prompt_template))\n",
        "if prompt_len > 180:\n",
        "    raise ValueError(f\"Your prompt is too long! Please reduce the length from {prompt_len} to 180 tokens\")\n",
        "print(f\"Prompt Length: {prompt_len} tokens\")"
      ],
      "metadata": {
        "id": "ujW2fAlue6NZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "sample_idx = 0\n",
        "sample_input = format_dataset({k: [v] for k, v in dataset['train'][sample_idx].items()})[0].split('Output: ')[0]\n",
        "input_ids = tokenizer.encode(sample_input, return_tensors='pt').to('cuda')\n",
        "with torch.no_grad():\n",
        "  output = model.generate(input_ids=input_ids,\n",
        "                            max_length=512,\n",
        "                            use_cache=True,\n",
        "                            temperature=0.,\n",
        "                            eos_token_id=tokenizer.eos_token_id\n",
        "  )\n",
        "print(tokenizer.decode(output.to('cpu')[0], skip_special_tokens=True))"
      ],
      "metadata": {
        "id": "70z9VMYUfUyE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Then, let's define dataset.\n",
        "response_template = \"Output:\"\n",
        "collator = DataCollatorForCompletionOnlyLM(response_template, tokenizer=tokenizer)\n",
        "\n",
        "train_dataset = dataset['train']"
      ],
      "metadata": {
        "id": "ApZOSbzafWP4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# SFTTrainer Do everything else for you!\n",
        "\n",
        "lora_config=LoraConfig(\n",
        "    r=4,\n",
        "    task_type=\"CAUSAL_LM\",\n",
        "    target_modules= [\"Wqkv\", \"fc1\", \"fc2\" ]\n",
        ")\n",
        "\n",
        "training_args = TrainingArguments(\n",
        "    output_dir=\"./results\",\n",
        "    num_train_epochs=3,\n",
        "    fp16=True,\n",
        "    per_device_train_batch_size=2,\n",
        "    gradient_accumulation_steps=8,\n",
        "    learning_rate=1e-4,\n",
        "    optim=\"paged_adamw_32bit\",\n",
        "    save_strategy=\"no\",\n",
        "    warmup_ratio=0.03,\n",
        "    logging_steps=1,\n",
        "    lr_scheduler_type=\"cosine\",\n",
        "    report_to=None,\n",
        "    gradient_checkpointing=True\n",
        ")\n",
        "\n",
        "trainer = SFTTrainer(\n",
        "    model,\n",
        "    training_args,\n",
        "    train_dataset=train_dataset,\n",
        "    formatting_func=format_dataset,\n",
        "    data_collator=collator,\n",
        "    peft_config=lora_config,\n",
        "    max_seq_length=512,\n",
        "    tokenizer=tokenizer,\n",
        ")"
      ],
      "metadata": {
        "id": "lGJAv8hLfXuR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Run Training\n",
        "trainer.train()"
      ],
      "metadata": {
        "id": "pPceawKOfZcc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Wrap-up Training\n",
        "model = trainer.model\n",
        "model.eval()\n",
        "\n",
        "note_samples = dataset['test']['note']\n",
        "\n",
        "def inference(note, question, model):\n",
        "    prompt = prompt_template.format(note=note, question=question, answer=\"\")\n",
        "    tokens = tokenizer.encode(prompt, return_tensors=\"pt\").to('cuda')\n",
        "    outs = model.generate(input_ids=tokens,\n",
        "                          max_length=512,\n",
        "                          use_cache=True,\n",
        "                          temperature=1.,\n",
        "                          eos_token_id=tokenizer.eos_token_id\n",
        "                          )\n",
        "    output_text = tokenizer.decode(outs.to('cpu')[0], skip_special_tokens=True)\n",
        "    return output_text[len(prompt):]\n",
        "\n",
        "\n",
        "def compare_models(note, question):\n",
        "    with torch.no_grad():\n",
        "        asc_answer = inference(note, question, trainer.model)\n",
        "        with model.disable_adapter():\n",
        "            phi_answer = inference(note, question, trainer.model)\n",
        "    return asc_answer, phi_answer\n",
        "\n",
        "demo = gr.Interface(fn=compare_models, inputs=[gr.Dropdown(note_samples), \"text\"], outputs=[gr.Textbox(label=\"Trained\"), gr.Textbox(label=\"Phi-2\")])\n",
        "demo.launch()"
      ],
      "metadata": {
        "id": "Nz_Z_Gj1favb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "ogayHIqhHKrU"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}