examples/nlp_models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 📖 NLP Models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial will demonstrate how to fine-tune a pretrained HuggingFace transformer using the composer library! Composer provides a highly optimized and functional training loop and the ability to compose several methods that can accelerate training.\n",
    "\n",
    "We will focus on fine-tuning a pretrained BERT-base model on the Stanford Sentiment Treebank v2 (SST-2) dataset. After fine-tuning, the BERT model should be able to determine if a setence has positive or negative sentiment.\n",
    "\n",
    "Let's do this 🚀"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install Composer\n",
    "\n",
    "To develop NLP models with Composer, we'll need to install Composer with the NLP dependencies. If you haven't already, run: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install 'mosaicml[nlp]'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To log our results using Tensorboard, we will also need to install Tensorboard. If you haven't already, run:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install 'mosaicml[tensorboard]'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining a Composer Model "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first task is to create a composer model. A composer model defines four components for the Composer trainer:\n",
    "- `forward()` - parses the dataloader output for the model's forward function and extracts the necessary components of the model's output for loss calculation.\n",
    "- `loss()` - computes the loss for the current batch using the model and dataloader outputs.\n",
    "- `validate()` - parses the dataloader and model output for torchmetrics.\n",
    "- `metrics()` - defines the torchmetrics to use during training/validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import transformers\n",
    "from torchmetrics import Accuracy\n",
    "from torchmetrics.collections import MetricCollection\n",
    "from composer.models.base import ComposerModel\n",
    "from composer.metrics import CrossEntropy\n",
    "\n",
    "# Define a Composer Model\n",
    "class ComposerBERT(ComposerModel):\n",
    "    def __init__(self, model):\n",
    "        super().__init__()\n",
    "        self.module = model\n",
    "\n",
    "        # Metrics\n",
    "        self.train_loss = CrossEntropy()\n",
    "        self.val_loss = CrossEntropy()\n",
    "\n",
    "        self.train_acc = Accuracy()\n",
    "        self.val_acc = Accuracy()\n",
    "\n",
    "    def forward(self, batch):\n",
    "        output = self.module(**batch)\n",
    "        return output\n",
    "\n",
    "    def loss(self, outputs, batch):\n",
    "        return outputs['loss']\n",
    "\n",
    "    def validate(self, batch):\n",
    "        labels = batch.pop('labels')\n",
    "        output = self.forward(batch)\n",
    "        output = output['logits']\n",
    "        return (output, labels)\n",
    "  \n",
    "    def metrics(self, train: bool = False):\n",
    "        return MetricCollection([self.train_loss, self.train_acc]) \\\n",
    "     if train else MetricCollection([self.val_loss, self.val_acc])\n",
    "\n",
    "# Create a BERT sequence classification model using HuggingFace transformsers\n",
    "model = transformers.AutoModelForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2) # in BERT hparams\n",
    "\n",
    "# Package as a composer model\n",
    "composer_model = ComposerBERT(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating dataloaders"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we will download and tokenize the SST-2 datasets. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "from multiprocessing import cpu_count\n",
    "\n",
    "# Create BERT tokenizer\n",
    "tokenizer = transformers.AutoTokenizer.from_pretrained('bert-base-uncased') # from transformer_shared\n",
    "def tokenize_function(sample):\n",
    "    return tokenizer(\n",
    "        text=sample['sentence'],\n",
    "        padding=\"max_length\",\n",
    "        max_length=256,\n",
    "        truncation=True\n",
    "    )\n",
    "\n",
    "# Tokenize SST-2\n",
    "sst2_dataset = datasets.load_dataset(\"glue\", \"sst2\")\n",
    "tokenized_sst2_dataset = sst2_dataset.map(tokenize_function,\n",
    "                                          batched=True, \n",
    "                                          num_proc=cpu_count(),\n",
    "                                          batch_size=1000,\n",
    "                                          remove_columns=['idx', 'sentence'])\n",
    "\n",
    "# Split dataset into train and validation sets\n",
    "train_dataset = tokenized_sst2_dataset[\"train\"]\n",
    "eval_dataset = tokenized_sst2_dataset[\"validation\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we will create a PyTorch `DataLoader` for each of the datasets generated in the previous block."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "data_collator = transformers.data.data_collator.default_data_collator\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=16, shuffle=False, drop_last=False, collate_fn=data_collator)\n",
    "eval_dataloader = DataLoader(eval_dataset,batch_size=16, shuffle=False, drop_last=False, collate_fn=data_collator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use the composer `Trainer`, we need to define a `split_batch` function. This function defines how to split the dataloader output into several \"microbatches\". Microbatchs are chunks of the batch that were divided based on the amount of gradient accumulation used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from composer.core import DataSpec\n",
    "\n",
    "def split_batch_dict(batch, n_microbatches: int):\n",
    "    chunked = {k: v.chunk(n_microbatches) for k, v in batch.items()}\n",
    "    num_chunks = len(list(chunked.values())[0])\n",
    "    return [{k: v[idx] for k, v in chunked.items()} for idx in range(num_chunks)]\n",
    "\n",
    "train_dataspec = DataSpec(dataloader=train_dataloader,\n",
    "                          split_batch=split_batch_dict)\n",
    "eval_dataspec = DataSpec(dataloader=eval_dataloader,\n",
    "                         split_batch=split_batch_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimizers and Learning Rate Schedulers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last setup step is to create an optimizer and a learning rate scheduler. We will use PyTorch's AdamW optimizer and linear learning rate scheduler since these are typically used to fine-tune BERT on tasks such as SST-2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.optim import AdamW\n",
    "from torch.optim.lr_scheduler import LinearLR\n",
    "\n",
    "optimizer = AdamW(\n",
    "    params=composer_model.parameters(),\n",
    "    lr=3e-5, betas=(0.9, 0.98),\n",
    "    eps=1e-6, weight_decay=3e-6\n",
    ")\n",
    "linear_lr_decay = LinearLR(\n",
    "    optimizer, start_factor=1.0,\n",
    "    end_factor=0, total_iters=150\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Logging to Tensorboard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last step is to set up the ability to log our metrics to [Tensorboard](https://www.tensorflow.org/tensorboard). Tensorboard is nice tool that allows you to track and visualize metrics. To do this we need to create a TensorboardLogger object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from composer.loggers import TensorboardLogger\n",
    "\n",
    "# Here we specify that we want our tensorboard logs to be saved to ./tensorboard_logs.\n",
    "tb_logger = TensorboardLogger(log_dir=\"./tensorboard_logs\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Composer Trainer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now specify a Composer `Trainer` object and run our training! `Trainer` has many arguments that are described in our [documentation](https://docs.mosaicml.com/en/stable/api_reference/composer.trainer.trainer.html), but let's discuss the less obvious arguments used below:\n",
    "- `max_duration` - a string specifying how long to train, either in terms of batches (e.g. '10ba' is 10 batches) or epochs (e.g. '1ep' is 1 epoch).\n",
    "- `schedulers` - a list of PyTorch learning rate schedulers that will be composed together.\n",
    "- `device` - specifies if the training will be done on CPU or GPU by using 'cpu' or 'gpu', respectively.\n",
    "- `train_subset_num_batches` - specifies the number of training batches to use for each epoch. This is not a necessary argument but is useful for quickly testing code.\n",
    "- `precision` - whether to do the training in full precision 'fp32' or mixed precision 'amp'. Mixed precision provides an almost 2x speedup in training time on certain hardware. If you get a P100, try `precision='amp'`!\n",
    "- `seed` - sets the random seed for the training run, so the results are reproducible!\n",
    "- `loggers` - sets all the loggers we want to use during our training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from composer import Trainer\n",
    "\n",
    "# Create Trainer Object\n",
    "trainer = Trainer(\n",
    "    model=composer_model, \n",
    "    train_dataloader=train_dataspec,\n",
    "    eval_dataloader=eval_dataspec,\n",
    "    max_duration=\"1ep\",\n",
    "    optimizers=optimizer,\n",
    "    schedulers=[linear_lr_decay],\n",
    "    device='gpu' if torch.cuda.is_available() else 'cpu',\n",
    "    train_subset_num_batches=150,\n",
    "    precision='fp32',\n",
    "    seed=17,\n",
    "    loggers=[tb_logger],\n",
    ")\n",
    "# Start training\n",
    "trainer.fit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing Results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our model reaches almost 86% accuracy with only 100 iterations of training! Let's visualize a few samples from the validation set to see how our model performs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_batch = next(iter(eval_dataloader))\n",
    "\n",
    "# Move batch to gpu\n",
    "eval_batch = {k: v.cuda() if torch.cuda.is_available() else v for k, v in eval_batch.items()}\n",
    "with torch.no_grad():\n",
    "    predictions = composer_model(eval_batch)[\"logits\"].argmax(dim=1)\n",
    "\n",
    "# Visualize only 5 samples\n",
    "predictions = predictions[:6]\n",
    "\n",
    "label = ['negative', 'positive']\n",
    "for i, prediction in enumerate(predictions[:6]):\n",
    "    sentence = sst2_dataset[\"validation\"][i][\"sentence\"]\n",
    "    correct_label = label[sst2_dataset[\"validation\"][i][\"label\"]]\n",
    "    prediction_label = label[prediction]\n",
    "    print(f\"Sample: {sentence}\")\n",
    "    print(f\"Label: {correct_label}\")\n",
    "    print(f\"Prediction: {prediction_label}\")\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial showed how to use the Composer `Trainer` to fine-tune a pre-trained BERT on a subset of the SST-2 dataset. We focused on the Composer's basic functionality, but there are many more tools such as easy to use gradient accumulation and multi-GPU training! Check out many other features at our [documentation](https://docs.mosaicml.com/en/latest)."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}