Simple_CBOW_in_TF_2.0.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The **Continuous Bag-of-Words model (CBOW)** is frequently used in NLP deep\n",
    "learning. It is a model that tries to predict words given the context of\n",
    "a few words before and a few words after the target word. This is\n",
    "distinct from language modeling, since CBOW is not sequential and does\n",
    "not have to be probabilistic. Typcially, CBOW is used to quickly train\n",
    "word embeddings, and these embeddings are used to initialize the\n",
    "embeddings of some more complicated model. Usually, this is referred to\n",
    "as *pretraining embeddings*. It almost always helps performance a couple\n",
    "of percent.\n",
    "\n",
    "The CBOW model is as follows. Given a target word $w_i$ and an\n",
    "$N$ context window on each side, $w_{i-1}, \\dots, w_{i-N}$\n",
    "and $w_{i+1}, \\dots, w_{i+N}$, referring to all context words\n",
    "collectively as $C$, CBOW tries to minimize\n",
    "\n",
    "\\begin{align}-\\log p(w_i | C) = -\\log \\text{Softmax}(A(\\sum_{w \\in C} q_w) + b)\\end{align}\n",
    "\n",
    "where $q_w$ is the embedding for word $w$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, I am going to use `tensorflow 2.0.0-beta0` and take the advantage of its flexible custom model definitions to create a COBW model and will train it to get the embeddings of a toy dataset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.0.0-beta0\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.layers import *\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "print(tf.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2 words to the left, 2 to the right\n",
    "CONTEXT_SIZE = 2  \n",
    "# Embedding dimension on left\n",
    "EMBED_LEFT = 10\n",
    "# Embedding dimension on right\n",
    "EMBED_RIGHT = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(['When', 'forty', 'shall', 'besiege'], 'winters'), (['forty', 'winters', 'besiege', 'thy'], 'shall'), (['winters', 'shall', 'thy', 'brow,'], 'besiege'), (['shall', 'besiege', 'brow,', 'And'], 'thy'), (['besiege', 'thy', 'And', 'dig'], 'brow,')]\n"
     ]
    }
   ],
   "source": [
    "# Row test from Shakespeare Sonnet 2\n",
    "raw_text = \"\"\"When forty winters shall besiege thy brow,\n",
    "And dig deep trenches in thy beauty's field,\n",
    "Thy youth's proud livery so gazed on now,\n",
    "Will be a totter'd weed of small worth held:\n",
    "Then being asked, where all thy beauty lies,\n",
    "Where all the treasure of thy lusty days;\n",
    "To say, within thine own deep sunken eyes,\n",
    "Were an all-eating shame, and thriftless praise.\n",
    "How much more praise deserv'd thy beauty's use,\n",
    "If thou couldst answer 'This fair child of mine\n",
    "Shall sum my count, and make my old excuse,'\n",
    "Proving his beauty by succession thine!\n",
    "This were to be new made when thou art old,\n",
    "And see thy blood warm when thou feel'st it cold.\"\"\".split()\n",
    "\n",
    "# By deriving a set from `raw_text`, we deduplicate the array\n",
    "vocab = set(raw_text)\n",
    "vocab_size = len(vocab)\n",
    "\n",
    "word_to_ix = {word: i for i, word in enumerate(vocab)}\n",
    "data = []\n",
    "for i in range(2, len(raw_text) - 2):\n",
    "    context = [raw_text[i - 2], raw_text[i - 1],\n",
    "               raw_text[i + 1], raw_text[i + 2]]\n",
    "    target = raw_text[i]\n",
    "    data.append((context, target))\n",
    "print(data[:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The toy dataset is ready. Let's now specify the model definition. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CBOW(Model):\n",
    "\n",
    "    def __init__(self, vocab_size, embedding_dim_left, embedding_dim_right, context_size):\n",
    "        super(CBOW, self).__init__()\n",
    "        self.embeddings_left = Embedding(vocab_size, embedding_dim_left)\n",
    "        self.embeddings_right = Embedding(vocab_size, embedding_dim_right)\n",
    "        self.linear1 = Dense(128, activation='linear', input_dim=(context_size * embedding_dim_left))\n",
    "        self.linear2 = Dense(vocab_size, activation='linear', input_dim=128)\n",
    "\n",
    "    def call(self, inputs):\n",
    "        left_embeds = tf.reshape(self.embeddings_left(inputs[:2]), [1,20])\n",
    "        right_embeds = tf.reshape(self.embeddings_left(inputs[:2]), [1,20])\n",
    "        left_out = tf.keras.activations.relu(self.linear1(left_embeds))\n",
    "        right_out = tf.keras.activations.relu(self.linear1(right_embeds))\n",
    "        out = self.linear2(left_out + right_out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's instantiate the model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 175,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = CBOW(vocab_size, EMBED_LEFT, EMBED_RIGHT, CONTEXT_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define our loss function and optimizer. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 176,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_func = tf.keras.losses.SparseCategoricalCrossentropy()\n",
    "optimizer = tf.keras.optimizers.SGD(learning_rate=0.001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function which will train our model. We use `GradientTape` to make use of automatic differentiation. Makes the process of Backpropagation much more flexible. Makes it also easy to chart dependencies inside the training loop. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 177,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def model_train(features, labels):\n",
    "    # Define the GradientTape context\n",
    "    with tf.GradientTape() as tape:\n",
    "        # Get the probabilities\n",
    "        predictions = model(features)\n",
    "        # Calculate the loss\n",
    "        loss = loss_func(labels, predictions)\n",
    "    # Get the gradients\n",
    "    gradients = tape.gradient(loss, model.trainable_variables)\n",
    "    # Update the weights\n",
    "    optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we train the model for 10 epochs. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 178,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:01<00:00,  8.50it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[<tf.Tensor: id=66937, shape=(), dtype=float32, numpy=1550.9991>, <tf.Tensor: id=67382, shape=(), dtype=float32, numpy=1634.5785>, <tf.Tensor: id=67827, shape=(), dtype=float32, numpy=1486.1724>, <tf.Tensor: id=68272, shape=(), dtype=float32, numpy=1476.1912>, <tf.Tensor: id=68717, shape=(), dtype=float32, numpy=1470.2799>, <tf.Tensor: id=69162, shape=(), dtype=float32, numpy=1466.611>, <tf.Tensor: id=69607, shape=(), dtype=float32, numpy=1464.6365>, <tf.Tensor: id=70052, shape=(), dtype=float32, numpy=1462.1602>, <tf.Tensor: id=70497, shape=(), dtype=float32, numpy=1458.745>, <tf.Tensor: id=70942, shape=(), dtype=float32, numpy=1455.2542>]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "running_loss = []\n",
    "for epoch in tqdm(range(10)):\n",
    "    total_loss = 0\n",
    "    for context, target in data:\n",
    "        context_idxs = np.array([word_to_ix[w] for w in context])\n",
    "        target_idxs = np.array([word_to_ix[target]])\n",
    "        loss = model_train(context_idxs, target_idxs)\n",
    "        total_loss += loss\n",
    "    running_loss.append(total_loss)\n",
    "print(running_loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The loss is pretty high but it does decrease anyway. The purpose of this notebook is to show how we can combine the low level TF ops with the high-level APIs to aid utmost flexibility in model definition. The example is inspired from [this tutorial](https://pytorch.org/tutorials/beginner/nlp/word_embeddings_tutorial.html). "
   ]
  }
 ],
 "metadata": {
  "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.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}