1-1.NNLM/NNLM.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# code by Tae Hwan Jung @graykode\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "\n",
        "def make_batch():\n",
        "    input_batch = []\n",
        "    target_batch = []\n",
        "\n",
        "    for sen in sentences:\n",
        "        word = sen.split() # space tokenizer\n",
        "        input = [word_dict[n] for n in word[:-1]] # create (1~n-1) as input\n",
        "        target = word_dict[word[-1]] # create (n) as target, We usually call this 'casual language model'\n",
        "\n",
        "        input_batch.append(input)\n",
        "        target_batch.append(target)\n",
        "\n",
        "    return input_batch, target_batch\n",
        "\n",
        "# Model\n",
        "class NNLM(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(NNLM, self).__init__()\n",
        "        self.C = nn.Embedding(n_class, m)\n",
        "        self.H = nn.Linear(n_step * m, n_hidden, bias=False)\n",
        "        self.d = nn.Parameter(torch.ones(n_hidden))\n",
        "        self.U = nn.Linear(n_hidden, n_class, bias=False)\n",
        "        self.W = nn.Linear(n_step * m, n_class, bias=False)\n",
        "        self.b = nn.Parameter(torch.ones(n_class))\n",
        "\n",
        "    def forward(self, X):\n",
        "        X = self.C(X) # X : [batch_size, n_step, m]\n",
        "        X = X.view(-1, n_step * m) # [batch_size, n_step * m]\n",
        "        tanh = torch.tanh(self.d + self.H(X)) # [batch_size, n_hidden]\n",
        "        output = self.b + self.W(X) + self.U(tanh) # [batch_size, n_class]\n",
        "        return output\n",
        "\n",
        "if __name__ == '__main__':\n",
        "    n_step = 2 # number of steps, n-1 in paper\n",
        "    n_hidden = 2 # number of hidden size, h in paper\n",
        "    m = 2 # embedding size, m in paper\n",
        "\n",
        "    sentences = [\"i like dog\", \"i love coffee\", \"i hate milk\"]\n",
        "\n",
        "    word_list = \" \".join(sentences).split()\n",
        "    word_list = list(set(word_list))\n",
        "    word_dict = {w: i for i, w in enumerate(word_list)}\n",
        "    number_dict = {i: w for i, w in enumerate(word_list)}\n",
        "    n_class = len(word_dict)  # number of Vocabulary\n",
        "\n",
        "    model = NNLM()\n",
        "\n",
        "    criterion = nn.CrossEntropyLoss()\n",
        "    optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
        "\n",
        "    input_batch, target_batch = make_batch()\n",
        "    input_batch = torch.LongTensor(input_batch)\n",
        "    target_batch = torch.LongTensor(target_batch)\n",
        "\n",
        "    # Training\n",
        "    for epoch in range(5000):\n",
        "        optimizer.zero_grad()\n",
        "        output = model(input_batch)\n",
        "\n",
        "        # output : [batch_size, n_class], target_batch : [batch_size]\n",
        "        loss = criterion(output, target_batch)\n",
        "        if (epoch + 1) % 1000 == 0:\n",
        "            print('Epoch:', '%04d' % (epoch + 1), 'cost =', '{:.6f}'.format(loss))\n",
        "\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "    # Predict\n",
        "    predict = model(input_batch).data.max(1, keepdim=True)[1]\n",
        "\n",
        "    # Test\n",
        "    print([sen.split()[:2] for sen in sentences], '->', [number_dict[n.item()] for n in predict.squeeze()])"
      ],
      "outputs": [],
      "execution_count": null
    }
  ],
  "metadata": {
    "anaconda-cloud": {},
    "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.6.1"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 4
}