text_classify.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "text_classify.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "S1vZIU-yVlxY"
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Prepare data"
      ],
      "metadata": {
        "id": "jR3kTwlAW5tE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from torchtext.datasets import AG_NEWS\n",
        "train_iter = AG_NEWS(split='train')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RQvmVPQ4VpEv",
        "outputId": "bff216fc-cca7-4ece-972d-da669bfd1178"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "29.5MB [00:00, 44.2MB/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from torchtext.data.utils import get_tokenizer\n",
        "from torchtext.vocab import build_vocab_from_iterator\n",
        "\n",
        "tokenizer = get_tokenizer('basic_english')\n",
        "train_iter = AG_NEWS(split='train')\n",
        "\n",
        "def yield_tokens(data_iter):\n",
        "    for _, text in data_iter:\n",
        "        yield tokenizer(text)\n",
        "\n",
        "vocab = build_vocab_from_iterator(yield_tokens(train_iter), specials=[\"<unk>\"])\n",
        "vocab.set_default_index(vocab[\"<unk>\"])"
      ],
      "metadata": {
        "id": "VLFMFj-rVqPs"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "text_pipeline = lambda x: vocab(tokenizer(x))\n",
        "label_pipeline = lambda x: int(x) - 1"
      ],
      "metadata": {
        "id": "q74baD6jVspQ"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Generate data batch and iterator"
      ],
      "metadata": {
        "id": "fpBnJCgqW9O7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from torch.utils.data import DataLoader\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "def collate_batch(batch):\n",
        "    label_list, text_list, offsets = [], [], [0]\n",
        "    for (_label, _text) in batch:\n",
        "         label_list.append(label_pipeline(_label))\n",
        "         processed_text = torch.tensor(text_pipeline(_text), dtype=torch.int64)\n",
        "         text_list.append(processed_text)\n",
        "         offsets.append(processed_text.size(0))\n",
        "    label_list = torch.tensor(label_list, dtype=torch.int64)\n",
        "    offsets = torch.tensor(offsets[:-1]).cumsum(dim=0)\n",
        "    text_list = torch.cat(text_list)\n",
        "    return label_list.to(device), text_list.to(device), offsets.to(device)    \n",
        "\n",
        "train_iter = AG_NEWS(split='train')\n",
        "dataloader = DataLoader(train_iter, batch_size=8, shuffle=False, collate_fn=collate_batch)"
      ],
      "metadata": {
        "id": "quYvsy8vVuBo"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from torch import nn\n",
        "\n",
        "class TextClassificationModel(nn.Module):\n",
        "\n",
        "    def __init__(self, vocab_size, embed_dim, num_class):\n",
        "        super(TextClassificationModel, self).__init__()\n",
        "        self.embedding = nn.EmbeddingBag(vocab_size, embed_dim, sparse=True)\n",
        "        self.fc = nn.Linear(embed_dim, num_class)\n",
        "        self.init_weights()\n",
        "\n",
        "    def init_weights(self):\n",
        "        initrange = 0.5\n",
        "        self.embedding.weight.data.uniform_(-initrange, initrange)\n",
        "        self.fc.weight.data.uniform_(-initrange, initrange)\n",
        "        self.fc.bias.data.zero_()\n",
        "\n",
        "    def forward(self, text, offsets):\n",
        "        embedded = self.embedding(text, offsets)\n",
        "        return self.fc(embedded)"
      ],
      "metadata": {
        "id": "VpR1rR0CVvrE"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Initiate an instance"
      ],
      "metadata": {
        "id": "USPdQPASW__k"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "num_class = len(set([label for (label, text) in train_iter]))\n",
        "vocab_size = len(vocab)\n",
        "emsize = 64\n",
        "model = TextClassificationModel(vocab_size, emsize, num_class).to(device)"
      ],
      "metadata": {
        "id": "68CuJ4M7VxF5"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Function to train and evaluate the results"
      ],
      "metadata": {
        "id": "haToiroRXC4n"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import time\n",
        "\n",
        "def train(dataloader):\n",
        "    model.train()\n",
        "    total_acc, total_count = 0, 0\n",
        "    log_interval = 500\n",
        "    start_time = time.time()\n",
        "\n",
        "    for idx, (label, text, offsets) in enumerate(dataloader):\n",
        "        optimizer.zero_grad()\n",
        "        predicted_label = model(text, offsets)\n",
        "        loss = criterion(predicted_label, label)\n",
        "        loss.backward()\n",
        "        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1)\n",
        "        optimizer.step()\n",
        "        total_acc += (predicted_label.argmax(1) == label).sum().item()\n",
        "        total_count += label.size(0)\n",
        "        if idx % log_interval == 0 and idx > 0:\n",
        "            elapsed = time.time() - start_time\n",
        "            print('| epoch {:3d} | {:5d}/{:5d} batches '\n",
        "                  '| accuracy {:8.3f}'.format(epoch, idx, len(dataloader),\n",
        "                                              total_acc/total_count))\n",
        "            total_acc, total_count = 0, 0\n",
        "            start_time = time.time()\n",
        "\n",
        "def evaluate(dataloader):\n",
        "    model.eval()\n",
        "    total_acc, total_count = 0, 0\n",
        "\n",
        "    with torch.no_grad():\n",
        "        for idx, (label, text, offsets) in enumerate(dataloader):\n",
        "            predicted_label = model(text, offsets)\n",
        "            loss = criterion(predicted_label, label)\n",
        "            total_acc += (predicted_label.argmax(1) == label).sum().item()\n",
        "            total_count += label.size(0)\n",
        "    return total_acc/total_count"
      ],
      "metadata": {
        "id": "9JD5rLyZVzA9"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Split the dataset and run the model"
      ],
      "metadata": {
        "id": "PrWwj1-kXGWB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from torch.utils.data.dataset import random_split\n",
        "from torchtext.data.functional import to_map_style_dataset\n",
        "# Hyperparameters\n",
        "EPOCHS = 10 # epoch\n",
        "LR = 5  # learning rate\n",
        "BATCH_SIZE = 64 # batch size for training\n",
        "  \n",
        "criterion = torch.nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr=LR)\n",
        "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.1)\n",
        "total_accu = None\n",
        "train_iter, test_iter = AG_NEWS()\n",
        "train_dataset = to_map_style_dataset(train_iter)\n",
        "test_dataset = to_map_style_dataset(test_iter)\n",
        "num_train = int(len(train_dataset) * 0.95)\n",
        "split_train_, split_valid_ = \\\n",
        "    random_split(train_dataset, [num_train, len(train_dataset) - num_train])\n",
        "\n",
        "train_dataloader = DataLoader(split_train_, batch_size=BATCH_SIZE,\n",
        "                              shuffle=True, collate_fn=collate_batch)\n",
        "valid_dataloader = DataLoader(split_valid_, batch_size=BATCH_SIZE,\n",
        "                              shuffle=True, collate_fn=collate_batch)\n",
        "test_dataloader = DataLoader(test_dataset, batch_size=BATCH_SIZE,\n",
        "                             shuffle=True, collate_fn=collate_batch)\n",
        "\n",
        "for epoch in range(1, EPOCHS + 1):\n",
        "    epoch_start_time = time.time()\n",
        "    train(train_dataloader)\n",
        "    accu_val = evaluate(valid_dataloader)\n",
        "    if total_accu is not None and total_accu > accu_val:\n",
        "      scheduler.step()\n",
        "    else:\n",
        "       total_accu = accu_val\n",
        "    print('-' * 59)\n",
        "    print('| end of epoch {:3d} | time: {:5.2f}s | '\n",
        "          'valid accuracy {:8.3f} '.format(epoch,\n",
        "                                           time.time() - epoch_start_time,\n",
        "                                           accu_val))\n",
        "    print('-' * 59)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jmTMZA9fV0t6",
        "outputId": "4414b008-6631-4b57-b506-b0a08392939e"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "1.86MB [00:00, 29.2MB/s]                  \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "| epoch   1 |   500/ 1782 batches | accuracy    0.682\n",
            "| epoch   1 |  1000/ 1782 batches | accuracy    0.856\n",
            "| epoch   1 |  1500/ 1782 batches | accuracy    0.873\n",
            "-----------------------------------------------------------\n",
            "| end of epoch   1 | time: 16.59s | valid accuracy    0.879 \n",
            "-----------------------------------------------------------\n",
            "| epoch   2 |   500/ 1782 batches | accuracy    0.896\n",
            "| epoch   2 |  1000/ 1782 batches | accuracy    0.900\n",
            "| epoch   2 |  1500/ 1782 batches | accuracy    0.903\n",
            "-----------------------------------------------------------\n",
            "| end of epoch   2 | time: 15.16s | valid accuracy    0.894 \n",
            "-----------------------------------------------------------\n",
            "| epoch   3 |   500/ 1782 batches | accuracy    0.913\n",
            "| epoch   3 |  1000/ 1782 batches | accuracy    0.913\n",
            "| epoch   3 |  1500/ 1782 batches | accuracy    0.915\n",
            "-----------------------------------------------------------\n",
            "| end of epoch   3 | time: 15.16s | valid accuracy    0.894 \n",
            "-----------------------------------------------------------\n",
            "| epoch   4 |   500/ 1782 batches | accuracy    0.928\n",
            "| epoch   4 |  1000/ 1782 batches | accuracy    0.929\n",
            "| epoch   4 |  1500/ 1782 batches | accuracy    0.929\n",
            "-----------------------------------------------------------\n",
            "| end of epoch   4 | time: 15.20s | valid accuracy    0.901 \n",
            "-----------------------------------------------------------\n",
            "| epoch   5 |   500/ 1782 batches | accuracy    0.932\n",
            "| epoch   5 |  1000/ 1782 batches | accuracy    0.930\n",
            "| epoch   5 |  1500/ 1782 batches | accuracy    0.930\n",
            "-----------------------------------------------------------\n",
            "| end of epoch   5 | time: 14.93s | valid accuracy    0.903 \n",
            "-----------------------------------------------------------\n",
            "| epoch   6 |   500/ 1782 batches | accuracy    0.933\n",
            "| epoch   6 |  1000/ 1782 batches | accuracy    0.933\n",
            "| epoch   6 |  1500/ 1782 batches | accuracy    0.932\n",
            "-----------------------------------------------------------\n",
            "| end of epoch   6 | time: 14.84s | valid accuracy    0.905 \n",
            "-----------------------------------------------------------\n",
            "| epoch   7 |   500/ 1782 batches | accuracy    0.932\n",
            "| epoch   7 |  1000/ 1782 batches | accuracy    0.935\n",
            "| epoch   7 |  1500/ 1782 batches | accuracy    0.932\n",
            "-----------------------------------------------------------\n",
            "| end of epoch   7 | time: 15.08s | valid accuracy    0.902 \n",
            "-----------------------------------------------------------\n",
            "| epoch   8 |   500/ 1782 batches | accuracy    0.935\n",
            "| epoch   8 |  1000/ 1782 batches | accuracy    0.934\n",
            "| epoch   8 |  1500/ 1782 batches | accuracy    0.936\n",
            "-----------------------------------------------------------\n",
            "| end of epoch   8 | time: 15.16s | valid accuracy    0.904 \n",
            "-----------------------------------------------------------\n",
            "| epoch   9 |   500/ 1782 batches | accuracy    0.935\n",
            "| epoch   9 |  1000/ 1782 batches | accuracy    0.933\n",
            "| epoch   9 |  1500/ 1782 batches | accuracy    0.936\n",
            "-----------------------------------------------------------\n",
            "| end of epoch   9 | time: 14.70s | valid accuracy    0.904 \n",
            "-----------------------------------------------------------\n",
            "| epoch  10 |   500/ 1782 batches | accuracy    0.935\n",
            "| epoch  10 |  1000/ 1782 batches | accuracy    0.934\n",
            "| epoch  10 |  1500/ 1782 batches | accuracy    0.935\n",
            "-----------------------------------------------------------\n",
            "| end of epoch  10 | time: 14.73s | valid accuracy    0.904 \n",
            "-----------------------------------------------------------\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Evaluate on test set"
      ],
      "metadata": {
        "id": "M3VbMtU1XLTG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print('Checking the results of test dataset.')\n",
        "accu_test = evaluate(test_dataloader)\n",
        "print('test accuracy {:8.3f}'.format(accu_test))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "j8DUMCi0V75i",
        "outputId": "60d809e6-a227-4271-887b-a1c8e9549189"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Checking the results of test dataset.\n",
            "test accuracy    0.904\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Test on random news"
      ],
      "metadata": {
        "id": "QaUadtHPXMKB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "ag_news_label = {1: \"World\",\n",
        "                 2: \"Sports\",\n",
        "                 3: \"Business\",\n",
        "                 4: \"Sci/Tec\"}\n",
        "\n",
        "def predict(text, text_pipeline):\n",
        "    with torch.no_grad():\n",
        "        text = torch.tensor(text_pipeline(text))\n",
        "        output = model(text, torch.tensor([0]))\n",
        "        return output.argmax(1).item() + 1\n",
        "\n",
        "ex_text_str = \"MEMPHIS, Tenn. – Four days ago, Jon Rahm was \\\n",
        "    enduring the season’s worst weather conditions on Sunday at The \\\n",
        "    Open on his way to a closing 75 at Royal Portrush, which \\\n",
        "    considering the wind and the rain was a respectable showing. \\\n",
        "    Thursday’s first round at the WGC-FedEx St. Jude Invitational \\\n",
        "    was another story. With temperatures in the mid-80s and hardly any \\\n",
        "    wind, the Spaniard was 13 strokes better in a flawless round. \\\n",
        "    Thanks to his best putting performance on the PGA Tour, Rahm \\\n",
        "    finished with an 8-under 62 for a three-stroke lead, which \\\n",
        "    was even more impressive considering he’d never played the \\\n",
        "    front nine at TPC Southwind.\"\n",
        "\n",
        "model = model.to(\"cpu\")\n",
        "\n",
        "print(\"This is a %s news\" %ag_news_label[predict(ex_text_str, text_pipeline)])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dBYagTQgV9cU",
        "outputId": "231ba9fc-21bc-4bd0-b7c8-3b5df53b77c1"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "This is a Sports news\n"
          ]
        }
      ]
    }
  ]
}