{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LSTM-arithmetic\n",
    "\n",
    "## Dataset\n",
    "- [Arithmetic dataset](https://drive.google.com/file/d/1cMuL3hF9jefka9RyF4gEBIGGeFGZYHE-/view?usp=sharing)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! pip install seaborn\n",
    "# ! pip install opencc\n",
    "# ! pip install -U scikit-learn\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn\n",
    "import torch.nn.utils.rnn\n",
    "import torch.utils.data\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import opencc\n",
    "import os\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "data_path = './data'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>src</th>\n",
       "      <th>tgt</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0+0=</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0-0=</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0*0=</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>(0+0)*0=</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0+0*0=</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        src  tgt\n",
       "0      0+0=    0\n",
       "1      0-0=    0\n",
       "2      0*0=    0\n",
       "3  (0+0)*0=    0\n",
       "4    0+0*0=    0"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_train = pd.read_csv(os.path.join(data_path, 'arithmetic_train.csv'))\n",
    "df_eval = pd.read_csv(os.path.join(data_path, 'arithmetic_eval.csv'))\n",
    "df_train.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# transform the input data to string\n",
    "df_train['tgt'] = df_train['tgt'].apply(lambda x: str(x))\n",
    "df_train['src'] = df_train['src'].add(df_train['tgt'])\n",
    "df_train['len'] = df_train['src'].apply(lambda x: len(x))\n",
    "\n",
    "df_eval['tgt'] = df_eval['tgt'].apply(lambda x: str(x))\n",
    "df_eval['src'] = df_eval['src'].add(df_eval['tgt'])\n",
    "df_eval['len'] = df_eval['src'].apply(lambda x: len(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build Dictionary\n",
    " - The model cannot perform calculations directly with plain text.\n",
    " - Convert all text (numbers/symbols) into numerical representations.\n",
    " - Special tokens\n",
    "    - '&lt;pad&gt;'\n",
    "        - Each sentence within a batch may have different lengths.\n",
    "        - The length is padded with '&lt;pad&gt;' to match the longest sentence in the batch.\n",
    "    - '&lt;eos&gt;'\n",
    "        - Specifies the end of the generated sequence.\n",
    "        - Without '&lt;eos&gt;', the model will not know when to stop generating."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "字典大小: 18\n"
     ]
    }
   ],
   "source": [
    "char_to_id = {}\n",
    "id_to_char = {}\n",
    "\n",
    "# write your code here\n",
    "# Build a dictionary and give every token in the train dataset an id\n",
    "# The dictionary should contain <eos> and <pad>\n",
    "# char_to_id is to conver charactors to ids, while id_to_char is the opposite\n",
    "\n",
    "vocab_size = len(char_to_id)\n",
    "print('Vocab size{}'.format(vocab_size))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Preprocessing\n",
    " - The data is processed into the format required for the model's input and output.\n",
    " - Example: 1+2-3=0\n",
    "     - Model input: 1 + 2 - 3 = 0\n",
    "     - Model output: / / / / / 0 &lt;eos&gt;  (the '/' can be replaced with &lt;pad&gt;)\n",
    "     - The key for the model's output is that the model does not need to predict the next character of the previous part. What matters is that once the model sees '=', it should start generating the answer, which is '0'. After generating the answer, it should also generate&lt;eos&gt;\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>src</th>\n",
       "      <th>tgt</th>\n",
       "      <th>len</th>\n",
       "      <th>char_id_list</th>\n",
       "      <th>label_id_list</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0+0=0</td>\n",
       "      <td>0</td>\n",
       "      <td>5</td>\n",
       "      <td>[15, 3, 15, 17, 15, 1]</td>\n",
       "      <td>[0, 0, 0, 0, 15, 1]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0-0=0</td>\n",
       "      <td>0</td>\n",
       "      <td>5</td>\n",
       "      <td>[15, 7, 15, 17, 15, 1]</td>\n",
       "      <td>[0, 0, 0, 0, 15, 1]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0*0=0</td>\n",
       "      <td>0</td>\n",
       "      <td>5</td>\n",
       "      <td>[15, 13, 15, 17, 15, 1]</td>\n",
       "      <td>[0, 0, 0, 0, 15, 1]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>(0+0)*0=0</td>\n",
       "      <td>0</td>\n",
       "      <td>9</td>\n",
       "      <td>[14, 15, 3, 15, 10, 13, 15, 17, 15, 1]</td>\n",
       "      <td>[0, 0, 0, 0, 0, 0, 0, 0, 15, 1]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0+0*0=0</td>\n",
       "      <td>0</td>\n",
       "      <td>7</td>\n",
       "      <td>[15, 3, 15, 13, 15, 17, 15, 1]</td>\n",
       "      <td>[0, 0, 0, 0, 0, 0, 15, 1]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "         src tgt  len                            char_id_list  \\\n",
       "0      0+0=0   0    5                  [15, 3, 15, 17, 15, 1]   \n",
       "1      0-0=0   0    5                  [15, 7, 15, 17, 15, 1]   \n",
       "2      0*0=0   0    5                 [15, 13, 15, 17, 15, 1]   \n",
       "3  (0+0)*0=0   0    9  [14, 15, 3, 15, 10, 13, 15, 17, 15, 1]   \n",
       "4    0+0*0=0   0    7          [15, 3, 15, 13, 15, 17, 15, 1]   \n",
       "\n",
       "                     label_id_list  \n",
       "0              [0, 0, 0, 0, 15, 1]  \n",
       "1              [0, 0, 0, 0, 15, 1]  \n",
       "2              [0, 0, 0, 0, 15, 1]  \n",
       "3  [0, 0, 0, 0, 0, 0, 0, 0, 15, 1]  \n",
       "4        [0, 0, 0, 0, 0, 0, 15, 1]  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Write your code here\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hyper Parameters\n",
    "\n",
    "|Hyperparameter|Meaning|Value|\n",
    "|-|-|-|\n",
    "|`batch_size`|Number of data samples in a single batch|64|\n",
    "|`epochs`|Total number of epochs to train|10|\n",
    "|`embed_dim`|Dimension of the word embeddings|256|\n",
    "|`hidden_dim`|Dimension of the hidden state in each timestep of the LSTM|256|\n",
    "|`lr`|Learning Rate|0.001|\n",
    "|`grad_clip`|To prevent gradient explosion in RNNs, restrict the gradient range|1|"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 64\n",
    "epochs = 2\n",
    "embed_dim = 256\n",
    "hidden_dim = 256\n",
    "lr = 0.001\n",
    "grad_clip = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Batching\n",
    "- Use `torch.utils.data.Dataset` to create a data generation tool called  `dataset`.\n",
    "- The, use `torch.utils.data.DataLoader` to randomly sample from the `dataset` and group the samples into batches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Dataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, sequences):\n",
    "        self.sequences = sequences\n",
    "    \n",
    "    def __len__(self):\n",
    "        # return the amount of data\n",
    "        return # Write your code here\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        # Extract the input data x and the ground truth y from the data\n",
    "        x = # Write your code here\n",
    "        y = # Write your code here\n",
    "        return x, y\n",
    "\n",
    "# collate function, used to build dataloader\n",
    "def collate_fn(batch):\n",
    "    batch_x = [torch.tensor(data[0]) for data in batch]\n",
    "    batch_y = [torch.tensor(data[1]) for data in batch]\n",
    "    batch_x_lens = torch.LongTensor([len(x) for x in batch_x])\n",
    "    batch_y_lens = torch.LongTensor([len(y) for y in batch_y])\n",
    "    \n",
    "    # Pad the input sequence\n",
    "    pad_batch_x = torch.nn.utils.rnn.pad_sequence(batch_x,\n",
    "                                                  batch_first=True,\n",
    "                                                  padding_value=char_to_id['<pad>'])\n",
    "    \n",
    "    pad_batch_y = torch.nn.utils.rnn.pad_sequence(batch_y,\n",
    "                                                  batch_first=True,\n",
    "                                                  padding_value=char_to_id['<pad>'])\n",
    "    \n",
    "    return pad_batch_x, pad_batch_y, batch_x_lens, batch_y_lens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_train = Dataset(df_train[['char_id_list', 'label_id_list']])\n",
    "ds_eval = Dataset(df_eval[['char_id_list', 'label_id_list']])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build dataloader of train set and eval set, collate_fn is the collate function\n",
    "dl_train = # Write your code here\n",
    "dl_eval = # Write your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Design\n",
    "\n",
    "## Execution Flow\n",
    "1. Convert all characters in the sentence into embeddings.\n",
    "2. Pass the embeddings through an LSTM sequentially.\n",
    "3. The output of the LSTM is passed into another LSTM, and additional layers can be added.\n",
    "4. The output from all time steps of the final LSTM is passed through a Fully Connected layer.\n",
    "5. The character corresponding to the maximum value across all output dimensions is selected as the next character.\n",
    "\n",
    "## Loss Function\n",
    "Since this is a classification task, Cross Entropy is used as the loss function.\n",
    "\n",
    "## Gradient Update\n",
    "Adam algorithm is used for gradient updates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CharRNN(torch.nn.Module):\n",
    "    def __init__(self, vocab_size, embed_dim, hidden_dim):\n",
    "        super(CharRNN, self).__init__()\n",
    "        \n",
    "        self.embedding = torch.nn.Embedding(num_embeddings=vocab_size,\n",
    "                                            embedding_dim=embed_dim,\n",
    "                                            padding_idx=char_to_id['<pad>'])\n",
    "        \n",
    "        self.rnn_layer1 = torch.nn.LSTM(input_size=embed_dim,\n",
    "                                        hidden_size=hidden_dim,\n",
    "                                        batch_first=True)\n",
    "        \n",
    "        self.rnn_layer2 = torch.nn.LSTM(input_size=hidden_dim,\n",
    "                                        hidden_size=hidden_dim,\n",
    "                                        batch_first=True)\n",
    "        \n",
    "        self.linear = torch.nn.Sequential(torch.nn.Linear(in_features=hidden_dim,\n",
    "                                                          out_features=hidden_dim),\n",
    "                                          torch.nn.ReLU(),\n",
    "                                          torch.nn.Linear(in_features=hidden_dim,\n",
    "                                                          out_features=vocab_size))\n",
    "        \n",
    "    def forward(self, batch_x, batch_x_lens):\n",
    "        return self.encoder(batch_x, batch_x_lens)\n",
    "    \n",
    "    # The forward pass of the model\n",
    "    def encoder(self, batch_x, batch_x_lens):\n",
    "        batch_x = self.embedding(batch_x)\n",
    "        \n",
    "        batch_x = torch.nn.utils.rnn.pack_padded_sequence(batch_x,\n",
    "                                                          batch_x_lens,\n",
    "                                                          batch_first=True,\n",
    "                                                          enforce_sorted=False)\n",
    "        \n",
    "        batch_x, _ = self.rnn_layer1(batch_x)\n",
    "        batch_x, _ = self.rnn_layer2(batch_x)\n",
    "        \n",
    "        batch_x, _ = torch.nn.utils.rnn.pad_packed_sequence(batch_x,\n",
    "                                                            batch_first=True)\n",
    "        \n",
    "        batch_x = self.linear(batch_x)\n",
    "        \n",
    "        return batch_x\n",
    "    \n",
    "    def generator(self, start_char, max_len=200):\n",
    "        \n",
    "        char_list = [char_to_id[c] for c in start_char]\n",
    "        \n",
    "        next_char = None\n",
    "        \n",
    "        while len(char_list) < max_len: \n",
    "            # Write your code here \n",
    "            # Pack the char_list to tensor\n",
    "            # Input the tensor to the embedding layer, LSTM layers, linear respectively\n",
    "            y = # Obtain the next token prediction y\n",
    "            \n",
    "            next_char = # Use argmax function to get the next token prediction\n",
    "            \n",
    "            if next_char == char_to_id['<eos>']:\n",
    "                break\n",
    "            \n",
    "            char_list.append(next_char)\n",
    "            \n",
    "        return [id_to_char[ch_id] for ch_id in char_list]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(2)\n",
    "\n",
    "\n",
    "device = # Write your code here. Specify a device (cuda or cpu)\n",
    "\n",
    "model = CharRNN(vocab_size,\n",
    "                embed_dim,\n",
    "                hidden_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ikm/miniconda3/envs/py3917/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "criterion = # Write your code here. Cross-entropy loss function. The loss function should ignore <pad>\n",
    "optimizer = # Write your code here. Use Adam or AdamW for Optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training\n",
    "1. The outer `for` loop controls the `epoch`\n",
    "    1. The inner `for` loop uses `data_loader` to retrieve batches.\n",
    "        1. Pass the batch to the `model` for training.\n",
    "        2. Compare the predicted results `batch_pred_y` with the true labels `batch_y` using Cross Entropy to calculate the loss `loss`\n",
    "        3. Use `loss.backward` to automatically compute the gradients.\n",
    "        4. Use `torch.nn.utils.clip_grad_value_` to limit the gradient values between `-grad_clip` &lt; and &lt; `grad_clip`.\n",
    "        5. Use `optimizer.step()` to update the model (backpropagation).\n",
    "2.  After every `1000` batches, output the current loss to monitor whether it is converging."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train epoch 1: 100%|██████████| 37020/37020 [04:53<00:00, 125.96it/s, loss=0.326]\n",
      "Validation epoch 1: 100%|██████████| 4114/4114 [02:32<00:00, 26.94it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.5898499525166192\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train epoch 2: 100%|██████████| 37020/37020 [04:59<00:00, 123.80it/s, loss=0.242]\n",
      "Validation epoch 2: 100%|██████████| 4114/4114 [02:31<00:00, 27.10it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.6699905033238367\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "from copy import deepcopy\n",
    "model = model.to(device)\n",
    "model.train()\n",
    "i = 0\n",
    "for epoch in range(1, epochs+1):\n",
    "    # The process bar\n",
    "    bar = tqdm(dl_train, desc=f\"Train epoch {epoch}\")\n",
    "    for batch_x, batch_y, batch_x_lens, batch_y_lens in bar:\n",
    "        # Write your code here\n",
    "        # Clear the gradient\n",
    "    \n",
    "        batch_pred_y = model(batch_x.to(device), batch_x_lens)\n",
    "        \n",
    "        # Write your code here\n",
    "        # Input the prediction and ground truths to loss function\n",
    "        # Back propagation\n",
    "\n",
    "        torch.nn.utils.clip_grad_value_(model.parameters(), grad_clip) # gradient clipping\n",
    "\n",
    "        # Write your code here\n",
    "        # Optimize parameters in the model\n",
    "\n",
    "        i+=1\n",
    "        if i%50==0:\n",
    "            bar.set_postfix(loss = loss.item())\n",
    "    \n",
    "    # Evaluate your model\n",
    "    bar = tqdm(dl_eval, desc=f\"Validation epoch {epoch}\")\n",
    "    matched = 0\n",
    "    total = 0\n",
    "    for batch_x, batch_y, batch_x_lens, batch_y_lens in bar:\n",
    "        \n",
    "        predictions = # Write your code here. Input the batch_x to the model and generate the predictions\n",
    "        \n",
    "        # Write your code here.\n",
    "        # Check whether the prediction match the ground truths\n",
    "        # Compute exact match (EM) on the eval dataset\n",
    "        # EM = correct/total\n",
    "\n",
    "        \n",
    "    print(matched/total)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generation\n",
    "Use `model.generator` and provide an initial character to automatically generate a sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1+1=2\n"
     ]
    }
   ],
   "source": [
    "model = model.to(\"cpu\")\n",
    "print(\"\".join(model.generator('1+1=')))"
   ]
  }
 ],
 "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
