examples/sequence_tagging.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sequence Tagging With an RNN Using Poutyne\n",
    "\n",
    "In this notebook, we will do sequence tagging with RNNs using Poutyne.\n",
    "\n",
    "Let's install the latest version of Poutyne (if it is not already), and import all the needed packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --upgrade fasttext\n",
    "%pip install --upgrade poutyne\n",
    "%matplotlib inline\n",
    "import contextlib\n",
    "import os\n",
    "import pickle\n",
    "import re\n",
    "import sys\n",
    "from io import TextIOBase\n",
    "\n",
    "import fasttext\n",
    "import fasttext.util\n",
    "import requests\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.nn.utils.rnn import pad_packed_sequence, pack_padded_sequence, pad_sequence, PackedSequence\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from poutyne import set_seeds, Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Also, we need to set Python's, NumPy's and PyTorch's seeds by using Poutyne function so that our training is (almost) reproducible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "set_seeds(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Train a Recurrent Neural Network (RNN)\n",
    "\n",
    "We train an RNN, or more precisely, an LSTM, to predict the sequence of tags associated with a given address, which is also known as address parsing.\n",
    "\n",
    "This task consists of detecting (by tagging) the different parts of an address such as the civic number, the street name or the postal code (or zip code). The following figure shows an example of such a tagging.\n",
    "\n",
    "![address parsing canada](https://poutyne.org/_images/address_parsing.png)\n",
    "\n",
    "Since addresses are written in a predetermined sequence, RNN is the best way to crack this problem. For our architecture, we will use two components, an RNN and a fully-connected layer.\n",
    "\n",
    "### Training Constants"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's set our training constants. We first have the CUDA device used for training if one is present. Secondly, we set the batch size (i.e. the number of elements to see before updating the model) and the learning rate for the optimizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "cuda_device = 0\n",
    "device = torch.device(\"cuda:%d\" % cuda_device if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "batch_size = 32\n",
    "lr = 0.1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### RNN\n",
    "\n",
    "For the first components, instead of using a vanilla RNN, we use a variant of it, know as a long short-term memory (LSTM) (to learn more about [LSTM](https://colah.github.io/posts/2015-08-Understanding-LSTMs/). For now, we use a single layer unidirectional LSTM. \n",
    "\n",
    "Also, since our data is textual, we will use the well-known word embeddings to encode the textual information. So the LSTM input and hidden state dimensions will be of the same size. This size corresponds to the word embeddings dimension, which in our case will be the [French pre trained](https://fasttext.cc/docs/en/crawl-vectors.html) fastText embeddings of dimension `300`.\n",
    "\n",
    "> See [here](https://discuss.pytorch.org/t/could-someone-explain-batch-first-true-in-lstm/15402) the explanation why we use the `batch_first` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "dimension = 300\n",
    "num_layer = 1\n",
    "bidirectional = False\n",
    "\n",
    "lstm_network = nn.LSTM(\n",
    "    input_size=dimension,\n",
    "    hidden_size=dimension,\n",
    "    num_layers=num_layer,\n",
    "    bidirectional=bidirectional,\n",
    "    batch_first=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Fully-connected Layer\n",
    "\n",
    "We use this layer to map the representation of the LSTM (`300`) to the tag space (8, the number of tags) and predict the most likely tag using a softmax."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "input_dim = dimension  # the output of the LSTM\n",
    "tag_dimension = 8\n",
    "\n",
    "fully_connected_network = nn.Linear(input_dim, tag_dimension)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### The Dataset\n",
    "\n",
    "Now let's download our dataset; it already split into a train, valid and test set using the following."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def download_data(saving_dir, data_type):\n",
    "    \"\"\"\n",
    "    Function to download the dataset using data_type to specify if we want the train, valid or test.\n",
    "    \"\"\"\n",
    "\n",
    "    root_url = \"https://graal-research.github.io/poutyne-external-assets/tips_and_tricks_assets/{}.p\"\n",
    "\n",
    "    url = root_url.format(data_type)\n",
    "    r = requests.get(url)\n",
    "    os.makedirs(saving_dir, exist_ok=True)\n",
    "\n",
    "    open(os.path.join(saving_dir, f\"{data_type}.p\"), 'wb').write(r.content)\n",
    "\n",
    "\n",
    "download_data('./datasets/addresses/', \"train\")\n",
    "download_data('./datasets/addresses/', \"valid\")\n",
    "download_data('./datasets/addresses/', \"test\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now let's load in memory the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the data\n",
    "\n",
    "train_data = pickle.load(open(\"./datasets/addresses/train.p\", \"rb\"))  # 80,000 examples\n",
    "valid_data = pickle.load(open(\"./datasets/addresses/valid.p\", \"rb\"))  # 20,000 examples\n",
    "test_data = pickle.load(open(\"./datasets/addresses/test.p\", \"rb\"))  # 30,000 examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "If we take a look at the training dataset, it's a list of `80,000` tuples where the first element is the full address, and the second element is a list of the tag (the ground truth)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('33 harnesworth crescent city of hamilton ontario l8b0j3',\n",
       "  ['StreetNumber',\n",
       "   'StreetName',\n",
       "   'StreetName',\n",
       "   'Municipality',\n",
       "   'Municipality',\n",
       "   'Municipality',\n",
       "   'Province',\n",
       "   'PostalCode']),\n",
       " ('1449 mouettes longueuil quebec j4j5k4',\n",
       "  ['StreetNumber', 'StreetName', 'Municipality', 'Province', 'PostalCode'])]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data[0:2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Since the address is a text, we need to *convert* it into categorical value, such as word embeddings, for that we will use a vectorizer. This embedding vectorizer will be able to extract for every word embedding value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Warning : `load_model` does not return WordVectorModel or SupervisedModel any more, but a `FastText` object which is very similar.\n"
     ]
    }
   ],
   "source": [
    "# We use this class so that the download templating of the fasttext\n",
    "# script be not buggy as hell in notebooks.\n",
    "class LookForProgress(TextIOBase):\n",
    "    def __init__(self, stdout):\n",
    "        self.stdout = stdout\n",
    "        self.regex = re.compile(r'([0-9]+(\\.[0-9]+)?%)', re.IGNORECASE)\n",
    "\n",
    "    def write(self, o):\n",
    "        res = self.regex.findall(o)\n",
    "        if len(res) != 0:\n",
    "            print(f\"\\r{res[-1][0]}\", end='', file=self.stdout)\n",
    "\n",
    "\n",
    "class EmbeddingVectorizer:\n",
    "    def __init__(self):\n",
    "        \"\"\"\n",
    "        Embedding vectorizer\n",
    "        \"\"\"\n",
    "        with contextlib.redirect_stdout(LookForProgress(sys.stdout)):\n",
    "            fasttext.util.download_model('fr', if_exists='ignore')\n",
    "        self.embedding_model = fasttext.load_model(\"./cc.fr.300.bin\")\n",
    "\n",
    "    def __call__(self, address):\n",
    "        \"\"\"\n",
    "        Convert address to embedding vectors\n",
    "        :param address: The address to convert\n",
    "        :return: The embeddings vectors\n",
    "        \"\"\"\n",
    "        embeddings = []\n",
    "        for word in address.split():\n",
    "            embeddings.append(self.embedding_model[word])\n",
    "        return embeddings\n",
    "\n",
    "\n",
    "embedding_model = EmbeddingVectorizer()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We also need a vectorizer to convert the address tag (e.g. StreetNumber, StreetName) into categorical values. So we will use a Vectorizer class that can use the embedding vectorizer and convert the address tag. We will explain and use the argument `predict` later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "class Vectorizer:\n",
    "    def __init__(self, dataset, embedding_model, predict=False):\n",
    "        self.data = dataset\n",
    "        self.embedding_model = embedding_model\n",
    "        self.predict = predict\n",
    "        self.tags_set = {\n",
    "            \"StreetNumber\": 0,\n",
    "            \"StreetName\": 1,\n",
    "            \"Unit\": 2,\n",
    "            \"Municipality\": 3,\n",
    "            \"Province\": 4,\n",
    "            \"PostalCode\": 5,\n",
    "            \"Orientation\": 6,\n",
    "            \"GeneralDelivery\": 7,\n",
    "        }\n",
    "\n",
    "    def __len__(self):\n",
    "        # for the dataloader\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, item):\n",
    "        data = self.data[item]\n",
    "\n",
    "        if not self.predict:\n",
    "            address = data[0]\n",
    "            address_vector = self.embedding_model(address)\n",
    "\n",
    "            tags = data[1]\n",
    "            idx_tags = self._convert_tags_to_idx(tags)\n",
    "            return address_vector, idx_tags\n",
    "\n",
    "        address_vector = self.embedding_model(data)\n",
    "        return address_vector\n",
    "\n",
    "    def _convert_tags_to_idx(self, tags):\n",
    "        idx_tags = []\n",
    "        for tag in tags:\n",
    "            idx_tags.append(self.tags_set[tag])\n",
    "        return idx_tags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data_vectorized = Vectorizer(train_data, embedding_model)\n",
    "valid_data_vectorized = Vectorizer(valid_data, embedding_model)\n",
    "test_data_vectorized = Vectorizer(test_data, embedding_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "#### DataLoader\n",
    "\n",
    "Now, since all the addresses are not of the same size, it is impossible to batch them together since all elements of a tensor must have the same lengths. But there is a trick, padding!\n",
    "\n",
    "The idea is simple. We add *empty* tokens at the end of each sequence up to the longest one in a batch. For the word vectors, we add vectors of 0 as padding. For the tag indices, we pad with -100s. We do so because of the [cross-entropy loss](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html#torch.nn.CrossEntropyLoss), the accuracy metric and the [F1 metric](https://poutyne.org/metrics.html#poutyne.FBeta) all ignore targets with values of -100.\n",
    "\n",
    "To do this padding, we use the `collate_fn` argument of the [PyTorch `DataLoader`](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader), and on running time, that process will be done. One thing to take into account, since we pad the sequence, we need each sequence's lengths to unpad them in the forward pass. That way, we can pad and pack the sequence to minimize the training time (read [this good explanation](https://stackoverflow.com/questions/51030782/why-do-we-pack-the-sequences-in-pytorch) of why we pad and pack sequences)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def pad_collate_fn(batch):\n",
    "    \"\"\"\n",
    "    The collate_fn that can add padding to the sequences so all can have\n",
    "    the same length as the longest one.\n",
    "\n",
    "    Args:\n",
    "        batch (List[List, List]): The batch data, where the first element\n",
    "        of the tuple are the word idx and the second element are the target\n",
    "        label.\n",
    "\n",
    "    Returns:\n",
    "        A tuple (x, y). The element x is a tensor of packed sequence .\n",
    "        The element y is a tensor of padded tag indices. The word vectors are\n",
    "        padded with vectors of 0s and the tag indices are padded with -100s.\n",
    "        Padding with -100 is done because of the cross-entropy loss and the\n",
    "        accuracy metric ignores the targets with values -100.\n",
    "    \"\"\"\n",
    "\n",
    "    # This gets us two lists of tensors and a list of integer.\n",
    "    # Each tensor in the first list is a sequence of word vectors.\n",
    "    # Each tensor in the second list is a sequence of tag indices.\n",
    "    # The list of integer consist of the lengths of the sequences in order.\n",
    "    sequences_vectors, sequences_labels, lengths = zip(\n",
    "        *[\n",
    "            (torch.FloatTensor(np.stack(seq_vectors)), torch.LongTensor(labels), len(seq_vectors))\n",
    "            for (seq_vectors, labels) in sorted(batch, key=lambda x: len(x[0]), reverse=True)\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    lengths = torch.LongTensor(lengths)\n",
    "\n",
    "    padded_sequences_vectors = pad_sequence(sequences_vectors, batch_first=True, padding_value=0)\n",
    "    pack_padded_sequences_vectors = pack_padded_sequence(\n",
    "        padded_sequences_vectors, lengths.cpu(), batch_first=True\n",
    "    )  # We pack the padded sequence to improve the computational speed during training\n",
    "\n",
    "    padded_sequences_labels = pad_sequence(sequences_labels, batch_first=True, padding_value=-100)\n",
    "\n",
    "    return pack_padded_sequences_vectors, padded_sequences_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = DataLoader(\n",
    "    train_data_vectorized,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    "    collate_fn=pad_collate_fn,\n",
    "    num_workers=2,\n",
    ")\n",
    "valid_loader = DataLoader(valid_data_vectorized, batch_size=batch_size, collate_fn=pad_collate_fn, num_workers=2)\n",
    "test_loader = DataLoader(test_data_vectorized, batch_size=batch_size, collate_fn=pad_collate_fn, num_workers=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Full Network\n",
    "\n",
    "Since our sequences are of variable lengths and we want to be the most efficient possible by packing them, we cannot use the [PyTorch `nn.Sequential`](https://pytorch.org/docs/stable/generated/torch.nn.Sequential.html) class to define our model, so we define the forward pass for it to pack and unpack the sequences (again, you can read [this good explanation](https://stackoverflow.com/questions/51030782/why-do-we-pack-the-sequences-in-pytorch) of why we pad and pack sequences)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "class FullNetWork(nn.Module):\n",
    "    def __init__(self, lstm_network, fully_connected_network):\n",
    "        super().__init__()\n",
    "        self.hidden_state = None\n",
    "\n",
    "        self.lstm_network = lstm_network\n",
    "        self.fully_connected_network = fully_connected_network\n",
    "\n",
    "    def forward(self, pack_padded_sequences_vectors: PackedSequence):\n",
    "        \"\"\"\n",
    "        Defines the computation performed at every call.\n",
    "        \"\"\"\n",
    "        lstm_out, self.hidden_state = self.lstm_network(pack_padded_sequences_vectors)\n",
    "        lstm_out, _ = pad_packed_sequence(lstm_out, batch_first=True)\n",
    "\n",
    "        tag_space = self.fully_connected_network(lstm_out)\n",
    "        return tag_space.transpose(-1, 1)  # We need to transpose since it's a sequence\n",
    "\n",
    "\n",
    "full_network = FullNetWork(lstm_network, fully_connected_network)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Summary\n",
    "\n",
    "So we have created an LSTM network (`lstm_network`), a fully connected network (`fully_connected_network`), those two components are used in the full network. This full network used padded, packed sequences (defined in the forward pass), so we created the `pad_collate_fn` function to process the needed work. The DataLoader will conduct that process. Finally, when we load the data, this will be done using the vectorizer, so the address will be represented using word embeddings. Also, the address components will be converted into categorical value (from 0 to 7)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## The Training Loop\n",
    "\n",
    "Now that we have all the components for the network let's train our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[35mEpoch: \u001b[36m 1/10 \u001b[35mTrain steps: \u001b[36m2500 \u001b[35mVal steps: \u001b[36m625 \u001b[32m17.43s \u001b[35mloss:\u001b[94m 0.402440\u001b[35m acc:\u001b[94m 87.340841\u001b[35m val_loss:\u001b[94m 0.074837\u001b[35m val_acc:\u001b[94m 98.491375\u001b[0m\n",
      "\u001b[35mEpoch: \u001b[36m 2/10 \u001b[35mTrain steps: \u001b[36m2500 \u001b[35mVal steps: \u001b[36m625 \u001b[32m17.73s \u001b[35mloss:\u001b[94m 0.044888\u001b[35m acc:\u001b[94m 99.044959\u001b[35m val_loss:\u001b[94m 0.026150\u001b[35m val_acc:\u001b[94m 99.355461\u001b[0m\n",
      "\u001b[35mEpoch: \u001b[36m 3/10 \u001b[35mTrain steps: \u001b[36m2500 \u001b[35mVal steps: \u001b[36m625 \u001b[32m17.63s \u001b[35mloss:\u001b[94m 0.019844\u001b[35m acc:\u001b[94m 99.555795\u001b[35m val_loss:\u001b[94m 0.015576\u001b[35m val_acc:\u001b[94m 99.650330\u001b[0m\n",
      "\u001b[35mEpoch: \u001b[36m 4/10 \u001b[35mTrain steps: \u001b[36m2500 \u001b[35mVal steps: \u001b[36m625 \u001b[32m17.55s \u001b[35mloss:\u001b[94m 0.013548\u001b[35m acc:\u001b[94m 99.688221\u001b[35m val_loss:\u001b[94m 0.011891\u001b[35m val_acc:\u001b[94m 99.731806\u001b[0m\n",
      "\u001b[35mEpoch: \u001b[36m 5/10 \u001b[35mTrain steps: \u001b[36m2500 \u001b[35mVal steps: \u001b[36m625 \u001b[32m17.63s \u001b[35mloss:\u001b[94m 0.010814\u001b[35m acc:\u001b[94m 99.763319\u001b[35m val_loss:\u001b[94m 0.009818\u001b[35m val_acc:\u001b[94m 99.785883\u001b[0m\n",
      "\u001b[35mEpoch: \u001b[36m 6/10 \u001b[35mTrain steps: \u001b[36m2500 \u001b[35mVal steps: \u001b[36m625 \u001b[32m18.27s \u001b[35mloss:\u001b[94m 0.009196\u001b[35m acc:\u001b[94m 99.790861\u001b[35m val_loss:\u001b[94m 0.008578\u001b[35m val_acc:\u001b[94m 99.789669\u001b[0m\n",
      "\u001b[35mEpoch: \u001b[36m 7/10 \u001b[35mTrain steps: \u001b[36m2500 \u001b[35mVal steps: \u001b[36m625 \u001b[32m17.93s \u001b[35mloss:\u001b[94m 0.008007\u001b[35m acc:\u001b[94m 99.813415\u001b[35m val_loss:\u001b[94m 0.007688\u001b[35m val_acc:\u001b[94m 99.797362\u001b[0m\n",
      "\u001b[35mEpoch: \u001b[36m 8/10 \u001b[35mTrain steps: \u001b[36m2500 \u001b[35mVal steps: \u001b[36m625 \u001b[32m17.73s \u001b[35mloss:\u001b[94m 0.007109\u001b[35m acc:\u001b[94m 99.833434\u001b[35m val_loss:\u001b[94m 0.006799\u001b[35m val_acc:\u001b[94m 99.834102\u001b[0m\n",
      "\u001b[35mEpoch: \u001b[36m 9/10 \u001b[35mTrain steps: \u001b[36m2500 \u001b[35mVal steps: \u001b[36m625 \u001b[32m17.71s \u001b[35mloss:\u001b[94m 0.006367\u001b[35m acc:\u001b[94m 99.850936\u001b[35m val_loss:\u001b[94m 0.006059\u001b[35m val_acc:\u001b[94m 99.846338\u001b[0m\n",
      "\u001b[35mEpoch: \u001b[36m10/10 \u001b[35mTrain steps: \u001b[36m2500 \u001b[35mVal steps: \u001b[36m625 \u001b[32m17.54s \u001b[35mloss:\u001b[94m 0.005770\u001b[35m acc:\u001b[94m 99.863745\u001b[35m val_loss:\u001b[94m 0.005594\u001b[35m val_acc:\u001b[94m 99.855076\u001b[0m\n",
      "\u001b[35mTest steps: \u001b[36m938 \u001b[32m5.04s \u001b[35mtest_loss:\u001b[94m 0.005963\u001b[35m test_acc:\u001b[94m 99.842004\u001b[0m                                                   \n"
     ]
    }
   ],
   "source": [
    "optimizer = optim.SGD(full_network.parameters(), lr)\n",
    "loss_function = nn.CrossEntropyLoss()\n",
    "\n",
    "model = Model(\n",
    "    full_network,\n",
    "    optimizer,\n",
    "    loss_function,\n",
    "    batch_metrics=['accuracy'],\n",
    "    device=device,\n",
    ")\n",
    "\n",
    "model.fit_generator(train_loader, valid_loader, epochs=10)\n",
    "\n",
    "test_loss, test_acc = model.evaluate_generator(test_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predict on New Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let say we want to predict using our trained model. For these new addresses, we will only have the address and not the tags. Let us download this dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "download_data('./datasets/addresses/', \"predict\")\n",
    "predict_data = pickle.load(open(\"./datasets/addresses/predict.p\", \"rb\"))  # 30,000 examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we take a look at the predicted dataset, it is a list of 30,000 addresses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['15 major st london ontario n5z1e1',\n",
       " '1032 king st w 1202 kingston ontario k7m9g2']"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict_data[0:2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also need to reuse the vectorizer, but now with the `predict` argument set to `True` since we cannot parse the ground truth."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "predict_data_vectorized = Vectorizer(predict_data, embedding_model, predict=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to change a little the `pad_collate_fn` since we also pad the labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pad_collate_fn_predict(batch):\n",
    "    \"\"\"\n",
    "    The collate_fn add padding to the sequences so all can have\n",
    "    the same length as the longest one.\n",
    "\n",
    "    Args:\n",
    "        batch (List[List]): The batch data of the word idx.\n",
    "\n",
    "    Returns:\n",
    "        A tuple (x, y). The element x is a tensor of padded word vectors, and y\n",
    "        their respective lengths of the sequences. The word vectors are padded with vectors of 0s.\n",
    "    \"\"\"\n",
    "\n",
    "    # This gets us two lists of tensors and a list of integers.\n",
    "    # Each tensor in the first list is a sequence of word vectors.\n",
    "    # The list of integers consists of the lengths of the sequences in order.\n",
    "\n",
    "    sequences_vectors, lengths = zip(\n",
    "        *[\n",
    "            (torch.FloatTensor(seq_vectors), len(seq_vectors))\n",
    "            for seq_vectors in sorted(batch, key=lambda x: len(x), reverse=True)\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    lengths = torch.LongTensor(lengths)\n",
    "\n",
    "    padded_sequences_vectors = pad_sequence(sequences_vectors, batch_first=True, padding_value=0)\n",
    "    pack_padded_sequences_vectors = pack_padded_sequence(padded_sequences_vectors, lengths.cpu(), batch_first=True)\n",
    "\n",
    "    return pack_padded_sequences_vectors\n",
    "\n",
    "\n",
    "predict_loader = DataLoader(\n",
    "    predict_data_vectorized,\n",
    "    batch_size=batch_size,\n",
    "    collate_fn=pad_collate_fn_predict,\n",
    "    num_workers=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now, let's predict using the `predict_generator` method but without concatenating the returns since batches are not always the same size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([[[ 10.362392  ,   3.4235473 ,  -2.209073  , ..., -12.281605  ,\n",
       "          -14.648837  ,  -6.0009966 ],\n",
       "         [  3.0805492 ,  11.114887  ,  13.809945  , ...,   4.4265227 ,\n",
       "           -4.769767  , -11.348867  ],\n",
       "         [  2.4337537 ,   0.28365183,   1.1139965 , ...,   2.737724  ,\n",
       "           -0.5722164 ,   2.829398  ],\n",
       "         ...,\n",
       "         [ -2.9403193 , -10.5651045 , -13.855679  , ...,  -9.27598   ,\n",
       "            6.3255806 ,  21.084066  ],\n",
       "         [  1.1715801 ,   2.9763868 ,   3.6524632 , ...,  -0.72333026,\n",
       "           -3.7734149 ,  -7.133936  ],\n",
       "         [ -1.867637  ,  -3.8698483 ,  -5.643851  , ...,  -7.526398  ,\n",
       "           -6.2222004 ,  -4.0403423 ]],\n",
       " \n",
       "        [[ 11.595455  ,   3.7031417 ,  -2.3065462 , ..., -12.633009  ,\n",
       "          -14.922238  ,  -6.347353  ],\n",
       "         [  5.605796  ,  11.3620825 ,  13.836491  , ...,   3.9517322 ,\n",
       "           -4.875977  , -11.297518  ],\n",
       "         [  3.0809286 ,   1.3133283 ,   1.8499303 , ...,   3.0905602 ,\n",
       "           -0.50734043,   2.7491996 ],\n",
       "         ...,\n",
       "         [ -4.994604  , -10.714485  , -13.141624  , ...,  -9.527986  ,\n",
       "            6.1759133 ,  20.962942  ],\n",
       "         [  1.3510636 ,   2.8583338 ,   3.3109648 , ...,  -0.81182444,\n",
       "           -3.9521592 ,  -7.32445   ],\n",
       "         [ -2.0445874 ,  -4.0527463 ,  -5.774393  , ...,  -7.6246266 ,\n",
       "           -6.2599034 ,  -4.0781655 ]],\n",
       " \n",
       "        [[ 11.677182  ,   2.9273143 ,  -2.4910564 , ..., -12.611051  ,\n",
       "          -15.027013  ,  -6.554265  ],\n",
       "         [  3.4473739 ,  10.514426  ,  14.57351   , ...,   3.6238976 ,\n",
       "           -5.223489  , -11.982287  ],\n",
       "         [  3.0238602 ,   0.53050137,   1.3198154 , ...,   3.283758  ,\n",
       "           -0.4880113 ,   2.664648  ],\n",
       "         ...,\n",
       "         [ -4.0954447 , -10.94076   , -13.172517  , ...,  -9.57532   ,\n",
       "            6.4028664 ,  21.60267   ],\n",
       "         [  1.6864542 ,   2.9354289 ,   3.0842376 , ...,  -0.80098784,\n",
       "           -4.0576844 ,  -7.4614325 ],\n",
       "         [ -2.0126655 ,  -4.0368423 ,  -5.732325  , ...,  -7.6642017 ,\n",
       "           -6.267487  ,  -4.016219  ]],\n",
       " \n",
       "        ...,\n",
       " \n",
       "        [[  9.030349  ,   2.5147102 ,  -2.1758246 , ...,   2.8536558 ,\n",
       "            2.8536558 ,   2.8536558 ],\n",
       "         [  3.9188714 ,  10.553635  ,  17.176523  , ...,   0.4582528 ,\n",
       "            0.4582528 ,   0.4582528 ],\n",
       "         [  1.6585233 ,   0.24136674,   1.5541401 , ...,  -1.039604  ,\n",
       "           -1.039604  ,  -1.039604  ],\n",
       "         ...,\n",
       "         [ -3.2747264 ,  -8.869284  , -14.9407835 , ...,  -0.04295171,\n",
       "           -0.04295171,  -0.04295171],\n",
       "         [  0.9152781 ,   1.3144858 ,   3.6261466 , ...,  -0.41076294,\n",
       "           -0.41076294,  -0.41076294],\n",
       "         [ -1.9186795 ,  -3.814353  ,  -6.109922  , ...,  -1.0673251 ,\n",
       "           -1.0673251 ,  -1.0673251 ]],\n",
       " \n",
       "        [[ 12.884519  ,   4.024414  ,  -1.8850508 , ...,   2.8536558 ,\n",
       "            2.8536558 ,   2.8536558 ],\n",
       "         [  5.491664  ,  12.431355  ,  17.777035  , ...,   0.4582528 ,\n",
       "            0.4582528 ,   0.4582528 ],\n",
       "         [  3.3901393 ,   1.0163777 ,   2.2849805 , ...,  -1.039604  ,\n",
       "           -1.039604  ,  -1.039604  ],\n",
       "         ...,\n",
       "         [ -4.8606915 , -11.308829  , -16.548843  , ...,  -0.04295171,\n",
       "           -0.04295171,  -0.04295171],\n",
       "         [  1.6389235 ,   3.450939  ,   4.43767   , ...,  -0.41076294,\n",
       "           -0.41076294,  -0.41076294],\n",
       "         [ -1.9431152 ,  -4.075838  ,  -6.397602  , ...,  -1.0673251 ,\n",
       "           -1.0673251 ,  -1.0673251 ]],\n",
       " \n",
       "        [[  9.569925  ,   2.4876893 ,  -6.923353  , ...,   2.8536558 ,\n",
       "            2.8536558 ,   2.8536558 ],\n",
       "         [  3.3440304 ,  10.245568  ,   5.96705   , ...,   0.4582528 ,\n",
       "            0.4582528 ,   0.4582528 ],\n",
       "         [  2.1744015 ,   0.3859098 ,  -0.31355786, ...,  -1.039604  ,\n",
       "           -1.039604  ,  -1.039604  ],\n",
       "         ...,\n",
       "         [ -3.2376053 ,  -8.978806  ,  -7.538249  , ...,  -0.04295171,\n",
       "           -0.04295171,  -0.04295171],\n",
       "         [  0.64889824,   1.0244368 ,  -1.070083  , ...,  -0.41076294,\n",
       "           -0.41076294,  -0.41076294],\n",
       "         [ -1.8653474 ,  -3.8354082 ,  -5.3427086 , ...,  -1.0673251 ,\n",
       "           -1.0673251 ,  -1.0673251 ]]], dtype=float32),\n",
       " array([[[ 1.44394169e+01,  4.65035677e+00, -1.62813139e+00, ...,\n",
       "          -1.00203333e+01, -1.29954023e+01, -6.75655079e+00],\n",
       "         [ 3.95823622e+00,  1.10212898e+01,  1.24216394e+01, ...,\n",
       "          -3.94837761e+00, -7.04809093e+00, -1.23708277e+01],\n",
       "         [ 3.80071282e+00,  1.05917811e+00,  2.07201314e+00, ...,\n",
       "          -2.73017049e-01, -1.90600681e+00,  6.22655869e-01],\n",
       "         ...,\n",
       "         [-5.10277653e+00, -1.16044950e+01, -1.34266720e+01, ...,\n",
       "           4.64116144e+00,  9.04016113e+00,  1.99978733e+01],\n",
       "         [ 2.23618436e+00,  4.36150455e+00,  3.91225123e+00, ...,\n",
       "          -4.53451538e+00, -6.38264942e+00, -7.79837179e+00],\n",
       "         [-1.91084170e+00, -4.03150177e+00, -5.66816759e+00, ...,\n",
       "          -5.66546440e+00, -5.21284819e+00, -3.64689159e+00]],\n",
       " \n",
       "        [[ 1.02849340e+01,  4.15290070e+00, -1.56945372e+00, ...,\n",
       "          -1.24890566e+01, -1.48781157e+01, -6.53854179e+00],\n",
       "         [ 3.50871134e+00,  1.03837852e+01,  1.44168110e+01, ...,\n",
       "           3.91083145e+00, -4.90434790e+00, -1.15230646e+01],\n",
       "         [ 2.38652778e+00,  5.57615757e-01,  1.45711899e+00, ...,\n",
       "           3.25055194e+00, -4.36077595e-01,  2.68630886e+00],\n",
       "         ...,\n",
       "         [-3.35965657e+00, -1.10284996e+01, -1.25283394e+01, ...,\n",
       "          -9.55706024e+00,  6.16046715e+00,  2.11034908e+01],\n",
       "         [ 1.04619515e+00,  4.33481693e+00,  2.13645339e+00, ...,\n",
       "          -7.01354265e-01, -3.88502526e+00, -7.43781233e+00],\n",
       "         [-1.90720832e+00, -3.75380993e+00, -5.39478207e+00, ...,\n",
       "          -7.63562155e+00, -6.27026224e+00, -4.05933285e+00]],\n",
       " \n",
       "        [[ 9.61695099e+00,  3.93914938e+00,  1.98254383e+00, ...,\n",
       "          -1.23484688e+01, -6.60503864e+00,  2.85365582e+00],\n",
       "         [ 3.44057989e+00,  1.14006910e+01,  1.74034023e+01, ...,\n",
       "          -5.64836311e+00, -1.17811403e+01,  4.58252788e-01],\n",
       "         [ 2.09207201e+00,  8.10892820e-01,  4.28207684e+00, ...,\n",
       "          -1.09142005e+00,  1.22841072e+00, -1.03960395e+00],\n",
       "         ...,\n",
       "         [-3.50295281e+00, -1.02293396e+01, -1.51526909e+01, ...,\n",
       "           7.82587385e+00,  1.92622509e+01, -4.29517068e-02],\n",
       "         [ 7.83683419e-01,  2.16639543e+00,  4.49790287e+00, ...,\n",
       "          -5.42884779e+00, -7.08081627e+00, -4.10762936e-01],\n",
       "         [-1.95001125e+00, -3.83760834e+00, -6.00906038e+00, ...,\n",
       "          -5.37854481e+00, -3.84737134e+00, -1.06732512e+00]],\n",
       " \n",
       "        ...,\n",
       " \n",
       "        [[ 1.28347063e+01,  3.99542665e+00, -1.62145424e+00, ...,\n",
       "           2.85365582e+00,  2.85365582e+00,  2.85365582e+00],\n",
       "         [ 4.91154146e+00,  1.15387592e+01,  1.65973492e+01, ...,\n",
       "           4.58252788e-01,  4.58252788e-01,  4.58252788e-01],\n",
       "         [ 3.37475371e+00,  1.02676058e+00,  2.12015367e+00, ...,\n",
       "          -1.03960395e+00, -1.03960395e+00, -1.03960395e+00],\n",
       "         ...,\n",
       "         [-4.62063599e+00, -1.09527197e+01, -1.70621738e+01, ...,\n",
       "          -4.29517068e-02, -4.29517068e-02, -4.29517068e-02],\n",
       "         [ 1.83531797e+00,  3.35930157e+00,  5.53424644e+00, ...,\n",
       "          -4.10762936e-01, -4.10762936e-01, -4.10762936e-01],\n",
       "         [-1.91633105e+00, -4.00633430e+00, -6.20124960e+00, ...,\n",
       "          -1.06732512e+00, -1.06732512e+00, -1.06732512e+00]],\n",
       " \n",
       "        [[ 1.26212320e+01,  3.87416816e+00,  1.89483619e+00, ...,\n",
       "           2.85365582e+00,  2.85365582e+00,  2.85365582e+00],\n",
       "         [ 5.47309399e+00,  1.04094954e+01,  1.63003445e+01, ...,\n",
       "           4.58252788e-01,  4.58252788e-01,  4.58252788e-01],\n",
       "         [ 3.30542779e+00,  1.23913169e+00,  4.73438072e+00, ...,\n",
       "          -1.03960395e+00, -1.03960395e+00, -1.03960395e+00],\n",
       "         ...,\n",
       "         [-4.75571537e+00, -1.00918856e+01, -1.47780485e+01, ...,\n",
       "          -4.29517068e-02, -4.29517068e-02, -4.29517068e-02],\n",
       "         [ 1.60826671e+00,  2.94319534e+00,  4.59750843e+00, ...,\n",
       "          -4.10762936e-01, -4.10762936e-01, -4.10762936e-01],\n",
       "         [-1.94441748e+00, -3.97083282e+00, -6.11064291e+00, ...,\n",
       "          -1.06732512e+00, -1.06732512e+00, -1.06732512e+00]],\n",
       " \n",
       "        [[ 9.75561142e+00,  2.87291837e+00, -6.65499973e+00, ...,\n",
       "           2.85365582e+00,  2.85365582e+00,  2.85365582e+00],\n",
       "         [ 3.40803528e+00,  1.10199175e+01,  6.22449493e+00, ...,\n",
       "           4.58252788e-01,  4.58252788e-01,  4.58252788e-01],\n",
       "         [ 1.98146415e+00,  3.97124410e-01, -1.89704895e-02, ...,\n",
       "          -1.03960395e+00, -1.03960395e+00, -1.03960395e+00],\n",
       "         ...,\n",
       "         [-3.62741518e+00, -9.46802139e+00, -7.94740963e+00, ...,\n",
       "          -4.29517068e-02, -4.29517068e-02, -4.29517068e-02],\n",
       "         [ 1.24053550e+00,  1.79104197e+00, -6.40422463e-01, ...,\n",
       "          -4.10762936e-01, -4.10762936e-01, -4.10762936e-01],\n",
       "         [-1.93732297e+00, -3.88715410e+00, -5.39007044e+00, ...,\n",
       "          -1.06732512e+00, -1.06732512e+00, -1.06732512e+00]]],\n",
       "       dtype=float32)]"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predictions = model.predict_generator(predict_loader, concatenate_returns=False)\n",
    "predictions[0:2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or using the `predict_dataset` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "predictions = model.predict_dataset(\n",
    "    predict_data_vectorized,\n",
    "    batch_size=batch_size,\n",
    "    collate_fn=pad_collate_fn_predict,\n",
    "    num_workers=2,\n",
    "    concatenate_returns=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "But, as you can see, we have log probabilities of our model, so we need to take the highest probability per address element."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "idx_predictions = []\n",
    "for batch in predictions:\n",
    "    idx_predictions.extend(batch.argmax(axis=1).tolist())  # We take the highest argument (so the tag idx)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "To convert those indexes into tags, we only need to convert them back using the inverse of the previous `tags_to_idx` dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[['StreetNumber',\n",
       "  'StreetName',\n",
       "  'StreetName',\n",
       "  'StreetName',\n",
       "  'Orientation',\n",
       "  'Municipality',\n",
       "  'Municipality',\n",
       "  'Municipality',\n",
       "  'Province',\n",
       "  'PostalCode'],\n",
       " ['StreetNumber',\n",
       "  'StreetName',\n",
       "  'StreetName',\n",
       "  'Orientation',\n",
       "  'Unit',\n",
       "  'Municipality',\n",
       "  'Municipality',\n",
       "  'Municipality',\n",
       "  'Province',\n",
       "  'PostalCode']]"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "idx_to_tags = {\n",
    "    0: \"StreetNumber\",\n",
    "    1: \"StreetName\",\n",
    "    2: \"Unit\",\n",
    "    3: \"Municipality\",\n",
    "    4: \"Province\",\n",
    "    5: \"PostalCode\",\n",
    "    6: \"Orientation\",\n",
    "    7: \"GeneralDelivery\",\n",
    "}\n",
    "\n",
    "tags_predictions = []\n",
    "for address in idx_predictions:\n",
    "    tags_predictions.append([idx_to_tags.get(tag) for tag in address])\n",
    "\n",
    "tags_predictions[0:2]"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}