bert_icsl.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BERT for Intent Classification and Slot Labelling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "In this notebook, we will finetune a [BERT](https://arxiv.org/abs/1810.04805) model for the Intent Classification and Slot Labelling (ICSL) problem.\n",
    "\n",
    "Intent classification and slot labeling are two essential problems in Natural Language Understanding (NLU). In _intent classification_, the agent needs to detect the intention that the speaker's utterance conveys. For example, when the speaker says \"Book a flight from Long Beach to Seattle\", the intention is to book a flight ticket. In _slot labeling_, the agent needs to extract the semantic entities that are related to the intent. In our previous example, \"Long Beach\" and \"Seattle\" are two semantic constituents related to the flight, i.e., the origin and the destination.\n",
    "\n",
    "Essentially, _intent classification_ can be viewed as a sequence classification problem and _slot labeling_ can be viewed as a sequence tagging problem similar to Named-Entity Recognition (NER). Due to their inner correlation, these two tasks are usually trained jointly with a multi-task objective function.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "We first load the Airline Travel Information System (ATIS) dataset, which contains around 5000 utterance abouts travel plans and is a classical benchmark for ICSL."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded the ATIS dataset\n",
      "#Train/Dev/Test = 4478/500/893\n",
      "#Intent         = 18\n",
      "#Slot           = 127\n"
     ]
    }
   ],
   "source": [
    "from gluonnlp.data import ATISDataset\n",
    "\n",
    "train_data = ATISDataset('train')\n",
    "dev_data = ATISDataset('dev')\n",
    "test_data = ATISDataset('test')\n",
    "intent_vocab = train_data.intent_vocab\n",
    "slot_vocab = train_data.slot_vocab\n",
    "print('Loaded the ATIS dataset')\n",
    "print('#Train/Dev/Test = {}/{}/{}'.format(len(train_data), len(dev_data), len(test_data)))\n",
    "print('#Intent         = {}'.format(len(intent_vocab)))\n",
    "print('#Slot           = {}'.format(len(slot_vocab)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, let's display some samples from the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sentence: [\"i'm\", 'flying', 'from', 'boston', 'to', 'the', 'bay', 'area']\n",
      "    Tags: ['O', 'O', 'O', 'B-fromloc.city_name', 'O', 'O', 'B-toloc.city_name', 'I-toloc.city_name']\n",
      "   Label: atis_flight\n"
     ]
    }
   ],
   "source": [
    "print('Sentence:', dev_data[4][0])\n",
    "print('    Tags:', dev_data[4][1])\n",
    "print('   Label:', intent_vocab.idx_to_token[dev_data[4][2][0]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the BERT Model\n",
    "Next, we load the pretrained BERT model into the GPU. We load the BERT-base model trained on the Book Corpus + Wikipedia datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import mxnet as mx\n",
    "from mxnet import gluon\n",
    "from mxnet.gluon import nn, Block\n",
    "import gluonnlp as nlp\n",
    "import time\n",
    "from gluonnlp.data import BERTTokenizer\n",
    "\n",
    "dropout_prob = 0.1\n",
    "ctx = mx.gpu(0)\n",
    "\n",
    "bert_model, bert_vocab = nlp.model.get_model(name='bert_12_768_12',\n",
    "                                             dataset_name='book_corpus_wiki_en_cased',\n",
    "                                             pretrained=True,\n",
    "                                             ctx=ctx,\n",
    "                                             use_pooler=True,\n",
    "                                             use_decoder=False,\n",
    "                                             use_classifier=False,\n",
    "                                             dropout=dropout_prob,\n",
    "                                             embed_dropout=dropout_prob)\n",
    "tokenizer = BERTTokenizer(bert_vocab, lower=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "BERT uses the subword tokenization: e.g \"Sunnyvale\" will be tokenized as"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['sunny', '##vale']"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer('Sunnyvale')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thus, we will convert the original ATIS dataset to make sure that the state corresponds to the first subword is used to predict the slot label.\n",
    "\n",
    "<img src=\"explain_subword_tagging.png\" width=\"480\" align=\"left\"/>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class IDSLSubwordTransform(object):\n",
    "    \"\"\"Transform the dataset using the bert vocabulary and tokenizer\n",
    "    \"\"\"\n",
    "    def __init__(self, subword_vocab, subword_tokenizer, slot_vocab, cased=False):\n",
    "        \"\"\"\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        subword_vocab : Vocab\n",
    "        subword_tokenizer : Tokenizer\n",
    "        cased : bool\n",
    "            Whether to convert all characters to lower\n",
    "        \"\"\"\n",
    "        super(IDSLSubwordTransform, self).__init__()\n",
    "        self._subword_vocab = subword_vocab\n",
    "        self._subword_tokenizer = subword_tokenizer\n",
    "        self._slot_vocab = slot_vocab\n",
    "        self._cased = cased\n",
    "        self._slot_pad_id = self._slot_vocab['O']\n",
    "\n",
    "\n",
    "    def __call__(self, word_tokens, tags, intent_ids):\n",
    "        \"\"\" Transform the\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        word_tokens : List[str]\n",
    "        tags : List[str]\n",
    "        intent_ids : np.ndarray\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        subword_ids : np.ndarray\n",
    "        subword_mask : np.ndarray\n",
    "        selected : np.ndarray\n",
    "        padded_tag_ids : np.ndarray\n",
    "        intent_label : int\n",
    "        length : int\n",
    "        \"\"\"\n",
    "        subword_ids = []\n",
    "        subword_mask = []\n",
    "        selected = []\n",
    "        padded_tag_ids = []\n",
    "        intent_label = intent_ids[0]\n",
    "        ptr = 0\n",
    "        for token, tag in zip(word_tokens, tags):\n",
    "            if not self._cased:\n",
    "                token = token.lower()\n",
    "            token_sw_ids = self._subword_vocab[self._subword_tokenizer(token)]\n",
    "            subword_ids.extend(token_sw_ids)\n",
    "            subword_mask.extend([1] + [0] * (len(token_sw_ids) - 1))\n",
    "            selected.append(ptr)\n",
    "            padded_tag_ids.extend([self._slot_vocab[tag]] +\n",
    "                                  [self._slot_pad_id] * (len(token_sw_ids) - 1))\n",
    "            ptr += len(token_sw_ids)\n",
    "        length = len(subword_ids)\n",
    "        if len(subword_ids) != len(padded_tag_ids):\n",
    "            print(word_tokens)\n",
    "            print(tags)\n",
    "            print(subword_ids)\n",
    "            print(padded_tag_ids)\n",
    "        return np.array(subword_ids, dtype=np.int32),\\\n",
    "               np.array(subword_mask, dtype=np.int32),\\\n",
    "               np.array(selected, dtype=np.int32),\\\n",
    "               np.array(padded_tag_ids, dtype=np.int32),\\\n",
    "               intent_label,\\\n",
    "               length\n",
    "\n",
    "idsl_transform = IDSLSubwordTransform(subword_vocab=bert_vocab,\n",
    "                                      subword_tokenizer=tokenizer,\n",
    "                                      slot_vocab=slot_vocab,\n",
    "                                      cased=False)\n",
    "train_data_bert = train_data.transform(idsl_transform, lazy=False)\n",
    "dev_data_bert = dev_data.transform(idsl_transform, lazy=False)\n",
    "test_data_bert = test_data.transform(idsl_transform, lazy=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "token ids: [  178   112   182  3754  1121   171 15540  1320  1106  1103  5952  1298]\n",
      "mask: [1 0 0 1 1 1 0 0 1 1 1 1]\n",
      "index of the first subword: [ 0  3  4  5  8  9 10 11]\n",
      "slot label: [126 126 126 126 126  48 126 126 126 126  78 123]\n",
      "intent label: 10\n",
      "length: 10\n"
     ]
    }
   ],
   "source": [
    "print('token ids:', dev_data_bert[4][0])\n",
    "print('mask:', dev_data_bert[4][1])\n",
    "print('index of the first subword:', dev_data_bert[4][2])\n",
    "print('slot label:', dev_data_bert[4][3])\n",
    "print('intent label:', dev_data_bert[4][4])\n",
    "print('length:', dev_data_bert[4][4])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the Training Network\n",
    "We add two fully-connected layers on top of BERT to predict the slot labels and intent labels, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BERTForICSL(Block):\n",
    "    def __init__(self, bert, num_intent_classes, num_slot_classes, dropout_prob,\n",
    "                 prefix=None, params=None):\n",
    "        super(BERTForICSL, self).__init__(prefix=prefix, params=params)\n",
    "        self.bert = bert\n",
    "        with self.name_scope():\n",
    "            self.intent_classifier = nn.HybridSequential()\n",
    "            with self.intent_classifier.name_scope():\n",
    "                self.intent_classifier.add(nn.Dropout(rate=dropout_prob))\n",
    "                self.intent_classifier.add(nn.Dense(units=num_intent_classes, flatten=False))\n",
    "            self.slot_tagger = nn.HybridSequential()\n",
    "            with self.slot_tagger.name_scope():\n",
    "                self.slot_tagger.add(nn.Dropout(rate=dropout_prob))\n",
    "                self.slot_tagger.add(nn.Dense(units=num_slot_classes, flatten=False))\n",
    "\n",
    "    def forward(self, inputs, valid_length):\n",
    "        \"\"\"\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        inputs : NDArray\n",
    "            The input sentences, has shape (batch_size, seq_length)\n",
    "        valid_length : NDArray\n",
    "            The valid length of the sentences\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        intent_scores : NDArray\n",
    "            Shape (batch_size, num_classes)\n",
    "        slot_scores : NDArray\n",
    "            Shape (batch_size, seq_length, num_tag_types)\n",
    "        \"\"\"\n",
    "        token_types = mx.nd.zeros_like(inputs)\n",
    "        encoded_states, pooler_out = self.bert(inputs, token_types, valid_length)\n",
    "        intent_scores = self.intent_classifier(pooler_out)\n",
    "        slot_scores = self.slot_tagger(encoded_states)\n",
    "        return intent_scores, slot_scores\n",
    "net = BERTForICSL(bert_model, num_intent_classes=len(intent_vocab),\n",
    "                  num_slot_classes=len(slot_vocab), dropout_prob=dropout_prob)\n",
    "net.slot_tagger.initialize(ctx=ctx, init=mx.init.Normal(0.02))\n",
    "net.intent_classifier.initialize(ctx=ctx, init=mx.init.Normal(0.02))\n",
    "net.hybridize()\n",
    "intent_pred_loss = gluon.loss.SoftmaxCELoss()\n",
    "slot_pred_loss = gluon.loss.SoftmaxCELoss(batch_axis=[0, 1])\n",
    "intent_pred_loss.hybridize()\n",
    "slot_pred_loss.hybridize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BERTForICSL(\n",
      "  (intent_classifier): HybridSequential(\n",
      "    (0): Dropout(p = 0.1, axes=())\n",
      "    (1): Dense(None -> 18, linear)\n",
      "  )\n",
      "  (bert): BERTModel(\n",
      "    (pooler): Dense(768 -> 768, Activation(tanh))\n",
      "    (word_embed): HybridSequential(\n",
      "      (0): Embedding(28996 -> 768, float32)\n",
      "      (1): Dropout(p = 0.1, axes=())\n",
      "    )\n",
      "    (token_type_embed): HybridSequential(\n",
      "      (0): Embedding(2 -> 768, float32)\n",
      "      (1): Dropout(p = 0.1, axes=())\n",
      "    )\n",
      "    (encoder): BERTEncoder(\n",
      "      (transformer_cells): HybridSequential(\n",
      "        (0): BERTEncoderCell(\n",
      "          (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "          (attention_cell): MultiHeadAttentionCell(\n",
      "            (_base_cell): DotProductAttentionCell(\n",
      "              (_dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            )\n",
      "            (proj_query): Dense(768 -> 768, linear)\n",
      "            (proj_value): Dense(768 -> 768, linear)\n",
      "            (proj_key): Dense(768 -> 768, linear)\n",
      "          )\n",
      "          (proj): Dense(768 -> 768, linear)\n",
      "          (ffn): BERTPositionwiseFFN(\n",
      "            (activation): GELU()\n",
      "            (ffn_2): Dense(3072 -> 768, linear)\n",
      "            (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "            (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            (ffn_1): Dense(768 -> 3072, linear)\n",
      "          )\n",
      "          (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "        )\n",
      "        (1): BERTEncoderCell(\n",
      "          (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "          (attention_cell): MultiHeadAttentionCell(\n",
      "            (_base_cell): DotProductAttentionCell(\n",
      "              (_dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            )\n",
      "            (proj_query): Dense(768 -> 768, linear)\n",
      "            (proj_value): Dense(768 -> 768, linear)\n",
      "            (proj_key): Dense(768 -> 768, linear)\n",
      "          )\n",
      "          (proj): Dense(768 -> 768, linear)\n",
      "          (ffn): BERTPositionwiseFFN(\n",
      "            (activation): GELU()\n",
      "            (ffn_2): Dense(3072 -> 768, linear)\n",
      "            (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "            (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            (ffn_1): Dense(768 -> 3072, linear)\n",
      "          )\n",
      "          (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "        )\n",
      "        (2): BERTEncoderCell(\n",
      "          (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "          (attention_cell): MultiHeadAttentionCell(\n",
      "            (_base_cell): DotProductAttentionCell(\n",
      "              (_dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            )\n",
      "            (proj_query): Dense(768 -> 768, linear)\n",
      "            (proj_value): Dense(768 -> 768, linear)\n",
      "            (proj_key): Dense(768 -> 768, linear)\n",
      "          )\n",
      "          (proj): Dense(768 -> 768, linear)\n",
      "          (ffn): BERTPositionwiseFFN(\n",
      "            (activation): GELU()\n",
      "            (ffn_2): Dense(3072 -> 768, linear)\n",
      "            (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "            (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            (ffn_1): Dense(768 -> 3072, linear)\n",
      "          )\n",
      "          (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "        )\n",
      "        (3): BERTEncoderCell(\n",
      "          (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "          (attention_cell): MultiHeadAttentionCell(\n",
      "            (_base_cell): DotProductAttentionCell(\n",
      "              (_dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            )\n",
      "            (proj_query): Dense(768 -> 768, linear)\n",
      "            (proj_value): Dense(768 -> 768, linear)\n",
      "            (proj_key): Dense(768 -> 768, linear)\n",
      "          )\n",
      "          (proj): Dense(768 -> 768, linear)\n",
      "          (ffn): BERTPositionwiseFFN(\n",
      "            (activation): GELU()\n",
      "            (ffn_2): Dense(3072 -> 768, linear)\n",
      "            (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "            (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            (ffn_1): Dense(768 -> 3072, linear)\n",
      "          )\n",
      "          (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "        )\n",
      "        (4): BERTEncoderCell(\n",
      "          (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "          (attention_cell): MultiHeadAttentionCell(\n",
      "            (_base_cell): DotProductAttentionCell(\n",
      "              (_dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            )\n",
      "            (proj_query): Dense(768 -> 768, linear)\n",
      "            (proj_value): Dense(768 -> 768, linear)\n",
      "            (proj_key): Dense(768 -> 768, linear)\n",
      "          )\n",
      "          (proj): Dense(768 -> 768, linear)\n",
      "          (ffn): BERTPositionwiseFFN(\n",
      "            (activation): GELU()\n",
      "            (ffn_2): Dense(3072 -> 768, linear)\n",
      "            (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "            (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            (ffn_1): Dense(768 -> 3072, linear)\n",
      "          )\n",
      "          (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "        )\n",
      "        (5): BERTEncoderCell(\n",
      "          (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "          (attention_cell): MultiHeadAttentionCell(\n",
      "            (_base_cell): DotProductAttentionCell(\n",
      "              (_dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            )\n",
      "            (proj_query): Dense(768 -> 768, linear)\n",
      "            (proj_value): Dense(768 -> 768, linear)\n",
      "            (proj_key): Dense(768 -> 768, linear)\n",
      "          )\n",
      "          (proj): Dense(768 -> 768, linear)\n",
      "          (ffn): BERTPositionwiseFFN(\n",
      "            (activation): GELU()\n",
      "            (ffn_2): Dense(3072 -> 768, linear)\n",
      "            (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "            (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            (ffn_1): Dense(768 -> 3072, linear)\n",
      "          )\n",
      "          (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "        )\n",
      "        (6): BERTEncoderCell(\n",
      "          (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "          (attention_cell): MultiHeadAttentionCell(\n",
      "            (_base_cell): DotProductAttentionCell(\n",
      "              (_dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            )\n",
      "            (proj_query): Dense(768 -> 768, linear)\n",
      "            (proj_value): Dense(768 -> 768, linear)\n",
      "            (proj_key): Dense(768 -> 768, linear)\n",
      "          )\n",
      "          (proj): Dense(768 -> 768, linear)\n",
      "          (ffn): BERTPositionwiseFFN(\n",
      "            (activation): GELU()\n",
      "            (ffn_2): Dense(3072 -> 768, linear)\n",
      "            (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "            (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            (ffn_1): Dense(768 -> 3072, linear)\n",
      "          )\n",
      "          (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "        )\n",
      "        (7): BERTEncoderCell(\n",
      "          (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "          (attention_cell): MultiHeadAttentionCell(\n",
      "            (_base_cell): DotProductAttentionCell(\n",
      "              (_dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            )\n",
      "            (proj_query): Dense(768 -> 768, linear)\n",
      "            (proj_value): Dense(768 -> 768, linear)\n",
      "            (proj_key): Dense(768 -> 768, linear)\n",
      "          )\n",
      "          (proj): Dense(768 -> 768, linear)\n",
      "          (ffn): BERTPositionwiseFFN(\n",
      "            (activation): GELU()\n",
      "            (ffn_2): Dense(3072 -> 768, linear)\n",
      "            (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "            (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            (ffn_1): Dense(768 -> 3072, linear)\n",
      "          )\n",
      "          (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "        )\n",
      "        (8): BERTEncoderCell(\n",
      "          (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "          (attention_cell): MultiHeadAttentionCell(\n",
      "            (_base_cell): DotProductAttentionCell(\n",
      "              (_dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            )\n",
      "            (proj_query): Dense(768 -> 768, linear)\n",
      "            (proj_value): Dense(768 -> 768, linear)\n",
      "            (proj_key): Dense(768 -> 768, linear)\n",
      "          )\n",
      "          (proj): Dense(768 -> 768, linear)\n",
      "          (ffn): BERTPositionwiseFFN(\n",
      "            (activation): GELU()\n",
      "            (ffn_2): Dense(3072 -> 768, linear)\n",
      "            (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "            (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            (ffn_1): Dense(768 -> 3072, linear)\n",
      "          )\n",
      "          (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "        )\n",
      "        (9): BERTEncoderCell(\n",
      "          (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "          (attention_cell): MultiHeadAttentionCell(\n",
      "            (_base_cell): DotProductAttentionCell(\n",
      "              (_dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            )\n",
      "            (proj_query): Dense(768 -> 768, linear)\n",
      "            (proj_value): Dense(768 -> 768, linear)\n",
      "            (proj_key): Dense(768 -> 768, linear)\n",
      "          )\n",
      "          (proj): Dense(768 -> 768, linear)\n",
      "          (ffn): BERTPositionwiseFFN(\n",
      "            (activation): GELU()\n",
      "            (ffn_2): Dense(3072 -> 768, linear)\n",
      "            (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "            (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            (ffn_1): Dense(768 -> 3072, linear)\n",
      "          )\n",
      "          (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "        )\n",
      "        (10): BERTEncoderCell(\n",
      "          (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "          (attention_cell): MultiHeadAttentionCell(\n",
      "            (_base_cell): DotProductAttentionCell(\n",
      "              (_dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            )\n",
      "            (proj_query): Dense(768 -> 768, linear)\n",
      "            (proj_value): Dense(768 -> 768, linear)\n",
      "            (proj_key): Dense(768 -> 768, linear)\n",
      "          )\n",
      "          (proj): Dense(768 -> 768, linear)\n",
      "          (ffn): BERTPositionwiseFFN(\n",
      "            (activation): GELU()\n",
      "            (ffn_2): Dense(3072 -> 768, linear)\n",
      "            (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "            (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            (ffn_1): Dense(768 -> 3072, linear)\n",
      "          )\n",
      "          (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "        )\n",
      "        (11): BERTEncoderCell(\n",
      "          (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "          (attention_cell): MultiHeadAttentionCell(\n",
      "            (_base_cell): DotProductAttentionCell(\n",
      "              (_dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            )\n",
      "            (proj_query): Dense(768 -> 768, linear)\n",
      "            (proj_value): Dense(768 -> 768, linear)\n",
      "            (proj_key): Dense(768 -> 768, linear)\n",
      "          )\n",
      "          (proj): Dense(768 -> 768, linear)\n",
      "          (ffn): BERTPositionwiseFFN(\n",
      "            (activation): GELU()\n",
      "            (ffn_2): Dense(3072 -> 768, linear)\n",
      "            (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "            (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "            (ffn_1): Dense(768 -> 3072, linear)\n",
      "          )\n",
      "          (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "        )\n",
      "      )\n",
      "      (layer_norm): BERTLayerNorm(scale=True, axis=-1, eps=1e-12, center=True, in_channels=768)\n",
      "      (dropout_layer): Dropout(p = 0.1, axes=())\n",
      "    )\n",
      "  )\n",
      "  (slot_tagger): HybridSequential(\n",
      "    (0): Dropout(p = 0.1, axes=())\n",
      "    (1): Dense(None -> 127, linear)\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "print(net)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the DataLoader and Trainer for Training/Validation/Testing\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/gluon-nlp/src/gluonnlp/data/batchify/batchify.py:228: UserWarning: Padding value 0 is used in data.batchify.Pad(). Please check whether this is intended (e.g. value of padding index in the vocabulary).\n",
      "  'Padding value 0 is used in data.batchify.Pad(). '\n"
     ]
    }
   ],
   "source": [
    "batch_size = 32\n",
    "learning_rate = 5E-5\n",
    "\n",
    "trainer = gluon.Trainer(net.collect_params(), 'bertadam',\n",
    "                        {'learning_rate': learning_rate, 'wd': 0.0})\n",
    "batchify_fn = nlp.data.batchify.Tuple(nlp.data.batchify.Pad(),    # Subword ID\n",
    "                                      nlp.data.batchify.Pad(),    # Subword Mask\n",
    "                                      nlp.data.batchify.Pad(),    # Beginning of subword\n",
    "                                      nlp.data.batchify.Pad(),    # Tag IDs\n",
    "                                      nlp.data.batchify.Stack(),  # Intent Label\n",
    "                                      nlp.data.batchify.Stack())  # Valid Length\n",
    "train_batch_sampler = nlp.data.sampler.SortedBucketSampler(\n",
    "    [len(ele) for ele in train_data_bert],\n",
    "    batch_size=batch_size,\n",
    "    mult=20,\n",
    "    shuffle=True)\n",
    "train_loader = gluon.data.DataLoader(dataset=train_data_bert,\n",
    "                                     num_workers=4,\n",
    "                                     batch_sampler=train_batch_sampler,\n",
    "                                     batchify_fn=batchify_fn)\n",
    "dev_loader = gluon.data.DataLoader(dataset=dev_data_bert,\n",
    "                                   num_workers=4,\n",
    "                                   batch_size=batch_size,\n",
    "                                   batchify_fn=batchify_fn,\n",
    "                                   shuffle=False)\n",
    "test_loader = gluon.data.DataLoader(dataset=test_data_bert,\n",
    "                                    num_workers=4,\n",
    "                                    batch_size=batch_size,\n",
    "                                    batchify_fn=batchify_fn,\n",
    "                                    shuffle=False)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100%|██████████| 140/140 [00:13<00:00, 10.85it/s]\n",
      "[Epoch 0] train intent/slot = 0.827/1.207, #token per second=5465\n",
      "100%|██████████| 140/140 [00:12<00:00, 10.31it/s]\n",
      "[Epoch 1] train intent/slot = 0.183/0.229, #token per second=5940\n",
      "100%|██████████| 140/140 [00:12<00:00, 11.07it/s]\n",
      "[Epoch 2] train intent/slot = 0.089/0.124, #token per second=5968\n",
      "100%|██████████| 140/140 [00:12<00:00, 11.77it/s]\n",
      "[Epoch 3] train intent/slot = 0.056/0.088, #token per second=5955\n",
      "100%|██████████| 140/140 [00:12<00:00, 11.43it/s]\n",
      "[Epoch 4] train intent/slot = 0.039/0.071, #token per second=5964\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "import sys\n",
    "\n",
    "nepochs = 5\n",
    "warmup_ratio = 0.1\n",
    "step_num = 0\n",
    "num_train_steps = int(len(train_batch_sampler) * nepochs)\n",
    "num_warmup_steps = int(num_train_steps * warmup_ratio)\n",
    "best_dev_sf1 = -1\n",
    "for epoch_id in range(nepochs):\n",
    "    avg_train_intent_loss = 0.0\n",
    "    avg_train_slot_loss = 0.0\n",
    "    nsample = 0\n",
    "    nslot = 0\n",
    "    ntoken = 0\n",
    "    train_epoch_start = time.time()\n",
    "    for token_ids, mask, selected, slot_ids, intent_label, valid_length in tqdm(train_loader, file=sys.stdout):\n",
    "        # Copy data to the context, i.e., GPU in our example\n",
    "        token_ids = mx.nd.array(token_ids, ctx=ctx).astype(np.int32)\n",
    "        mask = mx.nd.array(mask, ctx=ctx).astype(np.float32)\n",
    "        slot_ids = mx.nd.array(slot_ids, ctx=ctx).astype(np.int32)\n",
    "        intent_label = mx.nd.array(intent_label, ctx=ctx).astype(np.int32)\n",
    "        valid_length = mx.nd.array(valid_length, ctx=ctx).astype(np.float32)\n",
    "        batch_nslots = mask.sum().asscalar()\n",
    "        batch_nsample = token_ids.shape[0]\n",
    "\n",
    "        # Set learning rate warm-up\n",
    "        step_num += 1\n",
    "        if step_num < num_warmup_steps:\n",
    "            new_lr = learning_rate * step_num / num_warmup_steps\n",
    "        else:\n",
    "            offset = ((step_num - num_warmup_steps) * learning_rate /\n",
    "                      (num_train_steps - num_warmup_steps))\n",
    "            new_lr = learning_rate - offset\n",
    "        trainer.set_learning_rate(new_lr)\n",
    "\n",
    "        # Begin to calculate the gradient\n",
    "        with mx.autograd.record():\n",
    "            intent_scores, slot_scores = net(token_ids, valid_length)\n",
    "            intent_loss = intent_pred_loss(intent_scores, intent_label)\n",
    "            slot_loss = slot_pred_loss(slot_scores, slot_ids, mask.expand_dims(axis=-1))\n",
    "            intent_loss = intent_loss.mean()\n",
    "            slot_loss = slot_loss.sum() / batch_nslots\n",
    "            loss = intent_loss + slot_loss\n",
    "            loss.backward()\n",
    "        trainer.update(1.0)\n",
    "        avg_train_intent_loss += intent_loss.asscalar() * batch_nsample\n",
    "        avg_train_slot_loss += slot_loss.asscalar() * batch_nslots\n",
    "        nsample += batch_nsample\n",
    "        nslot += batch_nslots\n",
    "        ntoken += valid_length.sum().asscalar()\n",
    "    train_epoch_end = time.time()\n",
    "    avg_train_intent_loss /= nsample\n",
    "    avg_train_slot_loss /= nslot\n",
    "    print('[Epoch {}] train intent/slot = {:.3f}/{:.3f}, #token per second={:.0f}'.format(\n",
    "        epoch_id, avg_train_intent_loss, avg_train_slot_loss,\n",
    "        ntoken / (train_epoch_end - train_epoch_start)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Epoch 4]    dev intent/slot = 0.128/0.092, slot f1 = 93.93, intent acc = 97.80\n",
      "[Epoch 4]    test intent/slot = 0.123/0.154, slot f1 = 91.92, intent acc = 98.21\n"
     ]
    }
   ],
   "source": [
    "from seqeval.metrics import f1_score as ner_f1_score\n",
    "\n",
    "def evaluation(ctx, data_loader, net, intent_pred_loss, slot_pred_loss, slot_vocab):\n",
    "    \"\"\"\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    ctx : Context\n",
    "    data_loader : DataLoader\n",
    "    net : Block\n",
    "    intent_pred_loss : Block\n",
    "    slot_pred_loss : Block\n",
    "    slot_vocab : Vocab\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    avg_intent_loss : float\n",
    "    avg_slot_loss : float\n",
    "    intent_acc : float\n",
    "    slot_f1 : float\n",
    "    pred_slots : list\n",
    "    gt_slots : list\n",
    "    \"\"\"\n",
    "    nsample = 0\n",
    "    nslot = 0\n",
    "    avg_intent_loss = 0\n",
    "    avg_slot_loss = 0\n",
    "    correct_intent = 0\n",
    "    pred_slots = []\n",
    "    gt_slots = []\n",
    "    for token_ids, mask, selected, slot_ids, intent_label, valid_length in data_loader:\n",
    "        token_ids = mx.nd.array(token_ids, ctx=ctx).astype(np.int32)\n",
    "        mask = mx.nd.array(mask, ctx=ctx).astype(np.float32)\n",
    "        slot_ids = mx.nd.array(slot_ids, ctx=ctx).astype(np.int32)\n",
    "        intent_label = mx.nd.array(intent_label, ctx=ctx).astype(np.int32)\n",
    "        valid_length = mx.nd.array(valid_length, ctx=ctx).astype(np.float32)\n",
    "        batch_nslot = mask.sum().asscalar()\n",
    "        batch_nsample = token_ids.shape[0]\n",
    "        # Forward network\n",
    "        intent_scores, slot_scores = net(token_ids, valid_length)\n",
    "        intent_loss = intent_pred_loss(intent_scores, intent_label)\n",
    "        slot_loss = slot_pred_loss(slot_scores, slot_ids, mask.expand_dims(axis=-1))\n",
    "        avg_intent_loss += intent_loss.sum().asscalar()\n",
    "        avg_slot_loss += slot_loss.sum().asscalar()\n",
    "        pred_slot_ids = mx.nd.argmax(slot_scores, axis=-1).astype(np.int32)\n",
    "        correct_intent += (mx.nd.argmax(intent_scores, axis=-1).astype(np.int32)\n",
    "                           == intent_label).sum().asscalar()\n",
    "        for i in range(batch_nsample):\n",
    "            ele_valid_length = int(valid_length[i].asscalar())\n",
    "            ele_sel = selected[i].asnumpy()[:ele_valid_length]\n",
    "            ele_gt_slot_ids = slot_ids[i].asnumpy()[ele_sel]\n",
    "            ele_pred_slot_ids = pred_slot_ids[i].asnumpy()[ele_sel]\n",
    "            ele_gt_slot_tokens = [slot_vocab.idx_to_token[v] for v in ele_gt_slot_ids]\n",
    "            ele_pred_slot_tokens = [slot_vocab.idx_to_token[v] for v in ele_pred_slot_ids]\n",
    "            gt_slots.append(ele_gt_slot_tokens)\n",
    "            pred_slots.append(ele_pred_slot_tokens)\n",
    "        nsample += batch_nsample\n",
    "        nslot += batch_nslot\n",
    "    avg_intent_loss /= nsample\n",
    "    avg_slot_loss /= nslot\n",
    "    intent_acc = correct_intent / float(nsample)\n",
    "    slot_f1 = ner_f1_score(pred_slots, gt_slots)\n",
    "    return avg_intent_loss, avg_slot_loss, intent_acc, slot_f1, pred_slots, gt_slots\n",
    "\n",
    "avg_dev_intent_loss, avg_dev_slot_loss, dev_intent_acc, dev_slot_f1, dev_pred_slots, dev_gt_slots\\\n",
    "    = evaluation(ctx, dev_loader, net, intent_pred_loss, slot_pred_loss, slot_vocab)\n",
    "print('[Epoch {}]    dev intent/slot = {:.3f}/{:.3f}, slot f1 = {:.2f}, intent acc = {:.2f}'.format(\n",
    "    epoch_id, avg_dev_intent_loss, avg_dev_slot_loss, dev_slot_f1 * 100, dev_intent_acc * 100))\n",
    "avg_test_intent_loss, avg_test_slot_loss, test_intent_acc, test_slot_f1, test_pred_slots, test_gt_slots \\\n",
    "    = evaluation(ctx, test_loader, net, intent_pred_loss, slot_pred_loss, slot_vocab)\n",
    "print('[Epoch {}]    test intent/slot = {:.3f}/{:.3f}, slot f1 = {:.2f}, intent acc = {:.2f}'.format(\n",
    "    epoch_id, avg_test_intent_loss, avg_test_slot_loss, test_slot_f1 * 100, test_intent_acc * 100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sentence:     ['show', 'me', 'all', 'round', 'trip', 'flights', 'between', 'houston', 'and', 'las', 'vegas']\n",
      "Ground Truth: ['O', 'O', 'O', 'B-round_trip', 'I-round_trip', 'O', 'O', 'B-fromloc.city_name', 'O', 'B-toloc.city_name', 'I-toloc.city_name', 'O', 'O', 'O']\n",
      "Prediction:   ['O', 'O', 'O', 'B-round_trip', 'I-round_trip', 'O', 'O', 'B-fromloc.city_name', 'O', 'B-toloc.city_name', 'I-toloc.city_name', 'O', 'O', 'O']\n"
     ]
    }
   ],
   "source": [
    "print('Sentence:    ', dev_data[1][0])\n",
    "print('Ground Truth:', dev_gt_slots[1])\n",
    "print('Prediction:  ', dev_gt_slots[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Full Script + Results\n",
    "\n",
    "You can run the experiments in https://github.com/dmlc/gluon-nlp/tree/master/scripts/intent_slot\n",
    "\n",
    "For ATIS\n",
    "\n",
    "| Models | Intent Acc (%) | Slot F1 (%) |\n",
    "| ------ | ------------------------ | ----------- |\n",
    "| [Intent Gating & self-attention, EMNLP 2018](https://www.aclweb.org/anthology/D18-1417) | 98.77 | 96.52 |\n",
    "| [BLSTM-CRF + ELMo, AAAI 2019](https://arxiv.org/abs/1811.05370) | 97.42 | 95.62 |\n",
    "| [Joint BERT, Arxiv 2019](https://arxiv.org/pdf/1902.10909.pdf) |  97.5 | 96.1 |\n",
    "| Ours | 98.66±0.00  | 95.88±0.04 |\n",
    "\n",
    "For SNIPS\n",
    "\n",
    "| Models | Intent Acc (%) | Slot F1 (%) |\n",
    "| ------ | ------------------------ | ----------- |\n",
    "| [BLSTM-CRF + ELMo, AAAI 2019](https://arxiv.org/abs/1811.05370) | 99.29 | 93.90 |\n",
    "| [Joint BERT, Arxiv 2019](https://arxiv.org/pdf/1902.10909.pdf) | 98.60 | 97.00 |\n",
    "| Ours | 98.81±0.13 | 95.94±0.10 |\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}