model_training_pipeline.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Language Model Training Pipeline"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Currently only tested for GPT2!* "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import utils\n",
    "import torch\n",
    "import evaluate\n",
    "import transformers\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from torch_lr_finder import LRFinder\n",
    "from pathlib import Path\n",
    "from tqdm.notebook import tqdm\n",
    "from datasets import Dataset\n",
    "from tokenizers import Tokenizer\n",
    "from tokenizers.normalizers import (Sequence, Lowercase, NFD, StripAccents)\n",
    "from tokenizers.pre_tokenizers import Whitespace\n",
    "from tokenizers.models import BPE\n",
    "from tokenizers.trainers import BpeTrainer\n",
    "from tokenizers.decoders import BPEDecoder\n",
    "from transformers import AutoConfig, \\\n",
    "    DataCollatorWithPadding, AutoModelForSequenceClassification, \\\n",
    "    Trainer, TrainingArguments, AutoTokenizer, GPT2Config\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "import importlib\n",
    "importlib.reload(utils)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = Path(\"\")\n",
    "'''Path to large data folder'''\n",
    "seed = 42\n",
    "'''Random seed: int'''\n",
    "\n",
    "# Tokenizer\n",
    "train_file = data_path/\"\"\n",
    "'''TXT file for training tokenizer'''\n",
    "tokenizer_path  = Path(\"\")\n",
    "'''Path to save tokenizer'''\n",
    "vocab_size = None   \n",
    "'''Vocab size for tokenizer: int'''\n",
    "special_tokens = None\n",
    "'''Special tokens for tokenizer: list of strings'''\n",
    "\n",
    "# Language Model Pretraining\n",
    "lm_train_file = data_path/\"\" \n",
    "'''TXT file for training language model'''\n",
    "lm_valid_file = data_path/\"\"\n",
    "'''TXT file for validating language model'''\n",
    "pretrained_output_model_path = Path(\"\")\n",
    "'''Path to save pretrained model'''\n",
    "cache_dir = data_path/\"\"\n",
    "'''Path to save cache'''\n",
    "config_class = GPT2Config\n",
    "'''Config class for language model: e.g. GPT2Config'''\n",
    "lm_config = {\n",
    "    'model_type': None, # e.g. 'gpt2',\n",
    "    'vocab_size': None, # e.g. 50257,\n",
    "    'n_positions': None, # e.g. 1024,\n",
    "    'n_layer': None, # e.g. 12,\n",
    "    # add more config here if needed\n",
    "}\n",
    "'''Config for language model: dict\n",
    "See https://huggingface.co/transformers/model_doc/gpt2.html#gpt2config for an example for GPT2'''\n",
    "\n",
    "# Classifier Finetuning\n",
    "load_pretrained_weights = None\n",
    "'''Whether to load pretrained weights: bool'''\n",
    "classifier_output_model_path = data_path/\"\"\n",
    "'''Path to save classifier model'''\n",
    "labeled_data = Path(\"\")\n",
    "'''Path to labeled data in any format, preprocessing is self-defined'''\n",
    "batch_size = 32\n",
    "'''Batch size for classifier training: int'''\n",
    "epochs = 12\n",
    "'''Epochs for classifier training: int'''\n",
    "learning_rate = 5e-5\n",
    "'''Learning rate for classifier training: float e.g. 5e-5'''"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train BPE Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer_path.parent.mkdir(exist_ok=True, parents=True)\n",
    "\n",
    "# Define tokenizer\n",
    "tokenizer = Tokenizer(BPE())\n",
    "tokenizer.normalizer = Sequence([NFD(),Lowercase(),StripAccents()])\n",
    "tokenizer.pre_tokenizer = Whitespace()\n",
    "tokenizer.decoder = BPEDecoder()\n",
    "trainer = BpeTrainer(vocab_size=vocab_size, special_tokens=special_tokens)\n",
    "tokenizer.train([str(train_file)], trainer=trainer) \n",
    "tokenizer.save(str(tokenizer_path))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Language Model Pretraining"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Run the cell below and it will output an LM training shell script in the output model directory you specify. Navigate to that directory and run the shell script in a persistent shell session (TMUX) with the corresponding Python environment*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fetch training script\n",
    "# !rm run_clm.py\n",
    "# !wget -O run_clm.py https://raw.githubusercontent.com/HMC-MIR/ExplorationWithLLMs/72d9d79c50e3639e215c0fb7e9b6b55748981064/source/02_arhan_temp_name/run_clm.py?token=GHSAT0AAAAAAB5TQDSTGPLS4OX6JTVBTSNSZEIWEXQ\n",
    "# copy file into same directory as this notebook: https://github.com/HMC-MIR/ExplorationWithLLMs/blob/main/source/02_arhan_temp_name/run_clm.py\n",
    "\n",
    "# use run_mlm.py for masked language modeling like RoBERTas\n",
    "\n",
    "# Language Model Config\n",
    "lm_config_dict = lm_config\n",
    "\n",
    "lm_config_str = \"\\\"\" + \",\".join(f\"{k}={v}\" for k,v in lm_config_dict.items()) + \"\\\"\"\n",
    "\n",
    "# Write out config to tokenizer directory for internal HuggingFace use\n",
    "with open(str(tokenizer_path.parent/\"config.json\"), 'w') as fp:\n",
    "    json.dump(lm_config_dict, fp)\n",
    "\n",
    "pretrained_output_model_path.mkdir(exist_ok=True, parents=True)\n",
    "\n",
    "# Create training shell script\n",
    "cmd = f''' \n",
    "torchrun  --nproc_per_node 2\n",
    "{Path.cwd()/\"run_clm.py\"}\n",
    "--model_type {lm_config_dict[\"model_type\"]}\n",
    "--tokenizer_name {tokenizer_path.parent.resolve()}\n",
    "--train_file {lm_train_file}\n",
    "--validation_file {lm_valid_file.resolve()}\n",
    "--output_dir {pretrained_output_model_path}\n",
    "--do_train\n",
    "--do_eval\n",
    "--evaluation_strategy steps\n",
    "--per_device_train_batch_size 8\n",
    "--per_gpu_eval_batch_size 8\n",
    "--learning_rate 1e-4\n",
    "--num_train_epochs 12\n",
    "--logging_steps 2000\n",
    "--save_steps 2000\n",
    "--seed {seed}\n",
    "--overwrite_output_dir\n",
    "--local_rank 0\n",
    "--cache_dir {cache_dir}\n",
    "--config_overrides=\"{lm_config_str}\"\n",
    "'''.replace(\"\\n\", \" \")\n",
    "\n",
    "# Write out training shell script to output model directory\n",
    "with open(str(pretrained_output_model_path/\"train_lm.sh\"), \"w\") as fout:\n",
    "    fout.write(cmd)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Language Model Pretraining Curves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr_loss = []\n",
    "val_loss = []\n",
    "accuracy = []\n",
    "step = []\n",
    "with open(\"/home/ajain/ttmp/EWLLMs/experiments/unpretrained_classifier/checkpoint-1314/trainer_state.json\", 'r') as f:\n",
    "    log_history = json.load(f)['log_history']\n",
    "    for entry in log_history:\n",
    "        if \"loss\" in entry:\n",
    "            step.append(int(entry[\"step\"]))\n",
    "            tr_loss.append(float(entry[\"loss\"]))\n",
    "        elif \"eval_loss\" in entry:\n",
    "            val_loss.append(float(entry[\"eval_loss\"]))\n",
    "\n",
    "        if \"eval_accuracy\" in entry:\n",
    "            accuracy.append(float(entry[\"eval_accuracy\"]))\n",
    "\n",
    "step, tr_loss, val_loss, acc = np.array(step), np.array(tr_loss), np.array(val_loss), np.array(accuracy)\n",
    "plt.plot(step, tr_loss, 'k-', label=\"Train\")\n",
    "plt.scatter(step, tr_loss, c='k')\n",
    "plt.plot(step, val_loss, 'g-', label=\"Validation\")\n",
    "plt.scatter(step, val_loss, c='g')\n",
    "plt.legend()\n",
    "plt.xlabel('Steps')\n",
    "plt.ylabel('Loss')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "plt.plot(step, acc, 'k-', label=\"Validation\")\n",
    "plt.scatter(step, acc, c='k')\n",
    "plt.legend()\n",
    "plt.xlabel('Steps')\n",
    "plt.ylabel('accuracy')\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Language Model Linear Probe Classifier"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fill in the cell below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def data_preparation(labeled_data):\n",
    "    \"\"\"Prepare data for training, validation, and testing.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    train_df : pd.DataFrame\n",
    "        Training data with 2 columns, \"text\" and \"label\".\n",
    "    val_df : pd.DataFrame\n",
    "        Validation data with 2 columns, \"text\" and \"label\".\n",
    "    test_df : pd.DataFrame\n",
    "        Testing data with 2 columns, \"text\" and \"label\".\n",
    "    \"\"\"\n",
    "\n",
    "    # FILL IN\n",
    "    train_X, train_y, val_X, val_y, test_X, test_y = utils.load_pkl(labeled_data)\n",
    "    \n",
    "\n",
    "    train_df = pd.DataFrame({\"text\": train_X, \"label\": train_y})\n",
    "    val_df = pd.DataFrame({\"text\": val_X, \"label\": val_y})\n",
    "    test_df = pd.DataFrame({\"text\": test_X, \"label\": test_y})\n",
    "    \n",
    "    return train_df, val_df, test_df"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prewritten Helper Functions & Classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenizer_function(examples, tokenizer):\n",
    "    return tokenizer(examples[\"text\"], padding='max_length', truncation=True)\n",
    "\n",
    "def label2id_function(examples, label2id):\n",
    "    return {\"label\": [label2id[label] for label in examples[\"label\"]]}\n",
    "\n",
    "def compute_metrics(eval_pred):\n",
    "    accuracy = evaluate.load(\"accuracy\")  \n",
    "    predictions, labels = eval_pred\n",
    "    predictions = np.argmax(predictions, axis=1)\n",
    "    return accuracy.compute(predictions=predictions, references=labels)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training Classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare data\n",
    "train_df, val_df, test_df = data_preparation(labeled_data)\n",
    "train_ds = Dataset.from_dict(train_df)\n",
    "val_ds = Dataset.from_dict(val_df)\n",
    "test_ds = Dataset.from_dict(test_df)\n",
    "\n",
    "# Define label map\n",
    "label2id = {label: i for i, label in enumerate(set(train_df['label']))}\n",
    "id2label = {i: label for label, i in label2id.items()}\n",
    "\n",
    "# Load Tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(pretrained_output_model_path)\n",
    "tokenizer.pad_token = '<pad>'\n",
    "\n",
    "\n",
    "# Define model\n",
    "if load_pretrained_weights:\n",
    "    config = AutoConfig.from_pretrained(pretrained_output_model_path)\n",
    "    config.num_labels = len(label2id)\n",
    "    config.pad_token_id = tokenizer.pad_token_id\n",
    "    model = AutoModelForSequenceClassification.from_pretrained(pretrained_output_model_path, config=config)\n",
    "else:\n",
    "    config = config_class(**lm_config)\n",
    "    config.num_labels = len(label2id)\n",
    "    config.pad_token_id = tokenizer.pad_token_id\n",
    "    model = AutoModelForSequenceClassification.from_config(config)\n",
    "tokenizer.model_max_length = config.n_positions\n",
    "\n",
    "# Tokenize and convert labels to ids\n",
    "train_ds = train_ds.map(tokenizer_function, batched=True, fn_kwargs={\"tokenizer\": tokenizer})\n",
    "val_ds = val_ds.map(tokenizer_function, batched=True, fn_kwargs={\"tokenizer\": tokenizer})\n",
    "train_ds = train_ds.map(label2id_function, batched=True, fn_kwargs={\"label2id\": label2id})\n",
    "val_ds = val_ds.map(label2id_function, batched=True, fn_kwargs={\"label2id\": label2id})\n",
    "\n",
    "data_collator = DataCollatorWithPadding(tokenizer=tokenizer, padding=\"longest\", max_length=1024)\n",
    "\n",
    "# Freeze all layers except the classifier\n",
    "for name, param in model.named_parameters():\n",
    "    param.requires_grad = False\n",
    "model.score.weight.requires_grad = True\n",
    "\n",
    "\n",
    "# Define training arguments\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=classifier_output_model_path,\n",
    "    learning_rate=learning_rate,\n",
    "    per_device_train_batch_size=batch_size,\n",
    "    per_device_eval_batch_size=batch_size,\n",
    "    num_train_epochs=epochs,\n",
    "    weight_decay=0.01,\n",
    "    logging_strategy=\"epoch\",\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    load_best_model_at_end=True,\n",
    "    push_to_hub=False,\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_ds,\n",
    "    eval_dataset=val_ds,\n",
    "    tokenizer=tokenizer,\n",
    "    data_collator=data_collator,\n",
    "    compute_metrics=compute_metrics,\n",
    ")\n",
    "\n",
    "# Train model\n",
    "trainer.train()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Language Model Classifier Finetuning Curves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Visualize training and validation losses\n",
    "with open(\"/home/abunn/ttmp/ExplorationWithLLMs/source/03_alec_temp_name/finetuned_models/dense_1_8/log_history.json\", \"r\") as fin:\n",
    "    log_history = json.load(fin)\n",
    "\n",
    "step, tr_loss, val_loss = [], [], []\n",
    "acc = []\n",
    "for epoch in log_history:\n",
    "    step.append(epoch[\"epoch\"])\n",
    "    tr_loss.append(epoch[\"train_loss\"])\n",
    "    val_loss.append(epoch[\"val_loss\"])\n",
    "    acc.append(epoch[\"accuracy\"])\n",
    "\n",
    "step, tr_loss, val_loss = np.array(step), np.array(tr_loss), np.array(val_loss)\n",
    "plt.plot(step, tr_loss, 'k-', label=\"Train\")\n",
    "plt.scatter(step, tr_loss, c='k')\n",
    "plt.plot(step, val_loss, 'g-', label=\"Validation\")\n",
    "plt.scatter(step, val_loss, c='g')\n",
    "plt.legend()\n",
    "plt.xlabel('Steps')\n",
    "plt.ylabel('Loss')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(step, acc, 'g-', label=\"Validation\")\n",
    "plt.scatter(step, acc, c='g')\n",
    "plt.legend()\n",
    "plt.xlabel('Steps')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "EWLLMs",
   "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.11"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}