notebooks/rnn_from_scratch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RNN from scratch with PyTorch\n",
    "A RNN ist just a normal NN.\n",
    "It's very easy to implement in PyTorch due to its dynamic nature.\n",
    "\n",
    "We'll build a very simple character based language model.\n",
    "\n",
    "Taken from http://www.fast.ai/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Init and helpers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:20.187029Z",
     "iopub.status.busy": "2020-12-06T08:47:20.186373Z",
     "iopub.status.idle": "2020-12-06T08:47:20.309644Z",
     "shell.execute_reply": "2020-12-06T08:47:20.310735Z"
    }
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:20.315957Z",
     "iopub.status.busy": "2020-12-06T08:47:20.314473Z",
     "iopub.status.idle": "2020-12-06T08:47:20.387279Z",
     "shell.execute_reply": "2020-12-06T08:47:20.386890Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I already have the data.\n"
     ]
    }
   ],
   "source": [
    "NIETSCHE_PATH = Path(\"../data/raw/nietzsche.txt\")\n",
    "if NIETSCHE_PATH.is_file():\n",
    "    print(\"I already have the data.\")\n",
    "else:\n",
    "    !wget -o ../data/raw/nietzsche.txt https://s3.amazonaws.com/text-datasets/nietzsche.txt\n",
    "        \n",
    "with NIETSCHE_PATH.open() as f:\n",
    "    data = f.read()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A tweet of Nietzsche:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:20.390460Z",
     "iopub.status.busy": "2020-12-06T08:47:20.390010Z",
     "iopub.status.idle": "2020-12-06T08:47:20.392264Z",
     "shell.execute_reply": "2020-12-06T08:47:20.392662Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PREFACE\n",
      "\n",
      "\n",
      "SUPPOSING that Truth is a woman--what then? Is there not ground\n",
      "for suspecting that all philosophers, in so far as they have been\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(data[:140])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to know the alphabet and we add a padding value \"\\0\" to the alphabet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:20.410193Z",
     "iopub.status.busy": "2020-12-06T08:47:20.409477Z",
     "iopub.status.idle": "2020-12-06T08:47:20.412698Z",
     "shell.execute_reply": "2020-12-06T08:47:20.412171Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "85"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "alphabet = [\"\\0\", *sorted(list(set(data)))]\n",
    "n_alphabet = len(alphabet)\n",
    "n_alphabet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:20.416803Z",
     "iopub.status.busy": "2020-12-06T08:47:20.416181Z",
     "iopub.status.idle": "2020-12-06T08:47:20.418426Z",
     "shell.execute_reply": "2020-12-06T08:47:20.418912Z"
    }
   },
   "outputs": [],
   "source": [
    "char2index = {c: i for i, c in enumerate(alphabet)}\n",
    "index2char = {i: c for i, c in enumerate(alphabet)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Convert the data into a list of integers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:20.459645Z",
     "iopub.status.busy": "2020-12-06T08:47:20.456579Z",
     "iopub.status.idle": "2020-12-06T08:47:20.461267Z",
     "shell.execute_reply": "2020-12-06T08:47:20.461520Z"
    }
   },
   "outputs": [],
   "source": [
    "index = [char2index[c] for c in data]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:20.463805Z",
     "iopub.status.busy": "2020-12-06T08:47:20.463476Z",
     "iopub.status.idle": "2020-12-06T08:47:20.465132Z",
     "shell.execute_reply": "2020-12-06T08:47:20.465383Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[40, 42, 29, 30, 25, 27, 29, 1, 1, 1, 43, 45, 40, 40, 39, 43, 33, 38, 31, 2, 73, 61, 54, 73, 2]\n",
      "PREFACE\n",
      "\n",
      "\n",
      "SUPPOSING that \n"
     ]
    }
   ],
   "source": [
    "print(index[:25])\n",
    "print(\"\".join(index2char[i] for i in index[:25]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:20.467519Z",
     "iopub.status.busy": "2020-12-06T08:47:20.467164Z",
     "iopub.status.idle": "2020-12-06T08:47:20.468936Z",
     "shell.execute_reply": "2020-12-06T08:47:20.469190Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[40, 42, 29]"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index[0: 3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:20.531092Z",
     "iopub.status.busy": "2020-12-06T08:47:20.521077Z",
     "iopub.status.idle": "2020-12-06T08:47:22.563153Z",
     "shell.execute_reply": "2020-12-06T08:47:22.563421Z"
    }
   },
   "outputs": [],
   "source": [
    "X, y = [], []\n",
    "for i in range(len(index) - 4):\n",
    "    X.append(index[i : i + 3])\n",
    "    y.append(index[i + 3])\n",
    "    \n",
    "X = np.stack(X)\n",
    "y = np.stack(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:22.565676Z",
     "iopub.status.busy": "2020-12-06T08:47:22.565356Z",
     "iopub.status.idle": "2020-12-06T08:47:22.567072Z",
     "shell.execute_reply": "2020-12-06T08:47:22.567331Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((600889, 3), (600889,))"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X.shape, y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:22.569823Z",
     "iopub.status.busy": "2020-12-06T08:47:22.569358Z",
     "iopub.status.idle": "2020-12-06T08:47:22.571051Z",
     "shell.execute_reply": "2020-12-06T08:47:22.571308Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([40, 42, 29]), 30)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X[0], y[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:22.573504Z",
     "iopub.status.busy": "2020-12-06T08:47:22.573159Z",
     "iopub.status.idle": "2020-12-06T08:47:22.574945Z",
     "shell.execute_reply": "2020-12-06T08:47:22.575186Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "numpy.ndarray"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:22.577419Z",
     "iopub.status.busy": "2020-12-06T08:47:22.577102Z",
     "iopub.status.idle": "2020-12-06T08:47:22.791865Z",
     "shell.execute_reply": "2020-12-06T08:47:22.792148Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader, Dataset, TensorDataset\n",
    "\n",
    "\n",
    "train_ds = TensorDataset(torch.from_numpy(X), torch.from_numpy(y))\n",
    "train_dl = DataLoader(train_ds, batch_size=500)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:22.794452Z",
     "iopub.status.busy": "2020-12-06T08:47:22.794136Z",
     "iopub.status.idle": "2020-12-06T08:47:22.795818Z",
     "shell.execute_reply": "2020-12-06T08:47:22.795504Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:22.799526Z",
     "iopub.status.busy": "2020-12-06T08:47:22.799090Z",
     "iopub.status.idle": "2020-12-06T08:47:22.801673Z",
     "shell.execute_reply": "2020-12-06T08:47:22.801383Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/stefan/projects/pytorch_tutorial/.venv/lib/python3.8/site-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx (Triggered internally at  /pytorch/c10/cuda/CUDAFunctions.cpp:100.)\n",
      "  return torch._C._cuda_getDeviceCount() > 0\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "device(type='cpu')"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:22.805642Z",
     "iopub.status.busy": "2020-12-06T08:47:22.805321Z",
     "iopub.status.idle": "2020-12-06T08:47:22.807151Z",
     "shell.execute_reply": "2020-12-06T08:47:22.806837Z"
    }
   },
   "outputs": [],
   "source": [
    "class CharModel(nn.Module):\n",
    "    def __init__(self, n_vocab, n_embedding, n_hidden):\n",
    "        super().__init__()\n",
    "        self.emb = nn.Embedding(n_vocab, n_embedding)\n",
    "        self.lin_in = nn.Linear(n_embedding, n_hidden)\n",
    "        \n",
    "        self.lin_hidden = nn.Linear(n_hidden, n_hidden)\n",
    "        self.lin_out = nn.Linear(n_hidden, n_vocab)\n",
    "        \n",
    "    def forward(self, X):\n",
    "        c1, c2, c3 = X[:, 0], X[:, 1], X[:, 2]\n",
    "        \n",
    "        in1 = F.relu(self.lin_in(self.emb(c1)))\n",
    "        h = F.tanh(self.lin_hidden(in1))\n",
    "                   \n",
    "        in2 = F.relu(self.lin_in(self.emb(c2)))\n",
    "        h = F.tanh(self.lin_hidden(h + in2))\n",
    "        \n",
    "        in3 = F.relu(self.lin_in(self.emb(c3)))\n",
    "        h = F.tanh(self.lin_hidden(h + in3))\n",
    "        \n",
    "        return F.log_softmax(self.lin_out(h), dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:22.809256Z",
     "iopub.status.busy": "2020-12-06T08:47:22.808949Z",
     "iopub.status.idle": "2020-12-06T08:47:22.811552Z",
     "shell.execute_reply": "2020-12-06T08:47:22.811233Z"
    }
   },
   "outputs": [],
   "source": [
    "n_embedding = 40\n",
    "n_hidden = 256\n",
    "\n",
    "model = CharModel(n_alphabet, n_embedding=40, n_hidden=128)\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:22.813741Z",
     "iopub.status.busy": "2020-12-06T08:47:22.813436Z",
     "iopub.status.idle": "2020-12-06T08:47:22.815254Z",
     "shell.execute_reply": "2020-12-06T08:47:22.814935Z"
    }
   },
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(model.parameters(), 0.001)\n",
    "#criterion = nn.CrossEntropyLoss()\n",
    "criterion = F.nll_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:22.818691Z",
     "iopub.status.busy": "2020-12-06T08:47:22.818372Z",
     "iopub.status.idle": "2020-12-06T08:47:22.819647Z",
     "shell.execute_reply": "2020-12-06T08:47:22.819884Z"
    }
   },
   "outputs": [],
   "source": [
    "def fit(model, n_epoch=2):\n",
    "    optimizer = optim.Adam(model.parameters(), 0.001)\n",
    "    \n",
    "    for epoch in range(n_epoch):\n",
    "        print(f\"Epoch {epoch}:\")\n",
    "        running_loss, correct = 0.0, 0\n",
    "\n",
    "        model.train()\n",
    "        for X, y in train_dl:\n",
    "            X, y = X.to(device), y.to(device)\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            y_ = model(X)\n",
    "            loss = criterion(y_, y)\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            _, y_label_ = torch.max(y_, 1)\n",
    "            correct += (y_label_ == y).sum().item()\n",
    "            running_loss += loss.item() * X.shape[0]\n",
    "\n",
    "        print(f\"  Train Loss: {running_loss / len(train_dl.dataset):0.4f}\")\n",
    "        print(f\"  Train Acc:  {correct / len(train_dl.dataset):0.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:22.891629Z",
     "iopub.status.busy": "2020-12-06T08:47:22.885064Z",
     "iopub.status.idle": "2020-12-06T08:47:44.559163Z",
     "shell.execute_reply": "2020-12-06T08:47:44.558554Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/stefan/projects/pytorch_tutorial/.venv/lib/python3.8/site-packages/torch/nn/functional.py:1628: UserWarning: nn.functional.tanh is deprecated. Use torch.tanh instead.\n",
      "  warnings.warn(\"nn.functional.tanh is deprecated. Use torch.tanh instead.\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:\n",
      "  Train Loss: 2.2354\n",
      "  Train Acc:  0.37\n",
      "Epoch 1:\n",
      "  Train Loss: 1.9220\n",
      "  Train Acc:  0.44\n"
     ]
    }
   ],
   "source": [
    "fit(model, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:44.564413Z",
     "iopub.status.busy": "2020-12-06T08:47:44.563821Z",
     "iopub.status.idle": "2020-12-06T08:47:44.565659Z",
     "shell.execute_reply": "2020-12-06T08:47:44.566157Z"
    }
   },
   "outputs": [],
   "source": [
    "def predict(word):\n",
    "    word_idx = [char2index[c] for c in word]\n",
    "    word_idx\n",
    "    with torch.no_grad():\n",
    "        X = torch.tensor(word_idx).unsqueeze(0).to(device)\n",
    "        model.eval()\n",
    "        y_ = model(X).cpu()\n",
    "    pred = index2char[torch.argmax(y_).item()]\n",
    "    print(f\"{word} --> '{pred}'\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:44.569376Z",
     "iopub.status.busy": "2020-12-06T08:47:44.568799Z",
     "iopub.status.idle": "2020-12-06T08:47:44.573254Z",
     "shell.execute_reply": "2020-12-06T08:47:44.572657Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "the --> ' '\n"
     ]
    }
   ],
   "source": [
    "predict(\"the\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:44.578272Z",
     "iopub.status.busy": "2020-12-06T08:47:44.577584Z",
     "iopub.status.idle": "2020-12-06T08:47:44.586743Z",
     "shell.execute_reply": "2020-12-06T08:47:44.587254Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "wom --> 'e'\n"
     ]
    }
   ],
   "source": [
    "predict(\"wom\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:44.591272Z",
     "iopub.status.busy": "2020-12-06T08:47:44.590478Z",
     "iopub.status.idle": "2020-12-06T08:47:44.594301Z",
     "shell.execute_reply": "2020-12-06T08:47:44.594786Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "man --> ' '\n"
     ]
    }
   ],
   "source": [
    "predict(\"man\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:44.599994Z",
     "iopub.status.busy": "2020-12-06T08:47:44.599344Z",
     "iopub.status.idle": "2020-12-06T08:47:44.602202Z",
     "shell.execute_reply": "2020-12-06T08:47:44.602759Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hum --> 'a'\n"
     ]
    }
   ],
   "source": [
    "predict(\"hum\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:44.613497Z",
     "iopub.status.busy": "2020-12-06T08:47:44.612587Z",
     "iopub.status.idle": "2020-12-06T08:47:44.614830Z",
     "shell.execute_reply": "2020-12-06T08:47:44.615544Z"
    }
   },
   "outputs": [],
   "source": [
    "class CharModel(nn.Module):\n",
    "    def __init__(self, n_vocab, n_embedding, n_hidden):\n",
    "        super().__init__()\n",
    "        self.emb = nn.Embedding(n_vocab, n_embedding)\n",
    "        self.lin_in = nn.Linear(n_embedding, n_hidden)\n",
    "        self.lin_hidden = nn.Linear(n_hidden, n_hidden)\n",
    "        self.lin_out = nn.Linear(n_hidden, n_vocab)\n",
    "        \n",
    "    def forward(self, X):\n",
    "        c1, c2, c3 = X[:, 0], X[:, 1], X[:, 2]\n",
    "        \n",
    "        in1 = F.relu(self.lin_in(self.emb(c1)))       \n",
    "        in2 = F.relu(self.lin_in(self.emb(c2)))\n",
    "        in3 = F.relu(self.lin_in(self.emb(c3)))\n",
    "\n",
    "        h = F.tanh(self.lin_hidden(in1))\n",
    "        h = F.tanh(self.lin_hidden(h + in2))\n",
    "        h = F.tanh(self.lin_hidden(h + in3))\n",
    "        \n",
    "        return F.log_softmax(self.lin_out(h), dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:47:44.620975Z",
     "iopub.status.busy": "2020-12-06T08:47:44.620431Z",
     "iopub.status.idle": "2020-12-06T08:48:12.736384Z",
     "shell.execute_reply": "2020-12-06T08:48:12.736856Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:\n",
      "  Train Loss: 2.2230\n",
      "  Train Acc:  0.37\n",
      "Epoch 1:\n",
      "  Train Loss: 1.9139\n",
      "  Train Acc:  0.44\n",
      "\n",
      "the --> ' '\n",
      "wom --> 'e'\n",
      "man --> ' '\n",
      "hum --> 'a'\n"
     ]
    }
   ],
   "source": [
    "model = CharModel(n_alphabet, n_embedding=n_embedding, n_hidden=128).to(device)\n",
    "fit(model)\n",
    "\n",
    "print()\n",
    "predict(\"the\")\n",
    "predict(\"wom\")\n",
    "predict(\"man\")\n",
    "predict(\"hum\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:48:12.743238Z",
     "iopub.status.busy": "2020-12-06T08:48:12.742679Z",
     "iopub.status.idle": "2020-12-06T08:48:12.744508Z",
     "shell.execute_reply": "2020-12-06T08:48:12.744957Z"
    }
   },
   "outputs": [],
   "source": [
    "class CharModel(nn.Module):\n",
    "    def __init__(self, n_vocab, n_embedding, n_hidden):\n",
    "        super().__init__()\n",
    "        self.emb = nn.Embedding(n_vocab, n_embedding)\n",
    "        self.lin_in = nn.Linear(n_embedding, n_hidden)\n",
    "        self.lin_hidden = nn.Linear(n_hidden, n_hidden)\n",
    "        self.lin_out = nn.Linear(n_hidden, n_vocab)\n",
    "        \n",
    "        self.n_hidden = n_hidden\n",
    "        \n",
    "    def forward(self, X):\n",
    "        c1, c2, c3 = X[:, 0], X[:, 1], X[:, 2]\n",
    "        \n",
    "        in1 = F.relu(self.lin_in(self.emb(c1)))       \n",
    "        in2 = F.relu(self.lin_in(self.emb(c2)))\n",
    "        in3 = F.relu(self.lin_in(self.emb(c3)))\n",
    "        \n",
    "        h = torch.zeros(X.shape[0], n_hidden, requires_grad=True).to(device)\n",
    "        h = F.tanh(self.lin_hidden(h + in1))\n",
    "        h = F.tanh(self.lin_hidden(h + in2))\n",
    "        h = F.tanh(self.lin_hidden(h + in3))\n",
    "        \n",
    "        return F.log_softmax(self.lin_out(h), dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:48:12.748353Z",
     "iopub.status.busy": "2020-12-06T08:48:12.747800Z",
     "iopub.status.idle": "2020-12-06T08:48:50.701859Z",
     "shell.execute_reply": "2020-12-06T08:48:50.701338Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:\n",
      "  Train Loss: 2.0938\n",
      "  Train Acc:  0.40\n",
      "Epoch 1:\n",
      "  Train Loss: 1.8023\n",
      "  Train Acc:  0.47\n",
      "\n",
      "the --> ' '\n",
      "wom --> 'e'\n",
      "man --> ' '\n",
      "hum --> 'a'\n"
     ]
    }
   ],
   "source": [
    "model = CharModel(n_alphabet, n_embedding=n_embedding, n_hidden=n_hidden).to(device)\n",
    "fit(model)\n",
    "\n",
    "print()\n",
    "predict(\"the\")\n",
    "predict(\"wom\")\n",
    "predict(\"man\")\n",
    "predict(\"hum\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:48:50.707961Z",
     "iopub.status.busy": "2020-12-06T08:48:50.707177Z",
     "iopub.status.idle": "2020-12-06T08:48:50.709589Z",
     "shell.execute_reply": "2020-12-06T08:48:50.709092Z"
    }
   },
   "outputs": [],
   "source": [
    "class CharModel(nn.Module):\n",
    "    def __init__(self, n_vocab, n_embedding, n_hidden):\n",
    "        super().__init__()\n",
    "        self.emb = nn.Embedding(n_vocab, n_embedding)\n",
    "        self.lin_in = nn.Linear(n_embedding, n_hidden)\n",
    "        self.lin_hidden = nn.Linear(n_hidden, n_hidden)\n",
    "        self.lin_out = nn.Linear(n_hidden, n_vocab)\n",
    "        \n",
    "        self.n_hidden = n_hidden\n",
    "        \n",
    "    def forward(self, X):\n",
    "        h = torch.zeros(X.shape[0], n_hidden, requires_grad=True).to(device)\n",
    "        for i in range(X.shape[1]):\n",
    "            c = X[:, i]\n",
    "            in_ = F.relu(self.lin_in(self.emb(c)))\n",
    "            h = F.tanh(self.lin_hidden(h + in_))\n",
    "\n",
    "        return F.log_softmax(self.lin_out(h), dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:48:50.713113Z",
     "iopub.status.busy": "2020-12-06T08:48:50.712537Z",
     "iopub.status.idle": "2020-12-06T08:49:28.649837Z",
     "shell.execute_reply": "2020-12-06T08:49:28.650337Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:\n",
      "  Train Loss: 2.0920\n",
      "  Train Acc:  0.40\n",
      "Epoch 1:\n",
      "  Train Loss: 1.7984\n",
      "  Train Acc:  0.47\n",
      "\n",
      "the --> ' '\n",
      "wom --> 'a'\n",
      "man --> ' '\n",
      "hum --> 'a'\n"
     ]
    }
   ],
   "source": [
    "model = CharModel(n_alphabet, n_embedding=n_embedding, n_hidden=n_hidden).to(device)\n",
    "fit(model)\n",
    "\n",
    "print()\n",
    "predict(\"the\")\n",
    "predict(\"wom\")\n",
    "predict(\"man\")\n",
    "predict(\"hum\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:49:28.653466Z",
     "iopub.status.busy": "2020-12-06T08:49:28.652870Z",
     "iopub.status.idle": "2020-12-06T08:49:28.656600Z",
     "shell.execute_reply": "2020-12-06T08:49:28.657108Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "the huma --> 'n'\n"
     ]
    }
   ],
   "source": [
    "predict(\"the huma\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:49:28.660105Z",
     "iopub.status.busy": "2020-12-06T08:49:28.659516Z",
     "iopub.status.idle": "2020-12-06T08:49:28.664202Z",
     "shell.execute_reply": "2020-12-06T08:49:28.663749Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "those  --> 'o'\n"
     ]
    }
   ],
   "source": [
    "predict(\"those \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:49:28.667736Z",
     "iopub.status.busy": "2020-12-06T08:49:28.666915Z",
     "iopub.status.idle": "2020-12-06T08:49:28.671387Z",
     "shell.execute_reply": "2020-12-06T08:49:28.670940Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "those o --> 'f'\n"
     ]
    }
   ],
   "source": [
    "predict(\"those o\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:49:28.674900Z",
     "iopub.status.busy": "2020-12-06T08:49:28.674088Z",
     "iopub.status.idle": "2020-12-06T08:49:28.678752Z",
     "shell.execute_reply": "2020-12-06T08:49:28.678304Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "those of  --> 's'\n"
     ]
    }
   ],
   "source": [
    "predict(\"those of \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:49:28.682279Z",
     "iopub.status.busy": "2020-12-06T08:49:28.681453Z",
     "iopub.status.idle": "2020-12-06T08:49:28.685718Z",
     "shell.execute_reply": "2020-12-06T08:49:28.686154Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "those of u --> 'p'\n"
     ]
    }
   ],
   "source": [
    "predict(\"those of u\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can use `nn.Sequential` to make it a bit more readable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:49:28.692174Z",
     "iopub.status.busy": "2020-12-06T08:49:28.691515Z",
     "iopub.status.idle": "2020-12-06T08:49:28.693952Z",
     "shell.execute_reply": "2020-12-06T08:49:28.693423Z"
    }
   },
   "outputs": [],
   "source": [
    "class CharModel(nn.Module):\n",
    "    def __init__(self, n_vocab, n_embedding, n_hidden):\n",
    "        super().__init__()\n",
    "        self.i2e = nn.Sequential(\n",
    "            nn.Embedding(n_vocab, n_embedding),\n",
    "            nn.Linear(n_embedding, n_hidden),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.h2h = nn.Sequential(\n",
    "            nn.Linear(n_hidden, n_hidden),\n",
    "            nn.Tanh(),\n",
    "        )\n",
    "        self.h2out = nn.Linear(n_hidden, n_vocab)\n",
    "        \n",
    "        self.n_hidden = n_hidden\n",
    "        \n",
    "    def forward(self, X):\n",
    "        h = torch.zeros(X.shape[0], n_hidden, requires_grad=True).to(device)\n",
    "        for i in range(X.shape[1]):\n",
    "            c = X[:, i]\n",
    "            h = self.h2h(h + self.i2e(c))\n",
    "\n",
    "        return F.log_softmax(self.h2out(h), dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-12-06T08:49:28.697759Z",
     "iopub.status.busy": "2020-12-06T08:49:28.697090Z",
     "iopub.status.idle": "2020-12-06T08:50:06.780817Z",
     "shell.execute_reply": "2020-12-06T08:50:06.781311Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:\n",
      "  Train Loss: 2.0896\n",
      "  Train Acc:  0.40\n",
      "Epoch 1:\n",
      "  Train Loss: 1.7967\n",
      "  Train Acc:  0.47\n",
      "\n",
      "the --> ' '\n",
      "wom --> 'a'\n",
      "man --> ' '\n",
      "hum --> 'a'\n"
     ]
    }
   ],
   "source": [
    "model = CharModel(n_alphabet, n_embedding=n_embedding, n_hidden=n_hidden).to(device)\n",
    "fit(model)\n",
    "\n",
    "print()\n",
    "predict(\"the\")\n",
    "predict(\"wom\")\n",
    "predict(\"man\")\n",
    "predict(\"hum\")"
   ]
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}