{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-rnn-md-1 fw=all -->\n",
    "\n",
    "# Sentiment Analysis: Using Recurrent Neural Networks\n",
    "\n",
    "\n",
    "Like word similarity and analogy tasks,\n",
    "we can also apply pretrained word vectors\n",
    "to sentiment analysis.\n",
    "Since the IMDb review dataset\n",
    "in that section\n",
    "is not very big,\n",
    "using text representations\n",
    "that were pretrained\n",
    "on large-scale corpora\n",
    "may reduce overfitting of the model.\n",
    "As a specific example\n",
    "illustrated in the figure,\n",
    "we will represent each token\n",
    "using the pretrained GloVe model,\n",
    "and feed these token representations\n",
    "into a multilayer bidirectional RNN\n",
    "to obtain the text sequence representation,\n",
    "which will\n",
    "be transformed into \n",
    "sentiment analysis outputs [@Maas.Daly.Pham.ea.2011].\n",
    "For the same downstream application,\n",
    "we will consider a different architectural\n",
    "choice later.\n",
    "\n",
    "![This section feeds pretrained GloVe to an RNN-based architecture for sentiment analysis.](../img/nlp-map-sa-rnn.svg)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-rnn-sentiment-analysis-using-recurrent-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 = 128\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-rnn-md-2 fw=all -->\n",
    "\n",
    "## Representing Single Text with RNNs\n",
    "\n",
    "In text classification tasks,\n",
    "such as sentiment analysis,\n",
    "a varying-length text sequence \n",
    "will be transformed into fixed-length categories.\n",
    "In the following `BiRNN` class,\n",
    "while each token of a text sequence\n",
    "gets its individual\n",
    "pretrained GloVe\n",
    "representation via the embedding layer\n",
    "(`self.embedding`),\n",
    "the entire sequence\n",
    "is encoded by a bidirectional RNN (`self.encoder`).\n",
    "More concretely,\n",
    "the hidden states (at the last layer)\n",
    "of the bidirectional LSTM\n",
    "at both the initial and final time steps\n",
    "are concatenated \n",
    "as the representation of the text sequence.\n",
    "This single text representation\n",
    "is then transformed into output categories\n",
    "by a fully connected layer (`self.decoder`)\n",
    "with two outputs (\"positive\" and \"negative\")."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-rnn-representing-single-text-with-rnns-1"
   },
   "source": [
    "class BiRNN(nnx.Module):\n",
    "    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,\n",
    "                 rngs=None):\n",
    "        rngs = nnx.Rngs(params=0, carry=1) if rngs is None else rngs\n",
    "        self.embedding = nnx.Embed(vocab_size, embed_size, rngs=rngs)\n",
    "        self.forward_rnns = nnx.List([])\n",
    "        self.backward_rnns = nnx.List([])\n",
    "        for i in range(num_layers):\n",
    "            num_inputs = embed_size if i == 0 else 2 * num_hiddens\n",
    "            self.forward_rnns.append(nnx.RNN(\n",
    "                nnx.LSTMCell(num_inputs, num_hiddens, rngs=rngs), rngs=rngs))\n",
    "            self.backward_rnns.append(nnx.RNN(\n",
    "                nnx.LSTMCell(num_inputs, num_hiddens, rngs=rngs),\n",
    "                reverse=True, keep_order=True, rngs=rngs))\n",
    "        self.decoder = nnx.Linear(4 * num_hiddens, 2, rngs=rngs)\n",
    "\n",
    "    def __call__(self, inputs):\n",
    "        # The shape of `inputs` is (batch size, no. of time steps)\n",
    "        embeddings = self.embedding(inputs)\n",
    "        outputs = embeddings\n",
    "        for forward_rnn, backward_rnn in zip(\n",
    "                self.forward_rnns, self.backward_rnns):\n",
    "            outputs = jnp.concatenate(\n",
    "                [forward_rnn(outputs), backward_rnn(outputs)], axis=-1)\n",
    "        # Each endpoint contains both directions, so concatenating the first\n",
    "        # and last time steps produces 4 * num_hiddens features.\n",
    "        encoding = jnp.concatenate([outputs[:, 0, :], outputs[:, -1, :]],\n",
    "                                   axis=1)\n",
    "        outs = self.decoder(encoding)\n",
    "        return outs"
   ],
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-rnn-md-3 fw=all -->\n",
    "\n",
    "Let's construct a bidirectional RNN with two hidden layers to represent single text for sentiment analysis."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-rnn-representing-single-text-with-rnns-2"
   },
   "source": [
    "embed_size, num_hiddens, num_layers, devices = 100, 100, 2, d2l.try_all_gpus()\n",
    "net = BiRNN(len(vocab), embed_size, num_hiddens, num_layers)"
   ],
   "execution_count": 3,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-rnn-representing-single-text-with-rnns-3"
   },
   "source": [
    "# NNX modules create their parameters in the constructor.\n",
    "d2l.check_shape(net(jnp.ones((1, 500), dtype=jnp.int32)), (1, 2))"
   ],
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-rnn-md-4 fw=all -->\n",
    "\n",
    "## Loading Pretrained Word Vectors\n",
    "\n",
    "Below we load the pretrained 100-dimensional (needs to be consistent with `embed_size`) GloVe embeddings for tokens in the vocabulary."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-rnn-loading-pretrained-word-vectors-1"
   },
   "source": [
    "glove_embedding = d2l.TokenEmbedding('glove.6b.100d')"
   ],
   "execution_count": 5,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-rnn-md-5 fw=all -->\n",
    "\n",
    "Print the shape of the vectors\n",
    "for all the tokens in the vocabulary."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-rnn-loading-pretrained-word-vectors-2"
   },
   "source": [
    "embeds = glove_embedding[vocab.idx_to_token]\n",
    "embeds.shape"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(49346, 100)"
     },
     "metadata": {},
     "execution_count": 6
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-rnn-md-6 fw=all -->\n",
    "\n",
    "We use these pretrained\n",
    "word vectors\n",
    "to represent tokens in the reviews\n",
    "and will not update\n",
    "these vectors during training."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-rnn-loading-pretrained-word-vectors-3"
   },
   "source": [
    "# Store the pretrained table as non-trainable NNX data.\n",
    "net.embedding.embedding = nnx.data(jnp.array(embeds))"
   ],
   "execution_count": 7,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-rnn-md-7 fw=all -->\n",
    "\n",
    "## Training and Evaluating the Model\n",
    "\n",
    "Now we can train the bidirectional RNN for sentiment analysis."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-rnn-training-and-evaluating-the-model-1"
   },
   "source": [
    "lr, num_epochs = 0.01, 4\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",
    "@nnx.jit\n",
    "def eval_step(net, X):\n",
    "    return net(X)\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    loss_terms, train_correct_terms, num_train = [], [], 0\n",
    "    for X, y in train_iter:\n",
    "        l, logits = train_step(net, optimizer, X, y)\n",
    "        loss_terms.append(l * len(y))\n",
    "        train_correct_terms.append((logits.argmax(axis=-1) == y).sum())\n",
    "        num_train += len(y)\n",
    "    # Evaluate\n",
    "    correct_terms, total = [], 0\n",
    "    for X, y in test_iter:\n",
    "        logits = eval_step(net, X)\n",
    "        correct_terms.append((logits.argmax(axis=-1) == y).sum())\n",
    "        total += len(y)\n",
    "    loss_sum = float(jnp.stack(loss_terms).sum())\n",
    "    train_correct = int(jnp.stack(train_correct_terms).sum())\n",
    "    correct = int(jnp.stack(correct_terms).sum())\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.627, train acc 0.632, test acc 0.796\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "epoch 2, loss 0.418, train acc 0.814, test acc 0.825\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "epoch 3, loss 0.364, train acc 0.844, test acc 0.847\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "epoch 4, loss 0.333, train acc 0.857, test acc 0.846\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-rnn-md-8 fw=all -->\n",
    "\n",
    "We define the following function to predict the sentiment of a text sequence using the trained model `net`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-rnn-training-and-evaluating-the-model-2"
   },
   "source": [
    "\n",
    "def predict_sentiment(net, vocab, sequence):\n",
    "    \"\"\"Predict the sentiment of a text sequence.\"\"\"\n",
    "    sequence = jnp.array(vocab[sequence.split()])\n",
    "    label = jnp.argmax(net(sequence.reshape(1, -1)), axis=1)\n",
    "    return 'positive' if label == 1 else 'negative'"
   ],
   "execution_count": 9,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-rnn-md-9 fw=all -->\n",
    "\n",
    "Finally, let's use the trained model to predict the sentiment for two simple sentences."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-rnn-training-and-evaluating-the-model-3"
   },
   "source": [
    "predict_sentiment(net, vocab, 'this movie is so great')"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'positive'"
     },
     "metadata": {},
     "execution_count": 10
    }
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sentiment-analysis-rnn-training-and-evaluating-the-model-4"
   },
   "source": [
    "predict_sentiment(net, vocab, 'this movie is so bad')"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'negative'"
     },
     "metadata": {},
     "execution_count": 11
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-rnn-md-10 fw=all -->\n",
    "\n",
    "## Summary\n",
    "\n",
    "* Pretrained word vectors can represent individual tokens in a text sequence.\n",
    "* Bidirectional RNNs can represent a text sequence, such as via the concatenation of its hidden states at the initial and final time steps. This single text representation can be transformed into categories using a fully connected layer.\n",
    "\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Increase the number of epochs. Can you improve the training and testing accuracies? How about tuning other hyperparameters?\n",
    "1. Use larger pretrained word vectors, such as 300-dimensional GloVe embeddings. Does it improve classification accuracy?\n",
    "1. Can we improve the classification accuracy by using the spaCy tokenization? You need to install spaCy (`pip install spacy`) and install the English package (`python -m spacy download en_core_web_sm`). In the code, first, import spaCy (`import spacy`). Then, load the spaCy English package (`spacy_en = spacy.load('en_core_web_sm')`). Finally, define the function `def tokenizer(text): return [tok.text for tok in spacy_en.tokenizer(text)]` and replace the original `tokenizer` function. Note the different forms of phrase tokens in GloVe and spaCy. For example, the phrase token \"new york\" takes the form of \"new-york\" in GloVe and the form of \"new york\" after the spaCy tokenization."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=sentiment-analysis-rnn-md-11 fw=jax -->\n",
    "\n",
    "[Discussions](https://d2l.discourse.group/t/1424)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
