{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-cnn-md-1 fw=all -->\n",
    "\n",
    "# Sentiment Analysis: Using Convolutional Neural Networks\n",
    "\n",
    "\n",
    "In that section,\n",
    "we investigated mechanisms\n",
    "for processing\n",
    "two-dimensional image data\n",
    "with two-dimensional CNNs,\n",
    "which were applied to\n",
    "local features such as adjacent pixels.\n",
    "Though originally\n",
    "designed for computer vision,\n",
    "CNNs are also widely used\n",
    "for natural language processing.\n",
    "Simply put,\n",
    "just think of any text sequence\n",
    "as a one-dimensional image.\n",
    "In this way,\n",
    "one-dimensional CNNs\n",
    "can process local features\n",
    "such as $n$-grams in text.\n",
    "\n",
    "In this section,\n",
    "we will use the *textCNN* model\n",
    "to demonstrate\n",
    "how to design a CNN architecture\n",
    "for representing single text [@Kim.2014].\n",
    "Compared with\n",
    "the figure\n",
    "that uses an RNN architecture with GloVe pretraining\n",
    "for sentiment analysis,\n",
    "the only difference in the figure\n",
    "lies in\n",
    "the choice of the architecture.\n",
    "\n",
    "\n",
    "![This section feeds pretrained GloVe to a CNN-based architecture for sentiment analysis.](../img/nlp-map-sa-cnn.svg)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-cnn-sentiment-analysis-using-convolutional-neural-networks"
   },
   "source": [
    "from d2l import jax as d2l\n",
    "import jax\n",
    "from jax import numpy as jnp\n",
    "from flax import nnx\n",
    "import optax\n",
    "import numpy as np\n",
    "\n",
    "batch_size = 256\n",
    "train_iter, test_iter, vocab = d2l.load_data_imdb(batch_size)"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-cnn-md-2 fw=all -->\n",
    "\n",
    "## One-Dimensional Convolutions\n",
    "\n",
    "Before introducing the model,\n",
    "let's see how a one-dimensional convolution works.\n",
    "Bear in mind that it is just a special case\n",
    "of a two-dimensional convolution\n",
    "based on the cross-correlation operation.\n",
    "\n",
    "![One-dimensional cross-correlation operation. The shaded portions are the first output element as well as the input and kernel tensor elements used for the output computation: $0\\times1+1\\times2=2$.](../img/conv1d.svg)\n",
    "\n",
    "As shown in the figure,\n",
    "in the one-dimensional case,\n",
    "the convolution window\n",
    "slides from left to right\n",
    "across the input tensor.\n",
    "During sliding,\n",
    "the input subtensor (e.g., $0$ and $1$ in the figure) contained in the convolution window\n",
    "at a certain position\n",
    "and the kernel tensor (e.g., $1$ and $2$ in the figure) are multiplied elementwise.\n",
    "The sum of these multiplications\n",
    "gives the single scalar value (e.g., $0\\times1+1\\times2=2$ in the figure)\n",
    "at the corresponding position of the output tensor.\n",
    "\n",
    "We implement one-dimensional cross-correlation in the following `corr1d` function.\n",
    "Given an input tensor `X`\n",
    "and a kernel tensor `K`,\n",
    "it returns the output tensor `Y`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-cnn-one-dimensional-convolutions-1"
   },
   "source": [
    "def corr1d(X, K):\n",
    "    w = K.shape[0]\n",
    "    Y = d2l.zeros((X.shape[0] - w + 1))\n",
    "    for i in range(Y.shape[0]):\n",
    "        Y = Y.at[i].set((X[i: i + w] * K).sum())\n",
    "    return Y"
   ],
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-cnn-md-3 fw=all -->\n",
    "\n",
    "We can construct the input tensor `X` and the kernel tensor `K` from the figure to validate the output of the above one-dimensional cross-correlation implementation."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-cnn-one-dimensional-convolutions-2"
   },
   "source": [
    "X, K = d2l.tensor([0, 1, 2, 3, 4, 5, 6]), d2l.tensor([1, 2])\n",
    "corr1d(X, K)"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([ 2.,  5.,  8., 11., 14., 17.], dtype=float32)"
     },
     "metadata": {},
     "execution_count": 3
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-cnn-md-4 fw=all -->\n",
    "\n",
    "For any\n",
    "one-dimensional input with multiple channels,\n",
    "the convolution kernel\n",
    "needs to have the same number of input channels.\n",
    "Then for each channel,\n",
    "perform a cross-correlation operation on the one-dimensional tensor of the input and the one-dimensional tensor of the convolution kernel,\n",
    "summing the results over all the channels\n",
    "to produce the one-dimensional output tensor.\n",
    "the figure shows a one-dimensional cross-correlation operation with 3 input channels.\n",
    "\n",
    "![One-dimensional cross-correlation operation with 3 input channels. The shaded portions are the first output element as well as the input and kernel tensor elements used for the output computation: $0\\times1+1\\times2+1\\times3+2\\times4+2\\times(-1)+3\\times(-3)=2$.](../img/conv1d-channel.svg)\n",
    "\n",
    "\n",
    "We can implement the one-dimensional cross-correlation operation for multiple input channels\n",
    "and validate the results in the figure."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-cnn-one-dimensional-convolutions-3"
   },
   "source": [
    "def corr1d_multi_in(X, K):\n",
    "    # First, iterate through the 0th dimension (channel dimension) of `X` and\n",
    "    # `K`. Then, add them together\n",
    "    return sum(corr1d(x, k) for x, k in zip(X, K))\n",
    "\n",
    "X = d2l.tensor([[0, 1, 2, 3, 4, 5, 6],\n",
    "              [1, 2, 3, 4, 5, 6, 7],\n",
    "              [2, 3, 4, 5, 6, 7, 8]])\n",
    "K = d2l.tensor([[1, 2], [3, 4], [-1, -3]])\n",
    "corr1d_multi_in(X, K)"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([ 2.,  8., 14., 20., 26., 32.], dtype=float32)"
     },
     "metadata": {},
     "execution_count": 4
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-cnn-md-5 fw=all -->\n",
    "\n",
    "Note that\n",
    "multi-input-channel one-dimensional cross-correlations\n",
    "are equivalent\n",
    "to\n",
    "single-input-channel\n",
    "two-dimensional cross-correlations.\n",
    "To illustrate,\n",
    "an equivalent form of\n",
    "the multi-input-channel one-dimensional cross-correlation\n",
    "in the figure\n",
    "is\n",
    "the\n",
    "single-input-channel\n",
    "two-dimensional cross-correlation\n",
    "in the figure,\n",
    "where the height of the convolution kernel\n",
    "has to be the same as that of the input tensor.\n",
    "\n",
    "\n",
    "![Two-dimensional cross-correlation operation with a single input channel. The shaded portions are the first output element as well as the input and kernel tensor elements used for the output computation: $2\\times(-1)+3\\times(-3)+1\\times3+2\\times4+0\\times1+1\\times2=2$.](../img/conv1d-2d.svg)\n",
    "\n",
    "Both the outputs in the figure and the figure have only one channel.\n",
    "Same as two-dimensional convolutions with multiple output channels described in that section,\n",
    "we can also specify multiple output channels\n",
    "for one-dimensional convolutions.\n",
    "\n",
    "## Max-Over-Time Pooling\n",
    "\n",
    "Similarly, we can use pooling\n",
    "to extract the highest value\n",
    "from sequence representations\n",
    "as the most important feature\n",
    "across time steps.\n",
    "The *max-over-time pooling* used in textCNN\n",
    "works like\n",
    "the one-dimensional global max-pooling\n",
    "[@Collobert.Weston.Bottou.ea.2011].\n",
    "For a multi-channel input\n",
    "where each channel stores values\n",
    "at different time steps,\n",
    "the output at each channel\n",
    "is the maximum value\n",
    "for that channel.\n",
    "Note that\n",
    "the max-over-time pooling\n",
    "allows different numbers of time steps\n",
    "at different channels.\n",
    "\n",
    "## The textCNN Model\n",
    "\n",
    "Using the one-dimensional convolution\n",
    "and max-over-time pooling,\n",
    "the textCNN model\n",
    "takes individual pretrained token representations\n",
    "as input,\n",
    "then obtains and transforms sequence representations\n",
    "for the downstream application.\n",
    "\n",
    "For a single text sequence\n",
    "with $n$ tokens represented by\n",
    "$d$-dimensional vectors,\n",
    "the width, height, and number of channels\n",
    "of the input tensor\n",
    "are $n$, $1$, and $d$, respectively.\n",
    "The textCNN model transforms the input\n",
    "into the output as follows:\n",
    "\n",
    "1. Define multiple one-dimensional convolution kernels and perform convolution operations separately on the inputs. Convolution kernels with different widths may capture local features among different numbers of adjacent tokens.\n",
    "1. Perform max-over-time pooling on all the output channels, and then concatenate all the scalar pooling outputs as a vector.\n",
    "1. Transform the concatenated vector into the output categories using the fully connected layer. Dropout can be used for reducing overfitting.\n",
    "\n",
    "![The model architecture of textCNN.](../img/textcnn.svg)\n",
    "\n",
    "the figure\n",
    "illustrates the model architecture of textCNN\n",
    "with a concrete example.\n",
    "The input is a sentence with 11 tokens,\n",
    "where\n",
    "each token is represented by a 6-dimensional vector.\n",
    "So we have a 6-channel input with width 11.\n",
    "Define\n",
    "two one-dimensional convolution kernels\n",
    "of widths 2 and 4,\n",
    "with 4 and 5 output channels, respectively.\n",
    "They produce\n",
    "4 output channels with width $11-2+1=10$\n",
    "and 5 output channels with width $11-4+1=8$.\n",
    "Despite different widths of these 9 channels,\n",
    "the max-over-time pooling\n",
    "gives a concatenated 9-dimensional vector,\n",
    "which is finally transformed\n",
    "into a 2-dimensional output vector\n",
    "for binary sentiment predictions.\n",
    "\n",
    "\n",
    "### Defining the Model\n",
    "\n",
    "We implement the textCNN model in the following class.\n",
    "Compared with the bidirectional RNN model in\n",
    "that section,\n",
    "besides\n",
    "replacing recurrent layers with convolutional layers,\n",
    "we also use two embedding layers:\n",
    "one with trainable weights and the other\n",
    "with fixed weights."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-cnn-defining-the-model-1"
   },
   "source": [
    "class TextCNN(nnx.Module):\n",
    "    def __init__(self, vocab_size, embed_size, kernel_sizes, num_channels,\n",
    "                 rngs=None):\n",
    "        rngs = nnx.Rngs(params=0, dropout=1) if rngs is None else rngs\n",
    "        self.embedding = nnx.Embed(vocab_size, embed_size, rngs=rngs)\n",
    "        # The embedding layer not to be trained\n",
    "        self.constant_embedding = nnx.Embed(vocab_size, embed_size, rngs=rngs)\n",
    "        self.dropout = nnx.Dropout(0.5, rngs=rngs)\n",
    "        self.decoder = nnx.Linear(sum(num_channels), 2, rngs=rngs)\n",
    "        # Create multiple one-dimensional convolutional layers\n",
    "        self.convs = nnx.List([\n",
    "            nnx.Conv(2 * embed_size, c, kernel_size=(k,), rngs=rngs)\n",
    "            for c, k in zip(num_channels, kernel_sizes)])\n",
    "\n",
    "    def __call__(self, inputs):\n",
    "        # Concatenate two embedding layer outputs with shape (batch size, no.\n",
    "        # of tokens, token vector dimension) along vectors\n",
    "        embeddings = jnp.concatenate((\n",
    "            self.embedding(inputs), self.constant_embedding(inputs)), axis=2)\n",
    "        # For Flax Conv, input shape is (batch, length, channels) which is\n",
    "        # already the shape of embeddings\n",
    "        # For each one-dimensional convolutional layer, after max-over-time\n",
    "        # pooling, a tensor of shape (batch size, no. of channels) is obtained.\n",
    "        # Concatenate along channels\n",
    "        encoding = jnp.concatenate([\n",
    "            jnp.max(nnx.relu(conv(embeddings)), axis=1)\n",
    "            for conv in self.convs], axis=1)\n",
    "        outputs = self.decoder(self.dropout(encoding))\n",
    "        return outputs"
   ],
   "execution_count": 5,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-cnn-md-6 fw=all -->\n",
    "\n",
    "Let's create a textCNN instance.\n",
    "It has 3 convolutional layers with kernel widths of 3, 4, and 5, all with 100 output channels."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-cnn-defining-the-model-2"
   },
   "source": [
    "embed_size, kernel_sizes, nums_channels = 100, [3, 4, 5], [100, 100, 100]\n",
    "devices = d2l.try_all_gpus()\n",
    "net = TextCNN(len(vocab), embed_size, kernel_sizes, nums_channels)\n",
    "d2l.check_shape(nnx.view(net, deterministic=True)(\n",
    "    jnp.ones((1, 500), dtype=jnp.int32)), (1, 2))"
   ],
   "execution_count": 6,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-cnn-md-7 fw=all -->\n",
    "\n",
    "### Loading Pretrained Word Vectors\n",
    "\n",
    "Same as that section,\n",
    "we load pretrained 100-dimensional GloVe embeddings\n",
    "as the initialized token representations.\n",
    "These token representations (embedding weights)\n",
    "will be trained in `embedding`\n",
    "and fixed in `constant_embedding`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-cnn-loading-pretrained-word-vectors"
   },
   "source": [
    "glove_embedding = d2l.TokenEmbedding('glove.6b.100d')\n",
    "embeds = glove_embedding[vocab.idx_to_token]\n",
    "# Set pretrained embedding weights in the parameters\n",
    "net.embedding.embedding[...] = jnp.array(embeds)\n",
    "net.constant_embedding.embedding = nnx.data(jnp.array(embeds))"
   ],
   "execution_count": 7,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-cnn-md-8 fw=all -->\n",
    "\n",
    "### Training and Evaluating the Model\n",
    "\n",
    "Now we can train the textCNN model for sentiment analysis."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-cnn-training-and-evaluating-the-model-1"
   },
   "source": [
    "lr, num_epochs = 0.001, 5\n",
    "optimizer = nnx.Optimizer(net, optax.adam(lr), wrt=nnx.Param)\n",
    "loss_fn = optax.softmax_cross_entropy_with_integer_labels\n",
    "\n",
    "@nnx.jit\n",
    "def train_step(net, optimizer, X, y):\n",
    "    def compute_loss(model):\n",
    "        logits = model(X)\n",
    "        return loss_fn(logits, y).mean(), logits\n",
    "    (loss, logits), grads = nnx.value_and_grad(\n",
    "        compute_loss, has_aux=True)(net)\n",
    "    optimizer.update(net, grads)\n",
    "    return loss, logits\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    loss_sum, train_correct, num_train = (\n",
    "        jnp.array(0.0), jnp.array(0), 0)\n",
    "    for X, y in train_iter:\n",
    "        l, logits = train_step(net, optimizer, X, y)\n",
    "        loss_sum += l * len(y)\n",
    "        train_correct += (logits.argmax(axis=-1) == y).sum()\n",
    "        num_train += len(y)\n",
    "    # Evaluate\n",
    "    correct, total = jnp.array(0), 0\n",
    "    for X, y in test_iter:\n",
    "        logits = nnx.view(net, deterministic=True)(X)\n",
    "        correct += (logits.argmax(axis=-1) == y).sum()\n",
    "        total += len(y)\n",
    "    loss_sum, train_correct, correct = (\n",
    "        float(loss_sum), int(train_correct), int(correct))\n",
    "    print(f'epoch {epoch + 1}, loss {loss_sum / num_train:.3f}, '\n",
    "          f'train acc {train_correct / num_train:.3f}, '\n",
    "          f'test acc {correct / total:.3f}')"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "epoch 1, loss 0.644, train acc 0.664, test acc 0.812\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "epoch 2, loss 0.413, train acc 0.811, test acc 0.842\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "epoch 3, loss 0.339, train acc 0.851, test acc 0.856\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "epoch 4, loss 0.280, train acc 0.886, test acc 0.866\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "epoch 5, loss 0.219, train acc 0.916, test acc 0.875\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-cnn-md-9 fw=all -->\n",
    "\n",
    "Below we use the trained model to predict the sentiment for two simple sentences."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-cnn-training-and-evaluating-the-model-2"
   },
   "source": [
    "tokens = jnp.array(vocab['this movie is so great'.split()])\n",
    "logits = nnx.view(net, deterministic=True)(tokens.reshape(1, -1))\n",
    "'positive' if int(jnp.argmax(logits, axis=1)[0]) == 1 else 'negative'"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'positive'"
     },
     "metadata": {},
     "execution_count": 9
    }
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-cnn-training-and-evaluating-the-model-3"
   },
   "source": [
    "tokens = jnp.array(vocab['this movie is so bad'.split()])\n",
    "logits = nnx.view(net, deterministic=True)(tokens.reshape(1, -1))\n",
    "'positive' if int(jnp.argmax(logits, axis=1)[0]) == 1 else 'negative'"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'negative'"
     },
     "metadata": {},
     "execution_count": 10
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-cnn-md-10 fw=all -->\n",
    "\n",
    "## Summary\n",
    "\n",
    "* One-dimensional CNNs can process local features such as $n$-grams in text.\n",
    "* Multi-input-channel one-dimensional cross-correlations are equivalent to single-input-channel two-dimensional cross-correlations.\n",
    "* The max-over-time pooling allows different numbers of time steps at different channels.\n",
    "* The textCNN model transforms individual token representations into downstream application outputs using one-dimensional convolutional layers and max-over-time pooling layers.\n",
    "\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Tune hyperparameters and compare the two architectures for sentiment analysis in that section and in this section, such as in classification accuracy and computational efficiency.\n",
    "1. Can you further improve the classification accuracy of the model by using the methods introduced in the exercises of that section?\n",
    "1. Add positional encoding in the input representations. Does it improve the classification accuracy?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-cnn-md-11 fw=jax -->\n",
    "\n",
    "[Discussions](https://d2l.discourse.group/t/1425)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
