{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-attention-md-1 fw=all -->\n",
    "\n",
    "# Natural Language Inference: Using Attention\n",
    "\n",
    "We introduced the natural language inference task and the SNLI dataset in that section. In view of many models that are based on complex and deep architectures, @Parikh.Tackstrom.Das.ea.2016 proposed to address natural language inference with attention mechanisms and called it a \"decomposable attention model\".\n",
    "This results in a model without recurrent or convolutional layers, achieving the best result at the time on the SNLI dataset with much fewer parameters.\n",
    "In this section, we will describe and implement this attention-based method (with MLPs) for natural language inference, as depicted in the figure.\n",
    "\n",
    "![This section feeds pretrained GloVe to an architecture based on attention and MLPs for natural language inference.](../img/nlp-map-nli-attention.svg)\n",
    "\n",
    "\n",
    "## The Model\n",
    "\n",
    "Rather than preserving the order of tokens in premises and hypotheses,\n",
    "we can simply align tokens in one text sequence to every token in the other, and vice versa,\n",
    "then compare and aggregate such information to predict the logical relationships\n",
    "between premises and hypotheses.\n",
    "Similar to alignment of tokens between source and target sentences in machine translation,\n",
    "the alignment of tokens between premises and hypotheses\n",
    "can be neatly accomplished by attention mechanisms.\n",
    "\n",
    "![Natural language inference using attention mechanisms.](../img/nli-attention.svg)\n",
    "\n",
    "the figure depicts the natural language inference method using attention mechanisms.\n",
    "At a high level, it consists of three jointly trained steps: attending, comparing, and aggregating.\n",
    "We will illustrate them step by step in the following."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-attention-the-model"
   },
   "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"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-attention-md-2 fw=all -->\n",
    "\n",
    "### Attending\n",
    "\n",
    "The first step is to align tokens in one text sequence to each token in the other sequence.\n",
    "Suppose that the premise is \"i do need sleep\" and the hypothesis is \"i am tired\".\n",
    "Due to semantical similarity,\n",
    "we may wish to align \"i\" in the hypothesis with \"i\" in the premise,\n",
    "and align \"tired\" in the hypothesis with \"sleep\" in the premise.\n",
    "Likewise, we may wish to align \"i\" in the premise with \"i\" in the hypothesis,\n",
    "and align \"need\" and \"sleep\" in the premise with \"tired\" in the hypothesis.\n",
    "Note that such alignment is *soft* using weighted average,\n",
    "where ideally large weights are associated with the tokens to be aligned.\n",
    "For ease of demonstration, the figure shows such alignment in a *hard* way.\n",
    "\n",
    "Now we describe the soft alignment using attention mechanisms in more detail.\n",
    "Denote by $\\mathbf{A} = (\\mathbf{a}_1, \\ldots, \\mathbf{a}_m)$\n",
    "and $\\mathbf{B} = (\\mathbf{b}_1, \\ldots, \\mathbf{b}_n)$ the premise and hypothesis,\n",
    "whose number of tokens are $m$ and $n$, respectively,\n",
    "where $\\mathbf{a}_i, \\mathbf{b}_j \\in \\mathbb{R}^{d}$ ($i = 1, \\ldots, m, j = 1, \\ldots, n$) is a $d$-dimensional word vector.\n",
    "For soft alignment, we compute the attention weights $e_{ij} \\in \\mathbb{R}$ as\n",
    "\n",
    "$$e_{ij} = f(\\mathbf{a}_i)^\\top f(\\mathbf{b}_j),$$\n",
    "\n",
    "where the function $f$ is an MLP defined in the following `mlp` function.\n",
    "The output dimension of $f$ is specified by the `num_hiddens` argument of `mlp`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-attention-attending-1"
   },
   "source": [
    "class MLP(nnx.Module):\n",
    "    def __init__(self, num_inputs, num_hiddens, flatten, rngs=None):\n",
    "        rngs = nnx.Rngs(params=0, dropout=1) if rngs is None else rngs\n",
    "        self.dropout1 = nnx.Dropout(0.2, rngs=rngs)\n",
    "        self.dense1 = nnx.Linear(num_inputs, num_hiddens, rngs=rngs)\n",
    "        self.dropout2 = nnx.Dropout(0.2, rngs=rngs)\n",
    "        self.dense2 = nnx.Linear(num_hiddens, num_hiddens, rngs=rngs)\n",
    "        self.flatten = flatten\n",
    "\n",
    "    def __call__(self, x):\n",
    "        x = nnx.relu(self.dense1(self.dropout1(x)))\n",
    "        if self.flatten:\n",
    "            x = x.reshape((x.shape[0], -1))\n",
    "        x = nnx.relu(self.dense2(self.dropout2(x)))\n",
    "        if self.flatten:\n",
    "            x = x.reshape((x.shape[0], -1))\n",
    "        return x"
   ],
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-attention-md-3 fw=all -->\n",
    "\n",
    "It should be highlighted that, in the equation\n",
    "$f$ takes inputs $\\mathbf{a}_i$ and $\\mathbf{b}_j$ separately rather than taking a pair of them together as input.\n",
    "This *decomposition* trick leads to only $m + n$ applications (linear complexity) of $f$ rather than $mn$ applications\n",
    "(quadratic complexity).\n",
    "\n",
    "\n",
    "Normalizing the attention weights in the equation,\n",
    "we compute the weighted average of all the token vectors in the hypothesis\n",
    "to obtain representation of the hypothesis that is softly aligned with the token indexed by $i$ in the premise:\n",
    "\n",
    "$$\n",
    "\\boldsymbol{\\beta}_i = \\sum_{j=1}^{n}\\frac{\\exp(e_{ij})}{ \\sum_{k=1}^{n} \\exp(e_{ik})} \\mathbf{b}_j.\n",
    "$$\n",
    "\n",
    "Likewise, we compute soft alignment of premise tokens for each token indexed by $j$ in the hypothesis:\n",
    "\n",
    "$$\n",
    "\\boldsymbol{\\alpha}_j = \\sum_{i=1}^{m}\\frac{\\exp(e_{ij})}{ \\sum_{k=1}^{m} \\exp(e_{kj})} \\mathbf{a}_i.\n",
    "$$\n",
    "\n",
    "Below we define the `Attend` class to compute the soft alignment of hypotheses (`beta`) with input premises `A` and soft alignment of premises (`alpha`) with input hypotheses `B`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-attention-attending-2"
   },
   "source": [
    "class Attend(nnx.Module):\n",
    "    def __init__(self, embed_size, num_hiddens, rngs=None):\n",
    "        self.f = MLP(embed_size, num_hiddens, flatten=False, rngs=rngs)\n",
    "\n",
    "    def __call__(self, A, B):\n",
    "        # Shape of `A`/`B`: (`batch_size`, no. of tokens in sequence A/B,\n",
    "        # `embed_size`)\n",
    "        # Shape of `f_A`/`f_B`: (`batch_size`, no. of tokens in sequence A/B,\n",
    "        # `num_hiddens`)\n",
    "        f_A = self.f(A)\n",
    "        f_B = self.f(B)\n",
    "        # Shape of `e`: (`batch_size`, no. of tokens in sequence A,\n",
    "        # no. of tokens in sequence B)\n",
    "        e = jnp.matmul(f_A, f_B.transpose(0, 2, 1))\n",
    "        # Shape of `beta`: (`batch_size`, no. of tokens in sequence A,\n",
    "        # `embed_size`), where sequence B is softly aligned with each token\n",
    "        # (axis 1 of `beta`) in sequence A\n",
    "        beta = jnp.matmul(jax.nn.softmax(e, axis=-1), B)\n",
    "        # Shape of `alpha`: (`batch_size`, no. of tokens in sequence B,\n",
    "        # `embed_size`), where sequence A is softly aligned with each token\n",
    "        # (axis 1 of `alpha`) in sequence B\n",
    "        alpha = jnp.matmul(jax.nn.softmax(e.transpose(0, 2, 1), axis=-1), A)\n",
    "        return beta, alpha"
   ],
   "execution_count": 3,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-attention-md-4 fw=all -->\n",
    "\n",
    "### Comparing\n",
    "\n",
    "In the next step, we compare a token in one sequence with the other sequence that is softly aligned with that token.\n",
    "Note that in soft alignment, all the tokens from one sequence, though with probably different attention weights, will be compared with a token in the other sequence.\n",
    "For ease of demonstration, the figure pairs tokens with aligned tokens in a *hard* way.\n",
    "For example, suppose that the attending step determines that \"need\" and \"sleep\" in the premise are both aligned with \"tired\" in the hypothesis, the pair \"tired--need sleep\" will be compared.\n",
    "\n",
    "In the comparing step, we feed the concatenation (operator $[\\cdot, \\cdot]$) of tokens from one sequence and aligned tokens from the other sequence into a function $g$ (an MLP):\n",
    "\n",
    "$$\\mathbf{v}_{A,i} = g([\\mathbf{a}_i, \\boldsymbol{\\beta}_i]), i = 1, \\ldots, m\\\\ \\mathbf{v}_{B,j} = g([\\mathbf{b}_j, \\boldsymbol{\\alpha}_j]), j = 1, \\ldots, n.$$\n",
    "\n",
    "\n",
    "In the equation, $\\mathbf{v}_{A,i}$ is the comparison between token $i$ in the premise and all the hypothesis tokens that are softly aligned with token $i$;\n",
    "while $\\mathbf{v}_{B,j}$ is the comparison between token $j$ in the hypothesis and all the premise tokens that are softly aligned with token $j$.\n",
    "The following `Compare` class defines such a comparing step."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-attention-comparing"
   },
   "source": [
    "class Compare(nnx.Module):\n",
    "    def __init__(self, embed_size, num_hiddens, rngs=None):\n",
    "        self.g = MLP(2 * embed_size, num_hiddens, flatten=False, rngs=rngs)\n",
    "\n",
    "    def __call__(self, A, B, beta, alpha):\n",
    "        V_A = self.g(jnp.concatenate([A, beta], axis=2))\n",
    "        V_B = self.g(jnp.concatenate([B, alpha], axis=2))\n",
    "        return V_A, V_B"
   ],
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-attention-md-5 fw=all -->\n",
    "\n",
    "### Aggregating\n",
    "\n",
    "With two sets of comparison vectors $\\mathbf{v}_{A,i}$ ($i = 1, \\ldots, m$) and $\\mathbf{v}_{B,j}$ ($j = 1, \\ldots, n$) on hand,\n",
    "in the last step we will aggregate such information to infer the logical relationship.\n",
    "We begin by summing up both sets:\n",
    "\n",
    "$$\n",
    "\\mathbf{v}_A = \\sum_{i=1}^{m} \\mathbf{v}_{A,i}, \\quad \\mathbf{v}_B = \\sum_{j=1}^{n}\\mathbf{v}_{B,j}.\n",
    "$$\n",
    "\n",
    "Next we feed the concatenation of both summarization results into function $h$ (an MLP) to obtain the classification result of the logical relationship:\n",
    "\n",
    "$$\n",
    "\\hat{\\mathbf{y}} = h([\\mathbf{v}_A, \\mathbf{v}_B]).\n",
    "$$\n",
    "\n",
    "The aggregation step is defined in the following `Aggregate` class."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-attention-aggregating"
   },
   "source": [
    "class Aggregate(nnx.Module):\n",
    "    def __init__(self, num_hiddens, num_outputs, rngs=None):\n",
    "        rngs = nnx.Rngs(params=0, dropout=1) if rngs is None else rngs\n",
    "        self.h = MLP(2 * num_hiddens, num_hiddens, flatten=True, rngs=rngs)\n",
    "        self.output = nnx.Linear(num_hiddens, num_outputs, rngs=rngs)\n",
    "\n",
    "    def __call__(self, V_A, V_B):\n",
    "        # Sum up both sets of comparison vectors\n",
    "        V_A = V_A.sum(axis=1)\n",
    "        V_B = V_B.sum(axis=1)\n",
    "        # Feed the concatenation of both summarization results into an MLP\n",
    "        Y_hat = self.output(self.h(jnp.concatenate([V_A, V_B], axis=1)))\n",
    "        return Y_hat"
   ],
   "execution_count": 5,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-attention-md-6 fw=all -->\n",
    "\n",
    "### Putting It All Together\n",
    "\n",
    "By putting the attending, comparing, and aggregating steps together,\n",
    "we define the decomposable attention model to jointly train these three steps."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-attention-putting-it-all-together"
   },
   "source": [
    "class DecomposableAttention(nnx.Module):\n",
    "    def __init__(self, vocab_size, embed_size, num_hiddens, 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",
    "        self.attend = Attend(embed_size, num_hiddens, rngs=rngs)\n",
    "        self.compare = Compare(embed_size, num_hiddens, rngs=rngs)\n",
    "        self.aggregate = Aggregate(num_hiddens, 3, rngs=rngs)\n",
    "\n",
    "    def __call__(self, premises, hypotheses):\n",
    "        A = self.embedding(premises)\n",
    "        B = self.embedding(hypotheses)\n",
    "        beta, alpha = self.attend(A, B)\n",
    "        V_A, V_B = self.compare(A, B, beta, alpha)\n",
    "        # There are 3 possible outputs: entailment, contradiction, and neutral\n",
    "        Y_hat = self.aggregate(V_A, V_B)\n",
    "        return Y_hat"
   ],
   "execution_count": 6,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-attention-md-7 fw=all -->\n",
    "\n",
    "## Training and Evaluating the Model\n",
    "\n",
    "Now we will train and evaluate the defined decomposable attention model on the SNLI dataset.\n",
    "We begin by reading the dataset.\n",
    "\n",
    "\n",
    "### Reading the dataset\n",
    "\n",
    "We download and read the SNLI dataset using the function defined in\n",
    "that section. Most implementations use\n",
    "minibatches of $256$ examples padded or truncated to $50$ tokens. The JAX\n",
    "implementation uses minibatches of $512$ and length $32$ to avoid spending most\n",
    "of its time on padding. With the tokenization performed by `read_snli`, about\n",
    "1.2% of the labeled training premises and 0.02% of the hypotheses exceed 32\n",
    "tokens. This substantially reduces computation while introducing little\n",
    "additional truncation."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-attention-reading-the-dataset"
   },
   "source": [
    "batch_size, num_steps = 512, 32\n",
    "train_iter, test_iter, vocab = d2l.load_data_snli(batch_size, num_steps)"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "read 549367 examples\nread 9824 examples\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-attention-md-8 fw=all -->\n",
    "\n",
    "### Creating the Model\n",
    "\n",
    "We use the pretrained 100-dimensional GloVe embedding to represent the input tokens.\n",
    "Thus, we predefine the dimension of vectors $\\mathbf{a}_i$ and $\\mathbf{b}_j$ in the equation as 100.\n",
    "The output dimension of functions $f$ in the equation and $g$ in the equation is set to 200.\n",
    "Then we create a model instance, initialize its parameters,\n",
    "and load the GloVe embedding to initialize vectors of input tokens."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-attention-creating-the-model"
   },
   "source": [
    "embed_size, num_hiddens, devices = 100, 200, d2l.try_all_gpus()\n",
    "net = DecomposableAttention(len(vocab), embed_size, num_hiddens)\n",
    "glove_embedding = d2l.TokenEmbedding('glove.6b.100d')\n",
    "embeds = glove_embedding[vocab.idx_to_token]"
   ],
   "execution_count": 8,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-attention-md-9 fw=all -->\n",
    "\n",
    "### Training and Evaluating the Model\n",
    "\n",
    "In contrast to the `split_batch` function in that section that takes single inputs such as text sequences (or images),\n",
    "we define a `split_batch_multi_inputs` function to take multiple inputs such as premises and hypotheses in minibatches."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-attention-md-10 fw=all -->\n",
    "\n",
    "Now we can train and evaluate the model on the SNLI dataset."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-attention-training-and-evaluating-the-model-2-2"
   },
   "source": [
    "lr, num_epochs = 0.001, 4\n",
    "net.embedding.embedding[...] = jnp.array(embeds)\n",
    "optimizer = nnx.Optimizer(net, optax.adam(lr), wrt=nnx.Param)\n",
    "train_net = nnx.view(net, deterministic=False)\n",
    "eval_net = nnx.view(net, deterministic=True)\n",
    "\n",
    "@nnx.jit\n",
    "def train_step(net, optimizer, premises, hypotheses, labels):\n",
    "    def loss_fn(model):\n",
    "        logits = model(premises, hypotheses)\n",
    "        return optax.softmax_cross_entropy_with_integer_labels(\n",
    "            logits, labels).mean()\n",
    "    loss, grads = nnx.value_and_grad(loss_fn)(net)\n",
    "    optimizer.update(net, grads)\n",
    "    return loss\n",
    "\n",
    "@nnx.jit\n",
    "def eval_step(net, premises, hypotheses, labels):\n",
    "    logits = net(premises, hypotheses)\n",
    "    return (logits.argmax(axis=-1) == labels).sum()\n",
    "\n",
    "timer = d2l.Timer()\n",
    "for epoch in range(num_epochs):\n",
    "    batch_losses, n_train = [], 0\n",
    "    for batch in train_iter:\n",
    "        premises, hypotheses, labels = batch[0], batch[1], batch[2]\n",
    "        batch_losses.append(train_step(\n",
    "            train_net, optimizer, premises, hypotheses, labels))\n",
    "        n_train += len(labels)\n",
    "    # Synchronize once per epoch, rather than once per minibatch.\n",
    "    train_loss = float(jnp.stack(batch_losses).mean())\n",
    "    # Evaluate on test set\n",
    "    batch_correct, n_test = [], 0\n",
    "    for batch in test_iter:\n",
    "        premises, hypotheses, labels = batch[0], batch[1], batch[2]\n",
    "        batch_correct.append(eval_step(\n",
    "            eval_net, premises, hypotheses, labels))\n",
    "        n_test += len(labels)\n",
    "    n_correct = int(jnp.stack(batch_correct).sum())\n",
    "    print(f'epoch {epoch + 1}, loss {train_loss:.4f}, '\n",
    "          f'test acc {n_correct / n_test:.4f}')\n",
    "print(f'{num_epochs * n_train / timer.stop():.1f} examples/sec')"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "epoch 1, loss 0.8543, test acc 0.7414\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "epoch 2, loss 0.6290, test acc 0.7910\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "epoch 3, loss 0.5628, test acc 0.8029\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "epoch 4, loss 0.5267, test acc 0.8184\n937.5 examples/sec\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-attention-md-11 fw=all -->\n",
    "\n",
    "### Using the Model\n",
    "\n",
    "Finally, define the prediction function to output the logical relationship between a pair of premise and hypothesis."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-attention-using-the-model-1"
   },
   "source": [
    "\n",
    "def predict_snli(net, vocab, premise, hypothesis):\n",
    "    \"\"\"Predict the logical relationship between the premise and hypothesis.\"\"\"\n",
    "    premise = jnp.array(vocab[premise]).reshape((1, -1))\n",
    "    hypothesis = jnp.array(vocab[hypothesis]).reshape((1, -1))\n",
    "    label = jnp.argmax(nnx.view(net, deterministic=True)(premise, hypothesis),\n",
    "                       axis=1)\n",
    "    return 'entailment' if label == 0 else 'contradiction' if label == 1 \\\n",
    "            else 'neutral'"
   ],
   "execution_count": 10,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-attention-md-12 fw=all -->\n",
    "\n",
    "We can use the trained model to obtain the natural language inference result for a sample pair of sentences."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-attention-using-the-model-2"
   },
   "source": [
    "predict_snli(net, vocab, ['he', 'is', 'good', '.'],\n",
    "             ['he', 'is', 'bad', '.'])"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'contradiction'"
     },
     "metadata": {},
     "execution_count": 11
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-attention-md-13 fw=all -->\n",
    "\n",
    "## Summary\n",
    "\n",
    "* The decomposable attention model consists of three steps for predicting the logical relationships between premises and hypotheses: attending, comparing, and aggregating.\n",
    "* With attention mechanisms, we can align tokens in one text sequence to every token in the other, and vice versa. Such alignment is soft using weighted average, where ideally large weights are associated with the tokens to be aligned.\n",
    "* The decomposition trick leads to a more desirable linear complexity than quadratic complexity when computing attention weights.\n",
    "* We can use pretrained word vectors as the input representation for downstream natural language processing tasks such as natural language inference.\n",
    "\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Train the model with other combinations of hyperparameters. Can you get better accuracy on the test set?\n",
    "1. What are major drawbacks of the decomposable attention model for natural language inference?\n",
    "1. Suppose that we want to get the level of semantical similarity (e.g., a continuous value between 0 and 1) for any pair of sentences. How shall we collect and label the dataset? Can you design a model with attention mechanisms?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-attention-md-14 fw=jax -->\n",
    "\n",
    "[Discussions](https://d2l.discourse.group/t/1530)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
