Transformer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "012f19fa-d88b-4356-b1cb-fd2dce3f7bbb",
   "metadata": {},
   "source": [
    "# Transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "359954e4-ec88-4389-ad02-998016bb0311",
   "metadata": {},
   "source": [
    "## Config\n",
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f4be0056-d6ee-4416-aa13-69ef8593a357",
   "metadata": {},
   "outputs": [],
   "source": [
    "# basic\n",
    "import os\n",
    "import math\n",
    "\n",
    "# display\n",
    "from tqdm.notebook import tqdm\n",
    "from bertviz import head_view, model_view\n",
    "\n",
    "# Data\n",
    "from dataclasses import dataclass\n",
    "from datasets import load_dataset, Dataset\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Machine leaning\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.checkpoint import checkpoint\n",
    "import tiktoken"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bab98150-3662-4ef7-8701-dbe53fdbc17d",
   "metadata": {},
   "source": [
    "### Global variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "997eff3a-807b-4131-8848-26ae66f1854c",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset_size = 500000"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9e811e4-75c2-4774-ba77-5b5347c6313d",
   "metadata": {},
   "source": [
    "## Utils Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38525540-00ba-4690-9dcd-d66d74b3ee3b",
   "metadata": {},
   "source": [
    "### Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "887222e1-02e2-457e-af0b-5616dcbef351",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Tokenizer():\n",
    "    \"\"\"A tokenizer class for encoding/decoding text sequences.\"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        \"\"\"\n",
    "        Constructor method to initialize special token indices and tokenizer encoding. \n",
    "        \"\"\"\n",
    "        # Initialize special token indices\n",
    "        self.BOS_IDX: int = 100264  # Index for the Beginning of Sentence token\n",
    "        self.EOS_IDX: int = 100265  # Index for the End of Sentence token\n",
    "        self.PAD_IDX: int = 100266  # Index for the Padding token\n",
    "\n",
    "        # Initialize base encoding from tiktoken\n",
    "        cl100k_base = tiktoken.get_encoding(\"cl100k_base\")\n",
    "\n",
    "        # Initialize the tokenizer's encoding with special tokens added\n",
    "        self.encoder = tiktoken.Encoding(\n",
    "            name=\"cl100k_bep\", # Name for the encoder with BOS, EOS, and PAD tokens added\n",
    "            pat_str=cl100k_base._pat_str, # Pattern string from the base encoding\n",
    "            mergeable_ranks=cl100k_base._mergeable_ranks, # Mergeable ranks from the base encoding\n",
    "            special_tokens={\n",
    "                **cl100k_base._special_tokens, # Special tokens from the base encoding\n",
    "                \"<|bos|>\": self.BOS_IDX,  # BOS token\n",
    "                \"<|eos|>\": self.EOS_IDX,  # EOS token\n",
    "                \"<|pad|>\": self.PAD_IDX,  # PAD token\n",
    "            }\n",
    "        )\n",
    "        \n",
    "    def vocab_size(self) -> int:\n",
    "        \"\"\"\n",
    "        Method to return the size of the vocabulary in the tokenizer's encoding.\n",
    "\n",
    "        Returns:\n",
    "            int: The size of the vocabulary.\n",
    "        \"\"\"\n",
    "        return self.encoder.n_vocab\n",
    "\n",
    "\n",
    "    def sequence_padding(self, sequence, max_size: int = 512, device: str = \"cpu\") -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Method to add BOS/PAD/EOS special tokens and ensure the sequence length is within the maximum size.\n",
    "\n",
    "        Args:\n",
    "            sequence (torch.Tensor or list): The input sequence.\n",
    "            max_size (int, optional): The maximum allowed size for the sequence. Defaults to 512.\n",
    "            device (str, optional): The device where the tensors will be allocated. Defaults to \"cpu\".\n",
    "\n",
    "        Returns:\n",
    "            torch.Tensor: The processed sequence with special tokens added and length limited.\n",
    "        \"\"\"\n",
    "        assert max_size > 2, f\"[max_size]: {max_size} should be greater than 2\"\n",
    "        # Ensure the sequence is a torch tensor\n",
    "        tensor_sequence = torch.tensor(sequence, dtype=torch.long).to(device) if not torch.is_tensor(sequence) else sequence.to(device)\n",
    "\n",
    "        # Limit the sequence length within (max_size - 2) where 2 corresponding to bos and eos tags\n",
    "        cutted_sequence_size = max(0, min(max_size - 2, tensor_sequence.size()[0]))\n",
    "        tensor_sequence = tensor_sequence[:cutted_sequence_size]\n",
    "        \n",
    "        # Add BOS token\n",
    "        tensor_sequence = torch.cat([torch.tensor([self.BOS_IDX], dtype=torch.long, device=device), tensor_sequence], dim=0)\n",
    "\n",
    "        # Calculate the padding size\n",
    "        padding_size = max_size - tensor_sequence.size()[0] - 1 # expected size - current size - EOS tag\n",
    "\n",
    "        # Create PAD tensor\n",
    "        pad_tensor = torch.full((padding_size,), self.PAD_IDX, dtype=torch.long, device=device)\n",
    "\n",
    "        # Add PAD and EOS tokens\n",
    "        tensor_sequence = torch.cat([tensor_sequence, torch.tensor([self.EOS_IDX], dtype=torch.long, device=device), pad_tensor], dim=0)\n",
    "        \n",
    "        return tensor_sequence\n",
    "    \n",
    "    def sequence_cleaner(self, sequence):\n",
    "        \"\"\" Method used to remove BOS/PAD/EOS special tokens \"\"\"\n",
    "        # Checking tensor format\n",
    "        list_sequence = sequence.tolist() if torch.is_tensor(sequence) else sequence\n",
    "        def check_special(number):\n",
    "            return number not in [self.BOS_IDX, self.EOS_IDX, self.PAD_IDX]\n",
    "        return list(filter(check_special, list_sequence))\n",
    "\n",
    "    def generate_padding_mask(self, seq):\n",
    "        # seq shape is (B, T) where B is batch size and T is sequence length\n",
    "        # padding mask should be of size (B, 1, 1, T), mask should be True for padding tokens and False for others\n",
    "        return (seq != self.PAD_IDX).unsqueeze(0).unsqueeze(0)\n",
    "\n",
    "    def tokenize(self, sequence, device=\"cpu\") -> list:\n",
    "        \"\"\"\n",
    "        Method to generate a str list of separated tokens token.\n",
    "\n",
    "        Args:\n",
    "            sequence (torch.Tensor or list): The input sequence.\n",
    "            device (str, optional): The device where the tensors will be allocated. Defaults to \"cpu\".\n",
    "\n",
    "        Returns:\n",
    "            list: The processed sequence converted in a list of tokens in string format.\n",
    "        \"\"\"\n",
    "        # Ensure the sequence is a torch tensor\n",
    "        tensor_sequence = torch.tensor(sequence, dtype=torch.long).to(device) if not torch.is_tensor(sequence) else sequence.to(device)\n",
    "        # create batch of idx tokens\n",
    "        tensor_sequence = tensor_sequence.unsqueeze(0).T\n",
    "        # decode all batch to recreate list of separated tokens \n",
    "        tensor_sequence = self.encoder.decode_batch(tensor_sequence.detach().tolist())\n",
    "        return tensor_sequence\n",
    "\n",
    "    def tokenize_from_str(self, sequence, device=\"cpu\") -> list:\n",
    "        return self.tokenize(self.encoder.encode(sequence), device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e243e01a-8b46-4105-aea3-fa7246df6b3a",
   "metadata": {},
   "source": [
    "### TransformerConfig\n",
    "Data class that stores the configuration for a Transformer model, and training "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "2b7d7325-63c2-4c89-887d-1d28c6288293",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class TransformerConfig:\n",
    "    \"\"\"Data class that stores the configuration for a Transformer model.\n",
    "\n",
    "    Attributes:\n",
    "        - tokenizer: An instance of the Tokenizer class.\n",
    "        - block_size (int): Number of tokens in each sequence. Defaults to 512.\n",
    "        - batch_size (int): Number of sequences in each batch. Defaults to 12.\n",
    "        - vocab_size (int): Total size of the tokenizer vocabulary. It is set to the size of the tokenizer vocabulary.\n",
    "        - n_layer (int): Number of transformer encoder and decoder blocks (N). Defaults to 1.\n",
    "        - n_head (int): Number of heads in each attention block. Defaults to 2.\n",
    "        - n_embd (int): Token embedding size. This is from the original Transformer paper. Defaults to 128.\n",
    "        - dropout (float): Dropout rate to use in the Transformer model. Defaults to 0.1.\n",
    "        - bias (bool): Indicates whether to use bias in Linears and LayerNorms.\n",
    "            If True, bias is used similar to GPT-2.\n",
    "            If False, it is a bit better and faster. Defaults to False.\n",
    "        - device (str): The device to run the model on. Defaults to 'cpu'. 'cuda' is used if a GPU is available.\n",
    "        - learning_rate (float): Learning rate for the model optimization. Defaults to 3e-4.\n",
    "        - max_epochs (int): Number of training epochs.\n",
    "        - max_iters (int): Number of training steps. Defaults to 20.\n",
    "        - eval_iters (int): Number of validation epochs. Defaults to 20.\n",
    "        - train_data_size (int): Size of train data.\n",
    "        - visualize (bool): Define if we want to get the attention scores.\n",
    "    \"\"\"\n",
    "    tokenizer: any\n",
    "    block_size: int = 256 # 512\n",
    "    batch_size: int = 12\n",
    "    n_layer: int = 2 # 6\n",
    "    n_head: int = 4 # 8\n",
    "    n_embd: int = 256 # 512\n",
    "    dropout: float = 0.1\n",
    "    bias: bool = False # True:\n",
    "    device: str = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "    learning_rate = 3e-4\n",
    "    max_epochs: int = 100\n",
    "    max_iters: int = 2000\n",
    "    eval_iters: int = 20 # 200\n",
    "    train_data_size: int = 5000000\n",
    "    visualize: bool = False\n",
    "\n",
    "    @property\n",
    "    def vocab_size(self) -> int:\n",
    "        \"\"\"Returns the total size of the tokenizer vocabulary.\n",
    "\n",
    "        Returns:\n",
    "            int: The size of the tokenizer vocabulary.\n",
    "        \"\"\"\n",
    "        return self.tokenizer.vocab_size()\n",
    "\n",
    "    def __str__(self):\n",
    "        return (\n",
    "            f\"TransformerConfig(\\n\"\n",
    "            f\"\\t{self.tokenizer=},\\n\"\n",
    "            f\"\\t{self.block_size=},\\n\"\n",
    "            f\"\\t{self.batch_size=},\\n\"\n",
    "            f\"\\t{self.n_layer=},\\n\"\n",
    "            f\"\\t{self.n_head=},\\n\"\n",
    "            f\"\\t{self.n_embd=},\\n\"\n",
    "            f\"\\t{self.dropout=},\\n\"\n",
    "            f\"\\t{self.bias=},\\n\"\n",
    "            f\"\\t{self.device=},\\n\"\n",
    "            f\"\\t{self.learning_rate=},\\n\"\n",
    "            f\"\\t{self.max_epochs=},\\n\"\n",
    "            f\"\\t{self.max_iters=},\\n\"\n",
    "            f\"\\t{self.eval_iters=},\\n\"\n",
    "            f\"\\t{self.train_data_size=},\\n\"\n",
    "            f\"\\t{self.visualize=},\\n\"\n",
    "            f\")\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b678d53d-fe07-471e-8c54-c250c4e010be",
   "metadata": {},
   "source": [
    "### Dataset\n",
    "Classes used to managed train val and test dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "7bb57475-96fe-42c1-9568-06926c2c8f44",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TranslationDataset(Dataset):\n",
    "    \"\"\"\n",
    "    A PyTorch Dataset class to be used in a PyTorch DataLoader to create batches.\n",
    "\n",
    "    Args:\n",
    "        - dataset (Dataset): a dataset from HuggingFace datasets library.\n",
    "        - tokenizer (Tokenizer): The custom tiktoken tokenizer used to encode sequences.\n",
    "        - block_size (int): The maximum sequence length for tokenization.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, dataset, tokenizer, block_size):\n",
    "        self.dataset = dataset\n",
    "        self.tokenizer = tokenizer\n",
    "        self.block_size = block_size\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        \"\"\"\n",
    "        Get a tokenized example from the dataset at the specified index.\n",
    "\n",
    "        Args:\n",
    "            - index (int): the index of the example to fetch.\n",
    "\n",
    "        Returns:\n",
    "            - Dict: dictionary with keys 'inputs', 'inputs_mask', 'targets' , 'targets_mask' and 'translation', containing tokenized input,\n",
    "            target sequences and original translation.\n",
    "        \"\"\"\n",
    "        translation = self.dataset[index]['translation']\n",
    "        encode = self.tokenizer.encoder.encode\n",
    "        inputs = self.tokenizer.sequence_padding(encode(translation['en']), self.block_size) # source language\n",
    "        inputs_mask = self.tokenizer.generate_padding_mask(inputs)\n",
    "        targets = self.tokenizer.sequence_padding(encode(translation['fr']), self.block_size) # target language\n",
    "        targets_mask = self.tokenizer.generate_padding_mask(targets)\n",
    "        return {\n",
    "            'inputs': inputs,\n",
    "            'inputs_mask': inputs_mask,\n",
    "            'targets': targets,\n",
    "            'targets_mask': targets_mask,\n",
    "            'translation': translation\n",
    "        }\n",
    "\n",
    "    def __len__(self) -> int :\n",
    "        \"\"\"\n",
    "        Returns the length of the dataset.\n",
    "\n",
    "        Returns:\n",
    "            - int: the length of the dataset.\n",
    "        \"\"\"\n",
    "        return self.dataset.num_rows\n",
    "\n",
    "\n",
    "class DataLoaderFactory():\n",
    "    \"\"\"\n",
    "    A class to instantiate PyTorch DataLoaders for different splits of a HuggingFace Dataset.\n",
    "\n",
    "    Args:\n",
    "        - block_size (int): The maximum sequence length for tokenization.\n",
    "        - batch_size (int): The batch size for DataLoader.\n",
    "        - tokenizer (Tokenizer): a tokenizer that has an encode method.\n",
    "        - device (str): 'cpu' or 'cuda', depending on whether we use CPU or GPU.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, block_size, batch_size, tokenizer, device, train_dataset_size = 5_000_000):\n",
    "        self.train_data = TranslationDataset(load_dataset(\"wmt14\", \"fr-en\", split=f\"train[:{train_dataset_size}]\"), tokenizer, block_size)\n",
    "        self.val_data = TranslationDataset(load_dataset(\"wmt14\", \"fr-en\", split=\"validation\"), tokenizer, block_size)\n",
    "        self.test_data = TranslationDataset(load_dataset(\"wmt14\", \"fr-en\", split=\"test\"), tokenizer, block_size)\n",
    "\n",
    "        self.block_size = block_size\n",
    "        self.batch_size = batch_size\n",
    "        self.device = device\n",
    "\n",
    "        self.dataloader_train = DataLoader(self.train_data, batch_size=batch_size, shuffle=True)\n",
    "        self.dataloader_val = DataLoader(self.val_data, batch_size=batch_size, shuffle=True)\n",
    "        self.dataloader_test = DataLoader(self.test_data, batch_size=batch_size, shuffle=True)\n",
    "\n",
    "    \n",
    "    \n",
    "    def __len__(self) -> int :\n",
    "        \"\"\"\n",
    "        Print the length of each dataset and returns the length of all datasets.\n",
    "\n",
    "        Returns:\n",
    "            - int: the length of all dataset (train + val + test).\n",
    "        \"\"\"\n",
    "        print(\"\\033[95m\\033[1m\\033[4mNumber of data by datasets splits\\033[0m\")\n",
    "        print(f\"Train\\t\\t: {len(self.train_data)}\\t-> {len(self.train_data)/self.batch_size}\")\n",
    "        print(f\"Validation\\t: {len(self.val_data)}\\t\\t-> {len(self.val_data)/self.batch_size}\")\n",
    "        print(f\"Test\\t\\t: {len(self.test_data)}\\t\\t-> {len(self.test_data)/self.batch_size}\")\n",
    "        total = len(self.train_data) + len(self.val_data) + len(self.test_data)\n",
    "        print(f\"Total\\t\\t: {total}\")\n",
    "        return total\n",
    "\n",
    "    def get_batch(self, split):\n",
    "        \"\"\"\n",
    "        Choose the correct DataLoader and yield batches from it.\n",
    "\n",
    "        Args:\n",
    "            - split (str): 'train', 'val' or 'test'.\n",
    "\n",
    "        Yields:\n",
    "            - Dict: a dictionary with keys 'inputs', 'targets' and 'translation', containing a batch of tokenized input,\n",
    "            target sequences and original translation.\n",
    "        \"\"\"\n",
    "        # choose the correct dataloader\n",
    "        if split == 'train':\n",
    "            dataloader = self.dataloader_train\n",
    "        elif split == 'val':\n",
    "            dataloader = self.dataloader_val\n",
    "        else:\n",
    "            dataloader = self.dataloader_test\n",
    "\n",
    "        for batch in dataloader:\n",
    "            # Separate the 'translation' from the rest of the batch\n",
    "            translation = batch.pop('translation')\n",
    "    \n",
    "            # Move tensors to device\n",
    "            batch_on_device = {k: v.to(self.device) for k, v in batch.items()}\n",
    "    \n",
    "            # Add 'translation' back into the batch\n",
    "            batch_on_device['translation'] = translation\n",
    "    \n",
    "            yield batch_on_device"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f955159-1950-4102-a05a-a33737cd647c",
   "metadata": {},
   "source": [
    "## Transformer Classes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ef92dcf-f8a9-4405-98ac-32eb758e097e",
   "metadata": {},
   "source": [
    "<div style=\"text-align:center\">\n",
    "  <img src=\"./img/transformer.jpg\" width=\"25%\" alt=\"transformer.jpg\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55c728a3-1709-4ce6-a285-605bcb75ae68",
   "metadata": {},
   "source": [
    "### LayerNorm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2ce6e629-fa89-4931-83d0-e62727dac927",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LayerNorm(nn.Module):\n",
    "    \"\"\"A Layer Normalization module with optional bias.\n",
    "\n",
    "    This implementation of Layer Normalization allows turning off the bias term,\n",
    "    which is not directly supported by PyTorch's layer normalization function.\n",
    "\n",
    "    Attributes:\n",
    "        - weight: Learnable weights for the layer normalization. Initialized as an all ones tensor.\n",
    "        - bias: Learnable biases for the layer normalization. Initialized as an all zeros tensor if bias argument in constructor is True, otherwise it's set to None.\n",
    "\n",
    "    Args:\n",
    "        - ndim: An integer for the dimension of the input vectors.\n",
    "        - bias: A boolean which, if True, adds a learnable bias to the output.\n",
    "\n",
    "    \"\"\"\n",
    "    def __init__(self, ndim: int, bias: bool):\n",
    "        super().__init__()\n",
    "        self.weight = nn.Parameter(torch.ones(ndim))\n",
    "        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None\n",
    "\n",
    "    def forward(self, input):\n",
    "        \"\"\"Defines the computation performed at every call.\n",
    "\n",
    "        Args:\n",
    "            - input (tensor): The input tensor to be normalized.\n",
    "\n",
    "        Returns:\n",
    "            - tensor: The normalized input tensor.\n",
    "\n",
    "        \"\"\"\n",
    "        return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23ce34e4-7325-4431-a255-fa805eed3bf1",
   "metadata": {},
   "source": [
    "### Multi Head Self Attention"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86dbce89-9828-4cfd-a575-9f17254582de",
   "metadata": {},
   "source": [
    "<p align=\"middle\" style=\"text-align:center\">\n",
    "  <img src=\"./img/Multi-Head_Attention.jpg\" width=\"40%\" alt=\"Multi-Head_Attention.jpg\"/>\n",
    "  <img src=\"./img/Scale-Dot-Product.jpg\" width=\"27%\" alt=\"Scale-Dot-Product.jpg\"/>\n",
    "</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "357d6e13-1c75-485f-9625-29a174b44d77",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    \"\"\"\n",
    "    Multi-Head Attention module.\n",
    "    \n",
    "    This module applies multi-head attention mechanism on the input sequence. This implementation doesn't apply mask over the attention scores.\n",
    "    \n",
    "    Attributes:\n",
    "        - n_head (int): Number of attention heads.\n",
    "        - n_embd (int): Embedding dimensionality.\n",
    "        - dropout (float): Dropout rate.\n",
    "        - q_attn (nn.Linear): Linear layer for the query projection.\n",
    "        - k_attn (nn.Linear): Linear layer for the key projection.\n",
    "        - v_attn (nn.Linear): Linear layer for the value projection.\n",
    "        - c_proj (nn.Linear): Linear layer for the output projection.\n",
    "        - attn_dropout (nn.Dropout): Dropout layer for the attention scores.\n",
    "        - resid_dropout (nn.Dropout): Dropout layer for the residual connection.\n",
    "        - flash (bool): Flag indicating if flash attention is available.\n",
    "    \"\"\"\n",
    "    def __init__(self, config):\n",
    "        \"\"\"\n",
    "        Constructor for the MultiHeadAttention class.\n",
    "        \n",
    "        Args:\n",
    "            - config: The configuration object containing model parameters.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        assert config.n_embd % config.n_head == 0\n",
    "\n",
    "        # Params\n",
    "        self.n_head = config.n_head\n",
    "        self.n_embd = config.n_embd\n",
    "        self.dropout = config.dropout\n",
    "        self.visualize = config.visualize\n",
    "        \n",
    "        # INPUTS: query, key, value projections for all heads, but in a batch\n",
    "        self.q_attn = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)\n",
    "        self.k_attn = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)\n",
    "        self.v_attn = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)\n",
    "        \n",
    "        # OUTPUT: output projection\n",
    "        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)\n",
    "        \n",
    "        # regularization\n",
    "        self.attn_dropout = nn.Dropout(config.dropout)\n",
    "        self.resid_dropout = nn.Dropout(config.dropout)\n",
    "        \n",
    "        # flash attention make GPU go br but support is only in PyTorch >= 2.0\n",
    "        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')\n",
    "        if not self.flash:\n",
    "            print(\"WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0\")\n",
    "        # causal mask to ensure that attention is only applied to the left in the input sequence\n",
    "        self.register_buffer(\n",
    "            \"bias\",\n",
    "            torch.tril(\n",
    "                torch.ones(config.block_size, config.block_size)\n",
    "            ).view(1, 1, config.block_size, config.block_size)\n",
    "        )\n",
    "\n",
    "    def scaled_dot_product_attention(self, q, k, v, mask: bool = None, is_casual = False):\n",
    "        \"\"\"\n",
    "        Computes the scaled dot product attention.\n",
    "        \n",
    "        Args:\n",
    "            - q (Tensor): Query tensor of shape (batch_size, num_heads, seq_length, emb_dim).\n",
    "            - k (Tensor): Key tensor of shape (batch_size, num_heads, seq_length, emb_dim).\n",
    "            - v (Tensor): Value tensor of shape (batch_size, num_heads, seq_length, emb_dim).\n",
    "            - mask (bool, optional): Flag indicating whether to apply mask on the attention scores.\n",
    "\n",
    "        Returns:\n",
    "            - y (Tensor): Output tensor after applying attention.\n",
    "            - attn_weights (list): Attention weights usefull to visualized how attention work\n",
    "        \"\"\"\n",
    "        # manual implementation of attention\n",
    "        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) # Step 1 & 2: (MatMul) and (Scale)\n",
    "        if mask is not None:\n",
    "            att = att.masked_fill(mask == 0, float('-inf'))\n",
    "        att = F.softmax(att, dim=-1) # Step 3: Softmax\n",
    "        att_weights = att  # Save attention weights for visualization\n",
    "        if self.training:\n",
    "            att = self.attn_dropout(att)\n",
    "        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs) # Step 4: MatMul\n",
    "        return y, att_weights\n",
    "\n",
    "    def forward(self, q_x, k_x, v_x, is_causal = False, mask = None):\n",
    "        \"\"\"\n",
    "        Forward pass for the MultiHeadAttention module.\n",
    "        \n",
    "        Args:\n",
    "            - q_x (Tensor): Input query tensor of shape (batch_size, seq_length, emb_dim).\n",
    "            - k_x (Tensor): Input key tensor of shape (batch_size, seq_length, emb_dim).\n",
    "            - v_x (Tensor): Input value tensor of shape (batch_size, seq_length, emb_dim).\n",
    "            - is_causal (bool, optional): Flag indicating whether to apply mask on the attention scores.\n",
    "            - mask (torch.Tensor, optional): The mask tensor to ignore padding, size (B, 1, 1, T).\n",
    "\n",
    "        Returns:\n",
    "            - y (Tensor): Output tensor after applying multi-head attention.\n",
    "            - attn_weights (list): Attention weights usefull to visualized how attention work\n",
    "        \"\"\"\n",
    "        B_q, T_q, C_q = q_x.size() # batch size, sequence length, embedding dimensionality (n_embd)\n",
    "        B_kv, T_kv, C_kv = k_x.size() # batch size, sequence length, embedding dimensionality (n_embd)\n",
    "\n",
    "        # calculate query, key, values for all heads in batch and move head forward to be the batch dim\n",
    "        q, k, v = self.q_attn(q_x), self.k_attn(k_x), self.v_attn(v_x)\n",
    "        k = k.view(B_kv, T_kv, self.n_head, C_kv // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n",
    "        q = q.view(B_q, T_q, self.n_head, C_q // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n",
    "        v = v.view(B_kv, T_kv, self.n_head, C_kv // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n",
    "\n",
    "        # in case of is_casual (decoder) we are mixing triangular and padding mask\n",
    "        if mask is not None:\n",
    "            if is_causal:\n",
    "                attn_mask = mask * self.bias[:, :, :mask.size(-1), :mask.size(-1)]\n",
    "        \n",
    "        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)\n",
    "        if self.flash and not self.visualize:\n",
    "            \n",
    "            # efficient attention using Flash Attention CUDA kernels\n",
    "            y = torch.nn.functional.scaled_dot_product_attention(\n",
    "                q, k, v,\n",
    "                attn_mask=attn_mask if mask is not None and is_causal else mask,\n",
    "                dropout_p=self.dropout if self.training else 0,\n",
    "                is_causal=mask is None and is_causal == True\n",
    "            )\n",
    "            attn_weights = None\n",
    "        else:\n",
    "            # manual implementation of attention\n",
    "            y, attn_weights = self.scaled_dot_product_attention(\n",
    "                q, k, v,\n",
    "                attn_mask if mask is not None and is_causal else mask,\n",
    "                is_causal\n",
    "            )\n",
    "        y = y.transpose(1, 2).contiguous().view(B_q, T_q, C_q) # re-assemble all head outputs side by side # Step 5: Concatenate\n",
    "        # output projection\n",
    "        y = self.resid_dropout(self.c_proj(y)) # Step 6 : Linear\n",
    "        return y, attn_weights"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb4ec876-098a-4734-af01-61fffde39634",
   "metadata": {},
   "source": [
    "### Feed Forward layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "eb156ccb-b9fd-4b62-b2a4-d2762cb71a12",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForward(nn.Module):\n",
    "    \"\"\"\n",
    "    A position-wise Feed Forward Neural Network (FFNN) class for transformer models.\n",
    "    \n",
    "    The class implementing a position-wise FFNN.\n",
    "    The FFNN consists of two linear transformations with a GELU activation in between, \n",
    "    followed by a dropout for regularization.\n",
    "\n",
    "    Attributes:\n",
    "        - c_fc (nn.Linear): First fully connected layer.\n",
    "        - gelu (nn.GELU): GELU activation function layer.\n",
    "        - c_proj (nn.Linear): Second fully connected layer.\n",
    "        - dropout (nn.Dropout): Dropout layer for regularization.\n",
    "\n",
    "    Args:\n",
    "        - config (Config): A configuration object with attribute `n_embd`, `bias`, and `dropout`.\n",
    "    \"\"\"\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)\n",
    "        self.gelu    = nn.GELU()\n",
    "        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)\n",
    "        self.dropout = nn.Dropout(config.dropout)\n",
    "\n",
    "    def forward(self, x) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Define the computation performed at every call.\n",
    "\n",
    "        Args:\n",
    "            - x (torch.Tensor): The input tensor to the forward pass.\n",
    "\n",
    "        Returns:\n",
    "            - torch.Tensor: The output of the FFNN.\n",
    "        \"\"\"\n",
    "        x = self.c_fc(x)\n",
    "        x = self.gelu(x)\n",
    "        x = self.c_proj(x)\n",
    "        x = self.dropout(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47ac0eb2-b777-42d0-b02d-341930f62a59",
   "metadata": {},
   "source": [
    "### Blocks\n",
    "#### Encoder block"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54a4c170-1681-49f8-8c59-ca63f45d5145",
   "metadata": {},
   "source": [
    "<p align=\"middle\" style=\"text-align:center\">\n",
    "  <img src=\"./img/EncoderBlock.jpg\" width=\"15%\" alt=\"EncoderBlock.jpg\"/>\n",
    "</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "d648bfa6-2687-4d48-b9b1-be60061f48b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderBlock(nn.Module):\n",
    "    \"\"\"\n",
    "    A class that implements a single encoder block in the Transformer model.\n",
    "    \n",
    "    Each block consists of two sub-layers: a multi-head self-attention mechanism,\n",
    "    and a position-wise fully connected feed-forward network. There is a residual \n",
    "    connection around each of the two sub-layers, followed by layer normalization.\n",
    "\n",
    "    Attributes:\n",
    "        - ln_1 (LayerNorm): Layer normalization before the multi-head attention layer.\n",
    "        - attn (MultiHeadAttention): Multi-head attention layer.\n",
    "        - ln_2 (LayerNorm): Layer normalization before the feed-forward network.\n",
    "        - ffw (FeedForward): Position-wise feed-forward network.\n",
    "\n",
    "    Args:\n",
    "        - config (Config): A configuration object with attribute `n_embd` and `bias`.\n",
    "    \"\"\"\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)\n",
    "        self.attn = MultiHeadAttention(config)\n",
    "        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)\n",
    "        self.ffw = FeedForward(config)\n",
    "\n",
    "    def forward(self, x, mask=None) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Defines the computation performed at every call.\n",
    "\n",
    "        Args:\n",
    "            - x (torch.Tensor): The input tensor to the forward pass.\n",
    "            - mask (torch.Tensor, optional): The mask tensor to ignore padding, size (B, 1, 1, T).\n",
    "\n",
    "        Returns:\n",
    "            - torch.Tensor: The output tensor of the block.\n",
    "            - decoder_attn: The attention weight of the current block.\n",
    "        \"\"\"\n",
    "        # MultiHeadAttention\n",
    "        x = self.ln_1(x)\n",
    "        x_attn, decoder_attn = checkpoint(self.attn, x, x, x, False, mask)\n",
    "        x = x + x_attn\n",
    "        # FeedForward\n",
    "        x = x + checkpoint(self.ffw, self.ln_2(x))\n",
    "        return x, decoder_attn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7c18f22-c552-4270-b193-25999fc079d8",
   "metadata": {},
   "source": [
    "#### Decoder block"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cdaa130-b999-4bb2-9118-b303678a8765",
   "metadata": {},
   "source": [
    "<p align=\"middle\" style=\"text-align:center\">\n",
    "  <img src=\"./img/DecoderBlock.jpg\" width=\"15%\" alt=\"DecoderBlock.jpg\"/>\n",
    "</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "4b0c31f8-35a6-4f31-a0ca-39943bff97cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderBlock(nn.Module):\n",
    "    \"\"\"\n",
    "    A class that implements a single decoder block in the Transformer model.\n",
    "\n",
    "    Each block consists of three sub-layers: a multi-head self-attention mechanism,\n",
    "    a multi-head attention mechanism over the encoder's output, and a position-wise \n",
    "    fully connected feed-forward network. There is a residual connection around \n",
    "    each of the three sub-layers, followed by layer normalization.\n",
    "\n",
    "    Attributes:\n",
    "        - ln_1 (LayerNorm): Layer normalization before the first multi-head attention layer.\n",
    "        - attn1 (MultiHeadAttention): First multi-head attention layer, with self-attention.\n",
    "        - ln_2 (LayerNorm): Layer normalization before the second multi-head attention layer.\n",
    "        - attn2 (MultiHeadAttention): Second multi-head attention layer, attends to encoder outputs.\n",
    "        - ln_3 (LayerNorm): Layer normalization before the feed-forward network.\n",
    "        - ffw (FeedForward): Position-wise feed-forward network.\n",
    "\n",
    "    Args:\n",
    "        config (Config): A configuration object with attribute `n_embd` and `bias`.\n",
    "    \"\"\"\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)\n",
    "        self.attn1 = MultiHeadAttention(config)\n",
    "        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)\n",
    "        self.attn2 = MultiHeadAttention(config)\n",
    "        self.ln_3 = LayerNorm(config.n_embd, bias=config.bias)\n",
    "        self.ffw = FeedForward(config)\n",
    "\n",
    "    def forward(self, x, encoder_output, mask=None) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Defines the computation performed at every call.\n",
    "\n",
    "        Args:\n",
    "            - x (torch.Tensor): The input tensor to the forward pass.\n",
    "            - encoder_output (torch.Tensor): The output tensor from the last encoder block.\n",
    "            - mask (torch.Tensor, optional): The mask tensor to ignore padding, size (B, 1, 1, T).\n",
    "\n",
    "        Returns:\n",
    "            - torch.Tensor: The output tensor of the block.\n",
    "            - encoder_attn: The encoder attention weight of the current block.\n",
    "            - cross_attn: The cross attention weight of the current block.\n",
    "        \"\"\"\n",
    "        # Masked MultiHeadAttention\n",
    "        x = self.ln_1(x)\n",
    "        x_attn, encoder_attn = checkpoint(self.attn1, x, x, x, True, mask)\n",
    "        x = x + x_attn\n",
    "        # MultiHeadAttention with q, k from encoder and x from decoder\n",
    "        x = self.ln_2(x)\n",
    "        sliced_encoder_output = encoder_output[:, :-1] if x.size(-2) != encoder_output.size(-2) else encoder_output\n",
    "        x_attn, cross_attn = checkpoint(self.attn2, x, sliced_encoder_output, sliced_encoder_output, False, mask)\n",
    "        x = x + x_attn\n",
    "        # FeedForward\n",
    "        x = x + checkpoint(self.ffw, self.ln_3(x))\n",
    "        return x, encoder_attn, cross_attn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a4fd44b-c4cc-4d08-9e1d-b48b650b7801",
   "metadata": {},
   "source": [
    "### Encoder"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b12e3bba-ba56-4f79-b93f-ccfbe728b782",
   "metadata": {},
   "source": [
    "<p align=\"middle\" style=\"text-align:center\">\n",
    "  <img src=\"./img/encoder.jpg\" width=\"15%\" alt=\"encoder.jpg\"/>\n",
    "</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "d415dfae-bc35-4f10-9fa0-cf182bf30a3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    \"\"\"\n",
    "    A class that implements the encoder part of the Transformer model.\n",
    "\n",
    "    The encoder consists of several EncoderBlocks arranged in sequence.\n",
    "    The input first goes through an embedding layer followed by a positional encoding layer.\n",
    "    The output of this is then passed through each EncoderBlock in sequence.\n",
    "\n",
    "    Attributes:\n",
    "        - encoder (nn.ModuleDict): A dictionary of modules making up the transformer encoder.\n",
    "\n",
    "    Args:\n",
    "        - config (Config): A configuration object with attributes such as `vocab_size`, `block_size`, `n_embd`, `dropout`, `n_layer`, and `bias`.\n",
    "    \"\"\"\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        assert config.vocab_size is not None\n",
    "        assert config.block_size is not None\n",
    "        self.config = config\n",
    "\n",
    "        self.encoder = nn.ModuleDict(dict(\n",
    "            wte = nn.Embedding(config.vocab_size, config.n_embd),\n",
    "            # Learned positional encoding:\n",
    "            # In this case, instead of using a fixed function to determine positional encoding,\n",
    "            # we initialize a tensor of positional encodings which gets updated during training via backpropagation.\n",
    "            # This method may potentially capture more complex position-related patterns than fixed positional encoding,\n",
    "            # but it also introduces additional parameters to the model.\n",
    "            wpe = nn.Embedding(config.block_size, config.n_embd),\n",
    "            drop = nn.Dropout(config.dropout),\n",
    "            h = nn.ModuleList([EncoderBlock(config) for _ in range(config.n_layer)]),\n",
    "            ln_f = LayerNorm(config.n_embd, bias=config.bias),\n",
    "        ))\n",
    "        \n",
    "        # init all weights\n",
    "        self.apply(self._init_weights)\n",
    "        \n",
    "        # apply special scaled init to the residual projections, based on GPT-2 paper\n",
    "        for pn, p in self.named_parameters():\n",
    "            if pn.endswith('c_proj.weight'):\n",
    "                # This initialization is used to preventing the variance of the outputs of each layer from exploding or vanishing\n",
    "                # during the forward pass through the network.\n",
    "                # Preventing \"vanishing/exploding gradients\" problem\n",
    "                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))\n",
    "\n",
    "        # report number of parameters\n",
    "        print(\"Number of Encoder parameters: %.2fM\" % (self.get_num_params()/1e6,))\n",
    "\n",
    "    def get_num_params(self, non_embedding: bool = True):\n",
    "        \"\"\"\n",
    "        Returns the number of parameters in the model.\n",
    "        For non-embedding count (default), the position embeddings get subtracted.\n",
    "        The token embeddings would too, except due to the parameter sharing these\n",
    "        params are actually used as weights in the final layer, so we include them.\n",
    "\n",
    "        Args:\n",
    "            -non_embedding (bool, optional): If True, excludes the position embeddings count from the total (Default is True).\n",
    "\n",
    "        Returns:\n",
    "            - int: The number of parameters in the model.\n",
    "        \"\"\"\n",
    "        n_params = sum(p.numel() for p in self.parameters())\n",
    "        if non_embedding:\n",
    "            n_params -= self.encoder.wpe.weight.numel()\n",
    "        return n_params\n",
    "\n",
    "    def _init_weights(self, module):\n",
    "        \"\"\"\n",
    "        Initializes the weights of the model. Proper weight initialization can help speed up the training process and improve model performance.\n",
    "\n",
    "        Args:\n",
    "            - module (nn.Module): The module of the model to be initialized.\n",
    "        \"\"\"\n",
    "        if isinstance(module, nn.Linear):\n",
    "            # init Linear layers with normal distribution (Gaussian initialization)\n",
    "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "            if module.bias is not None:\n",
    "                # bias initialization if necessary\n",
    "                torch.nn.init.zeros_(module.bias)\n",
    "        elif isinstance(module, nn.Embedding):\n",
    "            # init Embedding layers with normal distribution (Gaussian initialization)\n",
    "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "\n",
    "    def forward(self, idx, mask=None):\n",
    "        \"\"\"\n",
    "        Defines the computation performed at every call.\n",
    "\n",
    "        Args:\n",
    "            - idx (torch.Tensor): The input tensor to the forward pass.\n",
    "            - mask (torch.Tensor, optional): The mask tensor to ignore padding, size (B, 1, 1, T).\n",
    "\n",
    "        Returns:\n",
    "            - torch.Tensor: The output tensor (logits) of the model.\n",
    "            - list: all encoder layers attentions weights.\n",
    "        \"\"\"\n",
    "        device = idx.device\n",
    "        b, t = idx.size()\n",
    "        assert t <= self.config.block_size, f\"Cannot forward sequence of length {t}, block size is only {self.config.block_size}\"\n",
    "        pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)\n",
    "\n",
    "        # pre-encoder block\n",
    "        tok_emb = self.encoder.wte(idx) # token embeddings of shape (b, t, n_embd)\n",
    "        pos_emb = self.encoder.wpe(pos) # position embeddings of shape (t, n_embd)\n",
    "        x = self.encoder.drop(tok_emb + pos_emb) # Addition of input embd + positional encoding\n",
    "\n",
    "        # encoders block\n",
    "        encoder_attn_all = []\n",
    "        for block in self.encoder.h:\n",
    "            x, encoder_attn = block(x, mask)\n",
    "            encoder_attn_all.append(encoder_attn)\n",
    "\n",
    "        return  self.encoder.ln_f(x), encoder_attn_all"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbbc7cdd-a078-4442-a3f0-2dceca89e256",
   "metadata": {},
   "source": [
    "### Decoder"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b06bb5c-b666-4d60-a154-84050836688c",
   "metadata": {},
   "source": [
    "<p align=\"middle\" style=\"text-align:center\">\n",
    "  <img src=\"./img/decoder.jpg\" width=\"15%\" alt=\"decoder.jpg\"/>\n",
    "</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "5a0264d2-09df-4222-9cc3-945f5b4272af",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    \"\"\"\n",
    "    This class implements the decoder part of the Transformer model.\n",
    "\n",
    "    The Decoder consists of several DecoderBlocks arranged in sequence. The input first goes through an embedding \n",
    "    layer followed by a positional encoding layer. The output of this is then passed through each DecoderBlock in \n",
    "    sequence.\n",
    "\n",
    "    Attributes:\n",
    "        - decoder (nn.ModuleDict): A dictionary of modules making up the transformer decoder.\n",
    "        - lm_head (nn.Linear): The final linear layer mapping from the embedding dimension to the vocabulary size.\n",
    "        - config (:obj:`Config`): The configuration object for the transformer model.\n",
    "\n",
    "    .. note:: The weight of the embedding layer and the linear layer are shared.\n",
    "\n",
    "    Args:\n",
    "        - config (:obj:`Config`): The configuration object with attributes such as `vocab_size`, `block_size`, `n_embd`, `dropout`, `n_layer`, and `bias`.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        assert config.vocab_size is not None\n",
    "        assert config.block_size is not None\n",
    "        self.config = config\n",
    "\n",
    "        self.decoder = nn.ModuleDict(dict(\n",
    "            wte = nn.Embedding(config.vocab_size, config.n_embd),\n",
    "            # Learned positional encoding:\n",
    "            # In this case, instead of using a fixed function to determine positional encoding,\n",
    "            # we initialize a tensor of positional encodings which gets updated during training via backpropagation.\n",
    "            # This method may potentially capture more complex position-related patterns than fixed positional encoding,\n",
    "            # but it also introduces additional parameters to the model.\n",
    "            wpe = nn.Embedding(config.block_size, config.n_embd),\n",
    "            drop = nn.Dropout(config.dropout),\n",
    "            h = nn.ModuleList([DecoderBlock(config) for _ in range(config.n_layer)]),\n",
    "            ln_f = LayerNorm(config.n_embd, bias=config.bias),\n",
    "        ))\n",
    "        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)\n",
    "        # with weight tying when using torch.compile() some warnings get generated:\n",
    "        # \"UserWarning: functional_call was passed multiple values for tied weights.\n",
    "        # This behavior is deprecated and will be an error in future versions\"\n",
    "        # not 100% sure what this is, so far seems to be harmless. TODO investigate\n",
    "        self.decoder.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying\n",
    "\n",
    "        # init all weights\n",
    "        self.apply(self._init_weights)\n",
    "        # apply special scaled init to the residual projections, per GPT-2 paper\n",
    "        for pn, p in self.named_parameters():\n",
    "            if pn.endswith('c_proj.weight'):\n",
    "                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))\n",
    "\n",
    "        # report number of parameters\n",
    "        print(\"number of Decoder parameters: %.2fM\" % (self.get_num_params()/1e6,))\n",
    "\n",
    "    def get_num_params(self, non_embedding: bool = True) -> int:\n",
    "        \"\"\"\n",
    "        Returns the number of parameters in the model.\n",
    "        For non-embedding count (default), the position embeddings get subtracted.\n",
    "        The token embeddings would too, except due to the parameter sharing these\n",
    "        params are actually used as weights in the final layer, so we include them.\n",
    "\n",
    "        Args:\n",
    "            - non_embedding (bool): If True, excludes the position embeddings count from the total. Default is True.\n",
    "\n",
    "        Returns:\n",
    "            - int: The number of parameters in the model.\n",
    "        \"\"\"\n",
    "        n_params = sum(p.numel() for p in self.parameters())\n",
    "        if non_embedding:\n",
    "            n_params -= self.decoder.wpe.weight.numel()\n",
    "        return n_params\n",
    "\n",
    "    def _init_weights(self, module):\n",
    "        \"\"\"\n",
    "        Initializes the weights of the model.\n",
    "\n",
    "        Args:\n",
    "            - module (torch.nn.Module): The module of the model to be initialized.\n",
    "        \"\"\"\n",
    "        if isinstance(module, nn.Linear):\n",
    "            # init Linear layers with normal distribution (Gaussian initialization)\n",
    "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "            if module.bias is not None:\n",
    "                # bias initialization if necessary\n",
    "                torch.nn.init.zeros_(module.bias)\n",
    "        elif isinstance(module, nn.Embedding):\n",
    "            # init Embedding layers with normal distribution (Gaussian initialization)\n",
    "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "\n",
    "    def forward(self, idx, enc_output=None, mask=None):\n",
    "        \"\"\"\n",
    "        Defines the computation performed at every call.\n",
    "\n",
    "        Args:\n",
    "            - idx (torch.Tensor): The input tensor to the forward pass.\n",
    "            - enc_output (torch.Tensor): The output tensor from the encoder.\n",
    "            - mask (torch.Tensor, optional): The mask tensor to ignore padding, size (B, 1, 1, T).\n",
    "\n",
    "        Returns:\n",
    "            - torch.Tensor: The output tensor (logits) of the model.\n",
    "            - list: all layers of decoder attentions weights.\n",
    "            - list: all layers cross attentions weights.\n",
    "        \"\"\"\n",
    "        device = idx.device\n",
    "        b, t = idx.size()\n",
    "        \n",
    "        pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)\n",
    "        tok_emb = self.decoder.wte(idx) # token embeddings of shape (b, t, n_embd)\n",
    "        pos_emb = self.decoder.wpe(pos) # position embeddings of shape (t, n_embd)\n",
    "        x = self.decoder.drop(tok_emb + pos_emb) # Addition of input embd + positional encoding\n",
    "\n",
    "        cross_attn_all = []\n",
    "        decoder_attn_all = []\n",
    "        for block in self.decoder.h:\n",
    "            x, encoder_attn, cross_attn = block(x, enc_output, mask)\n",
    "            decoder_attn_all.append(encoder_attn)\n",
    "            cross_attn_all.append(cross_attn)\n",
    "            \n",
    "        x = self.decoder.ln_f(x)\n",
    "        return self.lm_head(x), decoder_attn_all, cross_attn_all"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11ca5944-fe85-4edf-9b90-bbe9c04a8bff",
   "metadata": {},
   "source": [
    "### Transformer Main class"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83a83a93-b3b3-4488-8b05-a8add81aba6e",
   "metadata": {},
   "source": [
    "<p align=\"middle\" style=\"text-align:center\">\n",
    "  <img src=\"./img/transformer.jpg\" width=\"25%\" alt=\"transformer.jpg\"/>\n",
    "</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "7d62cda5-5bc1-4bd9-a3e8-134bab8da831",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "    \"\"\"\n",
    "    This class implements the Transformer model, which includes both the encoder and decoder.\n",
    "\n",
    "    The Transformer is a sequence transduction model that uses attention mechanisms.\n",
    "    It is primarily used in tasks that require understanding of context or relationships among words in a text.\n",
    "\n",
    "    Attributes:\n",
    "        - encoder (Encoder): The transformer encoder.\n",
    "        - decoder (Decoder): The transformer decoder.\n",
    "        - config (:obj:`Config`): The configuration object for the transformer model.\n",
    "\n",
    "    Args:\n",
    "        - config (:obj:`Config`): The configuration object with attributes such as `vocab_size`, `block_size`, `n_embd`, `dropout`, `n_layer`, and `bias`.\n",
    "    \"\"\"\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        assert config.vocab_size is not None\n",
    "        assert config.block_size is not None\n",
    "        self.config = config\n",
    "        self.encoder = Encoder(config)\n",
    "        self.decoder = Decoder(config)\n",
    "\n",
    "        \n",
    "        # report number of parameters\n",
    "        print(\"Total number of parameters: %.2fM\" % (self.encoder.get_num_params()/1e6 + self.decoder.get_num_params()/1e6,))\n",
    "\n",
    "    def forward(self, src, tgt, src_mask=None, tgt_mask=None):\n",
    "        \"\"\"\n",
    "        Defines the computation performed at every call.\n",
    "\n",
    "        Args:\n",
    "            - src (torch.Tensor): The input tensor to the encoder.\n",
    "            - tgt (torch.Tensor): The input tensor to the decoder.\n",
    "            - src_mask (torch.Tensor): The input_mask tensor to the encoder, size (B, 1, 1, T).\n",
    "            - tgt_mask (torch.Tensor): The target_masks tensor to the decoder, size (B, 1, 1, T).\n",
    "\n",
    "        Returns:\n",
    "            - torch.Tensor: The output tensor (logits) of the model.\n",
    "            - torch.Tensor: The loss tensor calculated on the basis of the decoder's output and target tensor.\n",
    "        \"\"\"\n",
    "        assert src.dim() == 2 and tgt.dim() == 2, \"src and tgt should be 2D (B, S)\"\n",
    "        if src_mask is not None:\n",
    "            assert src_mask.dim() == 4, \"src_mask should be 4D (B, H, S, S)\"\n",
    "        if tgt_mask is not None:\n",
    "            assert tgt_mask.dim() == 4, \"tgt_mask should be 4D (B, H, S, S)\"\n",
    "\n",
    "        enc_output, _ = self.encoder(src, src_mask)\n",
    "        tgt_shifted = tgt[:, :-1] # Shifted target\n",
    "        output, _, _ = self.decoder(tgt_shifted, enc_output, tgt_mask[:, :, :, :-1])\n",
    "\n",
    "        # Calculate the loss, using both the output and the target\n",
    "        loss_fct = nn.CrossEntropyLoss(ignore_index=self.config.tokenizer.PAD_IDX) # Ignore padding tokens\n",
    "        # The targets for the loss function are the input sequences shifted\n",
    "        tgt_tgt = tgt[:, 1:].contiguous()\n",
    "        loss = loss_fct(output.view(-1, output.size(-1)), tgt_tgt.view(-1))\n",
    "        return output, loss\n",
    "        \n",
    "    @torch.no_grad()\n",
    "    def translate_beam_search(self, src, temperature=1.0, top_k=None, src_mask=None):\n",
    "        \"\"\"\n",
    "        Generates translations of the source sequences using beam search.\n",
    "\n",
    "        Args:\n",
    "            - src (torch.Tensor): The source sequences to translate.\n",
    "            - temperature (float): control the randomness of predictions.\n",
    "            - src_mask (torch.Tensor): The input_mask tensor to the encoder, size (B, 1, 1, T).\n",
    "\n",
    "        Returns:\n",
    "            - Tuple[torch.Tensor, Dict[str, torch.Tensor]]: The best sequence found by beam search and a dictionary containing the attention weights.\n",
    "        \"\"\"\n",
    "        enc_output, encoder_attn = self.encoder(src, src_mask)\n",
    "        # initialize beam with start token\n",
    "        idx = torch.full((src.size(0), 1), self.config.tokenizer.BOS_IDX).long().to(src.device)\n",
    "\n",
    "        for iter in range(self.config.block_size):\n",
    "            print(f\"\\r{iter+1}/{self.config.block_size}\", end=\"\")\n",
    "            output, dec_attention, cross_attention = self.decoder(idx, enc_output)\n",
    "\n",
    "            # get top k tokens probabilities\n",
    "            topk_prob, topk_indices = torch.topk(F.softmax(output[:, -1, :] / temperature, dim=-1), top_k)\n",
    "    \n",
    "            # sample from the top k tokens\n",
    "            idx_next = topk_indices[:, torch.multinomial(topk_prob, num_samples=1).squeeze(1)]\n",
    "            idx = torch.cat((idx, idx_next), dim=1)\n",
    "    \n",
    "            # Stop generating when EOS token is sampled\n",
    "            if idx[0][-1].item() == self.config.tokenizer.EOS_IDX:\n",
    "                break\n",
    "\n",
    "        return idx, dict(encoder_attn=encoder_attn, decoder_attn=dec_attention, cross_attn=cross_attention)\n",
    "\n",
    "    def save_model(self, path: str):\n",
    "        \"\"\"\n",
    "        Saves the current state of the model to a file.\n",
    "\n",
    "        Args:\n",
    "            path (str): The path to the file where the model state should be saved.\n",
    "        \"\"\"\n",
    "        torch.save(self.state_dict(), path)\n",
    "\n",
    "    def load_model(self, path: str):\n",
    "        \"\"\"\n",
    "        Loads the model state from a file.\n",
    "\n",
    "        Args:\n",
    "            path (str): The path to the file from where the model state should be loaded.\n",
    "\n",
    "        Raises:\n",
    "            ValueError: If the specified file does not exist.\n",
    "        \"\"\"\n",
    "        if not os.path.exists(path):\n",
    "            raise ValueError(f\"{path} does not exist.\")\n",
    "        self.load_state_dict(torch.load(path))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "661e3d7e-cb08-4dce-aae1-41fc9b5aceca",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "420d3a0a-3f1f-4cc9-a3ae-b08570f2d0a6",
   "metadata": {},
   "source": [
    "### Loss estimation function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "882759cf-c512-4701-acce-0da8efdc6491",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def estimate_loss(model, dataset, config, splits = ['train', 'val']):\n",
    "    \"\"\"\n",
    "    This function estimates the loss of a model on specified data splits without performing backpropagation.\n",
    "    It sets the model to evaluation mode, iterates over the data splits and calculates the average loss.\n",
    "\n",
    "    Args:\n",
    "        model (Transformer): The model for which loss needs to be estimated.\n",
    "        dataset (CustomDataset): The dataset used for estimation. It should provide a 'get_batch' method.\n",
    "        config (Config): The configuration object defining the number of evaluation iterations.\n",
    "        splits (list[str]): List of the names of data splits to use for estimation.\n",
    "\n",
    "    Returns:\n",
    "        out (dict): A dictionary with split names as keys and corresponding average loss as values.\n",
    "    \"\"\"\n",
    "    # Create an empty dictionary to store the average loss for each split\n",
    "    out = {}\n",
    "\n",
    "    # Set the model to evaluation mode\n",
    "    model.eval()\n",
    "\n",
    "    # Loop over the data splits\n",
    "    for split in splits:\n",
    "        # Initialize a tensor to store the losses for each iteration in the current split\n",
    "        losses = torch.zeros(config.eval_iters)\n",
    "\n",
    "        # Get a batch iterator for the current split\n",
    "        batch = dataset.get_batch(split)\n",
    "        \n",
    "        # Initialize a progress bar for the inner loop\n",
    "        inner_loop = tqdm(range(config.eval_iters), desc=f\"Evaluation - {split}\", leave=False)\n",
    "\n",
    "        # Start the inner loop\n",
    "        for k in inner_loop:\n",
    "            # Sample a new batch of data\n",
    "            n_batch = next(batch)\n",
    "            X = n_batch['inputs']\n",
    "            Y = n_batch['targets']\n",
    "            \n",
    "            # Evaluate the loss for the current batch\n",
    "            logits, loss = model(X, Y)\n",
    "\n",
    "            # Store the current loss\n",
    "            losses[k] = loss.item()\n",
    "            \n",
    "        # Calculate and store the mean loss for the current split\n",
    "        out[split] = losses.mean()\n",
    "\n",
    "    # Return the dictionary with the average losses\n",
    "    return out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cd996f5-da49-4824-98c8-5c3e859ca7dd",
   "metadata": {},
   "source": [
    "### Init training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "5fdf6a48-2441-4e6b-9070-fcf67b8751ba",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset wmt14 (C:/Users/thiba/.cache/huggingface/datasets/wmt14/fr-en/1.0.0/2de185b074515e97618524d69f5e27ee7545dcbed4aa9bc1a4235710ffca33f4)\n",
      "Found cached dataset wmt14 (C:/Users/thiba/.cache/huggingface/datasets/wmt14/fr-en/1.0.0/2de185b074515e97618524d69f5e27ee7545dcbed4aa9bc1a4235710ffca33f4)\n",
      "Found cached dataset wmt14 (C:/Users/thiba/.cache/huggingface/datasets/wmt14/fr-en/1.0.0/2de185b074515e97618524d69f5e27ee7545dcbed4aa9bc1a4235710ffca33f4)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[95m\u001b[1m\u001b[4mNumber of data by datasets splits\u001b[0m\n",
      "Train\t\t: 5000000\t-> 416666.6666666667\n",
      "Validation\t: 3000\t\t-> 250.0\n",
      "Test\t\t: 3003\t\t-> 250.25\n",
      "Total\t\t: 5006003\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "5006003"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# init tokenizer\n",
    "tokenizer = Tokenizer()\n",
    "# init config\n",
    "config = TransformerConfig(tokenizer)\n",
    "# loading dataset\n",
    "dataset = DataLoaderFactory(config.block_size, config.batch_size, tokenizer, config.device)\n",
    "len(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0026ea4b-292c-444e-bbb4-e39ae391a59d",
   "metadata": {},
   "source": [
    "### Define model and optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "cd10ccbf-da9f-4704-a8dc-37e1b64218b7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of Encoder parameters: 27.25M\n",
      "number of Decoder parameters: 27.77M\n",
      "Total number of parameters: 55.01M\n"
     ]
    }
   ],
   "source": [
    "# Create model\n",
    "model = Transformer(config)\n",
    "model.train()\n",
    "model = model.to(config.device)\n",
    "\n",
    "# create a PyTorch optimizer\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate, betas=(0.9, 0.98), eps=1e-9)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93abe559-b946-4596-92a0-aa94124dadbc",
   "metadata": {},
   "source": [
    "### Training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "29636df1-1fa6-4b0e-b3fb-ed5f6fad567a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.optim.lr_scheduler import LambdaLR\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "from . import estimate_loss\n",
    "\n",
    "# learning rate warmup and then decay, which is a standard practice in Transformer training.\n",
    "def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps):\n",
    "    \"\"\" \"warm-up, then decay\" strategy. \"\"\"\n",
    "    def lr_lambda(current_step):\n",
    "        if current_step < num_warmup_steps:\n",
    "            return float(current_step) / float(max(1, num_warmup_steps))\n",
    "        return max(0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps)))\n",
    "\n",
    "    return LambdaLR(optimizer, lr_lambda)\n",
    "\n",
    "def training_loop(model, optimizer, dataset, config, saved_path = \"./out/transformer_state_dict.pth\"):\n",
    "    \"\"\"\n",
    "    This function performs the training loop for the given transformer model. It trains the model using the provided \n",
    "    optimizer and dataset according to the specified configuration. \n",
    "\n",
    "    Args:\n",
    "        - model (Transformer): The transformer model to be trained.\n",
    "        - optimizer (torch.optim.Optimizer): The optimizer used to update the model's parameters.\n",
    "        - dataset (CustomDataset): The dataset used for training and validation. It should provide a 'get_batch' method.\n",
    "        - config (Config): The configuration object that defines parameters like max_iters.\n",
    "\n",
    "    Returns:\n",
    "        - losses_list (dict): A dictionary that contains the training and validation losses per evaluation step.\n",
    "    \"\"\"\n",
    "    # This is the total number of training steps,\n",
    "    # which is typically the number of training examples times the number of epochs.\n",
    "    num_training_steps = config.train_data_size * config.max_epochs\n",
    "    # Choose warmup_steps such that it's 1% of total steps\n",
    "    num_warmup_steps = num_training_steps // 100\n",
    "    # You can add this after defining your optimizer\n",
    "    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps)\n",
    "\n",
    "    # Initialize a dictionary to keep track of training and validation losses\n",
    "    losses_list = {\n",
    "        'train': [],\n",
    "        'val': [],\n",
    "    }\n",
    "    \n",
    "    # Initialize minimum loss with a high value and the iteration number where the minimum was observed\n",
    "    iter_saved = 0\n",
    "\n",
    "    # init early stop\n",
    "    best_loss = float('inf')\n",
    "    val_loss = float('inf')\n",
    "    # This is the number of epochs with no improvement after which training will be stopped.\n",
    "    patience = 5\n",
    "    # This is used to keep track of the number of epochs without improvement.\n",
    "    patience_counter = 0\n",
    "\n",
    "    # Initialize a progress bar for the outer training loop\n",
    "    outer_loop = tqdm(range(config.max_epochs), desc=\"Train loss nan, Val loss nan, Saved nan\", leave=True)\n",
    "\n",
    "    # Start the training loop\n",
    "    for epochs in outer_loop:\n",
    "        # Initialize a batch of data from the 'train' part of the dataset\n",
    "        iter_loop = tqdm(range(config.max_iters), leave=False)\n",
    "        batch = dataset.get_batch('train')\n",
    "        for iter in iter_loop:\n",
    "            # Sample a new batch of data\n",
    "            n_batch = next(batch)\n",
    "            xb = n_batch['inputs']\n",
    "            xb_mask = n_batch['inputs_mask']\n",
    "            yb = n_batch['targets']\n",
    "            yb_mask = n_batch['targets_mask']\n",
    "        \n",
    "            # Evaluate the loss\n",
    "            logits, loss = model(xb, yb, xb_mask, yb_mask)\n",
    "            # Reset gradients\n",
    "            optimizer.zero_grad(set_to_none=True)\n",
    "            # Perform backpropagation\n",
    "            loss.backward()\n",
    "            # Update the model parameters\n",
    "            optimizer.step()\n",
    "            # update sheduler\n",
    "            scheduler.step()\n",
    "\n",
    "        ############\n",
    "        # Evaluation\n",
    "        ############\n",
    "        # Estimate the losses for both training and validation datasets\n",
    "        losses = estimate_loss(model, dataset, config)\n",
    "        # Return the model to training mode\n",
    "        model.train()\n",
    "        \n",
    "        # Record the estimated losses\n",
    "        losses_list['train'].append(losses['train'])\n",
    "        losses_list['val'].append(losses['val'])\n",
    "                    \n",
    "        # Get the latest losses\n",
    "        last_loss_train = losses_list['train'][-1]\n",
    "        last_loss_val = losses_list['val'][-1]\n",
    "        \n",
    "        ############\n",
    "        # early stop\n",
    "        ############\n",
    "        val_loss = losses['train'] # TODO switch to val\n",
    "        if val_loss < best_loss:\n",
    "            torch.save(model.module.state_dict(), saved_path)\n",
    "            best_loss = val_loss\n",
    "            patience_counter = 0\n",
    "            iter_saved = epochs\n",
    "        else:\n",
    "            patience_counter += 1\n",
    "\n",
    "        # Update the description of the progress bar\n",
    "        outer_loop.set_description(f\"Train loss {last_loss_train:.4f}, Val loss {last_loss_val:.4f}, Saved {iter_saved}\")\n",
    "        \n",
    "        if patience_counter >= patience:\n",
    "            print(\"Early stopping\")\n",
    "            break\n",
    "    \n",
    "    # Return the list of losses\n",
    "    return losses_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "8a8c6051-cbef-4107-ab9b-323215a4561f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "085fafe759df411d8475c49ef23b7c39",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Train loss nan, Val loss nan, Saved nan:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "AttributeError",
     "evalue": "'TransformerConfig' object has no attribute 'eval_interval'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[26], line 1\u001b[0m\n\u001b[1;32m----> 1\u001b[0m losses_list \u001b[38;5;241m=\u001b[39m \u001b[43mtraining_loop\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptimizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdataset\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[1;32mIn[25], line 35\u001b[0m, in \u001b[0;36mtraining_loop\u001b[1;34m(model, optimizer, dataset, config)\u001b[0m\n\u001b[0;32m     32\u001b[0m \u001b[38;5;66;03m# Start the training loop\u001b[39;00m\n\u001b[0;32m     33\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m \u001b[38;5;28miter\u001b[39m \u001b[38;5;129;01min\u001b[39;00m outer_loop:\n\u001b[0;32m     34\u001b[0m     \u001b[38;5;66;03m# Every eval_interval iterations, compute and log the losses on the training and validation sets\u001b[39;00m\n\u001b[1;32m---> 35\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28miter\u001b[39m \u001b[38;5;241m%\u001b[39m \u001b[43mconfig\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43meval_interval\u001b[49m \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[0;32m     36\u001b[0m         \u001b[38;5;66;03m# Estimate the losses for both training and validation datasets\u001b[39;00m\n\u001b[0;32m     37\u001b[0m         losses \u001b[38;5;241m=\u001b[39m estimate_loss(model, dataset, config)\n\u001b[0;32m     38\u001b[0m         \u001b[38;5;66;03m# Return the model to training mode\u001b[39;00m\n",
      "\u001b[1;31mAttributeError\u001b[0m: 'TransformerConfig' object has no attribute 'eval_interval'"
     ]
    }
   ],
   "source": [
    "losses_list = training_loop(model, optimizer, dataset, config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c114f72a-17c7-4d3a-9f6f-9002e8b9c002",
   "metadata": {},
   "source": [
    "### Plotting loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0a43ec6-a1e0-40af-ac96-5c9f29ff5284",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_losses(losses):\n",
    "    \"\"\"\n",
    "    Plots the losses for training and validation with more vibrant colors on a black background.\n",
    "\n",
    "    Args:\n",
    "        - losses (dict): A dictionary containing 'train' and 'val' lists with losses recorded for each epoch.\n",
    "    \"\"\"\n",
    "    train_losses = losses['train']\n",
    "    val_losses = losses['val']\n",
    "\n",
    "    plt.style.use('dark_background')\n",
    "    plt.figure(figsize=(10,5))\n",
    "    plt.title(\"Training and Validation Losses per epoch\", color='white')\n",
    "    plt.plot(train_losses, label='Training loss', color='pink')\n",
    "    plt.plot(val_losses, label='Validation loss', color='lime')\n",
    "    plt.xlabel(\"Epochs\", color='white')\n",
    "    plt.ylabel(\"Loss\", color='white')\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c22141ec-130b-4084-b718-e13685ff516b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Call the function\n",
    "plot_losses(losses_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af644c8b-c006-44c4-90d4-7f88666214f8",
   "metadata": {},
   "source": [
    "## Testing\n",
    "### Init testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f902cf9-6dc6-4c18-93f8-e1ee081128ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = TransformerConfig(tokenizer, visualize=True)\n",
    "# First, you need to initialize the model\n",
    "model2 = Transformer(config)\n",
    "# Then, load the state dict\n",
    "model2.load_model('./out/transformer_state_dict.pth')\n",
    "# If you are ready to perform inference (and not training), put the model in evaluation mode\n",
    "model2.eval()\n",
    "print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0daf645-f79e-497c-810b-75a45466c383",
   "metadata": {},
   "source": [
    "### Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5561094-15a6-4203-adb4-ceec39306307",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loss = estimate_loss(model2, dataset, config, ['test'])\n",
    "test_loss['test'].item()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4347fc4d-75de-4b62-afa5-ea2c4cb3ed00",
   "metadata": {},
   "source": [
    "## Visualization of attentions\n",
    "### Specific sentence translation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e538ea60-8b86-4a98-9c86-15cd7c4a9c13",
   "metadata": {},
   "outputs": [],
   "source": [
    "def translate(sentences, tokenizer, model, config):\n",
    "    \"\"\"\n",
    "    This function tokenizes input sentences, translates them using the provided model,\n",
    "    and decodes the output into human-readable text. It also returns the attention dictionary from the model.\n",
    "\n",
    "    Args:\n",
    "        - sentences (list[str]): List of sentences to be translated.\n",
    "        - tokenizer (Tokenizer): Tokenizer used for encoding and decoding sequences.\n",
    "        - model (Transformer): The model used for translation.\n",
    "        - config (Config): The configuration object that defines parameters like block_size.\n",
    "\n",
    "    Returns:\n",
    "        - decode_output (list[str]): List of translated sentences.\n",
    "        - attn (dict): Dictionary containing attention information from the last layer of the model.\n",
    "    \"\"\"\n",
    "    # Tokenize sentences\n",
    "    tknzr = tokenizer.encoder\n",
    "    sequences = []\n",
    "\n",
    "    # Encode each sentence and add it to the list of sequences\n",
    "    for sentence in sentences:\n",
    "        sequence = tokenizer.sequence_padding(tknzr.encode(sentence), config.block_size).unsqueeze(dim=0)\n",
    "        sequences.append(sequence)\n",
    "\n",
    "    # Concatenate the sequences into a tensor\n",
    "    sequences = torch.cat(sequences, dim=0)\n",
    "\n",
    "    # Set the model to evaluation mode and translate sentences\n",
    "    model.eval()\n",
    "    outputs, attn = model.translate_beam_search(sequences)\n",
    "\n",
    "    # Initialize a list to store the decoded sentences\n",
    "    decode_output = []\n",
    "\n",
    "    # Decode each output sequence and add it to the list of decoded outputs\n",
    "    for output in outputs:\n",
    "        output = tokenizer.sequence_cleaner(output)\n",
    "        decode_output += [tknzr.decode(output)]\n",
    "\n",
    "    # Return the decoded sentences and the attention dictionary\n",
    "    return decode_output, attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "828d48ee-0257-4131-bb93-273414d21410",
   "metadata": {},
   "outputs": [],
   "source": [
    "input = ['I am a teacher.']\n",
    "# expected_output = ['Je suis un professeur.']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35d8b57a-1410-4f39-a560-fd1668fe7af4",
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs, attentions = translate(input, tokenizer, model2, config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a7d8520-80e1-431a-bc3c-025dc8b7b965",
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcd57494-708c-4e6e-b41c-f09a4f83a161",
   "metadata": {},
   "source": [
    "### Formatting attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bf86860-2717-4ff7-a638-1a5db0a99019",
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_attn(input, output, attentions, batch: int = 0):\n",
    "    \"\"\"\n",
    "    This function formats the attention outputs and tokenized inputs and outputs for easier interpretation and visualization.\n",
    "\n",
    "    Args:\n",
    "        - input (str): The original input sentence.\n",
    "        - output (str): The translated output sentence.\n",
    "        - attentions (dict): A dictionary containing the attention information from the model.\n",
    "        - batch (int, optional): The batch index to format. Defaults to 0.\n",
    "\n",
    "    Returns:\n",
    "        - tokens_input (list[str]): The tokenized input sentence, padded to max_len.\n",
    "        - tokens_output (list[str]): The tokenized output sentence, padded to max_len.\n",
    "        - tensor_encoder_attn (torch.Tensor): The attention tensor for the encoder, trimmed and reshaped.\n",
    "        - tensor_cross_attn (torch.Tensor): The cross-attention tensor, trimmed and reshaped.\n",
    "        - tensor_decoder_attn (torch.Tensor): The attention tensor for the decoder, trimmed and reshaped.\n",
    "    \"\"\"\n",
    "\n",
    "    # Stack the attention tensors along a new dimension\n",
    "    tensor_encoder_attn = torch.stack(attentions['encoder_attn'], dim=0)\n",
    "    tensor_cross_attn = torch.stack(attentions['cross_attn'], dim=0)\n",
    "    tensor_decoder_attn = torch.stack(attentions['decoder_attn'], dim=0)\n",
    "\n",
    "    # Tokenize the input and output sentences\n",
    "    tokens_input = tokenizer.tokenize_from_str(input[batch])\n",
    "    tokens_output = tokenizer.tokenize_from_str(output[batch])\n",
    "\n",
    "    # Find the maximum length of the input and output tokens\n",
    "    max_len = max(len(tokens_input), len(tokens_output))\n",
    "\n",
    "    # If the input tokens are shorter than the max length, pad with empty strings\n",
    "    if len(tokens_input) < max_len:\n",
    "        tokens_input = tokens_input + [''] * (max_len - len(tokens_input))\n",
    "    # Otherwise, pad the output tokens with empty strings\n",
    "    else:\n",
    "        tokens_output = tokens_output + [''] * (max_len - len(tokens_output))\n",
    "\n",
    "    # Trim and reshape the attention tensors\n",
    "    tensor_encoder_attn = tensor_encoder_attn[:, batch:batch+1, :, 1:max_len+1, 1:max_len+1] # layers, batch, heads, seq_len, seq_len\n",
    "    tensor_cross_attn = tensor_cross_attn[:, batch:batch+1, :, 1:max_len+1, 1:max_len+1] # layers, batch, heads, seq_len, seq_len\n",
    "    tensor_decoder_attn = tensor_decoder_attn[:, batch:batch+1, :, 1:max_len+1, 1:max_len+1] # layers, batch, heads, seq_len, seq_len\n",
    "\n",
    "    # Return the formatted tokens and attention tensors\n",
    "    return tokens_input, tokens_output, tensor_encoder_attn, tensor_cross_attn, tensor_decoder_attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82bb706e-012e-4f49-9ed2-7e8661f2a203",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokens_input,\\\n",
    "tokens_output,\\\n",
    "tensor_encoder_attn,\\\n",
    "tensor_cross_attn,\\\n",
    "tensor_decoder_attn = format_attn(input, outputs, attentions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "563f8fe7-8eec-4ea9-bd00-58d810adb0f0",
   "metadata": {},
   "source": [
    "### Bertviz Model view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e62b6d9c-02e6-4e50-a610-1ea24492c19d",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_view(\n",
    "    encoder_attention=tensor_encoder_attn,\n",
    "    decoder_attention=tensor_decoder_attn,\n",
    "    cross_attention=tensor_cross_attn,\n",
    "    encoder_tokens=tokens_input,\n",
    "    decoder_tokens=tokens_output\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29162ff1-5b68-4331-af8f-650e3412db86",
   "metadata": {},
   "source": [
    "### Bertviz Head view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a14a8b2-fd6e-43b9-981b-eebae0dbc987",
   "metadata": {},
   "outputs": [],
   "source": [
    "head_view(\n",
    "    encoder_attention=tensor_encoder_attn,\n",
    "    decoder_attention=tensor_decoder_attn,\n",
    "    cross_attention=tensor_cross_attn,\n",
    "    encoder_tokens=tokens_input,\n",
    "    decoder_tokens=tokens_output\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2841b166-1814-4a7e-8787-faa663c0eddd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}