Next-Token-Prediction-checkpoint.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0f901b18",
   "metadata": {},
   "source": [
    "# Next-Token-Prediction\n",
    "This is based on the following blog posts: \n",
    "* Predicting Next Word — NLP & Deep Learning: https://medium.com/@vijay2340025/predicting-next-word-nlp-deep-learning-85010d966671 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b2471e25",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import nltk\n",
    "import pandas as pd\n",
    "import torch\n",
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c0febc54",
   "metadata": {},
   "outputs": [],
   "source": [
    "nltk.download('punkt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8a781aa1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dataset = \"\"\"\n",
    "    Is Antwerp a city?,\n",
    "    Is Antwerp a municipality?,\n",
    "    Is Antwerp in Belgium?,\n",
    "    What is Antwerp?,\n",
    "    What is the population of the city of Antwerp?,\n",
    "    Where is the city of Antwerp?,\n",
    "    Why is Antwerp important to fashion?,\n",
    "    Antwerp is to the east of what river?,\n",
    "    How many municipalities does Antwerp have?,\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "39255d82",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_all_possible_sequences(text):\n",
    "    seq = []\n",
    "    words = nltk.word_tokenize(text)\n",
    "    total_words = len(words)\n",
    "    for i in range(1, total_words):\n",
    "        for j in range(1, len(words)-i+1):\n",
    "            arr = words[j-1:j+i]\n",
    "            seq.append((arr[:-1], arr[-1]))\n",
    "    return seq\n",
    "def build_vocabulary(docs):\n",
    "    vocabulary = []\n",
    "    for doc in docs:\n",
    "        for w in nltk.word_tokenize(doc):\n",
    "            if w not in vocabulary:\n",
    "                vocabulary.append(w)\n",
    "    vocabulary.append('UNK')\n",
    "    return vocabulary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "edb54d0d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "docs = []\n",
    "for row in dataset.split(\",\"):\n",
    "    docs.append(row.lower())\n",
    "\n",
    "lst = []\n",
    "for doc in docs:\n",
    "    tmp_lst = get_all_possible_sequences(doc)\n",
    "    lst = lst + tmp_lst\n",
    "\n",
    "vocabulary = build_vocabulary(docs)\n",
    "id2word = {idx: w for (idx, w) in enumerate(vocabulary)}\n",
    "word2id = {w: idx for (idx, w) in enumerate(vocabulary)}\n",
    "def seq2id(arr):\n",
    "    return torch.tensor([word2id[i] for i in arr])\n",
    "def get_max_seq():\n",
    "    return len(list(set([len(i[0]) for i in lst])))\n",
    "MAX_SEQ_LEN = get_max_seq()\n",
    "def get_padded_x(data):\n",
    "    new_data = F.pad(input=data.view(1,-1), pad=(0, MAX_SEQ_LEN-data.shape[0], 0, 0), mode='constant', value=word2id['UNK'])\n",
    "    return new_data\n",
    "def get_xy_vector(arr):\n",
    "    x = seq2id(arr[0])\n",
    "    y = seq2id([arr[1]])\n",
    "    return x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "431fd558",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class NextWordModel(nn.Module):\n",
    "    \"\"\" Prediction of Next word based on the MAX_SEQ_LEN Sequence \"\"\"\n",
    "    def __init__(self, embedding_dim, hidden_dim, vocab_size):\n",
    "        super(NextWordModel, self).__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.word_embeddings = nn.Embedding(vocab_size, embedding_dim)\n",
    "        self.gru = nn.GRU(embedding_dim * MAX_SEQ_LEN, hidden_dim)\n",
    "        self.linear = nn.Linear(hidden_dim, vocab_size)\n",
    "\n",
    "    def forward(self, sentence):\n",
    "        embeds = self.word_embeddings(sentence)\n",
    "        lstm_out, _ = self.gru(embeds.view(1, 1, -1))\n",
    "        x = self.linear(lstm_out.view(1, -1))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d891422e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "  dev = \"cuda:0\"\n",
    "else:\n",
    "  dev = \"cpu\"\n",
    "print(f'Running on {dev}')\n",
    "# set the model to be copied on GPU\n",
    "device = torch.device(dev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7c231c92",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "EMBEDDING_DIM = 10\n",
    "NO_OF_EPOCHS = 300\n",
    "HIDDEN_DIM = len(vocabulary)\n",
    "model = NextWordModel(EMBEDDING_DIM, HIDDEN_DIM, len(vocabulary))\n",
    "loss_function = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.1)\n",
    "model.to(device)\n",
    "for epoch in range(NO_OF_EPOCHS):\n",
    "    running_loss = 0.0\n",
    "    i = 0\n",
    "    for data in lst:\n",
    "        model.zero_grad()\n",
    "        x, y = get_xy_vector(data)\n",
    "# convert to max seq length with padding\n",
    "        x = get_padded_x(x)\n",
    "\n",
    "        x = x.to(device)\n",
    "        y = y.to(device)\n",
    "\n",
    "        predicted = model(x)\n",
    "\n",
    "        loss = loss_function(predicted, y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        running_loss += loss\n",
    "        i += 1\n",
    "        if i % 100 == 0:\n",
    "            #print(f'Loss at iteration {i} and epoch {epoch} is {running_loss / 100}')\n",
    "            running_loss = 0\n",
    "print('Finished')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d0ff01d",
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    print('Type something here . . .')\n",
    "    while True:\n",
    "        inp = input(\"\")\n",
    "        inp = inp.strip()\n",
    "        if inp == \"q\":\n",
    "            break\n",
    "\n",
    "        tokens = nltk.word_tokenize(inp.lower())\n",
    "        x = seq2id(tokens)\n",
    "        x = get_padded_x(x)\n",
    "\n",
    "        x = x.to(device)\n",
    "        predicted = model(x).to(device)\n",
    "\n",
    "        predicted = predicted[0].cpu().numpy()\n",
    "\n",
    "        print(f'Answer: {inp} {id2word[np.argmax(predicted)]} ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1e828f0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}