flax_bert_demo_neurips_2020.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine-tuning BERT in Flax on GLUE\n",
    "\n",
    "This notebook fine-tunes a BERT model one of the [GLUE tasks](https://gluebenchmark.com/). It has the following features:\n",
    "\n",
    "*   Uses the [HuggingFace](https://github.com/huggingface/) datasets and tokenizers libraries.\n",
    "*   Loads the pre-trained BERT weights from HuggingFace.\n",
    "*   Model and training code is written in [Flax](http://www.github.com/google/flax).\n",
    "*   Can be configured to fine-tune on COLA, MRPC, SST2, STSB, QNLI, and RTE.\n",
    "\n",
    "Run-times on MRPC:\n",
    "\n",
    "*   Cloud TPU v3-8: 40s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# General imports.\n",
    "import os\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import flax\n",
    "\n",
    "# Huggingface datasets and transformers libraries.\n",
    "import datasets\n",
    "from transformers import BertTokenizerFast\n",
    "\n",
    "# flax_bert-specific imports.\n",
    "from flax import optim\n",
    "import data\n",
    "import modeling as flax_models\n",
    "import training\n",
    "from demo_lib import get_config, get_validation_splits, get_prefix, import_pretrained_params, create_model, create_optimizer, get_num_train_steps, get_learning_rate_fn\n",
    "os.environ['TOKENIZERS_PARALLELISM'] = 'true'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set your Training Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_settings = {\n",
    "    'train_batch_size': 32,\n",
    "    'eval_batch_size': 8,\n",
    "    'learning_rate': 5e-5,\n",
    "    'num_train_epochs': 3,\n",
    "    'dataset_path': 'glue',\n",
    "    'dataset_name': 'mrpc'  # ['cola', 'mrpc', 'sst2', 'stsb', 'qnli', 'rte']\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load dataset, tokenizers, and model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Reusing dataset glue (/home/marcvanzee/.cache/huggingface/datasets/glue/mrpc/1.0.0/7c99657241149a24692c402a5c3f34d4c9f1df5ac2e4c3759fadea38f6cb29c4)\n",
      "Loading cached processed dataset at /home/marcvanzee/.cache/huggingface/datasets/glue/mrpc/1.0.0/7c99657241149a24692c402a5c3f34d4c9f1df5ac2e4c3759fadea38f6cb29c4/cache-6ebbe5cc20e1150b.arrow\n",
      "Loading cached processed dataset at /home/marcvanzee/.cache/huggingface/datasets/glue/mrpc/1.0.0/7c99657241149a24692c402a5c3f34d4c9f1df5ac2e4c3759fadea38f6cb29c4/cache-ddcb7009256661ae.arrow\n",
      "Loading cached processed dataset at /home/marcvanzee/.cache/huggingface/datasets/glue/mrpc/1.0.0/7c99657241149a24692c402a5c3f34d4c9f1df5ac2e4c3759fadea38f6cb29c4/cache-51f953271952eb23.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading weights file https://huggingface.co/bert-base-uncased/resolve/main/pytorch_model.bin from cache at /home/marcvanzee/.cache/torch/transformers/a8041bf617d7f94ea26d15e218abd04afc2004805632abc0ed2066aa16d50d04.faf6ea826ae9c5867d12b22257f9877e6b8367890837bd60f7c54a29633f7f2f\n",
      "Num train examples: 3668\n"
     ]
    }
   ],
   "source": [
    "# Load the GLUE task.\n",
    "dataset = datasets.load_dataset('glue', train_settings['dataset_name'])\n",
    "\n",
    "# Get pre-trained config and update it with the train configuration.\n",
    "config = get_config('bert-base-uncased', dataset)\n",
    "config.update(train_settings)\n",
    "\n",
    "# Load HuggingFace tokenizer and data pipeline.\n",
    "tokenizer = BertTokenizerFast.from_pretrained(config.tokenizer)\n",
    "data_pipeline = data.ClassificationDataPipeline(dataset, tokenizer)\n",
    "\n",
    "# Create Flax model and optimizer.\n",
    "pretrained_params = import_pretrained_params(config)\n",
    "model = create_model(config, pretrained_params)\n",
    "optimizer = create_optimizer(config, model, pretrained_params)\n",
    "\n",
    "# Setup tokenizer, train step function and train iterator.\n",
    "tokenizer.model_max_length = config.max_seq_length\n",
    "num_train_steps = get_num_train_steps(config, data_pipeline)\n",
    "\n",
    "learning_rate_fn = get_learning_rate_fn(config, num_train_steps)\n",
    "train_history = training.TrainStateHistory(learning_rate_fn)\n",
    "train_state = train_history.initial_state()\n",
    "\n",
    "train_step_fn = training.create_train_step(clip_grad_norm=1.0)\n",
    "train_iter = data_pipeline.get_inputs(\n",
    "  split='train', batch_size=config.train_batch_size, training=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Starting training on mrpc for 343 steps (3 epochs)...\n",
      "\n",
      "Compiling train (takes about 20s)\n",
      "Step 0       grad_norm = 43.7472038269043\n",
      "             loss = 0.6523172855377197\n",
      "step 0/343\n",
      "step 10/343\n",
      "step 20/343\n",
      "step 30/343\n",
      "step 40/343\n",
      "step 50/343\n",
      "step 60/343\n",
      "step 70/343\n",
      "step 80/343\n",
      "step 90/343\n",
      "step 100/343\n",
      "step 110/343\n",
      "step 120/343\n",
      "step 130/343\n",
      "step 140/343\n",
      "step 150/343\n",
      "step 160/343\n",
      "step 170/343\n",
      "step 180/343\n",
      "step 190/343\n",
      "Step 200     grad_norm = 155.39651489257812\n",
      "             loss = 1.4091753959655762\n",
      "             seconds_per_step = 0.051291774958372116\n",
      "step 200/343\n",
      "step 210/343\n",
      "step 220/343\n",
      "step 230/343\n",
      "step 240/343\n",
      "step 250/343\n",
      "step 260/343\n",
      "step 270/343\n",
      "step 280/343\n",
      "step 290/343\n",
      "step 300/343\n",
      "step 310/343\n",
      "step 320/343\n",
      "step 330/343\n",
      "step 340/343\n",
      "\n",
      "Training finished! Running eval...\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(f'\\nStarting training on {config.dataset_name} for {num_train_steps} '\n",
    "      f'steps ({config.num_train_epochs:.0f} epochs)...\\n')\n",
    "\n",
    "for step, batch in zip(range(0, num_train_steps), train_iter):\n",
    "  optimizer, train_state = train_step_fn(optimizer, batch, train_state)\n",
    "  if step % 10 == 0:\n",
    "    print(f'step {step}/{num_train_steps}')\n",
    "\n",
    "print('\\nTraining finished!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run Evaluation\n",
    "\n",
    "The target eval_f1 for MRPC is 88.9 (variance of about 1.0)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "eval_accuracy = 0.877451\n",
      "eval_f1 = 0.914384\n"
     ]
    }
   ],
   "source": [
    "eval_step = training.create_eval_fn()\n",
    "\n",
    "for split in get_validation_splits(config.dataset_name):\n",
    "  eval_iter = data_pipeline.get_inputs(\n",
    "      split='validation', batch_size=config.eval_batch_size, training=False)\n",
    "  eval_stats = eval_step(optimizer, eval_iter)\n",
    "  eval_metric = datasets.load_metric(config.dataset_path, config.dataset_name)\n",
    "  eval_metric.add_batch(\n",
    "    predictions=eval_stats['prediction'],\n",
    "    references=eval_stats['label'])\n",
    "  eval_metrics = eval_metric.compute()\n",
    "  for name, val in sorted(eval_metrics.items()):\n",
    "    print(f'{get_prefix(split)}_{name} = {val:.06f}', flush=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}