medical_question_llm.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "A100",
      "authorship_tag": "ABX9TyP/e/kcAm9t7DY91mDSzwM8",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/e-gluzman/biomedical-ai/blob/main/medical_question_llm.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this notebook, I showcase finetuning a state of the art open source Large Language Model (Mistral) to solve Medical Questions using the Medical Question Dataset.\n",
        "\n",
        "I load the LLM from Hugging Face, create prompts from training data, fine-tune the model on the medical question dataset and then apply it to medical question answering.\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "Some useful resources: <br>\n",
        "https://www.datacamp.com/tutorial/mistral-7b-tutorial\n",
        "https://saankhya.medium.com/large-language-models-llms-a-comprehensive-guide-58ce825c8c0b"
      ],
      "metadata": {
        "id": "7iGspgmgSaJR"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kPzgxO0SH5Mx"
      },
      "outputs": [],
      "source": [
        "!pip install -q accelerate peft bitsandbytes\n",
        "!pip install -q git+https://github.com/huggingface/transformers\n",
        "!pip install -q trl py7zr auto-gptq optimum"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount(\"/content/drive\")"
      ],
      "metadata": {
        "id": "-d7FIqDHII8V"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import pandas as pd\n",
        "from datasets import load_dataset\n",
        "from transformers import AutoModelForSequenceClassification, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, GPTQConfig, TrainingArguments, GenerationConfig\n",
        "import torch\n",
        "from datasets import Dataset\n",
        "from peft import LoraConfig, AutoPeftModelForCausalLM, prepare_model_for_kbit_training, get_peft_model, PeftConfig\n",
        "from trl import SFTTrainer\n",
        "import os\n",
        "import tqdm\n",
        "import json\n",
        "from huggingface_hub import notebook_login\n"
      ],
      "metadata": {
        "id": "JB1_jXPQ1fp_"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "notebook_login()"
      ],
      "metadata": {
        "id": "dtb-gWg_OrZW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# lets load and view the traning data\n",
        "dataset = load_dataset(\"medmcqa\")"
      ],
      "metadata": {
        "id": "agQGR0mnIoMs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_df = pd.DataFrame(dataset['train'])"
      ],
      "metadata": {
        "id": "LLGwnoJJTYDm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_df.head()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 467
        },
        "id": "5D09M33lTZ__",
        "outputId": "ca6b1ab7-07ea-46f6-e408-14ecf2d0763f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "                                     id  \\\n",
              "0  e9ad821a-c438-4965-9f77-760819dfa155   \n",
              "1  e3d3c4e1-4fb2-45e7-9f88-247cc8f373b3   \n",
              "2  5c38bea6-787a-44a9-b2df-88f4218ab914   \n",
              "3  cdeedb04-fbe9-432c-937c-d53ac24475de   \n",
              "4  dc6794a3-b108-47c5-8b1b-3b4931577249   \n",
              "\n",
              "                                            question  \\\n",
              "0  Chronic urethral obstruction due to benign pri...   \n",
              "1  Which vitamin is supplied from only animal sou...   \n",
              "2  All of the following are surgical options for ...   \n",
              "3  Following endaerectomy on the right common car...   \n",
              "4   Growth hormone has its effect on growth through?   \n",
              "\n",
              "                          opa                        opb              opc  \\\n",
              "0                 Hyperplasia                  Hyperophy          Atrophy   \n",
              "1                   Vitamin C                 Vitamin B7      Vitamin B12   \n",
              "2  Adjustable gastric banding  Biliopancreatic diversion  Duodenal Switch   \n",
              "3  Central aery of the retina          Infraorbital aery    Lacrimal aery   \n",
              "4                    Directly                      IG1-1        Thyroxine   \n",
              "\n",
              "                          opd  cop choice_type  \\\n",
              "0                    Dyplasia    2      single   \n",
              "1                   Vitamin D    2      single   \n",
              "2  Roux en Y Duodenal By pass    3       multi   \n",
              "3          Nasociliary aretry    0       multi   \n",
              "4      Intranuclear receptors    1      single   \n",
              "\n",
              "                                                 exp   subject_name  \\\n",
              "0  Chronic urethral obstruction because of urinar...        Anatomy   \n",
              "1  Ans. (c) Vitamin B12 Ref: Harrison's 19th ed. ...   Biochemistry   \n",
              "2  Ans. is 'd' i.e., Roux en Y Duodenal Bypass Ba...        Surgery   \n",
              "3  The central aery of the retina is a branch of ...  Ophthalmology   \n",
              "4  Ans. is 'b' i.e., IGI-1GH has two major functi...     Physiology   \n",
              "\n",
              "                   topic_name  \\\n",
              "0               Urinary tract   \n",
              "1       Vitamins and Minerals   \n",
              "2  Surgical Treatment Obesity   \n",
              "3                        None   \n",
              "4                        None   \n",
              "\n",
              "                                                text  \n",
              "0  \\n    Question:\\n    Chronic urethral obstruct...  \n",
              "1  \\n    Question:\\n    Which vitamin is supplied...  \n",
              "2  \\n    Question:\\n    All of the following are ...  \n",
              "3  \\n    Question:\\n    Following endaerectomy on...  \n",
              "4  \\n    Question:\\n    Growth hormone has its ef...  "
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-71dc212c-d582-4e06-84e6-9419a11ba3d2\" class=\"colab-df-container\">\n",
              "    <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>id</th>\n",
              "      <th>question</th>\n",
              "      <th>opa</th>\n",
              "      <th>opb</th>\n",
              "      <th>opc</th>\n",
              "      <th>opd</th>\n",
              "      <th>cop</th>\n",
              "      <th>choice_type</th>\n",
              "      <th>exp</th>\n",
              "      <th>subject_name</th>\n",
              "      <th>topic_name</th>\n",
              "      <th>text</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>e9ad821a-c438-4965-9f77-760819dfa155</td>\n",
              "      <td>Chronic urethral obstruction due to benign pri...</td>\n",
              "      <td>Hyperplasia</td>\n",
              "      <td>Hyperophy</td>\n",
              "      <td>Atrophy</td>\n",
              "      <td>Dyplasia</td>\n",
              "      <td>2</td>\n",
              "      <td>single</td>\n",
              "      <td>Chronic urethral obstruction because of urinar...</td>\n",
              "      <td>Anatomy</td>\n",
              "      <td>Urinary tract</td>\n",
              "      <td>\\n    Question:\\n    Chronic urethral obstruct...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>e3d3c4e1-4fb2-45e7-9f88-247cc8f373b3</td>\n",
              "      <td>Which vitamin is supplied from only animal sou...</td>\n",
              "      <td>Vitamin C</td>\n",
              "      <td>Vitamin B7</td>\n",
              "      <td>Vitamin B12</td>\n",
              "      <td>Vitamin D</td>\n",
              "      <td>2</td>\n",
              "      <td>single</td>\n",
              "      <td>Ans. (c) Vitamin B12 Ref: Harrison's 19th ed. ...</td>\n",
              "      <td>Biochemistry</td>\n",
              "      <td>Vitamins and Minerals</td>\n",
              "      <td>\\n    Question:\\n    Which vitamin is supplied...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>5c38bea6-787a-44a9-b2df-88f4218ab914</td>\n",
              "      <td>All of the following are surgical options for ...</td>\n",
              "      <td>Adjustable gastric banding</td>\n",
              "      <td>Biliopancreatic diversion</td>\n",
              "      <td>Duodenal Switch</td>\n",
              "      <td>Roux en Y Duodenal By pass</td>\n",
              "      <td>3</td>\n",
              "      <td>multi</td>\n",
              "      <td>Ans. is 'd' i.e., Roux en Y Duodenal Bypass Ba...</td>\n",
              "      <td>Surgery</td>\n",
              "      <td>Surgical Treatment Obesity</td>\n",
              "      <td>\\n    Question:\\n    All of the following are ...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>cdeedb04-fbe9-432c-937c-d53ac24475de</td>\n",
              "      <td>Following endaerectomy on the right common car...</td>\n",
              "      <td>Central aery of the retina</td>\n",
              "      <td>Infraorbital aery</td>\n",
              "      <td>Lacrimal aery</td>\n",
              "      <td>Nasociliary aretry</td>\n",
              "      <td>0</td>\n",
              "      <td>multi</td>\n",
              "      <td>The central aery of the retina is a branch of ...</td>\n",
              "      <td>Ophthalmology</td>\n",
              "      <td>None</td>\n",
              "      <td>\\n    Question:\\n    Following endaerectomy on...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>dc6794a3-b108-47c5-8b1b-3b4931577249</td>\n",
              "      <td>Growth hormone has its effect on growth through?</td>\n",
              "      <td>Directly</td>\n",
              "      <td>IG1-1</td>\n",
              "      <td>Thyroxine</td>\n",
              "      <td>Intranuclear receptors</td>\n",
              "      <td>1</td>\n",
              "      <td>single</td>\n",
              "      <td>Ans. is 'b' i.e., IGI-1GH has two major functi...</td>\n",
              "      <td>Physiology</td>\n",
              "      <td>None</td>\n",
              "      <td>\\n    Question:\\n    Growth hormone has its ef...</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "    <div class=\"colab-df-buttons\">\n",
              "\n",
              "  <div class=\"colab-df-container\">\n",
              "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-71dc212c-d582-4e06-84e6-9419a11ba3d2')\"\n",
              "            title=\"Convert this dataframe to an interactive table.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
              "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "\n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    .colab-df-buttons div {\n",
              "      margin-bottom: 4px;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "    <script>\n",
              "      const buttonEl =\n",
              "        document.querySelector('#df-71dc212c-d582-4e06-84e6-9419a11ba3d2 button.colab-df-convert');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      async function convertToInteractive(key) {\n",
              "        const element = document.querySelector('#df-71dc212c-d582-4e06-84e6-9419a11ba3d2');\n",
              "        const dataTable =\n",
              "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                    [key], {});\n",
              "        if (!dataTable) return;\n",
              "\n",
              "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "          + ' to learn more about interactive tables.';\n",
              "        element.innerHTML = '';\n",
              "        dataTable['output_type'] = 'display_data';\n",
              "        await google.colab.output.renderOutput(dataTable, element);\n",
              "        const docLink = document.createElement('div');\n",
              "        docLink.innerHTML = docLinkHtml;\n",
              "        element.appendChild(docLink);\n",
              "      }\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "\n",
              "<div id=\"df-73a8ed57-d039-4a20-8ea1-caf4cc260ab9\">\n",
              "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-73a8ed57-d039-4a20-8ea1-caf4cc260ab9')\"\n",
              "            title=\"Suggest charts\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "     width=\"24px\">\n",
              "    <g>\n",
              "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
              "    </g>\n",
              "</svg>\n",
              "  </button>\n",
              "\n",
              "<style>\n",
              "  .colab-df-quickchart {\n",
              "      --bg-color: #E8F0FE;\n",
              "      --fill-color: #1967D2;\n",
              "      --hover-bg-color: #E2EBFA;\n",
              "      --hover-fill-color: #174EA6;\n",
              "      --disabled-fill-color: #AAA;\n",
              "      --disabled-bg-color: #DDD;\n",
              "  }\n",
              "\n",
              "  [theme=dark] .colab-df-quickchart {\n",
              "      --bg-color: #3B4455;\n",
              "      --fill-color: #D2E3FC;\n",
              "      --hover-bg-color: #434B5C;\n",
              "      --hover-fill-color: #FFFFFF;\n",
              "      --disabled-bg-color: #3B4455;\n",
              "      --disabled-fill-color: #666;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart {\n",
              "    background-color: var(--bg-color);\n",
              "    border: none;\n",
              "    border-radius: 50%;\n",
              "    cursor: pointer;\n",
              "    display: none;\n",
              "    fill: var(--fill-color);\n",
              "    height: 32px;\n",
              "    padding: 0;\n",
              "    width: 32px;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart:hover {\n",
              "    background-color: var(--hover-bg-color);\n",
              "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "    fill: var(--button-hover-fill-color);\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart-complete:disabled,\n",
              "  .colab-df-quickchart-complete:disabled:hover {\n",
              "    background-color: var(--disabled-bg-color);\n",
              "    fill: var(--disabled-fill-color);\n",
              "    box-shadow: none;\n",
              "  }\n",
              "\n",
              "  .colab-df-spinner {\n",
              "    border: 2px solid var(--fill-color);\n",
              "    border-color: transparent;\n",
              "    border-bottom-color: var(--fill-color);\n",
              "    animation:\n",
              "      spin 1s steps(1) infinite;\n",
              "  }\n",
              "\n",
              "  @keyframes spin {\n",
              "    0% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "      border-left-color: var(--fill-color);\n",
              "    }\n",
              "    20% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    30% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    40% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    60% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    80% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "    90% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "  }\n",
              "</style>\n",
              "\n",
              "  <script>\n",
              "    async function quickchart(key) {\n",
              "      const quickchartButtonEl =\n",
              "        document.querySelector('#' + key + ' button');\n",
              "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
              "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
              "      try {\n",
              "        const charts = await google.colab.kernel.invokeFunction(\n",
              "            'suggestCharts', [key], {});\n",
              "      } catch (error) {\n",
              "        console.error('Error during call to suggestCharts:', error);\n",
              "      }\n",
              "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
              "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
              "    }\n",
              "    (() => {\n",
              "      let quickchartButtonEl =\n",
              "        document.querySelector('#df-73a8ed57-d039-4a20-8ea1-caf4cc260ab9 button');\n",
              "      quickchartButtonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "    })();\n",
              "  </script>\n",
              "</div>\n",
              "\n",
              "    </div>\n",
              "  </div>\n"
            ]
          },
          "metadata": {},
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# lets generate instructions for model training and encode them into 'text' column\n",
        "def generate_prompt(x):\n",
        "    cop = 'Nothing'\n",
        "    if x['cop'] == 0:\n",
        "        cop = x['opa']\n",
        "    elif x['cop'] == 1:\n",
        "        cop = x['opb']\n",
        "    elif x['cop'] == 2:\n",
        "        cop = x['opc']\n",
        "    elif x['cop'] == 3:\n",
        "        cop = x['opd']\n",
        "    question = '{}\\nOptions:\\n1. {}\\n2. {}\\n3. {}\\n4. {}\\n'.format(x['question'], x['opa'], x['opb'], x['opc'], x['opd'])\n",
        "    answer = cop\n",
        "    prompt = f\"\"\"\n",
        "    Question:\n",
        "    {question}\n",
        "    [INST] Solve this post graduate medical entrance exam MCQ and provide the correct option. [/INST]\n",
        "    Answer: {answer} </s>\"\"\"\n",
        "    return prompt"
      ],
      "metadata": {
        "id": "jVxXCunXPVLP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_df = pd.DataFrame(dataset['train'])"
      ],
      "metadata": {
        "id": "Pd4-V8gBPW1J"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_df['text'] = train_df.apply(lambda x: generate_prompt(x),axis=1)"
      ],
      "metadata": {
        "id": "bMZ-Nwq7Ph9Y"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# here is an example prompt for our training\n",
        "train_df['text'][0]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 71
        },
        "id": "EDnA6W2ZSKRV",
        "outputId": "82c4c04f-6917-4dce-c248-c18b47ef3b83"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'\\n    Question:\\n    Chronic urethral obstruction due to benign prismatic hyperplasia can lead to the following change in kidney parenchyma\\nOptions:\\n1. Hyperplasia\\n2. Hyperophy\\n3. Atrophy\\n4. Dyplasia\\n\\n    [INST] Solve this post graduate medical entrance exam MCQ and provide the correct option. [/INST]\\n    Answer: Hyperophy </s>'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# this code loads the Mistral-7B LLM fine-tuned for instruction prompts\n",
        "model_id = \"mistralai/Mistral-7B-Instruct-v0.1\"\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
        "tokenizer.pad_token = tokenizer.unk_token\n",
        "tokenizer.pad_token_id =  tokenizer.unk_token_id\n",
        "tokenizer.padding_side = 'left'\n",
        "\n",
        "quantization_config_loading = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_quant_type=\"nf4\",\n",
        "    bnb_4bit_compute_dtype=torch.float16,\n",
        "    bnb_4bit_use_double_quant=False\n",
        ")\n",
        "\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "                                model_id,\n",
        "                                quantization_config=quantization_config_loading,\n",
        "                                device_map=\"auto\"\n",
        "                            )\n",
        "\n",
        "model.config.pad_token_id = tokenizer.pad_token_id\n",
        "model.config.use_cache=False\n",
        "model.config.pretraining_tp=1\n",
        "model.gradient_checkpointing_enable()\n",
        "model = prepare_model_for_kbit_training(model)"
      ],
      "metadata": {
        "id": "wL1GxpwoI0MS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Lora allows to us accelerate the fine-tuning of large models while consuming less memory.\n",
        "peft_config = LoraConfig(\n",
        "                    r=16,\n",
        "                    lora_alpha=32,\n",
        "                    lora_dropout=0.05,\n",
        "                    bias=\"none\",\n",
        "                    task_type=\"CAUSAL_LM\",\n",
        "                    target_modules=[\n",
        "                        \"q_proj\",\n",
        "                        \"k_proj\",\n",
        "                        \"v_proj\",\n",
        "                        \"o_proj\",\n",
        "                        \"gate_proj\",\n",
        "                        \"up_proj\",\n",
        "                        \"down_proj\",\n",
        "                        \"lm_head\",\n",
        "                    ]\n",
        "                )\n",
        "\n",
        "model = get_peft_model(model, peft_config)"
      ],
      "metadata": {
        "id": "CWtRU-7VKK3O"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "data = Dataset.from_pandas(train_df)"
      ],
      "metadata": {
        "id": "WfGJ5p46Q0j2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "training_arguments = TrainingArguments(\n",
        "                            output_dir=\"mistral-gptq-finetuned-medmcqa\",\n",
        "                            per_device_train_batch_size=8,\n",
        "                            gradient_accumulation_steps=1,\n",
        "                            optim=\"paged_adamw_32bit\",\n",
        "                            learning_rate=2e-4,\n",
        "                            lr_scheduler_type=\"cosine\",\n",
        "                            save_strategy=\"epoch\",\n",
        "                            logging_steps=50,\n",
        "                            num_train_epochs=1,\n",
        "                            max_steps=5000,\n",
        "                            fp16=True,\n",
        "                            push_to_hub=True\n",
        "                        )\n",
        "\n",
        "trainer = SFTTrainer(\n",
        "            model=model,\n",
        "            train_dataset=data,\n",
        "            peft_config=peft_config,\n",
        "            dataset_text_field=\"text\",\n",
        "            args=training_arguments,\n",
        "            tokenizer=tokenizer,\n",
        "            packing=False,\n",
        "            max_seq_length=512\n",
        "    )"
      ],
      "metadata": {
        "id": "1f775k2eKNTZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# train the model on our data and save to hub\n",
        "trainer.train()\n",
        "trainer.push_to_hub()\n"
      ],
      "metadata": {
        "id": "FNxw7rtBKQ5j"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "username = \"egluzman\"\n",
        "model_id = username + \"mistral-finetuned-medmcqa-2\"\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
        "\n",
        "model = AutoPeftModelForCausalLM.from_pretrained(\n",
        "    model_id,\n",
        "    low_cpu_mem_usage=True,\n",
        "    return_dict=True,\n",
        "    torch_dtype=torch.float16,\n",
        "    device_map=\"cuda\",\n",
        "    )"
      ],
      "metadata": {
        "id": "CmYGedr2lfr_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "generation_config = GenerationConfig(\n",
        "    do_sample=True,\n",
        "    top_k=1,\n",
        "    temperature=0.1,\n",
        "    max_new_tokens=25,\n",
        "    pad_token_id=tokenizer.pad_token_id\n",
        ")"
      ],
      "metadata": {
        "id": "cY7KpSMf2Ppd"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def generate_test_prompt(x):\n",
        "    question = '{}\\nOptions:\\n1. {}\\n2. {}\\n3. {}\\n4. {}\\n'.format(x['question'], x['opa'], x['opb'], x['opc'], x['opd'])\n",
        "    prompt = f\"\"\"\n",
        "    Question:\n",
        "    {question}\n",
        "    [INST] Solve this post graduate medical entrance exam MCQ and answer correctly. [/INST]\n",
        "    Answer: \"\"\"\n",
        "    return prompt\n"
      ],
      "metadata": {
        "id": "nfbprH3dP8Ez"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "val_data_df = pd.DataFrame(dataset['validation'])\n",
        "val_data_df['text'] = val_data_df.apply(lambda x: generate_test_prompt(x),axis=1)"
      ],
      "metadata": {
        "id": "HFW2Eh3TPExZ"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "example = val_data_df['text'][3]\n",
        "inputs = tokenizer(example, return_tensors=\"pt\").to(\"cuda\")\n",
        "\n",
        "outputs = model.generate(**inputs, generation_config=generation_config)\n",
        "answer = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
        "print(answer)"
      ],
      "metadata": {
        "id": "NkLdK0QKMprN",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "fefd9336-188f-438c-a1a2-8c8e8959a10f"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "    Question:\n",
            "    Axonal transport is:\n",
            "Options:\n",
            "1. Antegrade\n",
            "2. Retrograde\n",
            "3. Antegrade and retrograde\n",
            "4. None\n",
            "\n",
            "    [INST] Solve this post graduate medical entrance exam MCQ and answer correctly. [/INST]\n",
            "    Answer: 1. Antegrade\n",
            "\n",
            "    Axonal transport refers to the movement of vesicles and organelles along the ax\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Iterate through the set of questions and generate model answers\n",
        "def solve_question(question_prompt):\n",
        "    inputs = tokenizer(question_prompt, return_tensors=\"pt\", padding=True, truncation=True).to(\"cuda\")\n",
        "    outputs = model.generate(**inputs, generation_config=generation_config)\n",
        "    answer = tokenizer.batch_decode(outputs, skip_special_tokens=True)\n",
        "    return answer\n",
        "\n",
        "all_answers = []\n",
        "import re\n",
        "val_data_prompts = list(val_data_df['text'])\n",
        "for i in tqdm.tqdm(range(0, len(val_data_prompts), 16)):\n",
        "    question_prompts = val_data_prompts[i:i+16]\n",
        "    ans = solve_question(question_prompts)\n",
        "    ans_option = []\n",
        "    for text in ans:\n",
        "        ans_option.append(re.search(r'Answer: \\s*(.*)', text).group(1))\n",
        "    all_answers.extend(ans_option)"
      ],
      "metadata": {
        "id": "IMRTF9pJ1wH1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "correct_answers = []\n",
        "for i in range(len(val_data_df)):\n",
        "    if val_data_df['cop'][i] == 0:\n",
        "        correct_answers.append(val_data_df['opa'][i])\n",
        "    elif val_data_df['cop'][i] == 1:\n",
        "        correct_answers.append(val_data_df['opb'][i])\n",
        "    elif val_data_df['cop'][i] == 2:\n",
        "        correct_answers.append(val_data_df['opc'][i])\n",
        "    elif val_data_df['cop'][i] == 3:\n",
        "        correct_answers.append(val_data_df['opd'][i])\n",
        "correct_count = 0\n",
        "for i in range(len(val_data_df)):\n",
        "    correct_count += correct_answers[i] == all_answers[i]"
      ],
      "metadata": {
        "id": "UbdrWsAOQpWq"
      },
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Here is an example propmpt and the model's answers\n",
        "print(question_prompts[0])\n",
        "print('')\n",
        "print('The correct answer is:')\n",
        "print('')\n",
        "print(correct_answers[0])\n",
        "print('')\n",
        "print('The model answered:')\n",
        "print(all_answers[0])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jKyoOdmilD8T",
        "outputId": "d434529a-9407-4979-bd1c-5062c10cd4b2"
      },
      "execution_count": 43,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "    Question:\n",
            "    To remove centric interference, reduce:\n",
            "Options:\n",
            "1. Supporting cusps\n",
            "2. Central fossa\n",
            "3. Both of the above\n",
            "4. None\n",
            "\n",
            "    [INST] Solve this post graduate medical entrance exam MCQ and answer correctly. [/INST]\n",
            "    Answer: \n",
            "\n",
            "The correct answer is:\n",
            "\n",
            "Impulse through myelinated fibers is slower than non-myelinated fibers\n",
            "\n",
            "The model answered:\n",
            "4. Local anesthesia is effective only when the nerve is not covered by myelin sheath\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# calculate the model's accuracy\n",
        "print(f'The final score is {round(correct_count/len(val_data_df),2)} %!')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HDY7z7dZkL2c",
        "outputId": "a63cb336-418f-40bd-c483-4a2512c708d9"
      },
      "execution_count": 32,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The final score is 0.37 %!\n"
          ]
        }
      ]
    }
  ]
}