{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pk-jM1vivDzp"
      },
      "source": [
        "# Generating Names with Transformer\n",
        "\n",
        "In this lab, you will perform character-level text generation again, but this time using a Transformer instead of an RNN.\n",
        "We will work with a small dataset of names.\n",
        "\n",
        "You will start by tokenizing the data, then implement the core component of the Transformer architecture, multi-head self-attention. Finally, you will build the model by assembling the Transformer layers, train it and run inference.\n",
        "\n",
        "We recommend using Google Colab for this lab."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rjbfy3DVvI0o"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nhrsBW44Q70a",
        "outputId": "e851ed49-0193-4b84-e8ac-8606d6ceaf99"
      },
      "outputs": [],
      "source": [
        "# Load data\n",
        "!wget --quiet --show-progress \"https://gist.githubusercontent.com/saparina/c202be3b05f16375bd94644b802bc2a8/raw/91a0507776656b476499b55d3b12e3dd7cda4fa2/names\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a_s_Z5lbIaPG"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "import numpy as np\n",
        "import torch\n",
        "from torch import optim\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "from IPython.display import clear_output\n",
        "\n",
        "%matplotlib inline\n",
        "\n",
        "random.seed(42)\n",
        "np.random.seed(42)\n",
        "torch.manual_seed(42)\n",
        "if torch.cuda.is_available():\n",
        "    torch.cuda.manual_seed_all(42)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EMpB7jXVwQj2"
      },
      "source": [
        "## Data Processing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n6nXxU8WIaPM"
      },
      "source": [
        "The `names` file contains approximately 8,000 names in Latin script.\n",
        "\n",
        "We want the model to predict next character for a given prefix. To help the model learn to generate the full name from the beginning, we need to first predict a capital letter. We will add a special `_`  token at the beginning that will indicate a start of the sequence.\n",
        "\n",
        "We also need a stopping rule for generation (this could simply be a limit on the number of steps). Alternatively, we can add a special `<EOS>` token at the end of each training example. In this case, we will denote it as `#`.\n",
        "\n",
        "\n",
        "For example, to generate the name `Amandy`, the model will sequentially predict next character using the letters predicted in the previous steps as input. The generation starts with `_`:\n",
        "\n",
        "`_` -> `A`\n",
        "\n",
        "`_A` -> `m`\n",
        "\n",
        "`_Am` -> `a`\n",
        "\n",
        "...\n",
        "\n",
        "`_Amandy` -> `#`\n",
        "\n",
        "\n",
        "Generation can stop early if the `<EOS>` token (`#`) appears prematurely."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TFRHva2zIaPN",
        "outputId": "a1b9a0a7-f7da-406f-c12a-c596046b20cb"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "There are 7578 names: \n",
            "_Hall#\n",
            "_Zorina#\n",
            "_Jobina#\n",
            "_Orlando#\n",
            "_Domenic#\n",
            "_Jourdan#\n",
            "_Sileas#\n",
            "_Melessa#\n"
          ]
        }
      ],
      "source": [
        "import os\n",
        "start_token = \"_\"\n",
        "eos = '#'\n",
        "\n",
        "with open(\"names\") as f:\n",
        "    names = f.readlines()\n",
        "    names = [start_token + name.strip() + eos for name in names]\n",
        "\n",
        "names = list(set(names))  # filter duplicates\n",
        "print('There are {} names: '.format(len(names)))\n",
        "for x in names[::1000]:\n",
        "    print(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KWOHXx_DXSvz"
      },
      "source": [
        "**Task 1**: Select all unique tokens and create two dictionaries for token-to-index and index-to-token conversion. Ensure that the token \"_\" ( ```start_token```) has the index 0 in the dictionary."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zgB0VE9BIaPa",
        "outputId": "bc037e16-0567-4e3a-d2eb-e30a594073ae"
      },
      "outputs": [],
      "source": [
        "tokens = <create list of all tokens>\n",
        "tok2id = <create dict to map tokens to indices>\n",
        "id2tok = <create dict to map indices to tokens>\n",
        "\n",
        "n_tokens = len(tokens)\n",
        "print ('There are {} tokens',n_tokens)\n",
        "\n",
        "assert 50 < n_tokens < 60\n",
        "\n",
        "print('Vocabulary: ' + \"\".join(tokens))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jF_ukJotIaPd"
      },
      "outputs": [],
      "source": [
        "def to_matrix(names, max_len=None, pad=tok2id[' '], dtype=np.int64):\n",
        "    \"\"\"Casts a list of names into matrix\"\"\"\n",
        "\n",
        "    max_len = max_len or max(map(len, names))\n",
        "    names_ix = np.zeros([len(names), max_len], dtype) + pad\n",
        "\n",
        "    for i in range(len(names)):\n",
        "        try:\n",
        "            name_ix = list(map(tok2id.get, names[i]))\n",
        "            names_ix[i, :len(name_ix)] = name_ix\n",
        "        except Exception as e:\n",
        "            print(names[i])\n",
        "            raise e\n",
        "\n",
        "    return names_ix"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wmg_il5MIaPg",
        "outputId": "018702ea-24bc-44ed-86e1-fc5972efea70"
      },
      "outputs": [],
      "source": [
        "print('\\n'.join(names[:10]))\n",
        "print(to_matrix(names[:10]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yrMOyd0xXSv0"
      },
      "source": [
        "Let's split the data into two files, one for training your model and another one for testing:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "64C8xOqCIaPk",
        "outputId": "277b918e-b964-489b-92a6-560ae55a028c"
      },
      "outputs": [],
      "source": [
        "def split_data(names, train_ratio=0.8):\n",
        "    random.shuffle(names)  # Shuffle the data\n",
        "    split_idx = int(len(names) * train_ratio)\n",
        "    train_data = names[:split_idx]  # First 80% for training\n",
        "    val_data = names[split_idx:]  # Remaining 20% for validation\n",
        "    return train_data, val_data\n",
        "\n",
        "train_data, val_data = split_data(names)\n",
        "len(train_data), len(val_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NhgqoEHOIaPr"
      },
      "source": [
        "# Char-Transformer for Name Generation\n",
        "\n",
        "You need to build a character-based Transformer model that encodes the input sequence and predicts the next character.\n",
        "\n",
        "During training, the next character is known. Unlike RNNs, which process sequences step by step, the Transformer can process the entire sequence in a single pass. However, to prevent the model from \"looking into the future\" (i.e., using information from upcoming characters) you need to apply a mask.\n",
        "\n",
        "We will start by building the Transformer with its most important component: multi-head self-attention. Self-attention helps the model understand how different parts of a sequence relate to each other:\n",
        "\n",
        "![self-attention](https://jalammar.github.io/images/t/transformer_self-attention_visualization.png)\n",
        "\n",
        "Note that in our case, we are working with **characters**, not words.\n",
        "\n",
        "[The Illustrated Transformer](https://jalammar.github.io/illustrated-transformer/) is a very useful guide that explains Transformers in an intuitive way."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "noN540WdW6uc"
      },
      "source": [
        "## Multi-Head Self-Attention  \n",
        "\n",
        "We will start by implementing **Scaled-Dot Product Attention**. This mechanism calculates *attention scores* to determine which characters are most relevant to each other.  \n",
        "\n",
        "You need to implement the following formula to calculate attention:  \n",
        "\n",
        "   $$\n",
        "   \\text{Attention}(Q, K, V) = \\text{softmax} \\left( \\frac{QK^T}{\\sqrt{d_k}} \\right) V\n",
        "   $$  \n",
        "\n",
        "In self-attention, Q, K, and V (short for Queries, Keys, and Values) are different projections of the input sequence, and d_k is the key dimension, used for scaling to stabilize gradients.  \n",
        "\n",
        "By multiplying Q and K, we create a square matrix where the element at position *(i, j)* represents the *relevance score* between the *i-th* and *j-th* characters. We then weight V (which is also a projection of the input) using these attention scores to generate the final output.  \n",
        "\n",
        "Instead of using a single attention mechanism, **Multi-Head Self-Attention** applies multiple self-attention processes in parallel.\n",
        "Each attention head learns different relationships within the sequence.\n",
        "\n",
        "At the end, all attention heads are combined to create a better representation of the sequence.\n",
        "\n",
        "![multi_head_attetion](https://nlp.seas.harvard.edu/images/the-annotated-transformer_38_0.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RLjM7zv1hvrc"
      },
      "source": [
        "\n",
        "**Task 2**: Implement the scaled-dot product attention."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lUhkJ6VKcM32"
      },
      "outputs": [],
      "source": [
        "class ScaledDotProductAttention(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "\n",
        "    def forward(self, Q, K, V, mask=None):\n",
        "        \"\"\"\n",
        "        Q, K, V: (batch_size, n_heads, seq_len, head_dim)\n",
        "        mask:    (batch_size, seq_len, seq_len)\n",
        "        \"\"\"\n",
        "        d_k = Q.shape[-1]  # Head dimension\n",
        "        scores = <calculate dot product>  # Scaled dot-product\n",
        "\n",
        "        if mask is not None:\n",
        "            scores = <apply attention mask to scores>  # Apply mask for valid attention only>  # Apply mask if provided\n",
        "\n",
        "        attn_weights = <get probability distribution from scores>  # Normalize across last dimension\n",
        "        output = <get weighted sum of values>  # Compute attention-weighted values\n",
        "\n",
        "        return output"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "neqxpkxEiQ_h"
      },
      "source": [
        "**Task 3**: Implement Multi-Head Self-Attention."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tJ1JMcc5iNDL"
      },
      "outputs": [],
      "source": [
        "class MultiHeadSelfAttention(nn.Module):\n",
        "    def __init__(self, emb_size, n_head):\n",
        "        super().__init__()\n",
        "        assert emb_size % n_head == 0, \"Embedding size must be divisible by number of heads\"\n",
        "\n",
        "        self.n_head = n_head\n",
        "        self.head_dim = emb_size // n_head  # Dimension per head\n",
        "\n",
        "        # Linear layers for Q, K, V projections\n",
        "        self.W_q = nn.Linear(emb_size, emb_size, bias=False)\n",
        "        self.W_k = nn.Linear(emb_size, emb_size, bias=False)\n",
        "        self.W_v = nn.Linear(emb_size, emb_size, bias=False)\n",
        "\n",
        "        # Output projection layer\n",
        "        self.W_out = nn.Linear(emb_size, emb_size, bias=False)\n",
        "\n",
        "        self.scaled_dot_attn = ScaledDotProductAttention()\n",
        "\n",
        "    def forward(self, x, mask=None):\n",
        "        \"\"\"\n",
        "        x: (batch_size, seq_len, emb_size)\n",
        "        mask: Optional, used for causal masking (seq_len x seq_len)\n",
        "        \"\"\"\n",
        "        batch_size, seq_len, emb_size = x.shape\n",
        "\n",
        "        # Compute Q, K, V projections\n",
        "        Q = <get queries>\n",
        "        K = <get keys>\n",
        "        V = <get values>\n",
        "\n",
        "        # Reshape for multi-head attention\n",
        "        Q = Q.view(batch_size, seq_len, self.n_head, self.head_dim).transpose(1, 2)  # (batch, n_head, seq, head_dim)\n",
        "        K = K.view(batch_size, seq_len, self.n_head, self.head_dim).transpose(1, 2)\n",
        "        V = V.view(batch_size, seq_len, self.n_head, self.head_dim).transpose(1, 2)\n",
        "\n",
        "        # Compute scaled dot-product attention\n",
        "        attn_output = <apply attention mechanism>  # (batch, n_head, seq, head_dim)\n",
        "\n",
        "        # Merge heads\n",
        "        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, emb_size)\n",
        "\n",
        "        # Apply output projection\n",
        "        output = <project combined attention>\n",
        "\n",
        "        return output"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zay3rKf1TlMU"
      },
      "source": [
        "## Transformer Block\n",
        "\n",
        "Transformer block consists of two main components: Multi-Head Self-Attention and Feed-Forward Network. Each component is wrapped with Residual Connections and Layer Normalization for stable training and efficient gradient flow."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rNhZaxnRrHVj"
      },
      "outputs": [],
      "source": [
        "class TransformerBlock(nn.Module):\n",
        "    def __init__(self, emb_size, hidden_size, n_head, dropout=0.1):\n",
        "        super().__init__()\n",
        "        self.self_attn = MultiHeadSelfAttention(emb_size, n_head)\n",
        "        self.norm1 = nn.LayerNorm(emb_size)\n",
        "        self.norm2 = nn.LayerNorm(emb_size)\n",
        "\n",
        "        self.feed_forward = nn.Sequential(\n",
        "            nn.Linear(emb_size, hidden_size),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(hidden_size, emb_size)\n",
        "        )\n",
        "\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x, mask=None):\n",
        "        attn_output = self.self_attn(x, mask)\n",
        "        x = self.norm1(x + self.dropout(attn_output))  # Add & Norm\n",
        "\n",
        "        ff_output = self.feed_forward(x)\n",
        "        x = self.norm2(x + self.dropout(ff_output))  # Add & Norm\n",
        "\n",
        "        return x"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4YMuOxPuld5q"
      },
      "source": [
        "## Transformer\n",
        "\n",
        "Encoding a sequence with a Transformer follows a similar structure to RNNs:\n",
        "first, get input embeddings, then  process the sequence with Transformer Blocks and finally generates output scores.\n",
        "\n",
        "Unlike RNNs, Transformers process the entire sequence at once. To retain order, we must explicitly add *positional embeddings*.\n",
        "\n",
        "**Task 4**: Implement the forward pass."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KJCf0LYIIaPt"
      },
      "outputs": [],
      "source": [
        "MAX_LENGTH = 50  # Define maximum sequence length\n",
        "\n",
        "class NameTransformer(nn.Module):\n",
        "    def __init__(self, vocab_size, emb_size, hidden_size, n_layers=2, n_head=2, dropout=0.1):\n",
        "        super().__init__()\n",
        "        self.vocab_size = vocab_size\n",
        "\n",
        "        # Embedding layer\n",
        "        self.embedding = nn.Embedding(vocab_size, emb_size)\n",
        "\n",
        "        # Positional Encoding\n",
        "        self.position_embedding = nn.Embedding(MAX_LENGTH, emb_size)\n",
        "\n",
        "        # Transformer Blocks\n",
        "        self.layers = nn.ModuleList([\n",
        "            TransformerBlock(emb_size, hidden_size, n_head, dropout)\n",
        "            for _ in range(n_layers)\n",
        "        ])\n",
        "\n",
        "        # Output projection layer\n",
        "        self.fc_out = nn.Linear(emb_size, vocab_size)\n",
        "\n",
        "        # Register position ids\n",
        "        self.register_buffer(\"position_ids\", torch.arange(MAX_LENGTH).unsqueeze(1))\n",
        "\n",
        "    def _generate_square_subsequent_mask(self, seq_len):\n",
        "        \"\"\"\n",
        "        Generate a square mask of size (seq_len x seq_len) for self-attention.\n",
        "        \n",
        "        The mask ensures that position i can only attend to positions j where j <= i.\n",
        "        This prevents the model from \"cheating\" by looking at future tokens.\n",
        "\n",
        "        The mask values (i, j) are:\n",
        "        0.0 for positions that should be attended to (when j <= i)\n",
        "        -inf for positions that should be ignored (when j > i)\n",
        "        \"\"\"\n",
        "        mask = torch.triu(torch.full((seq_len, seq_len), float('-inf')), diagonal=1)\n",
        "        return mask\n",
        "\n",
        "    def forward(self, input):\n",
        "        \"\"\"\n",
        "        Forward pass through the transformer model.\n",
        "        \"\"\"\n",
        "        seq_len = input.size(1)\n",
        "        pos_ids = self.position_ids[:seq_len, :].squeeze(1)\n",
        "\n",
        "        # Embed input tokens and position\n",
        "        input_emb = <combine token and positional embeddings>\n",
        "\n",
        "        # Generate mask\n",
        "        mask = <create attention mask>\n",
        "\n",
        "        # Pass through Transformer Blocks\n",
        "        <apply multiple transformer layers>\n",
        "\n",
        "        # Project to vocabulary size\n",
        "        output = <transform final hidden states to logits>\n",
        "        return output"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BmkgMHc8IaPu"
      },
      "source": [
        "### Training\n",
        "\n",
        "**Task 5**: complete ```train_epoch``` function and train the model\n",
        "\n",
        "Note that train and test functions use `model.train()` and `model.eval()` to set the model in correct mode. This is needed because [`Dropout`](https://pytorch.org/docs/stable/generated/torch.nn.Dropout.html#torch.nn.Dropout) works differently during training and testing: during testing it simply computes the identity function, while during training it randomly zeros out some neurons and scales the remaining ones. \n",
        "Also, `model.eval()` can improve inference speed by disabling gradient tracking when not needed for testing."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "S164svO9IaPw"
      },
      "outputs": [],
      "source": [
        "# Function to train the model for one epoch\n",
        "def train_epoch(model, optimizer, train_batches, criterion=torch.nn.CrossEntropyLoss()):\n",
        "    loss_log = []\n",
        "    model.train() # train mode, compute grads, Dropout zeros out some neurons and scales the remaining ones\n",
        "\n",
        "    for batch in train_batches:\n",
        "        nums = to_matrix(batch)  # Convert batch to tensor matrix\n",
        "        nums = torch.tensor(nums, dtype=torch.long)\n",
        "\n",
        "        nums = nums.to(next(model.parameters()).device) # move tensor to the same device as model\n",
        "\n",
        "        # Forward pass\n",
        "        output = <forward pass>\n",
        "\n",
        "        # Compute loss (shift target left by 1 position)\n",
        "        target = nums[:, 1:].contiguous().view(-1)\n",
        "        output = output[:, :-1, :].contiguous().view(-1, model.vocab_size)\n",
        "\n",
        "        # Optimizer update\n",
        "        loss = <calculate cross-entropy loss>\n",
        "        <perform backpropagation and gradient descent step to update model parameters>\n",
        "\n",
        "        # Store loss\n",
        "        loss_log.append(loss.item())\n",
        "\n",
        "    return loss_log\n",
        "\n",
        "# Function to evaluate the model on the validation set\n",
        "def test(model, test_batches, criterion=torch.nn.CrossEntropyLoss()):\n",
        "    loss_log = []\n",
        "    model.eval() # eval mode, no grads, Dropout computes the identity function\n",
        "\n",
        "    with torch.no_grad():\n",
        "        for batch in test_batches:\n",
        "            nums = to_matrix(batch)\n",
        "            nums = torch.tensor(nums, dtype=torch.long).to(next(model.parameters()).device)\n",
        "\n",
        "            output = model(nums)\n",
        "\n",
        "            target = nums[:, 1:].contiguous().view(-1)  # Flatten target\n",
        "            output = output[:, :-1, :].contiguous().view(-1, model.vocab_size)  # Align with target\n",
        "\n",
        "            loss = criterion(output, target)\n",
        "            loss_log.append(loss.item())\n",
        "\n",
        "    return loss_log\n",
        "\n",
        "# Function to plot training history\n",
        "def plot_history(train_history, val_history, title='loss'):\n",
        "    plt.figure()\n",
        "    plt.title(title)\n",
        "    plt.plot(train_history, label='train', zorder=1)\n",
        "    points = np.array(val_history)\n",
        "    plt.scatter(points[:, 0], points[:, 1], marker='+', s=180, c='orange', label='val', zorder=2)\n",
        "    plt.xlabel('train steps')\n",
        "    plt.legend(loc='best')\n",
        "    plt.grid()\n",
        "    plt.show()\n",
        "\n",
        "# Full training function\n",
        "def train(model, opt, n_epochs):\n",
        "    train_log = []\n",
        "    val_log = []\n",
        "\n",
        "    bs = 32  # Batch size\n",
        "    total_steps = 0\n",
        "    train_batches = np.array_split(train_data, len(train_data) // bs)\n",
        "    test_batches = np.array_split(val_data, len(val_data) // bs)\n",
        "\n",
        "    for epoch in range(n_epochs):\n",
        "        print(f\"Epoch {epoch+1}/{n_epochs}\")\n",
        "\n",
        "        # Train for one epoch\n",
        "        train_loss = train_epoch(model, opt, train_batches)\n",
        "        train_log.extend(train_loss)\n",
        "        total_steps += len(train_batches)\n",
        "\n",
        "        # Evaluate on validation set\n",
        "        val_loss = test(model, test_batches)\n",
        "        val_log.append((len(train_log), np.mean(val_loss)))\n",
        "\n",
        "        # Clear output and plot loss curves\n",
        "        clear_output(wait=True)\n",
        "        plot_history(train_log, val_log, title='Training Loss')\n",
        "\n",
        "    print(\"Training complete!\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 490
        },
        "id": "Sxrc0a10IaPy",
        "outputId": "2cfcfb20-a18a-44f9-89a7-2ce87ea2ff88"
      },
      "outputs": [],
      "source": [
        "model = NameTransformer(len(tokens), 64, 64, n_layers=2, n_head=2, dropout=0.1)\n",
        "\n",
        "opt = torch.optim.Adam(model.parameters(), lr=1e-3)\n",
        "train(model, opt, 5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ffTAktWAIaP5"
      },
      "source": [
        "## Generation\n",
        "**Task 6**: complete ```gen_continuation``` function that run inference and generate a continuation of the input prefix. \n",
        "\n",
        "We'll use sampling with temperature. Experiment with different temperature values to see how they affect the quality and diversity of the generated names."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ugRlkX2ZIaP6",
        "outputId": "145cfd76-96de-401f-cddd-7873b3664a96"
      },
      "outputs": [],
      "source": [
        "def pick_by_distribution(logits, temperature=1.0):\n",
        "    \"\"\"\n",
        "    Selects a token based on probability distribution from logits.\n",
        "    \"\"\"\n",
        "    # Apply temperature to make distribution sharper or smoother\n",
        "    logits = logits / temperature\n",
        "\n",
        "    # Convert logits to probabilities\n",
        "    probs = torch.softmax(logits, dim=-1)\n",
        "    \n",
        "    # Sample from the probability distribution\n",
        "    return torch.multinomial(probs, 1).squeeze(-1)\n",
        "\n",
        "def ids2string(ids):\n",
        "    \"\"\"\n",
        "    Converts token indices back to a string.\n",
        "    \"\"\"\n",
        "    return \"\".join(id2tok[_] for _ in ids)\n",
        "\n",
        "def gen_continuation(model, prefix=\"_\", max_length=50, temperature=1.0):\n",
        "    \"\"\"\n",
        "    Generates a continuation of the input prefix using the trained model.\n",
        "    \"\"\"\n",
        "    model.eval()\n",
        "    with torch.no_grad():\n",
        "        nums = to_matrix([prefix])\n",
        "        nums = torch.tensor(nums, dtype=torch.long).to(next(model.parameters()).device)\n",
        "\n",
        "        outs = []  # Store generated tokens\n",
        "\n",
        "        for _ in range(max_length - len(prefix)):\n",
        "            output = <process current sequence through model>  # Forward pass\n",
        "\n",
        "            # Get the next token based on output logits\n",
        "            next_token = <sample the next token with the given temperature>\n",
        "\n",
        "            outs.append(next_token)  # Append token to generated sequence\n",
        "\n",
        "            # Stop if EOS token is reached\n",
        "            if next_token == tok2id.get(eos, None):\n",
        "                break\n",
        "\n",
        "            # Update model input by appending new token to prefix\n",
        "            nums = <concat the new token to the existing sequence>\n",
        "\n",
        "    print(prefix + \"|\" + ids2string(outs))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "gen_continuation(model, \" Ju\", temperature=0.7)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "fnlp",
      "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.10.16"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
