train_without_outputs.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#importing libs\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "import numpy as np\n",
    "import os, string, re, random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#selecting device to work on\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print(\"Working on\",device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#primitives\n",
    "\n",
    "chars = string.ascii_letters + string.digits + string.punctuation + ' ' + '\\n'\n",
    "char2int = {ch:i for i,ch in enumerate(chars)}\n",
    "int2char = {i:ch for i,ch in enumerate(chars)}\n",
    "char_len = len(chars)\n",
    "\n",
    "print(\"Characters include:\",repr(chars))\n",
    "print(\"Characters length:\",char_len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#loading data\n",
    "\n",
    "unwanted_characters = \"[àéñ¡]\"\n",
    "directory = 'data/'\n",
    "\n",
    "data = []\n",
    "try:\n",
    "    filenames = os.listdir(directory)\n",
    "    for fname in filenames:\n",
    "        txt = open(directory+fname, 'r', encoding='utf-8', errors='ignore').read()\n",
    "        txt = re.sub(unwanted_characters,\"\",txt)\n",
    "        data = np.append(data,txt)\n",
    "    print(\"Data loaded successfully\")\n",
    "\n",
    "except Exception as exception:\n",
    "    print(exception)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Supporting functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def preprocess(data, seq_len):\n",
    "    '''\n",
    "    Take data matrix and seqence length and converts it into\n",
    "    required format i.e. X - data and y - label\n",
    "    '''\n",
    "\n",
    "    y_t = data[0][1:]\n",
    "    num_of_seq = len(data[0]) // seq_len\n",
    "    X_t = np.array([char2int[ch] for ch in data[0][: seq_len * num_of_seq]]).reshape(-1,seq_len)\n",
    "    y_t = [char2int[ch] for ch in y_t[: seq_len * num_of_seq]]\n",
    "    for data_t in data[1:]:\n",
    "        data_y = data_t[1:] + \"\\n\"\n",
    "        num_of_seq = len(data_t) // seq_len\n",
    "        X_t = np.append(X_t, np.array([char2int[ch] for ch in data_t[: seq_len * num_of_seq]]).reshape(-1, seq_len), axis = 0)\n",
    "        y_t.extend([char2int[ch] for ch in data_y[: seq_len * num_of_seq]])\n",
    "\n",
    "    return (X_t,y_t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getBatches(X, y, batch_size, seq_len):\n",
    "    '''\n",
    "    Generates batches for training\n",
    "    '''\n",
    "\n",
    "    X_rows = X.shape[0]\n",
    "    target_size = seq_len * batch_size\n",
    "    y_ptr = 0\n",
    "    for batch_num in range(0,X_rows,batch_size):\n",
    "\n",
    "        X_tt = X[batch_num : batch_num + batch_size]\n",
    "\n",
    "        X_t = np.zeros((X_tt.shape[0], seq_len, char_len))\n",
    "\n",
    "        for batch_idx, seq_x in enumerate(X_tt):\n",
    "\n",
    "            for seq_idx, char_x in enumerate(seq_x):\n",
    "                X_t[batch_idx, seq_idx, char_x] = 1\n",
    "\n",
    "\n",
    "        yield X_t, y[y_ptr : y_ptr + target_size]\n",
    "\n",
    "        y_ptr += target_size  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RNNModel(nn.Module):\n",
    "    def __init__(self,input_size, hidden_size, rnn_layers, fc_hidden_dim):\n",
    "        super(RNNModel, self).__init__()\n",
    "        \n",
    "        self.rnn_layers = rnn_layers\n",
    "        self.hidden_size = hidden_size\n",
    "\n",
    "        self.rnn = nn.RNN(input_size=input_size, hidden_size=hidden_size, num_layers=rnn_layers, batch_first=True)\n",
    "        self.fcl = nn.Sequential(\n",
    "                    nn.Linear(in_features=hidden_size, out_features=fc_hidden_dim),\n",
    "                    nn.ReLU(),\n",
    "                    nn.Linear(in_features=fc_hidden_dim, out_features=input_size)\n",
    "                    )\n",
    "\n",
    "\n",
    "    def forward(self, x, prev_state):\n",
    "        out, state = self.rnn(x, prev_state)\n",
    "        fcl_out = self.fcl(out.reshape(-1,out.size(2)))\n",
    "        return fcl_out, state\n",
    "\n",
    "\n",
    "    def init_hidden(self, batch_size):\n",
    "        return torch.zeros((self.rnn_layers, batch_size, self.hidden_size), device=device)\n",
    "\n",
    "    def getModelName(self):\n",
    "        return \"RNN\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LSTMModel(nn.Module):\n",
    "    def __init__(self,input_size, hidden_size, lstm_layers, fc_hidden_dim):\n",
    "        super(LSTMModel, self).__init__()\n",
    "        \n",
    "        self.lstm_layers = lstm_layers\n",
    "        self.hidden_size = hidden_size\n",
    "\n",
    "        self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size, num_layers=lstm_layers, batch_first=True)\n",
    "        self.fcl = nn.Sequential(\n",
    "                    nn.Linear(in_features=hidden_size, out_features=fc_hidden_dim),\n",
    "                    nn.ReLU(),\n",
    "                    nn.Linear(in_features=fc_hidden_dim, out_features=input_size)\n",
    "                    )\n",
    "\n",
    "\n",
    "    def forward(self, x, prev_state):\n",
    "        out, state = self.lstm(x, prev_state)\n",
    "        fcl_out = self.fcl(out.reshape(-1,out.size(2)))\n",
    "        return fcl_out, state\n",
    "\n",
    "\n",
    "    def init_hidden(self, batch_size):\n",
    "        return [torch.zeros((self.lstm_layers, batch_size, self.hidden_size), device=device), \n",
    "                torch.zeros((self.lstm_layers, batch_size, self.hidden_size), device=device)]\n",
    "\n",
    "    def getModelName(self):\n",
    "        return \"LSTM\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GRUModel(nn.Module):\n",
    "    def __init__(self,input_size, hidden_size, gru_layers, fc_hidden_dim):\n",
    "        super(GRUModel, self).__init__()\n",
    "        \n",
    "        self.gru_layers = gru_layers\n",
    "        self.hidden_size = hidden_size\n",
    "\n",
    "        self.gru = nn.GRU(input_size=input_size, hidden_size=hidden_size, num_layers=gru_layers, batch_first=True)\n",
    "        self.fcl = nn.Sequential(\n",
    "                    nn.Linear(in_features=hidden_size, out_features=fc_hidden_dim),\n",
    "                    nn.ReLU(),\n",
    "                    nn.Linear(in_features=fc_hidden_dim, out_features=input_size)\n",
    "                    )\n",
    "\n",
    "\n",
    "    def forward(self, x, prev_state):\n",
    "        out, state = self.gru(x, prev_state)\n",
    "        fcl_out = self.fcl(out.reshape(-1,out.size(2)))\n",
    "        return fcl_out, state\n",
    "\n",
    "\n",
    "    def init_hidden(self, batch_size):\n",
    "        return torch.zeros((self.gru_layers, batch_size, self.hidden_size), device=device)\n",
    "\n",
    "\n",
    "    def getModelName(self):\n",
    "        return \"GRU\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hyperparameters\n",
    "\n",
    "batch_size = 128\n",
    "seq_len = 250\n",
    "\n",
    "EPOCHS = 500\n",
    "lr = 0.001\n",
    "\n",
    "loss_function = nn.CrossEntropyLoss()\n",
    "\n",
    "hidden_size = 512\n",
    "num_layers = 2\n",
    "fc_hidden_dim = 256\n",
    "\n",
    "initial_words = [\"I can do this all day\", \"I'm IRONMAN\", \"Reality is often disappointing\", \"I knew it!\", \"NoobMaster69\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def softmax(x):\n",
    "    '''Compute softmax values for each sets of scores in x.'''\n",
    "    return np.exp(x) / np.sum(np.exp(x), axis=0)\n",
    "    \n",
    "def decodeBasedOnTopK(one_hot_pred, top_k=1):\n",
    "    '''\n",
    "    Selects an integer from the topk distribution.\n",
    "    '''\n",
    "    choices = np.argpartition(one_hot_pred, -top_k)[-top_k:]\n",
    "    prob = one_hot_pred[choices]\n",
    "    prob = softmax(prob)\n",
    "    prob = prob / np.sum(prob)\n",
    "    choice = np.random.choice(choices, 1, p=prob)\n",
    "    return choice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def sample(model, initial_string = \"Dread it\", num_chars = 200, topk=1):\n",
    "    '''\n",
    "    Generates sample text.\n",
    "    '''\n",
    "    model.eval()\n",
    "    state = model.init_hidden(1)\n",
    "\n",
    "    #getting hidden states for starter string\n",
    "    for char in initial_string:\n",
    "        one_hot = torch.zeros((1, 1, char_len), dtype=torch.float, device=device)\n",
    "        one_hot[0,0,char2int[char]] = 1\n",
    "        _, state = model(one_hot, state)\n",
    "    \n",
    "    predict_letters = initial_string[-1]\n",
    "\n",
    "    for i in range(num_chars):\n",
    "        one_hot = torch.zeros((1, 1, char_len), dtype=torch.float, device=device)\n",
    "        one_hot[0,0,char2int[predict_letters[-1]]] = 1\n",
    "        out, state = model(one_hot, state)\n",
    "        predict_letters += int2char[decodeBasedOnTopK(out[0].cpu().data.numpy(), top_k=topk)]\n",
    "\n",
    "    print(\"{} -- {}\".format(initial_string, predict_letters[1:]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#initializations\n",
    "\n",
    "rnn_model = RNNModel(input_size=char_len, \n",
    "                       hidden_size=hidden_size, \n",
    "                       rnn_layers=num_layers, \n",
    "                       fc_hidden_dim=fc_hidden_dim).to(device)\n",
    "\n",
    "\n",
    "lstm_model = LSTMModel(input_size=char_len, \n",
    "                       hidden_size=hidden_size, \n",
    "                       lstm_layers=num_layers, \n",
    "                       fc_hidden_dim=fc_hidden_dim).to(device)\n",
    "\n",
    "                       \n",
    "gru_model = GRUModel(input_size=char_len, \n",
    "                       hidden_size=hidden_size, \n",
    "                       gru_layers=num_layers, \n",
    "                       fc_hidden_dim=fc_hidden_dim).to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train(model):\n",
    "\n",
    "    print(\"Working on {} Model\".format(model.getModelName()))\n",
    "    losses = []\n",
    "    opt = torch.optim.Adam(model.parameters(), lr=lr)\n",
    "    X,y = preprocess(data, seq_len)\n",
    "\n",
    "    for epoch in range(1, EPOCHS+1):\n",
    "        model.train()\n",
    "        epoch_loss = []\n",
    "        \n",
    "        for xx,yy in getBatches(X,y,batch_size,seq_len):\n",
    "            h_prev = model.init_hidden(xx.shape[0])\n",
    "\n",
    "            xx = torch.Tensor(xx).to(device)\n",
    "            yy = torch.LongTensor(yy).to(device)\n",
    "\n",
    "            opt.zero_grad()\n",
    "            out,_ = model(xx,h_prev)\n",
    "            loss = loss_function(out, yy)\n",
    "            epoch_loss.append(loss.item())\n",
    "\n",
    "        \n",
    "            loss.backward()\n",
    "            opt.step()\n",
    "\n",
    "\n",
    "        if(epoch % 10 == 0):\n",
    "            print(\"--\"*20)\n",
    "            print(\"Sample text generated at %dth epoch\"%(epoch))\n",
    "            sample(model=model, initial_string=random.choice(initial_words), num_chars=300, topk=5)\n",
    "            print(\"--\"*20)\n",
    "                \n",
    "        losses.append(np.mean(epoch_loss))\n",
    "        print('-------------------------')\n",
    "        print('Epoch: %d/%d ---- Epoch Loss: %f'%(epoch, EPOCHS, losses[-1]))\n",
    "\n",
    "    return losses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "rnn_losses = train(rnn_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "torch.save(rnn_model,\"model/{}.pth\".format(rnn_model.getModelName().lower()))\n",
    "print(\"Model Saved\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "lstm_losses = train(lstm_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "torch.save(lstm_model,\"model/{}.pth\".format(lstm_model.getModelName().lower()))\n",
    "print(\"Model Saved\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "gru_losses = train(gru_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "torch.save(gru_model,\"model/{}.pth\".format(gru_model.getModelName().lower()))\n",
    "print(\"Model Saved\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plotting\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.plot(range(1,EPOCHS+1),rnn_losses, label = \"RNN Loss\", color = 'red')    \n",
    "plt.plot(range(1,EPOCHS+1),lstm_losses, label = \"LSTM Loss\", color = 'green')\n",
    "plt.plot(range(1,EPOCHS+1),gru_losses, label = \"GRU Loss\", color = 'blue')  \n",
    "plt.title(\"Training Losses\")\n",
    "plt.xlabel(\"Epochs\")\n",
    "plt.ylabel(\"Loss Values\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.6.9 64-bit",
   "language": "python",
   "name": "python36964bit9b7bbc03f57b4d679e6428f48bb2d865"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9-final"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}