examples/text_summarization/abstractive_summarization_minilm_cnndm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright (c) Microsoft Corporation.  \n",
    "Licensed under the MIT License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Abstractive Summarization using MiniLM on CNN/DailyMails"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Before you start\n",
    "Set `QUICK_RUN = True` to run the notebook on a small subset of data and a smaller number of steps. If `QUICK_RUN = False`, the notebook takes about 2 hours to run on a VM with 4 16GB NVIDIA V100 GPUs. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "QUICK_RUN = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "This notebook demostrates how to fine-tune the [MiniLM](https://arxiv.org/abs/2002.10957) for abstractive summarization task. Utility functions and classes in the microsoft/nlp-recipes repo are used to facilitate data preprocessing, model training, model scoring, result postprocessing, and model evaluation.\n",
    "\n",
    "### Abstractive Summarization\n",
    "Abstractive summarization is the task of taking an input text and summarizing its content in a shorter output text. In contrast to extractive summarization, abstractive summarization doesn't take sentences directly from the input text, instead, rephrases the input text.\n",
    "\n",
    "### MiniLM\n",
    "[Unified Language Model](https://arxiv.org/abs/1905.03197) (UniLM) is a state of the art model developed by Microsoft Research Asia (MSRA). The model is first pre-trained on a large unlabeled natural language corpus (English Wikipedia and BookBorpus) and can be fine-tuned on different types of labeled data for various NLP tasks like text classification and abstractive summarization. For more information, please consult the notebook [Abstractive Summarization using MiniLM on CNN/DailyMails](./abstractive_summarization_unilm_cnndm.ipynb).\n",
    "\n",
    "Large pre-trained language models like BERT and UniLM usually consists of **hundreds** of millions of parameters and it's challleging to fine-tune such large models and also serve  real-life applications due to latency and capacity constraints.\n",
    "\n",
    "MiniLM is a small version of UniLM, which is trained to deelply mimic UniLM with  deep self-attention knowledge distillation. It only consits of **tens** of millions of parameters (33M), which is less than one third of BERT base model and only half of the size of [DistilBERT](https://arxiv.org/abs/1910.01108). Experimental results demonstrate that MiniLM retains most of the performance of UniLM on various NLP tasks with much less computation.  Our experiments show that to achieve the same performance, MiniLM funtuning on CNN/DailyMail dataset can be more than **ten times faster** and inferencing can be **six times faster** than UniLM's. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import os\n",
    "import shutil\n",
    "from tempfile import TemporaryDirectory\n",
    "import pprint\n",
    "import scrapbook as sb\n",
    "import sys\n",
    "import time\n",
    "import torch\n",
    "\n",
    "nlp_path = os.path.abspath(\"../../\")\n",
    "if nlp_path not in sys.path:\n",
    "    sys.path.insert(0, nlp_path)\n",
    "\n",
    "from utils_nlp.dataset.cnndm import CNNDMSummarizationDatasetOrg\n",
    "from utils_nlp.models.transformers.abstractive_summarization_seq2seq import S2SAbsSumProcessor, S2SAbstractiveSummarizer\n",
    "from utils_nlp.eval import compute_rouge_python\n",
    "\n",
    "from utils_nlp.models.transformers.datasets import SummarizationDataset\n",
    "from utils_nlp.dataset.cnndm import detokenize\n",
    "\n",
    "start_time = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# model parameters\n",
    "MODEL_NAME = \"minilm-l12-h384-uncased\" \n",
    "MAX_SEQ_LENGTH = 512 \n",
    "MAX_SOURCE_SEQ_LENGTH = 464 \n",
    "MAX_TARGET_SEQ_LENGTH = MAX_SEQ_LENGTH - MAX_SOURCE_SEQ_LENGTH \n",
    "\n",
    "# use 0 for CPU\n",
    "NUM_GPUS =  torch.cuda.device_count()\n",
    "\n",
    "# fine-tuning parameters\n",
    "TRAIN_PER_GPU_BATCH_SIZE = 4\n",
    "GRADIENT_ACCUMULATION_STEPS = 1\n",
    "LEARNING_RATE = 1e-4\n",
    "\n",
    "TOP_N = -1\n",
    "WARMUP_STEPS = 500\n",
    "MAX_STEPS = 5000\n",
    "BEAM_SIZE = 5\n",
    "if QUICK_RUN:\n",
    "    TOP_N = 1000\n",
    "    WARMUP_STEPS = 500\n",
    "    MAX_STEPS = 1000\n",
    "    BEAM_SIZE = 3\n",
    "    if NUM_GPUS == 0:\n",
    "        TOP_N = 5\n",
    "        MAX_STEPS = 10\n",
    "\n",
    "# inference parameters\n",
    "TEST_PER_GPU_BATCH_SIZE = 12\n",
    "FORBID_IGNORE_WORD = \".\"\n",
    "\n",
    "# mixed precision setting. To enable mixed precision training, follow instructions in SETUP.md. \n",
    "# You will be able to increase the batch sizes with mixed precision training.\n",
    "FP16 = False\n",
    "\n",
    "CLEANUP_RESULTS = False\n",
    "\n",
    "DATA_DIR = TemporaryDirectory().name\n",
    "CACHE_DIR = TemporaryDirectory().name\n",
    "\n",
    "MODEL_DIR = \"./minilm_cnndm_model\"\n",
    "RESULT_DIR = \"./minilm_cnndm_result\"\n",
    "os.makedirs(MODEL_DIR, exist_ok=True)\n",
    "os.makedirs(RESULT_DIR, exist_ok=True)\n",
    "OUTPUT_FILE = os.path.join(RESULT_DIR, 'nlp_cnndm_finetuning_results.txt')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the CNN/DailyMail dataset\n",
    "The [CNN/DailyMail dataset](https://cs.nyu.edu/~kcho/DMQA/) was original introduced for Q&A research. There are multiple versions of the dataset processed for summarization task available on the web. The `CNNDMSummarizationDatasetOrg` function downloads a version from the [UniLM repo](https://github.com/microsoft/unilm) with minimal processing. The function returns the training and testing dataset as `SummarizationDataset` which can be further processed for model training and testing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "train_ds, test_ds = CNNDMSummarizationDatasetOrg(local_path=DATA_DIR, top_n=TOP_N)\n",
    "print(len(train_ds))\n",
    "print(len(test_ds))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preprocessing\n",
    "The `S2SAbsSumProcessor` has multiple methods for converting input data in `SummarizationDataset`, `IterableSummarizationDataset` or json files into the format required for model training and testing. The preprocessing steps include\n",
    "- Tokenize input text\n",
    "- Convert tokens into token ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "processor = S2SAbsSumProcessor(model_name=MODEL_NAME,  cache_dir=CACHE_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = processor.s2s_dataset_from_sum_ds(train_ds, train_mode=True)\n",
    "test_dataset = processor.s2s_dataset_from_sum_ds(test_ds, train_mode=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# example code to load preprocessed xsum dataset from UniLM Repo\n",
    "# train_dataset = processor.s2s_dataset_from_json_or_file(\"/dadendev/unilm/data/xsum.train.uncased_tokenized.json\", train_mode=True, top_n=TOP_N)\n",
    "# test_dataset = processor.s2s_dataset_from_json_or_file(\"/dadendev/unilm/data/xsum.test.uncased_tokenized.json\", train_mode=False, top_n=TOP_N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fine tune model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `S2SAbstractiveSummarizer` loads a pre-trained UniLM model specified by `model_name`.  \n",
    "Call `S2SAbstractiveSummarizer.list_supported_models()` to see all the supported models.  \n",
    "If you want to use a model on the local disk, specify `load_model_from_dir` and `model_file_name`. This is particularly useful if you want to load a previously fine-tuned model and use it for inference directly without fine-tuning. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "S2SAbstractiveSummarizer.list_supported_models()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "abs_summarizer = S2SAbstractiveSummarizer(\n",
    "    model_name=MODEL_NAME,\n",
    "    max_seq_length=MAX_SEQ_LENGTH,\n",
    "    max_source_seq_length=MAX_SOURCE_SEQ_LENGTH,\n",
    "    max_target_seq_length=MAX_TARGET_SEQ_LENGTH,\n",
    "    cache_dir=CACHE_DIR\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "abs_summarizer.model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# example code to load the model from a saved checkpoint\n",
    "\"\"\"\n",
    "abs_summarizer = S2SAbstractiveSummarizer(\n",
    "     model_name=MODEL_NAME,\n",
    "     max_seq_length=MAX_SEQ_LENGTH,\n",
    "     max_source_seq_length=MAX_SOURCE_SEQ_LENGTH,\n",
    "    max_target_seq_length=MAX_TARGET_SEQ_LENGTH,\n",
    "     load_model_from_dir=RESULT_DIR,\n",
    "    model_file_name=\"model.5000.bin\",\n",
    " )\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "abs_summarizer.fit(\n",
    "    train_dataset=train_dataset,\n",
    "    num_gpus=NUM_GPUS,\n",
    "    per_gpu_batch_size=TRAIN_PER_GPU_BATCH_SIZE,\n",
    "    gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,\n",
    "    learning_rate=LEARNING_RATE,\n",
    "    warmup_steps=WARMUP_STEPS,\n",
    "    max_steps=MAX_STEPS,\n",
    "    fp16=FP16,\n",
    "    save_model_to_dir=MODEL_DIR\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the finetuned model\n",
    "# abs_summarizer.save_model(RESULT_DIR, 5000, False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate summaries on testing dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = abs_summarizer.predict(\n",
    "    test_dataset=test_dataset,\n",
    "    num_gpus=NUM_GPUS,\n",
    "    per_gpu_batch_size=TEST_PER_GPU_BATCH_SIZE,\n",
    "    beam_size=BEAM_SIZE,\n",
    "    max_tgt_length=MAX_TARGET_SEQ_LENGTH,\n",
    "    forbid_ignore_word=FORBID_IGNORE_WORD,\n",
    "    fp16=FP16\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for r in predictions[:5]:\n",
    "    print(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_ds.get_source()[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_ds.get_target()[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(OUTPUT_FILE, 'w', encoding=\"utf-8\") as f:\n",
    "    for line in predictions:\n",
    "        f.write(line + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prediction on a single input sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "source = \"\"\"\n",
    "But under the new rule, set to be announced in the next 48 hours, Border Patrol agents would immediately return anyone to Mexico — without any detainment and without any due process — who attempts to cross the southwestern border between the legal ports of entry. The person would not be held for any length of time in an American facility.\n",
    "\n",
    "Although they advised that details could change before the announcement, administration officials said the measure was needed to avert what they fear could be a systemwide outbreak of the coronavirus inside detention facilities along the border. Such an outbreak could spread quickly through the immigrant population and could infect large numbers of Border Patrol agents, leaving the southwestern border defenses weakened, the officials argued.\n",
    "The Trump administration plans to immediately turn back all asylum seekers and other foreigners attempting to enter the United States from Mexico illegally, saying the nation cannot risk allowing the coronavirus to spread through detention facilities and Border Patrol agents, four administration officials said.\n",
    "The administration officials said the ports of entry would remain open to American citizens, green-card holders and foreigners with proper documentation. Some foreigners would be blocked, including Europeans currently subject to earlier travel restrictions imposed by the administration. The points of entry will also be open to commercial traffic.\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "singel_test_ds = SummarizationDataset(\n",
    "    None, source=[source], source_preprocessing=[detokenize],\n",
    ")\n",
    "single_test_dataset = processor.s2s_dataset_from_sum_ds(singel_test_ds, train_mode=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "single_prediction = abs_summarizer.predict(\n",
    "    test_dataset=single_test_dataset,\n",
    "    num_gpus=NUM_GPUS,\n",
    "    per_gpu_batch_size=1,\n",
    "    beam_size=BEAM_SIZE,\n",
    "    forbid_ignore_word=FORBID_IGNORE_WORD,\n",
    "    fp16=FP16\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "single_prediction[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation\n",
    "We provide utility functions for evaluating summarization models and details can be found in the [summarization evaluation notebook](./summarization_evaluation.ipynb).  \n",
    "For the settings in this notebook with QUICK_RUN=False, you should get ROUGE scores close to the following numbers: <br />\n",
    "``\n",
    "{'rouge-1': {'f': 0.36208534811461,\n",
    "             'p': 0.4743143496862804,\n",
    "             'r': 0.30901813498597874},\n",
    " 'rouge-2': {'f': 0.1620935174111968,\n",
    "             'p': 0.2153396681546399,\n",
    "             'r': 0.13747476622638555},\n",
    " 'rouge-l': {'f': 0.2612394493528272,\n",
    "             'p': 0.3426511372716949,\n",
    "             'r': 0.22311445054693663}}\n",
    "``\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rouge_scores = compute_rouge_python(cand=predictions, ref=test_ds.get_target())\n",
    "pprint.pprint(rouge_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for testing\n",
    "sb.glue(\"rouge_1_f_score\", rouge_scores[\"rouge-1\"][\"f\"])\n",
    "sb.glue(\"rouge_2_f_score\", rouge_scores[\"rouge-2\"][\"f\"])\n",
    "sb.glue(\"rouge_l_f_score\", rouge_scores[\"rouge-l\"][\"f\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Distributed training with DistributedDataParallel (DDP)\n",
    "Please consult the notebook [Abstractive Summarization using MiniLM on CNN/DailyMails](./abstractive_summarization_unilm_cnndm.ipynb) for distributed training.    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Clean up "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.exists(DATA_DIR):\n",
    "    shutil.rmtree(DATA_DIR, ignore_errors=True)\n",
    "if os.path.exists(CACHE_DIR):\n",
    "    shutil.rmtree(CACHE_DIR, ignore_errors=True)\n",
    "    \n",
    "if CLEANUP_RESULTS:\n",
    "    if os.path.exists(MODEL_DIR):\n",
    "        shutil.rmtree(MODEL_DIR, ignore_errors=True)\n",
    "    if os.path.exists(RESULT_DIR):\n",
    "        shutil.rmtree(RESULT_DIR, ignore_errors=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Total notebook running time {}\".format(time.time() - start_time))"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python (nlp_gpu)",
   "language": "python",
   "name": "nlp_gpu"
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}