/4-3.Bi-LSTM(Attention)/Bi-LSTM(Attention).ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# code by Tae Hwan Jung(Jeff Jung) @graykode\n",
        "# Reference : https://github.com/prakashpandey9/Text-Classification-Pytorch/blob/master/models/LSTM_Attn.py\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "class BiLSTM_Attention(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(BiLSTM_Attention, self).__init__()\n",
        "\n",
        "        self.embedding = nn.Embedding(vocab_size, embedding_dim)\n",
        "        self.lstm = nn.LSTM(embedding_dim, n_hidden, bidirectional=True)\n",
        "        self.out = nn.Linear(n_hidden * 2, num_classes)\n",
        "\n",
        "    # lstm_output : [batch_size, n_step, n_hidden * num_directions(=2)], F matrix\n",
        "    def attention_net(self, lstm_output, final_state):\n",
        "        hidden = final_state.view(-1, n_hidden * 2, 1)   # hidden : [batch_size, n_hidden * num_directions(=2), 1(=n_layer)]\n",
        "        attn_weights = torch.bmm(lstm_output, hidden).squeeze(2) # attn_weights : [batch_size, n_step]\n",
        "        soft_attn_weights = F.softmax(attn_weights, 1)\n",
        "        # [batch_size, n_hidden * num_directions(=2), n_step] * [batch_size, n_step, 1] = [batch_size, n_hidden * num_directions(=2), 1]\n",
        "        context = torch.bmm(lstm_output.transpose(1, 2), soft_attn_weights.unsqueeze(2)).squeeze(2)\n",
        "        return context, soft_attn_weights.data.numpy() # context : [batch_size, n_hidden * num_directions(=2)]\n",
        "\n",
        "    def forward(self, X):\n",
        "        input = self.embedding(X) # input : [batch_size, len_seq, embedding_dim]\n",
        "        input = input.permute(1, 0, 2) # input : [len_seq, batch_size, embedding_dim]\n",
        "\n",
        "        hidden_state = torch.zeros(1*2, len(X), n_hidden) # [num_layers(=1) * num_directions(=2), batch_size, n_hidden]\n",
        "        cell_state = torch.zeros(1*2, len(X), n_hidden) # [num_layers(=1) * num_directions(=2), batch_size, n_hidden]\n",
        "\n",
        "        # final_hidden_state, final_cell_state : [num_layers(=1) * num_directions(=2), batch_size, n_hidden]\n",
        "        output, (final_hidden_state, final_cell_state) = self.lstm(input, (hidden_state, cell_state))\n",
        "        output = output.permute(1, 0, 2) # output : [batch_size, len_seq, n_hidden]\n",
        "        attn_output, attention = self.attention_net(output, final_hidden_state)\n",
        "        return self.out(attn_output), attention # model : [batch_size, num_classes], attention : [batch_size, n_step]\n",
        "\n",
        "if __name__ == '__main__':\n",
        "    embedding_dim = 2 # embedding size\n",
        "    n_hidden = 5  # number of hidden units in one cell\n",
        "    num_classes = 2  # 0 or 1\n",
        "\n",
        "    # 3 words sentences (=sequence_length is 3)\n",
        "    sentences = [\"i love you\", \"he loves me\", \"she likes baseball\", \"i hate you\", \"sorry for that\", \"this is awful\"]\n",
        "    labels = [1, 1, 1, 0, 0, 0]  # 1 is good, 0 is not good.\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",
        "    vocab_size = len(word_dict)\n",
        "\n",
        "    model = BiLSTM_Attention()\n",
        "\n",
        "    criterion = nn.CrossEntropyLoss()\n",
        "    optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
        "\n",
        "    inputs = torch.LongTensor([np.asarray([word_dict[n] for n in sen.split()]) for sen in sentences])\n",
        "    targets = torch.LongTensor([out for out in labels])  # To using Torch Softmax Loss function\n",
        "\n",
        "    # Training\n",
        "    for epoch in range(5000):\n",
        "        optimizer.zero_grad()\n",
        "        output, attention = model(inputs)\n",
        "        loss = criterion(output, targets)\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",
        "    # Test\n",
        "    test_text = 'sorry hate you'\n",
        "    tests = [np.asarray([word_dict[n] for n in test_text.split()])]\n",
        "    test_batch = torch.LongTensor(tests)\n",
        "\n",
        "    # Predict\n",
        "    predict, _ = model(test_batch)\n",
        "    predict = predict.data.max(1, keepdim=True)[1]\n",
        "    if predict[0][0] == 0:\n",
        "        print(test_text,\"is Bad Mean...\")\n",
        "    else:\n",
        "        print(test_text,\"is Good Mean!!\")\n",
        "\n",
        "    fig = plt.figure(figsize=(6, 3)) # [batch_size, n_step]\n",
        "    ax = fig.add_subplot(1, 1, 1)\n",
        "    ax.matshow(attention, cmap='viridis')\n",
        "    ax.set_xticklabels(['']+['first_word', 'second_word', 'third_word'], fontdict={'fontsize': 14}, rotation=90)\n",
        "    ax.set_yticklabels(['']+['batch_1', 'batch_2', 'batch_3', 'batch_4', 'batch_5', 'batch_6'], fontdict={'fontsize': 14})\n",
        "    plt.show()"
      ],
      "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
}