{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Emia6qd0MdJD"
      },
      "source": [
        "# LLM pretraining\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "owbKvxwhcMuO"
      },
      "source": [
        "In this notebook, we will explore the process of pretraining a large language model (LLM) from scratch, focusing on a decoder-only transformer architecture.\n",
        "\n",
        "The vast majority of current LLMs are trained with a causal auto-regressive objective, where the model predicts the next token given the previous tokens.\n",
        "\n",
        "We will follow this approach to train a character-level LLM on the works of Shakespeare."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "torZ8wDB_eGs"
      },
      "outputs": [],
      "source": [
        "import dataclasses\n",
        "import math\n",
        "import random\n",
        "import torch\n",
        "import torch.cuda\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim\n",
        "import torch.utils.data\n",
        "import time\n",
        "\n",
        "\n",
        "from typing import Optional"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yhmbmir2cuCn"
      },
      "source": [
        "## Part I: The training data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c7htNP9Jc0z9"
      },
      "source": [
        "We will start by preparing the data to train the model, available at [https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\n",
        "](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\n",
        ").\n",
        "\n",
        "This data was originally used in the following [blog post](https://karpathy.github.io/2015/05/21/rnn-effectiveness/), also worth reading."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jiy0tKQfMvUB",
        "outputId": "2d721bda-b8e6-4b9d-dfc2-3ec4601d7e26"
      },
      "outputs": [],
      "source": [
        "# Downloading the tiny shakespeare dataset\n",
        "!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\n",
        "\n",
        "# read it in to inspect it\n",
        "with open('input.txt', 'r', encoding='utf-8') as f:\n",
        "    text = f.read()\n",
        "\n",
        "print(\"Length in characters: \", len(text))\n",
        "\n",
        "# let's look at the first 100 characters\n",
        "print(\"First 100 characters\", \"-\" * 50)\n",
        "print(text[:100])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7dKUimzLevzp"
      },
      "source": [
        "From here, we can create the model vocabulary, which maps the characters into token ids.\n",
        "\n",
        "We can train the vocabulary on the whole training set and perform some light testing to check that everything is working as expected."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "W6E7HG-YNigp",
        "outputId": "23b6b40c-c6a7-4a26-d6f5-982094de7251"
      },
      "outputs": [],
      "source": [
        "class Vocabulary:\n",
        "  \"\"\"Character-level vocabulary\"\"\"\n",
        "\n",
        "  def __init__(self, text: str):\n",
        "    \"\"\"Initialize the vocabulary, mapping individual characters to integers.\"\"\"\n",
        "    # We create two dictionaries, ctoi and itoc: the former mapping characters\n",
        "    # to unique integers and the latter with the reverse mapping.\n",
        "    raise NotImplementedError(\"Implement me!\")\n",
        "\n",
        "  def encode(self, text: str) -> list[int]:\n",
        "    \"\"\"Encode a string as a list of token ids.\"\"\"\n",
        "    raise NotImplementedError(\"Implement me!\")\n",
        "\n",
        "  def decode(self, ids: list[int]) -> str:\n",
        "    \"\"\"Decode a list of token ids into a string.\"\"\"\n",
        "    raise NotImplementedError(\"Implement me!\")\n",
        "\n",
        "  def __len__(self):\n",
        "    return len(self.ctoi)\n",
        "\n",
        "  def __str__(self):\n",
        "    return f\"Vocabulary[size={len(self)}]{self.ctoi}\"\n",
        "\n",
        "vocab = Vocabulary(text)\n",
        "print(vocab)\n",
        "\n",
        "# Test if encoding/decoding work\n",
        "first_chars = text[:1000]\n",
        "encoded = vocab.encode(first_chars)\n",
        "decoded = vocab.decode(encoded)\n",
        "assert first_chars == decoded"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MCWHEgmgfpEg"
      },
      "source": [
        "With our vocabulary, we can encode the full text data as a tensor and split it into data for training and validation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jfmYqseQfoxs",
        "outputId": "86215e65-5b2d-49ee-a95a-d929461aa0f3"
      },
      "outputs": [],
      "source": [
        "# Encode all text\n",
        "encoded_text = torch.tensor(vocab.encode(text))\n",
        "print(\"Total tokens:\", encoded_text.size(0))\n",
        "# Split into train and validation\n",
        "val_size = 2000\n",
        "\n",
        "encoded_train = encoded_text[:-val_size]\n",
        "encoded_val = encoded_text[-val_size:]\n",
        "print(\"Train tokens:\", encoded_train.size(0))\n",
        "print(\"Val tokens: \", encoded_val.size(0))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "akQ2Ugf0gb66"
      },
      "source": [
        "From here, we can create the dataset class to iterate the data. Each record will be a chunk of adjacent tokens and the model will learn to predict the next token from the previous ones.\n",
        "\n",
        "We will train on sequences of 512 tokens, which offers a good trade-off between previous context length and efficiency (due to the quadractic behavior of attention)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7OYYZCc9bj3y",
        "outputId": "e2f7d306-41d3-4b7f-9b0c-44baee56a2da"
      },
      "outputs": [],
      "source": [
        "# Predict on block_size contiguous tokens\n",
        "seq_len = 512\n",
        "\n",
        "\n",
        "class ShakespeareDataset(torch.utils.data.Dataset):\n",
        "  \"\"\"Dataset class return input and target ids for a given idx\"\"\"\n",
        "\n",
        "  def __init__(self, data: torch.Tensor, seq_len: int):\n",
        "    self.data = data\n",
        "    self.seq_len = seq_len\n",
        "\n",
        "\n",
        "  def __getitem__(self, i: int) -> torch.Tensor:\n",
        "    \"\"\"Selects the seq_len+1 tokens starting at index i for training.\n",
        "\n",
        "    Returns:\n",
        "        (x, y) tuple, with the input and target ids.\n",
        "    \"\"\"\n",
        "    # Account for extra one in LM offset\n",
        "    raise NotImplementedError(\"Implement me!\")\n",
        "\n",
        "  def __len__(self):\n",
        "    \"\"\"The length of the dataset, only considering complete sequences.\"\"\"\n",
        "    return len(self.data) - self.seq_len - 1\n",
        "\n",
        "train = ShakespeareDataset(encoded_train, seq_len)\n",
        "val = ShakespeareDataset(encoded_val, seq_len)\n",
        "print(\"Train size:\", len(train), \", val size:\", len(val))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yOaZZGCICviC"
      },
      "source": [
        "## Part II: The Transformer Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AB7YL90ctBwg"
      },
      "source": [
        "In this section of the notebook, we will implement a standard decoder-only transformer model.\n",
        "\n",
        "There are many flavours of the transformer architecture, varying all components of the transformer network. In this notebook, we will consider the following configuration:\n",
        "\n",
        "* **Learnt Positional Encoddings** - In conjuntion with sinosoidal embeddings, these are one of the first approaches for encoding positional information in the transformer architecture, by learning position vectors to encode this information. Currently, however, most transformers adopt [RoPE embeddings](https://arxiv.org/abs/2104.09864).\n",
        "\n",
        "* **Pre layer normalization** - In contrast with the original transformer, we will add layer normalization *before* the attention and mlp blocks, which is shown to improve performance. Nevertheless, we will adopt the original layer normalization instead of the, nowadays more common, [RMS normalization](https://arxiv.org/abs/1910.07467).\n",
        "\n",
        "* **MLP with gelu activation** - Following the original transformer, we adopt a two layer MLP, but replace the ReLU activation with [GeLU](https://arxiv.org/abs/1606.08415). Recently, larger models also successfuly adopted variations of [GLUs](https://arxiv.org/abs/2002.05202).\n",
        "\n",
        "We start by defining the model configuration."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7NLfTPutY5ml"
      },
      "outputs": [],
      "source": [
        "@dataclasses.dataclass\n",
        "class Config:\n",
        "  seq_len: int = seq_len\n",
        "  n_layers: int = 6\n",
        "  n_heads: int = 6\n",
        "  hidden_size: int = 384\n",
        "  # Using the ffn hidden size as 4 times the\n",
        "  # network hidden size is a standard value.\n",
        "  ffn_hidden_size: int = 384 * 4\n",
        "  # Using a head size of hidden_size // n_heads\n",
        "  # is also a standard approach.\n",
        "  head_size: int = 384 // 6\n",
        "  vocab_size: int = len(vocab)\n",
        "  dropout = 0.3"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BCuB0WUYuHYd"
      },
      "source": [
        "Afterwards, we will define the self-attention mechanism. In order to make the implementation more efficient, we will:\n",
        "\n",
        "* Sharing projection matrices across heads;\n",
        "* Using torch.nn.functional.scaled_dot_product_attention for an efficient implementation of scaled dot-product attention. This implementation is based on [FlashAttention2](https://arxiv.org/abs/2307.08691)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cg-wikzjuDQr"
      },
      "outputs": [],
      "source": [
        "class MultiHeadAttention(nn.Module):\n",
        "\n",
        "  def __init__(self, config: Config, /, *, flash: bool = True):\n",
        "    super().__init__()\n",
        "    self.config = config\n",
        "    self.flash = True\n",
        "    # As is common in most transformer implementations, we will have single projection\n",
        "    # matrices to compute queries, keys and values for all heads.\n",
        "    # As such, each projection matrix will map an input of hidden_size to an output\n",
        "    # of n_heads * head_size.\n",
        "    self.queries_proj = nn.Linear(config.hidden_size, config.n_heads * config.head_size)\n",
        "    self.keys_proj = nn.Linear(config.hidden_size, config.n_heads * config.head_size)\n",
        "    self.values_proj = nn.Linear(config.hidden_size, config.n_heads * config.head_size)\n",
        "    self.out_proj = nn.Linear(config.n_heads * config.head_size, config.hidden_size)\n",
        "\n",
        "  def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
        "    # x shape: (batch, seq, hidden)\n",
        "    B, T, _ = x.size()\n",
        "    n_heads, head_dim = self.config.n_heads, self.config.head_size\n",
        "\n",
        "    # Compute projections and split by queries, keys and values.\n",
        "    # q, k, v shapes: (batch, seq, n_heads x head_size).\n",
        "\n",
        "    # Isolate heads\n",
        "    # q, k, v shapes: (batch, n_heads, seq, head_size)\n",
        "\n",
        "    # Compute scaled dot product attention:\n",
        "    # Hint 1: Use torch.nn.functional.scaled_dot_product_attention.\n",
        "    # Hint 2: Be careful to only add dropout during training and ensure a causal mask.\n",
        "    # out shape: (batch, n_heads, seq, head_size)\n",
        "\n",
        "    # Reorder to have all head values for token together\n",
        "    # out shape: (batch, seq, n_heads * head_size)\n",
        "\n",
        "    # Compute the final projection\n",
        "    # out shape: (batch, seq, hidden_size)\n",
        "    raise NotImplementedError(\"Implement me!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b7sLab5DepDK"
      },
      "source": [
        "Second, we will define the MLP layer of the transformer. In our implementation, we will adopt:\n",
        "* Up projection layer mapping hidden_size to ffn_hidden_size;\n",
        "* Activation function (GeLU);\n",
        "* Down projection layer mapping ffn_hidden_size to hidden_size;\n",
        "* Dropout."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gmJMINI_eoNN"
      },
      "outputs": [],
      "source": [
        "class MLP(nn.Module):\n",
        "\n",
        "  def __init__(self, config: Config):\n",
        "    raise NotImplementedError(\"Implement me!\")\n",
        "  \n",
        "  def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
        "    raise NotImplementedError(\"Implement me!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_lGdWenOxAlz"
      },
      "source": [
        "Now, we will define the transformer block. In contrast with the original transformer implementation, we will apply layer normalization before MHA and the MLP."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G-jXu95piwON"
      },
      "outputs": [],
      "source": [
        "class Block(nn.Module):\n",
        "\n",
        "  def __init__(self, config: Config):\n",
        "    raise NotImplementedError(\"Implement me!\")\n",
        "\n",
        "  def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
        "    raise NotImplementedError(\"Implement me!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iJM-iVehizyA"
      },
      "source": [
        "Finally, we will define the transformer decoder, using a causal language modeling objective. For this example, we use learnt embeddings to represent positional information in the transformer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SagYJ1_XsQ75"
      },
      "outputs": [],
      "source": [
        "class Transformer(nn.Module):\n",
        "\n",
        "  def __init__(self, config: Config):\n",
        "    super().__init__()\n",
        "    # Initialize both embedding layers, one with the size of the vocabulary,\n",
        "    # and another with the size of the sequence length.\n",
        "\n",
        "    # Initialize the transformer blocks\n",
        "\n",
        "    # It is also common to add a final layer norm before the lm_head\n",
        "    raise NotImplementedError(\"Implement me!\")\n",
        "\n",
        "  def forward(\n",
        "      self,\n",
        "      inputs: torch.Tensor,\n",
        "      *,\n",
        "      labels: Optional[torch.Tensor] = None\n",
        "  ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:\n",
        "    \"\"\"Computes the forward of a transformer model.\n",
        "\n",
        "    Returns a tuple with two elements. The first is the final logits,\n",
        "    and the second is the loss (only when labels are provided)\n",
        "    \"\"\"\n",
        "    raise NotImplementedError(\"Implement me!\")\n",
        "\n",
        "  def _get_pos_embeddings(self, tokens: torch.Tensor) -> torch.Tensor:\n",
        "    \"\"\"Computes the positional embeddings for the given tokens.\"\"\"\n",
        "    batch_size, seq_len = tokens.size()\n",
        "    pos_idx = torch.arange(seq_len, device=tokens.device)\n",
        "    pos_emb = self.pos_emb(pos_idx)\n",
        "    return pos_emb\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jfOy4OWJxPGa"
      },
      "source": [
        "## Part III: Training and Evaluation\n",
        "\n",
        "In the final part of this notebook, we will implement the training loop, as well as a generation function, and train our transformer.\n",
        "\n",
        "In order to mitigate expensive cost of training language models, we will implement [mixed precision training](https://arxiv.org/abs/1710.03740). In this approach, while the weights are stored in fp32, the operations inside the model are performed with fp16, levergaging faster gpu implementations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HcLQcMGaDP4g"
      },
      "outputs": [],
      "source": [
        "def train_step(\n",
        "    *,\n",
        "    model: Transformer,\n",
        "    optimizer: torch.optim.Optimizer,\n",
        "    scaler: torch.amp.GradScaler,\n",
        "    device: str,\n",
        "    x: torch.Tensor,\n",
        "    y: torch.Tensor,\n",
        "):\n",
        "  # Run the model in mixed precision, which improves overall speed\n",
        "  # with minimal performance degradation.\n",
        "  # https://pytorch.org/docs/stable/notes/amp_examples.html#typical-mixed-precision-training\n",
        "  raise NotImplementedError(\"Implement me!\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zlwQIoSuVOln"
      },
      "outputs": [],
      "source": [
        "@torch.no_grad\n",
        "def validate(\n",
        "    *,\n",
        "    model: Transformer,\n",
        "    val_dataloader: torch.utils.data.DataLoader,\n",
        "    device: str,\n",
        "):\n",
        "    model.eval()\n",
        "    total_loss = 0.0\n",
        "    for x, y in val_dataloader:\n",
        "        x, y = x.to(device), y.to(device)\n",
        "        with torch.autocast(device_type=device, dtype=torch.float16):\n",
        "          _, loss = model(x, labels=y)\n",
        "        total_loss += loss\n",
        "\n",
        "    total_loss /= len(val_dataloader)\n",
        "\n",
        "    model.train()\n",
        "    return total_loss.item()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ARtsnw91YoeX"
      },
      "outputs": [],
      "source": [
        "@torch.no_grad\n",
        "def generate(\n",
        "  model: nn.Module,\n",
        "  vocab: Vocabulary,\n",
        "  input: str,\n",
        "  new_tokens: int,\n",
        "  device,\n",
        "):\n",
        "  model.eval()\n",
        "  idx = vocab.encode(input)\n",
        "  idx = torch.tensor(idx, dtype=torch.long, device=device)\n",
        "  idx = idx.unsqueeze(0)\n",
        "\n",
        "  for t in range(new_tokens):\n",
        "    # The model only considers the previous seq len characters.\n",
        "    ctx = idx[:, -seq_len:]\n",
        "\n",
        "    with torch.autocast(device_type=device, dtype=torch.float16):\n",
        "      logits, _ = model(ctx)\n",
        "\n",
        "    next_logits = logits[:, -1, :]\n",
        "\n",
        "    next_probs = torch.softmax(next_logits, dim=-1)\n",
        "\n",
        "    next_idx = torch.multinomial(next_probs, num_samples=1)\n",
        "    # append sampled index to the running sequence and continue\n",
        "    idx = torch.cat((idx, next_idx), dim=1)\n",
        "\n",
        "  idx = idx.tolist()[0]\n",
        "  output = vocab.decode(idx)\n",
        "  model.train()\n",
        "  return output"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZkwRanJ1JN-o"
      },
      "outputs": [],
      "source": [
        "def run_validation(\n",
        "    *,\n",
        "    model,\n",
        "    val_dataloader,\n",
        "    device):\n",
        "\n",
        "  val_loss = validate(\n",
        "      model=model,\n",
        "      val_dataloader=val_dataloader,\n",
        "      device=device,\n",
        "  )\n",
        "  print(\"Validation\", \"-\" * 39)\n",
        "  print(f\"Total loss={val_loss:.3f}\")\n",
        "  test_input = \"\\n\" # Continue to use new line as if beginning a new paragraph\n",
        "  test_gen_len = 200\n",
        "  test_output = generate(\n",
        "      model,\n",
        "      vocab,\n",
        "      test_input,\n",
        "      test_gen_len,\n",
        "      device\n",
        "  )\n",
        "  print(test_output)\n",
        "  print(\"-\" * 50)\n",
        "  return val_loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ro7TCXQYJzym"
      },
      "outputs": [],
      "source": [
        "def set_seed(seed = 42):\n",
        "  random.seed(seed)\n",
        "  torch.manual_seed(seed)\n",
        "  torch.set_float32_matmul_precision(\"high\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W2caBQniwDqq"
      },
      "outputs": [],
      "source": [
        "def train_model(\n",
        "    *,\n",
        "    config = Config(),\n",
        "    num_steps=2000,\n",
        "    val_every=500,\n",
        "    log_every=50):\n",
        "\n",
        "  set_seed()\n",
        "\n",
        "  device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "  if device == \"cpu\":\n",
        "    print(\"WARNING: Using cpu device, training will be slow\")\n",
        "\n",
        "  train_dataloader = torch.utils.data.DataLoader(\n",
        "      train, batch_size=64, shuffle=True, pin_memory=True, num_workers=2, prefetch_factor=2,\n",
        "  )\n",
        "  val_dataloader = torch.utils.data.DataLoader(\n",
        "      val, batch_size=64, shuffle=False, pin_memory=True, num_workers=2, prefetch_factor=2,\n",
        "  )\n",
        "\n",
        "  model = Transformer(config)\n",
        "\n",
        "  model_params = sum(p.numel() for p in model.parameters())\n",
        "  print(f\"Total parameters: {model_params / 1e6:.2f}M\")\n",
        "\n",
        "  optimizer = torch.optim.AdamW(\n",
        "      model.parameters(), lr=1e-3, weight_decay=0.1, betas=(0.9, 0.95), fused=True\n",
        "  )\n",
        "  scaler = torch.amp.GradScaler()\n",
        "  model.to(device)\n",
        "\n",
        "  dataiter = iter(train_dataloader)\n",
        "  start_time = time.time()\n",
        "  start_step = time.time()\n",
        "\n",
        "  train_losses = []\n",
        "  val_losses = []\n",
        "\n",
        "  val_losses.append(run_validation(\n",
        "      model=model,\n",
        "      val_dataloader=val_dataloader,\n",
        "      device=device,\n",
        "  ))\n",
        "\n",
        "  for step in range(1, num_steps + 1):\n",
        "    x, y = next(dataiter)\n",
        "    x, y = x.to(device), y.to(device)\n",
        "    token_count = x.numel()\n",
        "    loss = train_step(\n",
        "        model=model,\n",
        "        optimizer=optimizer,\n",
        "        scaler=scaler,\n",
        "        device=device,\n",
        "        x=x,\n",
        "        y=y)\n",
        "\n",
        "    train_losses.append(loss)\n",
        "\n",
        "    step_time = time.time() - start_step\n",
        "    start_step = time.time()\n",
        "    if step % log_every == 0:\n",
        "      completed = step / num_steps * 100\n",
        "      ellapsed = time.time() - start_time\n",
        "      tok_per_sec = token_count / step_time\n",
        "      fmt_ellapsed = time.strftime(\"%H:%M:%S\", time.gmtime(ellapsed))\n",
        "      print(\n",
        "          f\"Step {step}/{num_steps} ({completed:.0f}%): \"\n",
        "          f\"TrainTime={fmt_ellapsed}, \"\n",
        "          f\"Loss={loss:.3f}, \"\n",
        "          f\"StepTime={step_time:.2}s\"\n",
        "      )\n",
        "    if step % val_every == 0:\n",
        "      val_losses.append(run_validation(\n",
        "          model=model,\n",
        "          val_dataloader=val_dataloader,\n",
        "          device=device\n",
        "      ))\n",
        "\n",
        "  test_input = \"\\n\" # Continue to use new line as if beginning a new paragraph\n",
        "  test_gen_len = 1000\n",
        "  test_output = generate(\n",
        "      model,\n",
        "      vocab,\n",
        "      test_input,\n",
        "      test_gen_len,\n",
        "      device\n",
        "  )\n",
        "  print(\"Final generation\", \"-\" * 30)\n",
        "  print(test_output)\n",
        "\n",
        "  return train_losses, val_losses"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DeOkj3UYgf87",
        "outputId": "5e3056c1-a0a5-4254-a2a7-c552ba1bd089"
      },
      "outputs": [],
      "source": [
        "train_losses, val_losses = train_model()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 472
        },
        "id": "F-o_51xWCIgh",
        "outputId": "d8dc4cf4-884d-47f0-86df-a3e5d5fa6c7f"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "\n",
        "# Create x-axis values for training and validation losses\n",
        "train_steps = list(range(len(train_losses)))\n",
        "num_val_steps = len(val_losses)\n",
        "num_steps = len(train_steps)\n",
        "val_steps = np.linspace(0, num_steps, num_val_steps, dtype=int).tolist()\n",
        "\n",
        "# Plot Cross Entropy Train losses\n",
        "plt.plot(train_steps, train_losses, label=\"Train\")\n",
        "plt.plot(val_steps, val_losses, label=\"Validation\")\n",
        "plt.xlabel(\"Step\")\n",
        "plt.ylabel(\"Loss\")\n",
        "plt.title(\"Cross entropy loss\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xsPbCbsEH0IP"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}