LAB 2 - Fine Tuning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "adbf511d-14d2-4b64-893c-de6234aff00e",
   "metadata": {},
   "source": [
    "<img src=\"static/event-image-small.jpg\" alt=\"drawing\" width=\"800\"/> \n",
    "\n",
    "# **LAB 2: FINE TUNING**\n",
    "\n",
    "Fine-tuning refers to the process of taking a pre-trained language model and \n",
    "training it further on a specific task or domain to improve its performance on that task.  \n",
    "<br />\n",
    "It is an important technique used to adapt LLMs to specific tasks and domains.  \n",
    "<br />\n",
    "In this lab we will explore basic ways to fine tune large language models using\n",
    "open soure tools. First we look at an example of doing this by hand with the open source 🤗 Transformers\n",
    "Python library. Familiarity with the 🤗 Transformers package is helpful once we\n",
    "introduce additional tools with more flexibility, such as H2O LLM Studio  \n",
    "<br />\n",
    "In this notebook, we will explore how do fine-tune a foundational large language\n",
    "model such that it can generate LinkedIn posts in the style of known influencers\n",
    "on the platform. \n",
    "\n",
    "The prepared data set from lab 1 can be found here: `s3://h2o-public-test-data/generative-ai/`\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc1b0bc2",
   "metadata": {},
   "source": [
    "# Using Hugging Face \n",
    "\n",
    "Among open source tools, Hugging Face provides some of the best to understnd \n",
    "how language modeling works. Before taking a look at expert fine tuning with H2O LLM Studio,\n",
    " we will look at a brief example of using the `transformers` and `datasets` python libraries. \n",
    "\n",
    "Let's load in the WNLI data set from the General Language Understanding Evaluation (GLUE)\n",
    "benchmark. (https://gluebenchmark.com/)\n",
    "\n",
    "From the paper, `The Winograd Schema Challenge (Levesque et al., 2011) is a reading comprehension task\n",
    "in which a system must read a sentence with a pronoun and select the referent of that pronoun from\n",
    "a list of choices.`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "709bbb6d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading builder script: 100%|██████████| 28.8k/28.8k [00:00<00:00, 48.7MB/s]\n",
      "Downloading metadata: 100%|██████████| 28.7k/28.7k [00:00<00:00, 67.7MB/s]\n",
      "Downloading readme: 100%|██████████| 27.9k/27.9k [00:00<00:00, 51.8MB/s]\n",
      "Downloading data: 100%|██████████| 29.0k/29.0k [00:00<00:00, 49.3MB/s]\n",
      "Generating train split: 100%|██████████| 635/635 [00:00<00:00, 24750.56 examples/s]\n",
      "Generating validation split: 100%|██████████| 71/71 [00:00<00:00, 18337.17 examples/s]\n",
      "Generating test split: 100%|██████████| 146/146 [00:00<00:00, 22774.78 examples/s]\n"
     ]
    }
   ],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "# set flag for training environment\n",
    "TRAINING = True\n",
    "\n",
    "from datasets import load_dataset\n",
    "from transformers import AutoTokenizer, DataCollatorWithPadding\n",
    "\n",
    "raw_datasets = load_dataset(\"glue\", \"wnli\")\n",
    "checkpoint = \"bert-base-uncased\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5e9b89c8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['sentence1', 'sentence2', 'label', 'idx'],\n",
       "        num_rows: 635\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['sentence1', 'sentence2', 'label', 'idx'],\n",
       "        num_rows: 71\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['sentence1', 'sentence2', 'label', 'idx'],\n",
       "        num_rows: 146\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "raw_datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "feedbd32",
   "metadata": {},
   "source": [
    "# Tokenizer\n",
    "\n",
    "We can automatically load the correct tokenizer used from the pretrained model\n",
    "via `AutoTokenizer`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "be2de095",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading (…)okenizer_config.json: 100%|██████████| 28.0/28.0 [00:00<00:00, 222kB/s]\n",
      "Downloading (…)lve/main/config.json: 100%|██████████| 570/570 [00:00<00:00, 5.00MB/s]\n",
      "Downloading (…)solve/main/vocab.txt: 100%|██████████| 232k/232k [00:00<00:00, 13.7MB/s]\n",
      "Downloading (…)/main/tokenizer.json: 100%|██████████| 466k/466k [00:00<00:00, 52.5MB/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['this', 'will', 'be', 'fun', '!']"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(checkpoint)\n",
    "\n",
    "sequence = 'this will be fun!'\n",
    "\n",
    "tokenizer.tokenize(sequence)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c46439e",
   "metadata": {},
   "source": [
    "# Tokenizer Output\n",
    "\n",
    "Let's take a look at the integers (input_ids) assigned to each token in the sequence\n",
    "as well as other information such as optional masks for any tokens that need to be\n",
    "masked from the attention mechanism - special tokens for truncating sequences for example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "575ab250",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': [101, 2023, 2097, 2022, 4569, 999, 102], 'token_type_ids': [0, 0, 0, 0, 0, 0, 0], 'attention_mask': [1, 1, 1, 1, 1, 1, 1]}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer(sequence)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a1cbf3c9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map: 100%|██████████| 635/635 [00:00<00:00, 28083.20 examples/s]\n",
      "Map: 100%|██████████| 71/71 [00:00<00:00, 12372.06 examples/s]\n",
      "Map: 100%|██████████| 146/146 [00:00<00:00, 16078.15 examples/s]\n"
     ]
    }
   ],
   "source": [
    "# function to crate\n",
    "def tokenize_function(example):\n",
    "    return tokenizer(example[\"sentence1\"], example[\"sentence2\"], truncation=True)\n",
    "\n",
    "tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)\n",
    "data_collator = DataCollatorWithPadding(tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "521096c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TrainingArguments\n",
    "\n",
    "training_args = TrainingArguments(\"test-trainer\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb96f253",
   "metadata": {},
   "source": [
    "# Load pretrained model weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a79a7c30",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading model.safetensors: 100%|██████████| 440M/440M [00:01<00:00, 383MB/s] \n",
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoModelForSequenceClassification\n",
    "\n",
    "model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ef50fa8",
   "metadata": {},
   "source": [
    "# Create a Trainer object to begin fine tuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7afdced6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import Trainer\n",
    "\n",
    "trainer = Trainer(\n",
    "    model,\n",
    "    training_args,\n",
    "    train_dataset=tokenized_datasets[\"train\"],\n",
    "    eval_dataset=tokenized_datasets[\"validation\"],\n",
    "    data_collator=data_collator,\n",
    "    tokenizer=tokenizer,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3f240084",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/240 [00:00<?, ?it/s]You're using a BertTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n",
      "100%|██████████| 240/240 [01:25<00:00,  2.82it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'train_runtime': 85.0038, 'train_samples_per_second': 22.411, 'train_steps_per_second': 2.823, 'train_loss': 0.7035834630330403, 'epoch': 3.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=240, training_loss=0.7035834630330403, metrics={'train_runtime': 85.0038, 'train_samples_per_second': 22.411, 'train_steps_per_second': 2.823, 'train_loss': 0.7035834630330403, 'epoch': 3.0})"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59dea71c",
   "metadata": {},
   "source": [
    "# Generate predictions on the validation data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "60ad3e34",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 9/9 [00:00<00:00, 13.60it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(71, 2) (71,)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "predictions = trainer.predict(tokenized_datasets[\"validation\"])\n",
    "print(predictions.predictions.shape, predictions.label_ids.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9b355bf",
   "metadata": {},
   "source": [
    "# Model Output\n",
    "\n",
    "As we can see, the transformer model outputs logits directly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ff0b3303",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0.09305581 0.06058019]\n",
      " [0.064982   0.07597283]\n",
      " [0.06600507 0.07268026]\n",
      " [0.11784149 0.07293674]\n",
      " [0.07580852 0.07072549]\n",
      " [0.07087269 0.07751691]\n",
      " [0.08125432 0.07708483]\n",
      " [0.10697782 0.06369373]\n",
      " [0.10485618 0.06580386]\n",
      " [0.10855884 0.073225  ]]\n"
     ]
    }
   ],
   "source": [
    "print(predictions.predictions[:10, :10])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "014a67c3",
   "metadata": {},
   "source": [
    "# Turn into label predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "60c214f8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1,\n",
       "       1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0,\n",
       "       0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0,\n",
       "       0, 0, 0, 0, 0])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "preds = np.argmax(predictions.predictions, axis=-1)\n",
    "preds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "569dcd70",
   "metadata": {},
   "source": [
    "<img src=\"static/llm-studio-logo.svg\" alt=\"drawing\" width=\"200\"/>  \n",
    "\n",
    "# Fine Tuning like a Kaggle Grandmaster using LLM Studio\n",
    "\n",
    "H2O LLM studio was created by some of the most successful kagglers in the world.\n",
    "Recently, it was used to win 1st place in the Kaggle LLM Science Exam competition, \n",
    "the first competitive LLM competion on the platform. \n",
    "\n",
    "LLM Studio is open source technology that allows anyone to fine tune their own\n",
    "large language model using their own data. \n",
    "\n",
    "It can be used as a UI, as well as via Python Command Line Interface (CLI). Please first\n",
    "complete the Aqarium lab found at https://aquarium.h2o.ai before trying to use\n",
    "the Python CLI. Examples of the CLI are found below as extended tasks.\n",
    "\n",
    "***\n",
    "\n",
    "<img src=\"static/aquarium-llm-studio.png\" alt=\"drawing\" width=\"600\"/> \n",
    "\n",
    "***\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b822d04",
   "metadata": {},
   "source": [
    "# 🎉 **CONGRATULATIONS!** You have completed this lab!\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "743c736b",
   "metadata": {},
   "source": [
    "\n",
    "# 📚 **EXTENDED TASKS**\n",
    "\n",
    "## Programmatic Fine Tuning with H2O LLM Studio\n",
    "\n",
    "H2O LLM Studio is open source technology located here: https://github.com/h2oai/h2o-llmstudio.\n",
    "\n",
    "It's recommended to have a modern GPU with sufficient VRAM to support the task\n",
    "of fine tuning some of these truly enormous neural network architectures. For best\n",
    "results, recommend at least one NVIDIA A100 with 80GB of GPU Memory\n",
    "\n",
    "Once up and running, you will be able to use LLM Studio pythonically through \n",
    "the use of a `cfg.yaml` file.\n",
    "\n",
    "You can pass this file to LLM Studio to launch your first experiment. For example,\n",
    "if you have a configuration file called `my_yaml_cfg.yaml` the Python command to \n",
    "launch the experiment would be:\n",
    "\n",
    "```python train.py -Y my_yaml_cfg.yaml```\n",
    "\n",
    "An example file can be found in this training repository as well: `example-config.yaml`)\n",
    "\n",
    "Let's now walk through the configurations within the YAML file"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61fc4c99",
   "metadata": {},
   "source": [
    "# Overall Configurations\n",
    "\n",
    "First of all, there are some overall configurations that are worth discussing:\n",
    "\n",
    "- `experiment_name` - the name you'd like to call your experiment\n",
    "- `output_directory` - where will you be storing your results\n",
    "- `problem_type` - the machine learning task you're fine tuning for (e.g. `text_causal_language_modeling`)\n",
    "- `llm_backbone` - the pretrained checkpoint you'll be using to undertake transfer learning\n",
    "\n",
    "The backbone is by far the most important configuration for your experiment. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd886e41",
   "metadata": {},
   "source": [
    "# Architecture\n",
    "\n",
    "The architecture settings allows the user to specify various aspects of the \n",
    "backbone architecture being used for fine tuning\n",
    "\n",
    "```\n",
    "architecture:\n",
    "    backbone_dtype: int4\n",
    "    force_embedding_gradients: false\n",
    "    gradient_checkpointing: true\n",
    "    intermediate_dropout: 0.0\n",
    "    pretrained: true\n",
    "    pretrained_weights:\n",
    "```\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af9c9d58",
   "metadata": {},
   "source": [
    "# Tokenizer Settings\n",
    "\n",
    "```\n",
    "tokenizer:\n",
    "    add_prefix_space: false\n",
    "    add_prompt_answer_tokens: false\n",
    "    max_length: 4032\n",
    "    max_length_answer: 2016\n",
    "    max_length_prompt: 2016\n",
    "    padding_quantile: 1.0\n",
    "    use_fast: true\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38539055",
   "metadata": {},
   "source": [
    "# Augmentation Settings\n",
    "\n",
    "```\n",
    "augmentation:\n",
    "    random_parent_probability: 0.0\n",
    "    skip_parent_probability: 0.0\n",
    "    token_mask_probability: 0.0\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5ee6e45",
   "metadata": {},
   "source": [
    "# Dataset Configurations\n",
    "\n",
    "```\n",
    "dataset:\n",
    "    add_eos_token_to_answer: true\n",
    "    add_eos_token_to_prompt: true\n",
    "    add_eos_token_to_system: true\n",
    "    answer_column: content\n",
    "    chatbot_author: H2O.ai\n",
    "    chatbot_name: h2oGPT\n",
    "    data_sample: 1.0\n",
    "    data_sample_choice:\n",
    "    - Train\n",
    "    - Validation\n",
    "    limit_chained_samples: false\n",
    "    mask_prompt_labels: true\n",
    "    parent_id_column: None\n",
    "    personalize: false\n",
    "    prompt_column:\n",
    "    - instruction\n",
    "    system_column: None\n",
    "    text_answer_separator: <|answer|>\n",
    "    text_prompt_start: <|prompt|>\n",
    "    text_system_start: <|system|>\n",
    "    train_dataframe: h2o_genai_world_training/influencers_data_cleaned_sample.csv\n",
    "    validation_dataframe: None\n",
    "    validation_size: 0.01\n",
    "    validation_strategy: automatic\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f83791a6",
   "metadata": {},
   "source": [
    "# Training Configurations\n",
    "\n",
    "```\n",
    "training:\n",
    "    batch_size: 2\n",
    "    differential_learning_rate: 1.0e-05\n",
    "    differential_learning_rate_layers: []\n",
    "    drop_last_batch: true\n",
    "    epochs: 1\n",
    "    evaluate_before_training: false\n",
    "    evaluation_epochs: 1.0\n",
    "    grad_accumulation: 1\n",
    "    gradient_clip: 0.0\n",
    "    learning_rate: 0.0001\n",
    "    lora: true\n",
    "    lora_alpha: 16\n",
    "    lora_dropout: 0.05\n",
    "    lora_r: 4\n",
    "    lora_target_modules: ''\n",
    "    loss_function: TokenAveragedCrossEntropy\n",
    "    optimizer: AdamW\n",
    "    save_best_checkpoint: false\n",
    "    schedule: Cosine\n",
    "    train_validation_data: false\n",
    "    warmup_epochs: 0.0\n",
    "    weight_decay: 0.0\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d9fa5d8",
   "metadata": {},
   "source": [
    "# "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python",
   "language": "python",
   "name": "python_kubernetes"
  },
  "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
}