5. Transformer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "64edd7ea",
   "metadata": {},
   "source": [
    "<div style=\"width: 100%; overflow: hidden;\">\n",
    "    <div style=\"width: 150px; float: left;\"> <img src=\"https://raw.githubusercontent.com/DataForScience/Networks/master/data/D4Sci_logo_ball.png\" alt=\"Data For Science, Inc\" align=\"left\" border=\"0\" width=150px> </div>\n",
    "    <div style=\"float: left; margin-left: 10px;\"> <h1>ChatGPT and Friends</h1>\n",
    "<h1>Transformer</h1>\n",
    "        <p>Bruno Gonçalves<br/>\n",
    "        <a href=\"http://www.data4sci.com/\">www.data4sci.com</a><br/>\n",
    "            @bgoncalves, @data4sci</p></div>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8790f6b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "\n",
    "from collections import Counter\n",
    "from pprint import pprint\n",
    "import time\n",
    "\n",
    "import math\n",
    "import copy\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt \n",
    "\n",
    "from ipywidgets import interact\n",
    "\n",
    "import os\n",
    "import gzip\n",
    "\n",
    "import tqdm as tq\n",
    "from tqdm.notebook import tqdm\n",
    "tqdm.pandas()\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import torchtext;\n",
    "from torchtext import data\n",
    "\n",
    "import spacy\n",
    "\n",
    "import networkx as nx\n",
    "\n",
    "import watermark\n",
    "\n",
    "%load_ext watermark\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f96cf266",
   "metadata": {},
   "source": [
    "We start by printing out the versions of the libraries we're using for future reference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f1fa0ecd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Python implementation: CPython\n",
      "Python version       : 3.11.7\n",
      "IPython version      : 8.12.3\n",
      "\n",
      "Compiler    : Clang 14.0.6 \n",
      "OS          : Darwin\n",
      "Release     : 23.6.0\n",
      "Machine     : arm64\n",
      "Processor   : arm\n",
      "CPU cores   : 16\n",
      "Architecture: 64bit\n",
      "\n",
      "Git hash: b44900a26f10de8fbaf559b307e69185828b77b4\n",
      "\n",
      "json      : 2.0.9\n",
      "spacy     : 3.7.4\n",
      "networkx  : 3.3\n",
      "numpy     : 1.26.4\n",
      "tqdm      : 4.66.4\n",
      "torch     : 2.3.0\n",
      "torchtext : 0.6.0\n",
      "watermark : 2.4.3\n",
      "pandas    : 2.2.3\n",
      "matplotlib: 3.8.0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%watermark -n -v -m -g -iv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45edbbc4",
   "metadata": {},
   "source": [
    "Load default figure style"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7da2e9ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.style.use('d4sci.mplstyle')\n",
    "colors = plt.rcParams['axes.prop_cycle'].by_key()['color']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd2b988f",
   "metadata": {},
   "source": [
    "# Code"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f39a5489",
   "metadata": {},
   "source": [
    "### Embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5a2e3d94",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Embedder(torch.nn.Module):\n",
    "    def __init__(self, vocab_size, d_model):\n",
    "        super().__init__()\n",
    "        self.embed = torch.nn.Embedding(vocab_size, d_model)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.embed(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcd13de8",
   "metadata": {},
   "source": [
    "### Positional encoding"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec6314e4",
   "metadata": {},
   "source": [
    "$$ PE_{(pos, 2i)} = sin(pos/10000^{2i/d_{model}}) $$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8f88e98",
   "metadata": {},
   "source": [
    "$$ PE_{(pos, 2i + 1)} = cos(pos/10000^{2i/d_{model}}) $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "518b7d6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEncoder(torch.nn.Module):\n",
    "    def __init__(self, d_model, max_seq_len=80):\n",
    "        super().__init__()\n",
    "        self.d_model = d_model\n",
    "        \n",
    "        # create constant positional encoding matrix\n",
    "        pe_matrix = torch.zeros(max_seq_len, d_model)\n",
    "        \n",
    "        for pos in range(max_seq_len):\n",
    "            for i in range(0, d_model, 2):\n",
    "                pe_matrix[pos, i] = math.sin(pos/10000**(2*i/d_model))\n",
    "                pe_matrix[pos, i+1] = math.cos(pos/10000**(2*i/d_model))\n",
    "        pe_matrix = pe_matrix.unsqueeze(0)     # Add one dimension for batch size\n",
    "        self.register_buffer('pe', pe_matrix)  # Register as persistent buffer\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # x is a sentence after embedding with dim (batch, number of words, vector dimension)\n",
    "        seq_len = x.size()[1]\n",
    "        x = x + self.pe[:, :seq_len]\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69f624ec",
   "metadata": {},
   "source": [
    "## Model layers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b7d0e4f",
   "metadata": {},
   "source": [
    "### Scaled Dot-Product Attention layer\n",
    "\n",
    "![](images/scaled_dot_product_attention.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "7cc28a09",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Given Query, Key, Value, calculate the final weighted value\n",
    "def scaled_dot_product_attention(q, k, v, mask=None, dropout=None):\n",
    "    # Shape of q and k are the same, both are (batch_size, seq_len, d_k)\n",
    "    # Shape of v is (batch_size, seq_len, d_v)\n",
    "    attention_scores = torch.matmul(q, k.transpose(-2, -1))/math.sqrt(q.shape[-1])  # size (batch_size, seq_len, seq_len)\n",
    "    \n",
    "    # Apply mask to scores\n",
    "    # <pad>\n",
    "    if mask is not None:\n",
    "        attention_scores = attention_scores.masked_fill(mask == 0, value=-1e9)\n",
    "        \n",
    "    # Softmax along the last dimension\n",
    "    attention_weights = F.softmax(attention_scores, dim=-1)\n",
    "    \n",
    "    if dropout is not None:\n",
    "        attention_weights = dropout(attention_weights)\n",
    "        \n",
    "    output = torch.matmul(attention_weights, v)\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfb3378e",
   "metadata": {},
   "source": [
    "### Multi-Head Attention layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6db322f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(torch.nn.Module):\n",
    "    def __init__(self, n_heads, d_model, dropout=0.1):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.n_heads = n_heads\n",
    "        self.d_model = d_model\n",
    "        self.d_k = self.d_v = d_model//n_heads\n",
    "        \n",
    "        # self attention linear layers\n",
    "        # Linear layers for q, k, v vectors generation in different heads\n",
    "        self.q_linear_layers = []\n",
    "        self.k_linear_layers = []\n",
    "        self.v_linear_layers = []\n",
    "        \n",
    "        for i in range(n_heads):\n",
    "            self.q_linear_layers.append(torch.nn.Linear(d_model, self.d_k))\n",
    "            self.k_linear_layers.append(torch.nn.Linear(d_model, self.d_k))\n",
    "            self.v_linear_layers.append(torch.nn.Linear(d_model, self.d_v))\n",
    "        \n",
    "        self.dropout = torch.nn.Dropout(dropout)\n",
    "        self.out = torch.nn.Linear(n_heads*self.d_v, d_model)\n",
    "        \n",
    "    def forward(self, q, k, v, mask=None):\n",
    "        multi_head_attention_outputs = []\n",
    "        for q_linear, k_linear, v_linear in zip(self.q_linear_layers,\n",
    "                                                self.k_linear_layers,\n",
    "                                                self.v_linear_layers):\n",
    "            new_q = q_linear(q)  # size: (batch_size, seq_len, d_k)\n",
    "            new_k = k_linear(k)  # size: (batch_size, seq_len, d_k)\n",
    "            new_v = v_linear(v)  # size: (batch_size, seq_len, d_v)\n",
    "            \n",
    "            # Scaled Dot-Product attention\n",
    "            head_v = scaled_dot_product_attention(new_q, new_k, new_v, mask, self.dropout)  # (batch_size, seq_len, d_v)\n",
    "            multi_head_attention_outputs.append(head_v)\n",
    "            \n",
    "        # Concat\n",
    "        #import pdb; pdb.set_trace()\n",
    "        concat = torch.cat(multi_head_attention_outputs, -1)  # (batch_size, seq_len, n_heads*d_v)\n",
    "        \n",
    "        # Linear layer to recover to original shap\n",
    "        output = self.out(concat)  # (batch_size, seq_len, d_model)\n",
    "        \n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6822e67",
   "metadata": {},
   "source": [
    "### Feed Forward layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "23cc0930",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForward(torch.nn.Module):\n",
    "    def __init__(self, d_model, d_ff=2048, dropout=0.1):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.linear_1 = torch.nn.Linear(d_model, d_ff)\n",
    "        self.dropout = torch.nn.Dropout(dropout)\n",
    "        self.linear_2 = torch.nn.Linear(d_ff, d_model)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.dropout(F.relu(self.linear_1(x)))\n",
    "        x = self.linear_2(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c0db77b",
   "metadata": {},
   "source": [
    "### Layer Normalization layer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39c98fac",
   "metadata": {},
   "source": [
    "#### Normalization\n",
    "\n",
    "$$\\mu = \\frac{1}{m} \\sum_{i=1}^{m}x_i$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5fdd3dc",
   "metadata": {},
   "source": [
    "$$\n",
    "\\sigma^{2} = \\frac{1}{m} \\sum^{m}_{i=1}(x_{i} - \\mu)^2\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "082cb48f",
   "metadata": {},
   "source": [
    "$$\n",
    "\\hat{Z}_i = \\frac{x_i - \\mu_i}{\\sqrt{\\sigma^{2}_{i} + \\epsilon}}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0ba6d0f",
   "metadata": {},
   "source": [
    "#### Add two learnable parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7cbd985",
   "metadata": {},
   "source": [
    "$$\n",
    "\\tilde{Z}_i = \\alpha_i * \\hat{Z}_i + \\beta_i\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "64adfc1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LayerNorm(torch.nn.Module):\n",
    "    def __init__(self, d_model, eps=1e-6):\n",
    "        super().__init__()\n",
    "        self.d_model = d_model\n",
    "        self.alpha = torch.nn.Parameter(torch.ones(self.d_model))\n",
    "        self.beta = torch.nn.Parameter(torch.zeros(self.d_model))\n",
    "        self.eps = eps\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # x size: (batch_size, seq_len, d_model)\n",
    "        x_hat = (x - x.mean(dim=-1, keepdim=True))/(x.std(dim=-1, keepdim=True) + self.eps)\n",
    "        x_tilde = self.alpha*x_hat + self.beta\n",
    "        return x_tilde"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1b8b1f0",
   "metadata": {},
   "source": [
    "## Encoder & Decoder layers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "244d8f51",
   "metadata": {},
   "source": [
    "### Encoder layer\n",
    "\n",
    "An encoder layer contains a multi-head attention layer and feed forward layer\n",
    "\n",
    "![](images/encoder.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5ca1aa94",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderLayer(torch.nn.Module):\n",
    "    def __init__(self, d_model, n_heads, dropout=0.1):\n",
    "        super().__init__()\n",
    "        self.d_model = d_model\n",
    "        self.norm_1 = LayerNorm(d_model)\n",
    "        self.norm_2 = LayerNorm(d_model)\n",
    "        self.multi_head_attention = MultiHeadAttention(n_heads, d_model)\n",
    "        self.feed_forward = FeedForward(d_model)\n",
    "        self.dropout_1 = torch.nn.Dropout(dropout)\n",
    "        self.dropout_2 = torch.nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, x, mask):\n",
    "        x = x + self.dropout_1(self.multi_head_attention(x, x, x, mask))\n",
    "        x = self.norm_1(x)\n",
    "        \n",
    "        x = x + self.dropout_2(self.feed_forward(x))\n",
    "        x = self.norm_2(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b857036",
   "metadata": {},
   "source": [
    "### Decoder layer\n",
    "\n",
    "An decoder layer contains two multi-head attention layers and one feed forward layer\n",
    "\n",
    "![](images/decoder.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9616fa16",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderLayer(torch.nn.Module):\n",
    "    def __init__(self, d_model, n_heads, dropout=0.1):\n",
    "        super().__init__()\n",
    "        self.norm_1 = LayerNorm(d_model)\n",
    "        self.norm_2 = LayerNorm(d_model)\n",
    "        self.norm_3 = LayerNorm(d_model)\n",
    "        \n",
    "        self.dropout_1 = torch.nn.Dropout(dropout)\n",
    "        self.dropout_2 = torch.nn.Dropout(dropout)\n",
    "        self.dropout_3 = torch.nn.Dropout(dropout)\n",
    "        \n",
    "        self.multi_head_attention_1 = MultiHeadAttention(n_heads, d_model)\n",
    "        self.multi_head_attention_2 = MultiHeadAttention(n_heads, d_model)\n",
    "        \n",
    "        self.feed_forward = FeedForward(d_model)\n",
    "        \n",
    "    def forward(self, x, encoder_output, src_mask, trg_mask):\n",
    "        x = self.dropout_1(self.multi_head_attention_1(x, x, x, trg_mask))\n",
    "        x = x + self.norm_1(x)\n",
    "        \n",
    "        x = self.dropout_2(self.multi_head_attention_2(x, encoder_output, encoder_output, src_mask))\n",
    "        x = x + self.norm_2(x)\n",
    "        \n",
    "        x = self.dropout_3(self.feed_forward(x))\n",
    "        x = x + self.norm_3(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c3e217f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def clone_layer(module, N):\n",
    "    return torch.nn.ModuleList([copy.deepcopy(module) for i in range(N)])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a72fa9f",
   "metadata": {},
   "source": [
    "## Encoder & Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "4d513e66",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(torch.nn.Module):\n",
    "    def __init__(self, vocab_size, d_model, N, n_heads):\n",
    "        super().__init__()\n",
    "        self.embed = Embedder(vocab_size, d_model)\n",
    "        self.pe = PositionalEncoder(d_model)\n",
    "        self.encoder_layers = clone_layer(EncoderLayer(d_model, n_heads), N)\n",
    "        self.norm = LayerNorm(d_model)\n",
    "        \n",
    "    def forward(self, src, mask):\n",
    "        x = self.embed(src)\n",
    "        x = self.pe(x)\n",
    "        for encoder in self.encoder_layers:\n",
    "            x = encoder(x, mask)\n",
    "        return self.norm(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "032849f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(torch.nn.Module):\n",
    "    def __init__(self, vocab_size, d_model, N, n_heads):\n",
    "        super().__init__()\n",
    "        self.embed = Embedder(vocab_size, d_model)\n",
    "        self.pe = PositionalEncoder(d_model)\n",
    "        self.decoder_layers = clone_layer(DecoderLayer(d_model, n_heads), N)\n",
    "        self.norm = LayerNorm(d_model)\n",
    "        \n",
    "    def forward(self, trg, encoder_output, src_mask, trg_mask):\n",
    "        x = self.embed(trg)\n",
    "        x = self.pe(x)\n",
    "        for decoder in self.decoder_layers:\n",
    "            x = decoder(x, encoder_output, src_mask, trg_mask)\n",
    "        return self.norm(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "093e7917",
   "metadata": {},
   "source": [
    "## Transformer\n",
    "\n",
    "![](images/transformer.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b78f211c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(torch.nn.Module):\n",
    "    def __init__(self, src_vocab_size, trg_vocab_size, d_model, N, n_heads):\n",
    "        super().__init__()\n",
    "        self.encoder = Encoder(src_vocab_size, d_model, N, n_heads)\n",
    "        self.decoder = Decoder(trg_vocab_size, d_model, N, n_heads)\n",
    "        self.linear = torch.nn.Linear(d_model, trg_vocab_size)\n",
    "        \n",
    "    def forward(self, src, trg, src_mask, trg_mask):\n",
    "        encoder_output = self.encoder(src, src_mask)\n",
    "        decoder_output = self.decoder(trg, encoder_output, src_mask, trg_mask)\n",
    "        output = self.linear(decoder_output)\n",
    "        \n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dc90152",
   "metadata": {},
   "source": [
    "# Data processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "fae31a73",
   "metadata": {},
   "outputs": [],
   "source": [
    "nlp = spacy.load('pt_core_news_sm')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "a1ef7833",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = lambda sentence: [tok.text for tok in nlp.tokenizer(sentence) if tok.text != \" \"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "34667518",
   "metadata": {},
   "outputs": [],
   "source": [
    "SRC = data.Field(lower=True, tokenize=tokenizer)\n",
    "TRG = data.Field(lower=True, tokenize=tokenizer, init_token='<sos>', eos_token='<eos>')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "92e6b379",
   "metadata": {},
   "outputs": [],
   "source": [
    "src_data = open('data/english.txt', 'r')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "e24d29b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "trg_data = open('data/french.txt', 'r')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "e90406cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_data = {'src': [line for line in src_data], 'trg': [line for line in trg_data]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "abe68f1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(raw_data, columns=['src', 'trg'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "a0d72149",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>src</th>\n",
       "      <th>trg</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Go.\\n</td>\n",
       "      <td>Va !\\n</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Run!\\n</td>\n",
       "      <td>Cours !\\n</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Run!\\n</td>\n",
       "      <td>Courez !\\n</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Fire!\\n</td>\n",
       "      <td>Au feu !\\n</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Help!\\n</td>\n",
       "      <td>À l'aide !\\n</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>154878</th>\n",
       "      <td>\"Top-down economics never works,\" said Obama. ...</td>\n",
       "      <td>« L'économie en partant du haut vers le bas, ç...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>154879</th>\n",
       "      <td>A carbon footprint is the amount of carbon dio...</td>\n",
       "      <td>Une empreinte carbone est la somme de pollutio...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>154880</th>\n",
       "      <td>Death is something that we're often discourage...</td>\n",
       "      <td>La mort est une chose qu'on nous décourage sou...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>154881</th>\n",
       "      <td>Since there are usually multiple websites on a...</td>\n",
       "      <td>Puisqu'il y a de multiples sites web sur chaqu...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>154882</th>\n",
       "      <td>If someone who doesn't know your background sa...</td>\n",
       "      <td>Si quelqu'un qui ne connaît pas vos antécédent...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>154883 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                      src  \\\n",
       "0                                                   Go.\\n   \n",
       "1                                                  Run!\\n   \n",
       "2                                                  Run!\\n   \n",
       "3                                                 Fire!\\n   \n",
       "4                                                 Help!\\n   \n",
       "...                                                   ...   \n",
       "154878  \"Top-down economics never works,\" said Obama. ...   \n",
       "154879  A carbon footprint is the amount of carbon dio...   \n",
       "154880  Death is something that we're often discourage...   \n",
       "154881  Since there are usually multiple websites on a...   \n",
       "154882  If someone who doesn't know your background sa...   \n",
       "\n",
       "                                                      trg  \n",
       "0                                                  Va !\\n  \n",
       "1                                               Cours !\\n  \n",
       "2                                              Courez !\\n  \n",
       "3                                              Au feu !\\n  \n",
       "4                                            À l'aide !\\n  \n",
       "...                                                   ...  \n",
       "154878  « L'économie en partant du haut vers le bas, ç...  \n",
       "154879  Une empreinte carbone est la somme de pollutio...  \n",
       "154880  La mort est une chose qu'on nous décourage sou...  \n",
       "154881  Puisqu'il y a de multiples sites web sur chaqu...  \n",
       "154882  Si quelqu'un qui ne connaît pas vos antécédent...  \n",
       "\n",
       "[154883 rows x 2 columns]"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "bd7de1b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "df.to_csv('data/en_to_fr.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "e03d40fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_fields = [('src', SRC), ('trg', TRG)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "73d3ee1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_set = data.TabularDataset('data/en_to_fr.csv', format='csv', fields=data_fields)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "bc0579fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "SRC.build_vocab(train_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "59fce309",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "14115"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(SRC.vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "ed11f91b",
   "metadata": {},
   "outputs": [],
   "source": [
    "TRG.build_vocab(train_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "3737a8aa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "28354"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(TRG.vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "dbdd08af",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torchtext.data.dataset.TabularDataset at 0x35cd0ebd0>"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_set"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f11ccfe9",
   "metadata": {},
   "source": [
    "# Train transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "017000e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set some parameters\n",
    "d_model = 512\n",
    "n_heads = 8\n",
    "N = 6\n",
    "src_vocab_size = len(SRC.vocab)\n",
    "trg_vocab_size = len(TRG.vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "50d166ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Transformer(src_vocab_size, trg_vocab_size, d_model, N, n_heads)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "105eeafd",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/lr/j1bs1q851k15cj5y777nxwph0000gn/T/ipykernel_17504/1058947583.py:3: UserWarning: nn.init.xavier_uniform is now deprecated in favor of nn.init.xavier_uniform_.\n",
      "  torch.nn.init.xavier_uniform(p)\n"
     ]
    }
   ],
   "source": [
    "for p in model.parameters():\n",
    "    if p.dim() > 1:\n",
    "        torch.nn.init.xavier_uniform(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "bb20010a",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "d95c5e69",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_iter = data.Iterator(train_set, batch_size=32, sort_key=lambda x: (len(x.src), len(x.trg)), shuffle=True, train=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "21f4ed6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_mask(src_input, trg_input):\n",
    "    # Source input mask\n",
    "    pad = SRC.vocab.stoi['<pad>']\n",
    "    src_mask = (src_input != pad).unsqueeze(1)\n",
    "    \n",
    "    # Target input mask\n",
    "    trg_mask = (trg_input != pad).unsqueeze(1)\n",
    "    \n",
    "    seq_len = trg_input.size(1)\n",
    "    nopeak_mask = np.tril(np.ones((1, seq_len, seq_len)), k=0).astype('uint8')\n",
    "    nopeak_mask = torch.from_numpy(nopeak_mask) != 0\n",
    "    trg_mask = trg_mask & nopeak_mask\n",
    "    \n",
    "    return src_mask, trg_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "20dec346",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(n_epochs, output_interval=100):\n",
    "    model.train()\n",
    "    start = time.time()\n",
    "    \n",
    "    for epoch in range(n_epochs):\n",
    "        total_loss = 0\n",
    "        for i, batch in tqdm(enumerate(train_iter)):\n",
    "            \n",
    "            src_input = batch.src.transpose(0, 1)  # size (batch_size, seq_len)\n",
    "            trg = batch.trg.transpose(0, 1)  # size (batch_size, seq_len)\n",
    "            \n",
    "            trg_input = trg[:, :-1]\n",
    "            ys = trg[:, 1:].contiguous().view(-1)\n",
    "            \n",
    "            # create src & trg masks\n",
    "            src_mask, trg_mask = create_mask(src_input, trg_input)\n",
    "            preds = model(src_input, trg_input, src_mask, trg_mask)\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "            loss = F.cross_entropy(preds.view(-1, preds.size(-1)), ys, ignore_index=1)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "            total_loss += loss.data.item()\n",
    "\n",
    "            if (i + 1) % output_interval == 0:\n",
    "                avg_loss = total_loss/output_interval\n",
    "                print('time = {}, epoch = {}, iter = {}, loss = {}'.format((time.time() - start)/60,\n",
    "                                                                           epoch + 1,\n",
    "                                                                           i + 1,\n",
    "                                                                           avg_loss))\n",
    "                total_loss = 0\n",
    "                start = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "d530a114",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "038299ac9d904ce29afdc2394bfb8838",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time = 0.0804331660270691, epoch = 1, iter = 10, loss = 9.759069347381592\n",
      "time = 0.08768448432286581, epoch = 1, iter = 20, loss = 9.148835849761962\n",
      "time = 0.07875199715296427, epoch = 1, iter = 30, loss = 8.58236484527588\n",
      "time = 0.09108409881591797, epoch = 1, iter = 40, loss = 8.055005025863647\n",
      "time = 0.09469588200251261, epoch = 1, iter = 50, loss = 7.459721231460572\n",
      "time = 0.10012988646825155, epoch = 1, iter = 60, loss = 6.964350795745849\n",
      "time = 0.08202158212661743, epoch = 1, iter = 70, loss = 6.53893551826477\n",
      "time = 0.09467024803161621, epoch = 1, iter = 80, loss = 6.396841096878052\n",
      "time = 0.08869445323944092, epoch = 1, iter = 90, loss = 6.188503885269165\n",
      "time = 0.0822692354520162, epoch = 1, iter = 100, loss = 6.092514514923096\n",
      "time = 0.10284863710403443, epoch = 1, iter = 110, loss = 6.092811155319214\n",
      "time = 0.09195224841435751, epoch = 1, iter = 120, loss = 6.04034423828125\n",
      "time = 0.0809195319811503, epoch = 1, iter = 130, loss = 6.0800868511199955\n",
      "time = 0.10584863026936848, epoch = 1, iter = 140, loss = 6.071158027648925\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[51], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m train_model(\u001b[38;5;241m3\u001b[39m, output_interval\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m10\u001b[39m)\n",
      "Cell \u001b[0;32mIn[48], line 21\u001b[0m, in \u001b[0;36mtrain_model\u001b[0;34m(n_epochs, output_interval)\u001b[0m\n\u001b[1;32m     19\u001b[0m optimizer\u001b[38;5;241m.\u001b[39mzero_grad()\n\u001b[1;32m     20\u001b[0m loss \u001b[38;5;241m=\u001b[39m F\u001b[38;5;241m.\u001b[39mcross_entropy(preds\u001b[38;5;241m.\u001b[39mview(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, preds\u001b[38;5;241m.\u001b[39msize(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m)), ys, ignore_index\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m)\n\u001b[0;32m---> 21\u001b[0m loss\u001b[38;5;241m.\u001b[39mbackward()\n\u001b[1;32m     22\u001b[0m optimizer\u001b[38;5;241m.\u001b[39mstep()\n\u001b[1;32m     24\u001b[0m total_loss \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m loss\u001b[38;5;241m.\u001b[39mdata\u001b[38;5;241m.\u001b[39mitem()\n",
      "File \u001b[0;32m/opt/anaconda3/lib/python3.11/site-packages/torch/_tensor.py:525\u001b[0m, in \u001b[0;36mTensor.backward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    515\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_unary(\u001b[38;5;28mself\u001b[39m):\n\u001b[1;32m    516\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[1;32m    517\u001b[0m         Tensor\u001b[38;5;241m.\u001b[39mbackward,\n\u001b[1;32m    518\u001b[0m         (\u001b[38;5;28mself\u001b[39m,),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    523\u001b[0m         inputs\u001b[38;5;241m=\u001b[39minputs,\n\u001b[1;32m    524\u001b[0m     )\n\u001b[0;32m--> 525\u001b[0m torch\u001b[38;5;241m.\u001b[39mautograd\u001b[38;5;241m.\u001b[39mbackward(\n\u001b[1;32m    526\u001b[0m     \u001b[38;5;28mself\u001b[39m, gradient, retain_graph, create_graph, inputs\u001b[38;5;241m=\u001b[39minputs\n\u001b[1;32m    527\u001b[0m )\n",
      "File \u001b[0;32m/opt/anaconda3/lib/python3.11/site-packages/torch/autograd/__init__.py:267\u001b[0m, in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    262\u001b[0m     retain_graph \u001b[38;5;241m=\u001b[39m create_graph\n\u001b[1;32m    264\u001b[0m \u001b[38;5;66;03m# The reason we repeat the same comment below is that\u001b[39;00m\n\u001b[1;32m    265\u001b[0m \u001b[38;5;66;03m# some Python versions print out the first line of a multi-line function\u001b[39;00m\n\u001b[1;32m    266\u001b[0m \u001b[38;5;66;03m# calls in the traceback and some print out the last line\u001b[39;00m\n\u001b[0;32m--> 267\u001b[0m _engine_run_backward(\n\u001b[1;32m    268\u001b[0m     tensors,\n\u001b[1;32m    269\u001b[0m     grad_tensors_,\n\u001b[1;32m    270\u001b[0m     retain_graph,\n\u001b[1;32m    271\u001b[0m     create_graph,\n\u001b[1;32m    272\u001b[0m     inputs,\n\u001b[1;32m    273\u001b[0m     allow_unreachable\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m    274\u001b[0m     accumulate_grad\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m    275\u001b[0m )\n",
      "File \u001b[0;32m/opt/anaconda3/lib/python3.11/site-packages/torch/autograd/graph.py:744\u001b[0m, in \u001b[0;36m_engine_run_backward\u001b[0;34m(t_outputs, *args, **kwargs)\u001b[0m\n\u001b[1;32m    742\u001b[0m     unregister_hooks \u001b[38;5;241m=\u001b[39m _register_logging_hooks_on_whole_graph(t_outputs)\n\u001b[1;32m    743\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 744\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m Variable\u001b[38;5;241m.\u001b[39m_execution_engine\u001b[38;5;241m.\u001b[39mrun_backward(  \u001b[38;5;66;03m# Calls into the C++ engine to run the backward pass\u001b[39;00m\n\u001b[1;32m    745\u001b[0m         t_outputs, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs\n\u001b[1;32m    746\u001b[0m     )  \u001b[38;5;66;03m# Calls into the C++ engine to run the backward pass\u001b[39;00m\n\u001b[1;32m    747\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    748\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m attach_logging_hooks:\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "train_model(3, output_interval=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ef9ec3b",
   "metadata": {},
   "source": [
    "<center>\n",
    "     <img src=\"https://raw.githubusercontent.com/DataForScience/Networks/master/data/D4Sci_logo_full.png\" alt=\"Data For Science, Inc\" align=\"center\" border=\"0\" width=300px> \n",
    "</center>"
   ]
  }
 ],
 "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.7"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}