Translator.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c1008ac5-b50d-48b7-a7bb-77654ea207a7",
   "metadata": {},
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec86144f",
   "metadata": {},
   "source": [
    "#### This class implements the Multi-Head Attention mechanism used in transformers.\n",
    "#### It allows the model to focus on different parts of the input sequence simultaneously by applying multiple attention heads. Each head independently computes attention, and the results are combined and passed through a linear layer to produce the final output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f540bd03",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from typing import Optional, Tuple\n",
    "\n",
    "class MultiHeadAttentionLayer(nn.Module):\n",
    "    def __init__(\n",
    "        self, hidden_size: int, n_heads: int, dropout: float, device: torch.device\n",
    "    ) -> None:\n",
    "        super().__init__()\n",
    "\n",
    "        assert (\n",
    "            hidden_size % n_heads == 0\n",
    "        ), \"Hidden size must be divisible by the number of heads.\"\n",
    "\n",
    "        self.hidden_size = hidden_size\n",
    "        self.n_heads = n_heads\n",
    "        self.head_size = hidden_size // n_heads\n",
    "\n",
    "        # Linear layers for query, key, and value projections\n",
    "        self.fc_query = nn.Linear(hidden_size, hidden_size)\n",
    "        self.fc_key = nn.Linear(hidden_size, hidden_size)\n",
    "        self.fc_value = nn.Linear(hidden_size, hidden_size)\n",
    "        self.fc_out = nn.Linear(hidden_size, hidden_size)\n",
    "\n",
    "        self.dp = nn.Dropout(dropout)\n",
    "\n",
    "        self.coefficient = torch.sqrt(torch.FloatTensor([self.head_size])).to(device)\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        query: torch.Tensor,\n",
    "        key: torch.Tensor,\n",
    "        value: torch.Tensor,\n",
    "        mask: Optional[torch.Tensor] = None,\n",
    "    ) -> Tuple[torch.Tensor, torch.Tensor]:\n",
    "        b_size = query.shape[0]\n",
    "\n",
    "        # Linear projections\n",
    "        query_output = self.fc_query(query)\n",
    "        key_output = self.fc_key(key)\n",
    "        value_output = self.fc_value(value)\n",
    "\n",
    "        # Reshape and permute for multi-head attention\n",
    "        query_output = query_output.view(\n",
    "            b_size, -1, self.n_heads, self.head_size\n",
    "        ).permute(0, 2, 1, 3)\n",
    "        key_output = key_output.view(b_size, -1, self.n_heads, self.head_size).permute(\n",
    "            0, 2, 1, 3\n",
    "        )\n",
    "        value_output = value_output.view(\n",
    "            b_size, -1, self.n_heads, self.head_size\n",
    "        ).permute(0, 2, 1, 3)\n",
    "\n",
    "        # Calculate attention scores\n",
    "        energy = (\n",
    "            torch.matmul(query_output, key_output.permute(0, 1, 3, 2))\n",
    "            / self.coefficient\n",
    "        )\n",
    "\n",
    "        if mask is not None:\n",
    "            energy = energy.masked_fill(mask == 0, -1e10)\n",
    "\n",
    "        # Apply softmax to get attention weights\n",
    "        attention = torch.softmax(energy, dim=-1)\n",
    "\n",
    "        # Calculate the weighted sum of values\n",
    "        output = torch.matmul(self.dp(attention), value_output)\n",
    "\n",
    "        # Concatenate heads and pass through the final linear layer\n",
    "        output = output.permute(0, 2, 1, 3).contiguous()\n",
    "        output = output.view(b_size, -1, self.hidden_size)\n",
    "        output = self.fc_out(output)\n",
    "\n",
    "        return output, attention"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7cdb573",
   "metadata": {},
   "source": [
    "#### The EncoderLayer class represents a single layer of the Transformer encoder, combining multi-head attention and a feed-forward neural network to process the input sequence.\n",
    "#### The Encoder class stacks multiple EncoderLayer instances to build the full encoder, producing a context-rich representation that captures relationships between tokens in the sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31aa2b5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderLayer(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        hidden_size: int,\n",
    "        n_heads: int,\n",
    "        ff_size: int,\n",
    "        dropout: float,\n",
    "        device: torch.device,\n",
    "    ) -> None:\n",
    "        super().__init__()\n",
    "\n",
    "        # Multi-head self-attention layer\n",
    "        self.self_atten = MultiHeadAttentionLayer(hidden_size, n_heads, dropout, device)\n",
    "        self.self_atten_norm = nn.LayerNorm(hidden_size)\n",
    "        self.ff_layer = FeedForwardLayer(hidden_size, ff_size, dropout)\n",
    "        self.dp = nn.Dropout(dropout)\n",
    "        self.ff_layer_norm = nn.LayerNorm(hidden_size)\n",
    "\n",
    "    def forward(self, input: torch.Tensor, input_mask: torch.Tensor) -> torch.Tensor:\n",
    "        # Self-attention\n",
    "        atten_result, _ = self.self_atten(input, input, input, input_mask)\n",
    "\n",
    "        # Add & norm\n",
    "        atten_norm = self.self_atten_norm(input + self.dp(atten_result))\n",
    "\n",
    "        # Feed-forward\n",
    "        ff_result = self.ff_layer(atten_norm)\n",
    "\n",
    "        # Add & norm\n",
    "        output = self.ff_layer_norm(atten_norm + self.dp(ff_result))\n",
    "\n",
    "        return output\n",
    "\n",
    "\n",
    "class Encoder(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_size: int,\n",
    "        hidden_size: int,\n",
    "        n_layers: int,\n",
    "        n_heads: int,\n",
    "        ff_size: int,\n",
    "        dropout: float,\n",
    "        device: torch.device,\n",
    "        MAX_LENGTH: int = 100,\n",
    "    ) -> None:\n",
    "        super().__init__()\n",
    "\n",
    "        self.device = device\n",
    "\n",
    "        # Embedding layers for tokens and positions\n",
    "        self.te = nn.Embedding(input_size, hidden_size)\n",
    "        self.pe = nn.Embedding(MAX_LENGTH, hidden_size)\n",
    "\n",
    "        # Stack of encoder layers\n",
    "        encoding_layers = [\n",
    "            EncoderLayer(hidden_size, n_heads, ff_size, dropout, device)\n",
    "            for _ in range(n_layers)\n",
    "        ]\n",
    "        self.encode_sequence = nn.Sequential(*encoding_layers)\n",
    "\n",
    "        self.dp = nn.Dropout(dropout)\n",
    "        self.coefficient = torch.sqrt(torch.FloatTensor([hidden_size])).to(device)\n",
    "\n",
    "    def forward(self, input: torch.Tensor, input_mask: torch.Tensor) -> torch.Tensor:\n",
    "        b_size, input_size = input.shape\n",
    "\n",
    "        # Create position tensor and add positional embeddings\n",
    "        pos = torch.arange(0, input_size).unsqueeze(0).repeat(b_size, 1).to(self.device)\n",
    "        input = self.dp((self.te(input) * self.coefficient) + self.pe(pos))\n",
    "\n",
    "        # Pass through each encoder layer\n",
    "        for layer in self.encode_sequence:\n",
    "            input = layer(input, input_mask)\n",
    "\n",
    "        return input"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7acb5b58",
   "metadata": {},
   "source": [
    "#### The DecoderLayer class represents a single layer of the Transformer decoder, combining multi-head attention with the encoder's output and a feed-forward neural network.\n",
    "#### The Decoder class stacks multiple DecoderLayer instances to build the full decoder, generating the output sequence by attending to both the encoded input and previously generated tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1b5191d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderLayer(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        hidden_size: int,\n",
    "        n_heads: int,\n",
    "        ff_size: int,\n",
    "        dropout: float,\n",
    "        device: torch.device,\n",
    "    ) -> None:\n",
    "        super().__init__()\n",
    "\n",
    "        # Self-attention and encoder-decoder attention layers\n",
    "        self.self_atten = MultiHeadAttentionLayer(hidden_size, n_heads, dropout, device)\n",
    "        self.self_atten_norm = nn.LayerNorm(hidden_size)\n",
    "        self.encoder_atten = MultiHeadAttentionLayer(\n",
    "            hidden_size, n_heads, dropout, device\n",
    "        )\n",
    "        self.encoder_atten_norm = nn.LayerNorm(hidden_size)\n",
    "        self.ff_layer = FeedForwardLayer(hidden_size, ff_size, dropout)\n",
    "        self.ff_layer_norm = nn.LayerNorm(hidden_size)\n",
    "        self.dp = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        target: torch.Tensor,\n",
    "        encoded_input: torch.Tensor,\n",
    "        target_mask: torch.Tensor,\n",
    "        input_mask: torch.Tensor,\n",
    "    ) -> Tuple[torch.Tensor, torch.Tensor]:\n",
    "        # Self-attention\n",
    "        atten_result, _ = self.self_atten(target, target, target, target_mask)\n",
    "        atten_norm = self.self_atten_norm(target + self.dp(atten_result))\n",
    "\n",
    "        # Encoder-decoder attention\n",
    "        atten_encoded, attention = self.encoder_atten(\n",
    "            atten_norm, encoded_input, encoded_input, input_mask\n",
    "        )\n",
    "        encoded_norm = self.encoder_atten_norm(atten_norm + self.dp(atten_encoded))\n",
    "\n",
    "        # Feed-forward\n",
    "        ff_result = self.ff_layer(encoded_norm)\n",
    "        output = self.ff_layer_norm(encoded_norm + self.dp(ff_result))\n",
    "\n",
    "        return output, attention\n",
    "\n",
    "\n",
    "class Decoder(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        output_size: int,\n",
    "        hidden_size: int,\n",
    "        n_layers: int,\n",
    "        n_heads: int,\n",
    "        ff_size: int,\n",
    "        dropout: float,\n",
    "        device: torch.device,\n",
    "        MAX_LENGTH: int = 100,\n",
    "    ) -> None:\n",
    "        super().__init__()\n",
    "\n",
    "        self.device = device\n",
    "\n",
    "        # Embedding layers for tokens and positions\n",
    "        self.te = nn.Embedding(output_size, hidden_size)\n",
    "        self.pe = nn.Embedding(MAX_LENGTH, hidden_size)\n",
    "\n",
    "        # Stack of decoder layers\n",
    "        decoding_layers = [\n",
    "            DecoderLayer(hidden_size, n_heads, ff_size, dropout, device)\n",
    "            for _ in range(n_layers)\n",
    "        ]\n",
    "        self.decode_sequence = nn.Sequential(*decoding_layers)\n",
    "\n",
    "        self.fc_out = nn.Linear(hidden_size, output_size)\n",
    "        self.dp = nn.Dropout(dropout)\n",
    "        self.coefficient = torch.sqrt(torch.FloatTensor([hidden_size])).to(device)\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        target: torch.Tensor,\n",
    "        encoded_input: torch.Tensor,\n",
    "        target_mask: torch.Tensor,\n",
    "        input_mask: torch.Tensor,\n",
    "    ) -> Tuple[torch.Tensor, torch.Tensor]:\n",
    "        b_size, target_size = target.shape\n",
    "\n",
    "        # Create position tensor and add positional embeddings\n",
    "        pos = (\n",
    "            torch.arange(0, target_size).unsqueeze(0).repeat(b_size, 1).to(self.device)\n",
    "        )\n",
    "        target = self.dp((self.te(target) * self.coefficient) + self.pe(pos))\n",
    "\n",
    "        # Pass through each decoder layer\n",
    "        for layer in self.decode_sequence:\n",
    "            target, attention = layer(target, encoded_input, target_mask, input_mask)\n",
    "\n",
    "        # Final linear layer to generate output predictions\n",
    "        output = self.fc_out(target)\n",
    "\n",
    "        return output, attention\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffb04dd3",
   "metadata": {},
   "source": [
    "#### This class implements the feed-forward neural network used in each layer of the Transformer.\n",
    "#### It consists of two linear transformations with a ReLU activation in between, applied independently to each position in the sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7993690",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForwardLayer(nn.Module):\n",
    "    def __init__(self, hidden_size: int, ff_size: int, dropout: float) -> None:\n",
    "        super().__init__()\n",
    "\n",
    "        # Feed-forward neural network with dropout and ReLU activation\n",
    "        self.ff_layer = nn.Sequential(\n",
    "            nn.Linear(hidden_size, ff_size),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(dropout),\n",
    "            nn.Linear(ff_size, hidden_size),\n",
    "        )\n",
    "\n",
    "    def forward(self, input: torch.Tensor) -> torch.Tensor:\n",
    "        return self.ff_layer(input)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75524622",
   "metadata": {},
   "source": [
    "#### This class defines the overall Transformer model, combining the encoder and decoder components.\n",
    "#### It processes input sequences through the encoder to generate context-rich representations, and then decodes these representations to produce the output sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c3489b07-2704-4c83-999a-4301ce4d299b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        encoder: Encoder,\n",
    "        decoder: Decoder,\n",
    "        device: torch.device,\n",
    "        padding_index: int = 0,\n",
    "    ) -> None:\n",
    "        super().__init__()\n",
    "\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.padding_index = padding_index\n",
    "        self.device = device\n",
    "\n",
    "    def make_input_mask(self, input: torch.Tensor) -> torch.Tensor:\n",
    "        # Create input mask to ignore padding tokens\n",
    "        input_mask = (input != self.padding_index).unsqueeze(1).unsqueeze(2)\n",
    "        return input_mask\n",
    "\n",
    "    def make_target_mask(self, target: torch.Tensor) -> torch.Tensor:\n",
    "        # Create target mask to ignore padding tokens and ensure autoregressive property\n",
    "        target_pad_mask = (target != self.padding_index).unsqueeze(1).unsqueeze(2)\n",
    "        target_sub_mask = torch.tril(\n",
    "            torch.ones((target.shape[1], target.shape[1]), device=self.device)\n",
    "        ).bool()\n",
    "        target_mask = target_pad_mask & target_sub_mask\n",
    "        return target_mask\n",
    "\n",
    "    def forward(\n",
    "        self, input: torch.Tensor, target: torch.Tensor\n",
    "    ) -> Tuple[torch.Tensor, torch.Tensor]:\n",
    "        input_mask = self.make_input_mask(input)\n",
    "        target_mask = self.make_target_mask(target)\n",
    "\n",
    "        # Encode input sequences\n",
    "        encoded_input = self.encoder(input, input_mask)\n",
    "\n",
    "        # Decode target sequences with encoded input\n",
    "        output, attention = self.decoder(target, encoded_input, target_mask, input_mask)\n",
    "\n",
    "        return output, attention\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecffa98b-f186-4f0e-860a-a608406a7d29",
   "metadata": {},
   "source": [
    "# Dictionary"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c46b5752",
   "metadata": {},
   "source": [
    "#### This class manages the vocabulary and token mappings for the Transformer model.\n",
    "#### It converts between words and their corresponding indices, facilitating encoding and decoding operations within the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a2dafe47-539f-4bf6-851d-0db37508780a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Dict\n",
    "\n",
    "# Special tokens\n",
    "PAD_TOKEN = 0\n",
    "SOS_TOKEN = 1\n",
    "EOS_TOKEN = 2\n",
    "UNK_TOKEN = 3\n",
    "\n",
    "\n",
    "class Dictionary:\n",
    "    def __init__(self, name: str) -> None:\n",
    "        \"\"\"\n",
    "        Initialize the dictionary with the given language name. The dictionary keeps track of words and their corresponding indices.\n",
    "\n",
    "        Args:\n",
    "            name (str): The name of the language.\n",
    "        \"\"\"\n",
    "        self.name = name\n",
    "        self.word2index: Dict[str, int] = {\n",
    "            \"<pad>\": PAD_TOKEN,\n",
    "            \"<sos>\": SOS_TOKEN,\n",
    "            \"<eos>\": EOS_TOKEN,\n",
    "            \"<unk>\": UNK_TOKEN,\n",
    "        }\n",
    "        self.word2count: Dict[str, int] = {}\n",
    "        self.index2word: Dict[int, str] = {\n",
    "            PAD_TOKEN: \"<pad>\",\n",
    "            SOS_TOKEN: \"<sos>\",\n",
    "            EOS_TOKEN: \"<eos>\",\n",
    "            UNK_TOKEN: \"<unk>\",\n",
    "        }\n",
    "        self.n_count: int = 4  # Count includes PAD, SOS, EOS, and UNK\n",
    "\n",
    "    def add_sentence(self, sentence: str) -> None:\n",
    "        \"\"\"\n",
    "        Add all words in a sentence to the dictionary.\n",
    "\n",
    "        Args:\n",
    "            sentence (str): The sentence whose words are to be added.\n",
    "        \"\"\"\n",
    "        for word in sentence.split(\" \"):\n",
    "            self.add_word(word)\n",
    "\n",
    "    def add_word(self, word: str) -> None:\n",
    "        \"\"\"\n",
    "        Add a word to the dictionary. If the word already exists, increment its count.\n",
    "\n",
    "        Args:\n",
    "            word (str): The word to be added.\n",
    "        \"\"\"\n",
    "        if word not in self.word2index:\n",
    "            self.word2index[word] = self.n_count\n",
    "            self.word2count[word] = 1\n",
    "            self.index2word[self.n_count] = word\n",
    "            self.n_count += 1\n",
    "        else:\n",
    "            self.word2count[word] += 1\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03d38e38-35ce-48da-991e-d77fc313a32f",
   "metadata": {},
   "source": [
    "# Utilities"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01fe5d43",
   "metadata": {},
   "source": [
    "#### This cell contains utility functions and helper methods that support various operations within the Transformer model, such as converting UNICODE TO ASCII, normalizing string, loading files & batches and tokenization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1c571b65-8b5f-4290-a146-907f73fb1dda",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import unicodedata\n",
    "import csv\n",
    "from typing import List, Tuple\n",
    "import torch\n",
    "\n",
    "PAD_TOKEN = 0\n",
    "SOS_TOKEN = 1\n",
    "EOS_TOKEN = 2\n",
    "\n",
    "\n",
    "def unicodeToAscii(s: str) -> str:\n",
    "    return \"\".join(\n",
    "        c for c in unicodedata.normalize(\"NFD\", s) if unicodedata.category(c) != \"Mn\"\n",
    "    )\n",
    "\n",
    "\n",
    "def normalizeString(s: str) -> str:\n",
    "    s = unicodeToAscii(s.lower().strip())\n",
    "    s = re.sub(r\"([.!?])\", r\" \\1\", s)\n",
    "    s = re.sub(r\"[^a-zA-Z.!?]+\", r\" \", s)\n",
    "    return s\n",
    "\n",
    "\n",
    "def load_files(\n",
    "    lang1: str,\n",
    "    lang2: str,\n",
    "    data_dir: str,\n",
    "    reverse: bool = True,\n",
    "    MAX_FILE_SIZE: int = 100000,\n",
    "    MAX_LENGTH: int = 60,\n",
    ") -> Tuple[Dictionary, Dictionary, List[str], List[str]]:\n",
    "    lang1_list = []\n",
    "    lang2_list = []\n",
    "\n",
    "    # Assume there's a single CSV file in the data_dir\n",
    "    csv_file_path = None\n",
    "    for root, _, files in os.walk(data_dir):\n",
    "        for file_name in files:\n",
    "            if file_name.endswith(\".csv\"):\n",
    "                csv_file_path = os.path.join(root, file_name)\n",
    "                break\n",
    "        if csv_file_path:\n",
    "            break\n",
    "\n",
    "    if not csv_file_path:\n",
    "        raise FileNotFoundError(f\"CSV file not found in {data_dir}\")\n",
    "\n",
    "    # Read all lines first to calculate interval\n",
    "    all_lang1_lines = []\n",
    "    all_lang2_lines = []\n",
    "\n",
    "    with open(csv_file_path, mode=\"r\", encoding=\"utf-8\") as csv_file:\n",
    "        reader = csv.DictReader(csv_file)\n",
    "        for row in reader:\n",
    "            all_lang1_lines.append(row[lang1.capitalize()].strip())\n",
    "            all_lang2_lines.append(row[lang2.capitalize()].strip())\n",
    "\n",
    "    # Ensure both files have the same number of lines\n",
    "    assert len(all_lang1_lines) == len(\n",
    "        all_lang2_lines\n",
    "    ), \"Mismatched number of lines in language files\"\n",
    "\n",
    "    # Determine the interval to sample lines\n",
    "    interval = max(1, len(all_lang1_lines) // MAX_FILE_SIZE)\n",
    "\n",
    "    # Select lines based on the interval\n",
    "    lang1_list = [all_lang1_lines[i] for i in range(0, len(all_lang1_lines), interval)]\n",
    "    lang2_list = [all_lang2_lines[i] for i in range(0, len(all_lang2_lines), interval)]\n",
    "\n",
    "    # Limit the number of selected lines to MAX_FILE_SIZE\n",
    "    lang1_list = lang1_list[:MAX_FILE_SIZE]\n",
    "    lang2_list = lang2_list[:MAX_FILE_SIZE]\n",
    "\n",
    "    # Debugging: Print the length of loaded lists\n",
    "    print(f\"Loaded {len(lang1_list)} sentences for {lang1}\")\n",
    "    print(f\"Loaded {len(lang2_list)} sentences for {lang2}\")\n",
    "\n",
    "    # Preprocess strings\n",
    "    lang1_normalized = list(map(normalizeString, lang1_list))\n",
    "    lang2_normalized = list(map(normalizeString, lang2_list))\n",
    "\n",
    "    lang1_sentences = []\n",
    "    lang2_sentences = []\n",
    "\n",
    "    for i in range(len(lang1_normalized)):\n",
    "        tokens1 = lang1_normalized[i].split(\" \")\n",
    "        tokens2 = lang2_normalized[i].split(\" \")\n",
    "        if len(tokens1) <= MAX_LENGTH and len(tokens2) <= MAX_LENGTH:\n",
    "            lang1_sentences.append(lang1_normalized[i])\n",
    "            lang2_sentences.append(lang2_normalized[i])\n",
    "\n",
    "    # Debugging: Print the number of sentences after filtering by length\n",
    "    print(f\"{len(lang1_sentences)} {lang1} sentences after length filtering\")\n",
    "    print(f\"{len(lang2_sentences)} {lang2} sentences after length filtering\")\n",
    "\n",
    "    if reverse:\n",
    "        input_dic = Dictionary(lang2)\n",
    "        output_dic = Dictionary(lang1)\n",
    "        return input_dic, output_dic, lang2_sentences, lang1_sentences\n",
    "    else:\n",
    "        input_dic = Dictionary(lang1)\n",
    "        output_dic = Dictionary(lang2)\n",
    "        return input_dic, output_dic, lang1_sentences, lang2_sentences\n",
    "\n",
    "\n",
    "\n",
    "def tokenize(sentence: str, dictionary: Dictionary, MAX_LENGTH: int = 60) -> List[int]:\n",
    "    split_sentence = [word for word in sentence.split(\" \")]\n",
    "    token = [SOS_TOKEN]\n",
    "    token += [\n",
    "        dictionary.word2index.get(word, dictionary.word2index[\"<unk>\"])\n",
    "        for word in sentence.split(\" \")\n",
    "    ]\n",
    "    token.append(EOS_TOKEN)\n",
    "    token += [PAD_TOKEN] * (MAX_LENGTH - len(split_sentence))\n",
    "    return token\n",
    "\n",
    "\n",
    "def load_batches(\n",
    "    input_lang: List[List[int]],\n",
    "    output_lang: List[List[int]],\n",
    "    batch_size: int,\n",
    "    device: torch.device,\n",
    ") -> List[Tuple[torch.Tensor, torch.Tensor]]:\n",
    "    data_loader = []\n",
    "    for i in range(0, len(input_lang), batch_size):\n",
    "        input_batch = input_lang[i : i + batch_size]\n",
    "        target_batch = output_lang[i : i + batch_size]\n",
    "\n",
    "        if len(input_batch) == 0 or len(target_batch) == 0:\n",
    "            continue\n",
    "\n",
    "        input_tensor = torch.LongTensor(input_batch).to(device)\n",
    "        target_tensor = torch.LongTensor(target_batch).to(device)\n",
    "        data_loader.append([input_tensor, target_tensor])\n",
    "    return data_loader\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1117eb5b-b7de-41ce-a62d-5f7831420600",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2abd065",
   "metadata": {},
   "source": [
    "#### This class handles the training process for the Transformer model, including the forward pass, loss computation, backpropagation, and optimization.\n",
    "#### It manages training epochs, tracks performance metrics, and saves model checkpoints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "409526bc-2b3b-4faf-93c7-c9a1f58946aa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded 200000 sentences for english\n",
      "Loaded 200000 sentences for french\n",
      "200000 english sentences after length filtering\n",
      "200000 french sentences after length filtering\n",
      "Sample tokenized input sentences:\n",
      "[1, 4, 5, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "[1, 4, 5, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "[1, 4, 5, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "[1, 4, 5, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "[1, 6, 5, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "Sample tokenized output sentences:\n",
      "[1, 4, 5, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "[1, 6, 7, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "[1, 8, 9, 5, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "[1, 10, 5, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "[1, 11, 5, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "Input vocabulary size: 13255\n",
      "Output vocabulary size: 21746\n",
      "Starting fold 1/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-08-08 07:38:40,061 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 07:38:40,063 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 07:38:40,065 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 07:38:40,184 - INFO - Fold 1/5, Epoch 1/10, Time: 85.1s, Estimated remaining time: 4168.2s\n",
      "2024-08-08 07:38:40,186 - INFO -   Training Loss: 3.6934, Perplexity: 40.1820, BLEU Score: 1.45\n",
      "2024-08-08 07:38:40,450 - INFO - Model saved to ./saved_models/english2french/transformer_epoch_0.pth\n",
      "2024-08-08 07:41:36,731 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 07:41:36,733 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 07:41:36,734 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 07:41:36,861 - INFO - Fold 1/5, Epoch 2/10, Time: 80.3s, Estimated remaining time: 3853.0s\n",
      "2024-08-08 07:41:36,862 - INFO -   Training Loss: 1.7606, Perplexity: 5.8162, BLEU Score: 3.84\n",
      "2024-08-08 07:44:35,609 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 07:44:35,613 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 07:44:35,614 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 07:44:35,735 - INFO - Fold 1/5, Epoch 3/10, Time: 80.4s, Estimated remaining time: 3777.5s\n",
      "2024-08-08 07:44:35,736 - INFO -   Training Loss: 1.1238, Perplexity: 3.0765, BLEU Score: 4.54\n",
      "2024-08-08 07:47:32,128 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 07:47:32,132 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 07:47:32,133 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 07:47:32,344 - INFO - Fold 1/5, Epoch 4/10, Time: 80.4s, Estimated remaining time: 3699.9s\n",
      "2024-08-08 07:47:32,346 - INFO -   Training Loss: 0.8613, Perplexity: 2.3662, BLEU Score: 4.77\n",
      "2024-08-08 07:50:27,724 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 07:50:27,726 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 07:50:27,728 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 07:50:27,934 - INFO - Fold 1/5, Epoch 5/10, Time: 80.7s, Estimated remaining time: 3631.2s\n",
      "2024-08-08 07:50:27,935 - INFO -   Training Loss: 0.7015, Perplexity: 2.0168, BLEU Score: 5.08\n",
      "2024-08-08 07:53:24,458 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 07:53:24,459 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 07:53:24,461 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 07:53:24,621 - INFO - Fold 1/5, Epoch 6/10, Time: 80.6s, Estimated remaining time: 3545.2s\n",
      "2024-08-08 07:53:24,622 - INFO -   Training Loss: 0.5986, Perplexity: 1.8196, BLEU Score: 5.08\n",
      "2024-08-08 07:53:24,868 - INFO - Model saved to ./saved_models/english2french/transformer_epoch_5.pth\n",
      "2024-08-08 07:56:20,243 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 07:56:20,245 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 07:56:20,246 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 07:56:20,350 - INFO - Fold 1/5, Epoch 7/10, Time: 80.6s, Estimated remaining time: 3464.2s\n",
      "2024-08-08 07:56:20,351 - INFO -   Training Loss: 0.5267, Perplexity: 1.6934, BLEU Score: 5.16\n",
      "2024-08-08 07:59:16,326 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 07:59:16,329 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 07:59:16,330 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 07:59:16,432 - INFO - Fold 1/5, Epoch 8/10, Time: 80.4s, Estimated remaining time: 3378.6s\n",
      "2024-08-08 07:59:16,433 - INFO -   Training Loss: 0.4744, Perplexity: 1.6070, BLEU Score: 5.34\n",
      "2024-08-08 08:02:12,520 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:02:12,522 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:02:12,524 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:02:12,701 - INFO - Fold 1/5, Epoch 9/10, Time: 80.3s, Estimated remaining time: 3292.6s\n",
      "2024-08-08 08:02:12,703 - INFO -   Training Loss: 0.4358, Perplexity: 1.5462, BLEU Score: 5.31\n",
      "2024-08-08 08:05:08,324 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:05:08,325 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:05:08,326 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:05:08,426 - INFO - Fold 1/5, Epoch 10/10, Time: 79.7s, Estimated remaining time: 3189.8s\n",
      "2024-08-08 08:05:08,427 - INFO -   Training Loss: 0.4049, Perplexity: 1.4991, BLEU Score: 5.34\n",
      "2024-08-08 08:05:08,689 - INFO - Model saved to ./saved_models/english2french/transformer_epoch_9.pth\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting fold 2/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-08-08 08:08:09,919 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:08:09,920 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:08:09,922 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:08:10,029 - INFO - Fold 2/5, Epoch 1/10, Time: 83.5s, Estimated remaining time: 3257.9s\n",
      "2024-08-08 08:08:10,030 - INFO -   Training Loss: 3.8082, Perplexity: 45.0687, BLEU Score: 0.93\n",
      "2024-08-08 08:08:10,249 - INFO - Model saved to ./saved_models/english2french/transformer_epoch_0.pth\n",
      "2024-08-08 08:11:07,598 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:11:07,599 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:11:07,600 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:11:07,701 - INFO - Fold 2/5, Epoch 2/10, Time: 80.6s, Estimated remaining time: 3062.4s\n",
      "2024-08-08 08:11:07,702 - INFO -   Training Loss: 1.9541, Perplexity: 7.0575, BLEU Score: 3.48\n",
      "2024-08-08 08:14:04,205 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:14:04,207 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:14:04,208 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:14:04,310 - INFO - Fold 2/5, Epoch 3/10, Time: 80.2s, Estimated remaining time: 2967.5s\n",
      "2024-08-08 08:14:04,311 - INFO -   Training Loss: 1.1939, Perplexity: 3.2998, BLEU Score: 4.42\n",
      "2024-08-08 08:17:00,316 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:17:00,317 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:17:00,319 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:17:00,478 - INFO - Fold 2/5, Epoch 4/10, Time: 80.2s, Estimated remaining time: 2886.9s\n",
      "2024-08-08 08:17:00,480 - INFO -   Training Loss: 0.8942, Perplexity: 2.4453, BLEU Score: 4.85\n",
      "2024-08-08 08:19:57,183 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:19:57,185 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:19:57,186 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:19:57,289 - INFO - Fold 2/5, Epoch 5/10, Time: 79.9s, Estimated remaining time: 2795.4s\n",
      "2024-08-08 08:19:57,290 - INFO -   Training Loss: 0.7305, Perplexity: 2.0760, BLEU Score: 4.94\n",
      "2024-08-08 08:22:54,076 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:22:54,077 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:22:54,078 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:22:54,184 - INFO - Fold 2/5, Epoch 6/10, Time: 80.1s, Estimated remaining time: 2721.8s\n",
      "2024-08-08 08:22:54,185 - INFO -   Training Loss: 0.6134, Perplexity: 1.8467, BLEU Score: 5.01\n",
      "2024-08-08 08:22:54,365 - INFO - Model saved to ./saved_models/english2french/transformer_epoch_5.pth\n",
      "2024-08-08 08:25:50,755 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:25:50,758 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:25:50,761 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:25:50,919 - INFO - Fold 2/5, Epoch 7/10, Time: 80.1s, Estimated remaining time: 2641.9s\n",
      "2024-08-08 08:25:50,920 - INFO -   Training Loss: 0.5367, Perplexity: 1.7104, BLEU Score: 5.12\n",
      "2024-08-08 08:28:47,598 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:28:47,600 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:28:47,601 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:28:47,702 - INFO - Fold 2/5, Epoch 8/10, Time: 80.1s, Estimated remaining time: 2561.6s\n",
      "2024-08-08 08:28:47,703 - INFO -   Training Loss: 0.4823, Perplexity: 1.6198, BLEU Score: 5.16\n",
      "2024-08-08 08:31:43,327 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:31:43,328 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:31:43,329 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:31:43,431 - INFO - Fold 2/5, Epoch 9/10, Time: 80.0s, Estimated remaining time: 2479.1s\n",
      "2024-08-08 08:31:43,431 - INFO -   Training Loss: 0.4412, Perplexity: 1.5546, BLEU Score: 5.27\n",
      "2024-08-08 08:34:39,465 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:34:39,467 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:34:39,469 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:34:39,626 - INFO - Fold 2/5, Epoch 10/10, Time: 79.9s, Estimated remaining time: 2397.8s\n",
      "2024-08-08 08:34:39,627 - INFO -   Training Loss: 0.4082, Perplexity: 1.5041, BLEU Score: 5.26\n",
      "2024-08-08 08:34:39,903 - INFO - Model saved to ./saved_models/english2french/transformer_epoch_9.pth\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting fold 3/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-08-08 08:37:41,280 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:37:41,281 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:37:41,283 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:37:41,437 - INFO - Fold 3/5, Epoch 1/10, Time: 83.2s, Estimated remaining time: 2412.9s\n",
      "2024-08-08 08:37:41,438 - INFO -   Training Loss: 3.8031, Perplexity: 44.8421, BLEU Score: 0.80\n",
      "2024-08-08 08:37:41,672 - INFO - Model saved to ./saved_models/english2french/transformer_epoch_0.pth\n",
      "2024-08-08 08:40:38,329 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:40:38,333 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:40:38,334 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:40:38,532 - INFO - Fold 3/5, Epoch 2/10, Time: 79.8s, Estimated remaining time: 2233.9s\n",
      "2024-08-08 08:40:38,533 - INFO -   Training Loss: 2.0648, Perplexity: 7.8834, BLEU Score: 3.33\n",
      "2024-08-08 08:43:32,773 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:43:32,774 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:43:32,775 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:43:32,877 - INFO - Fold 3/5, Epoch 3/10, Time: 80.3s, Estimated remaining time: 2168.4s\n",
      "2024-08-08 08:43:32,878 - INFO -   Training Loss: 1.2484, Perplexity: 3.4847, BLEU Score: 4.27\n",
      "2024-08-08 08:46:26,442 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:46:26,444 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:46:26,445 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:46:26,611 - INFO - Fold 3/5, Epoch 4/10, Time: 80.1s, Estimated remaining time: 2081.5s\n",
      "2024-08-08 08:46:26,612 - INFO -   Training Loss: 0.9301, Perplexity: 2.5347, BLEU Score: 4.66\n",
      "2024-08-08 08:49:23,147 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:49:23,152 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:49:23,153 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:49:23,353 - INFO - Fold 3/5, Epoch 5/10, Time: 80.2s, Estimated remaining time: 2005.7s\n",
      "2024-08-08 08:49:23,355 - INFO -   Training Loss: 0.7486, Perplexity: 2.1140, BLEU Score: 5.26\n",
      "2024-08-08 08:52:16,350 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:52:16,355 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:52:16,356 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:52:16,481 - INFO - Fold 3/5, Epoch 6/10, Time: 80.7s, Estimated remaining time: 1935.8s\n",
      "2024-08-08 08:52:16,482 - INFO -   Training Loss: 0.6317, Perplexity: 1.8808, BLEU Score: 4.97\n",
      "2024-08-08 08:52:16,686 - INFO - Model saved to ./saved_models/english2french/transformer_epoch_5.pth\n",
      "2024-08-08 08:55:05,034 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:55:05,036 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:55:05,037 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:55:05,186 - INFO - Fold 3/5, Epoch 7/10, Time: 80.3s, Estimated remaining time: 1847.2s\n",
      "2024-08-08 08:55:05,187 - INFO -   Training Loss: 0.5536, Perplexity: 1.7395, BLEU Score: 5.34\n",
      "2024-08-08 08:57:55,229 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 08:57:55,231 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 08:57:55,232 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 08:57:55,356 - INFO - Fold 3/5, Epoch 8/10, Time: 80.1s, Estimated remaining time: 1762.9s\n",
      "2024-08-08 08:57:55,357 - INFO -   Training Loss: 0.4954, Perplexity: 1.6412, BLEU Score: 5.18\n",
      "2024-08-08 09:00:45,767 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:00:45,770 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:00:45,771 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:00:45,925 - INFO - Fold 3/5, Epoch 9/10, Time: 80.4s, Estimated remaining time: 1688.8s\n",
      "2024-08-08 09:00:45,926 - INFO -   Training Loss: 0.4487, Perplexity: 1.5663, BLEU Score: 5.34\n",
      "2024-08-08 09:03:37,170 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:03:37,173 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:03:37,173 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:03:37,304 - INFO - Fold 3/5, Epoch 10/10, Time: 80.2s, Estimated remaining time: 1604.9s\n",
      "2024-08-08 09:03:37,304 - INFO -   Training Loss: 0.4180, Perplexity: 1.5189, BLEU Score: 5.16\n",
      "2024-08-08 09:03:37,550 - INFO - Model saved to ./saved_models/english2french/transformer_epoch_9.pth\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting fold 4/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-08-08 09:06:32,169 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:06:32,171 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:06:32,173 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:06:32,332 - INFO - Fold 4/5, Epoch 1/10, Time: 83.5s, Estimated remaining time: 1586.9s\n",
      "2024-08-08 09:06:32,333 - INFO -   Training Loss: 3.6860, Perplexity: 39.8858, BLEU Score: 1.19\n",
      "2024-08-08 09:06:32,699 - INFO - Model saved to ./saved_models/english2french/transformer_epoch_0.pth\n",
      "2024-08-08 09:09:22,131 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:09:22,133 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:09:22,134 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:09:22,243 - INFO - Fold 4/5, Epoch 2/10, Time: 80.4s, Estimated remaining time: 1446.5s\n",
      "2024-08-08 09:09:22,243 - INFO -   Training Loss: 1.8246, Perplexity: 6.2005, BLEU Score: 3.68\n",
      "2024-08-08 09:12:12,266 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:12:12,267 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:12:12,268 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:12:12,390 - INFO - Fold 4/5, Epoch 3/10, Time: 80.1s, Estimated remaining time: 1361.6s\n",
      "2024-08-08 09:12:12,391 - INFO -   Training Loss: 1.1463, Perplexity: 3.1466, BLEU Score: 4.42\n",
      "2024-08-08 09:15:01,920 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:15:01,922 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:15:01,923 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:15:02,036 - INFO - Fold 4/5, Epoch 4/10, Time: 80.1s, Estimated remaining time: 1281.2s\n",
      "2024-08-08 09:15:02,037 - INFO -   Training Loss: 0.8753, Perplexity: 2.3995, BLEU Score: 4.83\n",
      "2024-08-08 09:17:51,700 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:17:51,701 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:17:51,702 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:17:51,806 - INFO - Fold 4/5, Epoch 5/10, Time: 80.0s, Estimated remaining time: 1200.7s\n",
      "2024-08-08 09:17:51,807 - INFO -   Training Loss: 0.7124, Perplexity: 2.0390, BLEU Score: 4.99\n",
      "2024-08-08 09:20:41,668 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:20:41,669 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:20:41,670 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:20:41,783 - INFO - Fold 4/5, Epoch 6/10, Time: 80.1s, Estimated remaining time: 1121.4s\n",
      "2024-08-08 09:20:41,784 - INFO -   Training Loss: 0.6035, Perplexity: 1.8285, BLEU Score: 5.00\n",
      "2024-08-08 09:20:42,400 - INFO - Model saved to ./saved_models/english2french/transformer_epoch_5.pth\n",
      "2024-08-08 09:23:31,941 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:23:31,942 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:23:31,943 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:23:32,048 - INFO - Fold 4/5, Epoch 7/10, Time: 80.1s, Estimated remaining time: 1041.0s\n",
      "2024-08-08 09:23:32,049 - INFO -   Training Loss: 0.5319, Perplexity: 1.7021, BLEU Score: 5.07\n",
      "2024-08-08 09:26:21,731 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:26:21,732 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:26:21,733 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:26:21,848 - INFO - Fold 4/5, Epoch 8/10, Time: 80.1s, Estimated remaining time: 961.0s\n",
      "2024-08-08 09:26:21,849 - INFO -   Training Loss: 0.4818, Perplexity: 1.6190, BLEU Score: 5.19\n",
      "2024-08-08 09:29:12,107 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:29:12,111 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:29:12,112 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:29:12,253 - INFO - Fold 4/5, Epoch 9/10, Time: 80.3s, Estimated remaining time: 883.5s\n",
      "2024-08-08 09:29:12,254 - INFO -   Training Loss: 0.4393, Perplexity: 1.5516, BLEU Score: 5.19\n",
      "2024-08-08 09:32:01,817 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:32:01,819 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:32:01,820 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:32:01,924 - INFO - Fold 4/5, Epoch 10/10, Time: 80.2s, Estimated remaining time: 802.0s\n",
      "2024-08-08 09:32:01,925 - INFO -   Training Loss: 0.4075, Perplexity: 1.5030, BLEU Score: 5.28\n",
      "2024-08-08 09:32:02,248 - INFO - Model saved to ./saved_models/english2french/transformer_epoch_9.pth\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting fold 5/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-08-08 09:35:02,107 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:35:02,111 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:35:02,112 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:35:02,227 - INFO - Fold 5/5, Epoch 1/10, Time: 83.5s, Estimated remaining time: 751.2s\n",
      "2024-08-08 09:35:02,228 - INFO -   Training Loss: 3.7543, Perplexity: 42.7036, BLEU Score: 0.98\n",
      "2024-08-08 09:35:02,507 - INFO - Model saved to ./saved_models/english2french/transformer_epoch_0.pth\n",
      "2024-08-08 09:37:51,931 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:37:51,935 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:37:51,936 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:37:52,045 - INFO - Fold 5/5, Epoch 2/10, Time: 80.3s, Estimated remaining time: 642.4s\n",
      "2024-08-08 09:37:52,046 - INFO -   Training Loss: 1.8567, Perplexity: 6.4029, BLEU Score: 3.60\n",
      "2024-08-08 09:40:42,380 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:40:42,381 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:40:42,382 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:40:42,489 - INFO - Fold 5/5, Epoch 3/10, Time: 80.3s, Estimated remaining time: 561.9s\n",
      "2024-08-08 09:40:42,490 - INFO -   Training Loss: 1.1550, Perplexity: 3.1739, BLEU Score: 4.35\n",
      "2024-08-08 09:43:32,493 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:43:32,494 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:43:32,495 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:43:32,616 - INFO - Fold 5/5, Epoch 4/10, Time: 80.3s, Estimated remaining time: 481.8s\n",
      "2024-08-08 09:43:32,617 - INFO -   Training Loss: 0.8710, Perplexity: 2.3894, BLEU Score: 4.71\n",
      "2024-08-08 09:46:22,521 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:46:22,522 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:46:22,523 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:46:22,629 - INFO - Fold 5/5, Epoch 5/10, Time: 80.2s, Estimated remaining time: 400.9s\n",
      "2024-08-08 09:46:22,630 - INFO -   Training Loss: 0.7116, Perplexity: 2.0372, BLEU Score: 5.02\n",
      "2024-08-08 09:49:12,874 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:49:12,875 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:49:12,876 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:49:12,983 - INFO - Fold 5/5, Epoch 6/10, Time: 80.1s, Estimated remaining time: 320.5s\n",
      "2024-08-08 09:49:12,983 - INFO -   Training Loss: 0.5994, Perplexity: 1.8211, BLEU Score: 5.13\n",
      "2024-08-08 09:49:13,207 - INFO - Model saved to ./saved_models/english2french/transformer_epoch_5.pth\n",
      "2024-08-08 09:52:03,266 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:52:03,267 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:52:03,268 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:52:03,417 - INFO - Fold 5/5, Epoch 7/10, Time: 80.2s, Estimated remaining time: 240.7s\n",
      "2024-08-08 09:52:03,417 - INFO -   Training Loss: 0.5305, Perplexity: 1.6998, BLEU Score: 5.18\n",
      "2024-08-08 09:54:53,409 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:54:53,410 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:54:53,411 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:54:53,518 - INFO - Fold 5/5, Epoch 8/10, Time: 80.3s, Estimated remaining time: 160.7s\n",
      "2024-08-08 09:54:53,519 - INFO -   Training Loss: 0.4729, Perplexity: 1.6046, BLEU Score: 5.17\n",
      "2024-08-08 09:57:44,292 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 09:57:44,293 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 09:57:44,294 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 09:57:44,409 - INFO - Fold 5/5, Epoch 9/10, Time: 80.5s, Estimated remaining time: 80.5s\n",
      "2024-08-08 09:57:44,410 - INFO -   Training Loss: 0.4353, Perplexity: 1.5455, BLEU Score: 5.23\n",
      "2024-08-08 10:00:34,187 - WARNING - That's 100 lines that end in a tokenized period ('.')\n",
      "2024-08-08 10:00:34,189 - WARNING - It looks like you forgot to detokenize your test data, which may hurt your score.\n",
      "2024-08-08 10:00:34,190 - WARNING - If you insist your data is detokenized, or don't care, you can suppress this message with the `force` parameter.\n",
      "2024-08-08 10:00:34,335 - INFO - Fold 5/5, Epoch 10/10, Time: 80.1s, Estimated remaining time: 0.0s\n",
      "2024-08-08 10:00:34,336 - INFO -   Training Loss: 0.4039, Perplexity: 1.4977, BLEU Score: 5.18\n",
      "2024-08-08 10:00:34,559 - INFO - Model saved to ./saved_models/english2french/transformer_epoch_9.pth\n",
      "2024-08-08 10:00:34,561 - INFO - K-Fold Cross-Validation Results:\n",
      "2024-08-08 10:00:34,562 - INFO -   Average BLEU Score: 5.25\n",
      "2024-08-08 10:00:34,563 - INFO -   Average Loss: 0.4085\n",
      "2024-08-08 10:00:34,563 - INFO -   Average Perplexity: 1.5046\n",
      "2024-08-08 10:00:34,734 - INFO - Final model saved to ./saved_models/english2french/transformer_model.pt\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import numpy as np\n",
    "import time\n",
    "import argparse\n",
    "import pickle\n",
    "import os\n",
    "import logging\n",
    "from typing import List, Tuple\n",
    "from random import shuffle\n",
    "from sklearn.model_selection import KFold\n",
    "import sacrebleu\n",
    "\n",
    "# Configure logging for better output control and formatting\n",
    "logging.basicConfig(\n",
    "    level=logging.INFO, format=\"%(asctime)s - %(levelname)s - %(message)s\"\n",
    ")\n",
    "\n",
    "# Define special tokens\n",
    "PAD_TOKEN = 0  # Padding token\n",
    "SOS_TOKEN = 1  # Start of sentence token\n",
    "EOS_TOKEN = 2  # End of sentence token\n",
    "\n",
    "\n",
    "class Trainer:\n",
    "    def initialize_weights(self, model: nn.Module) -> None:\n",
    "        \"\"\"\n",
    "        Initialize model weights using Xavier uniform initialization.\n",
    "        \"\"\"\n",
    "        if hasattr(model, \"weight\") and model.weight.dim() > 1:\n",
    "            nn.init.xavier_uniform_(model.weight.data)\n",
    "\n",
    "    def save_dictionary(self, dictionary: dict, input: bool = True) -> None:\n",
    "        \"\"\"\n",
    "        Save the language dictionary to disk.\n",
    "        \"\"\"\n",
    "        # Create directory to save dictionaries if it doesn't exist\n",
    "        directory = (\n",
    "            f\"saved_models/{self.input_lang_dic.name}2{self.output_lang_dic.name}\"\n",
    "        )\n",
    "        if not os.path.exists(directory):\n",
    "            os.makedirs(directory)\n",
    "\n",
    "        # Determine the file path based on whether it's the input or output dictionary\n",
    "        file_path = f\"{directory}/{'input_dic.pkl' if input else 'output_dic.pkl'}\"\n",
    "        # Save the dictionary to a pickle file\n",
    "        with open(file_path, \"wb\") as f:\n",
    "            pickle.dump(dictionary, f, pickle.HIGHEST_PROTOCOL)\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        lang1: str,\n",
    "        lang2: str,\n",
    "        data_directory: str,\n",
    "        reverse: bool,\n",
    "        MAX_LENGTH: int,\n",
    "        MAX_FILE_SIZE: int,\n",
    "        batch_size: int,\n",
    "        lr: float = 0.0005,\n",
    "        hidden_size: int = 256,\n",
    "        encoder_layers: int = 3,\n",
    "        decoder_layers: int = 3,\n",
    "        encoder_heads: int = 8,\n",
    "        decoder_heads: int = 8,\n",
    "        encoder_ff_size: int = 512,\n",
    "        decoder_ff_size: int = 512,\n",
    "        encoder_dropout: float = 0.1,\n",
    "        decoder_dropout: float = 0.1,\n",
    "        device: str = \"cpu\",\n",
    "    ) -> None:\n",
    "        \"\"\"\n",
    "        Initialize the Trainer with various hyperparameters and configurations.\n",
    "        \"\"\"\n",
    "        self.MAX_LENGTH = MAX_LENGTH\n",
    "        self.MAX_FILE_SIZE = MAX_FILE_SIZE\n",
    "        self.device = device\n",
    "        self.batch_size = batch_size\n",
    "        self.lr = lr\n",
    "        self.hidden_size = hidden_size\n",
    "        self.encoder_layers = encoder_layers\n",
    "        self.decoder_layers = decoder_layers\n",
    "        self.encoder_heads = encoder_heads\n",
    "        self.decoder_heads = decoder_heads\n",
    "        self.encoder_ff_size = encoder_ff_size\n",
    "        self.decoder_ff_size = decoder_ff_size\n",
    "        self.encoder_dropout = encoder_dropout\n",
    "        self.decoder_dropout = decoder_dropout\n",
    "\n",
    "        # Load language data and create dictionaries\n",
    "        (\n",
    "            self.input_lang_dic,\n",
    "            self.output_lang_dic,\n",
    "            self.input_lang_list,\n",
    "            self.output_lang_list,\n",
    "        ) = load_files(\n",
    "            lang1, lang2, data_directory, reverse, self.MAX_FILE_SIZE, self.MAX_LENGTH\n",
    "        )\n",
    "\n",
    "        if self.input_lang_dic is None or self.output_lang_dic is None:\n",
    "            raise ValueError(\n",
    "                \"Loading language files failed due to mismatched line counts.\"\n",
    "            )\n",
    "\n",
    "        # Add sentences to input and output dictionaries\n",
    "        for sentence in self.input_lang_list:\n",
    "            self.input_lang_dic.add_sentence(sentence)\n",
    "        for sentence in self.output_lang_list:\n",
    "            self.output_lang_dic.add_sentence(sentence)\n",
    "\n",
    "        # Save the dictionaries to disk\n",
    "        self.save_dictionary(self.input_lang_dic, input=True)\n",
    "        self.save_dictionary(self.output_lang_dic, input=False)\n",
    "\n",
    "        # Tokenize sentences\n",
    "        self.tokenized_input_lang = [\n",
    "            tokenize(sentence, self.input_lang_dic, self.MAX_LENGTH)\n",
    "            for sentence in self.input_lang_list\n",
    "        ]\n",
    "        self.tokenized_output_lang = [\n",
    "            tokenize(sentence, self.output_lang_dic, self.MAX_LENGTH)\n",
    "            for sentence in self.output_lang_list\n",
    "        ]\n",
    "\n",
    "        # Debugging: Print a few tokenized sequences\n",
    "        print(\"Sample tokenized input sentences:\")\n",
    "        for i in range(5):\n",
    "            print(self.tokenized_input_lang[i])\n",
    "        print(\"Sample tokenized output sentences:\")\n",
    "        for i in range(5):\n",
    "            print(self.tokenized_output_lang[i])\n",
    "\n",
    "        self.batch_size = batch_size\n",
    "\n",
    "        # Create data loader for batching the training data\n",
    "        self.data_loader = load_batches(\n",
    "            self.tokenized_input_lang,\n",
    "            self.tokenized_output_lang,\n",
    "            self.batch_size,\n",
    "            self.device,\n",
    "        )\n",
    "\n",
    "        # Define sizes based on dictionaries\n",
    "        input_size = self.input_lang_dic.n_count\n",
    "        output_size = self.output_lang_dic.n_count\n",
    "\n",
    "        # Print vocabulary sizes for debugging\n",
    "        print(f\"Input vocabulary size: {input_size}\")\n",
    "        print(f\"Output vocabulary size: {output_size}\")\n",
    "\n",
    "        # Define encoder and decoder parts of the transformer\n",
    "        encoder_part = Encoder(\n",
    "            input_size,\n",
    "            hidden_size,\n",
    "            encoder_layers,\n",
    "            encoder_heads,\n",
    "            encoder_ff_size,\n",
    "            encoder_dropout,\n",
    "            self.device,\n",
    "        )\n",
    "        decoder_part = Decoder(\n",
    "            output_size,\n",
    "            hidden_size,\n",
    "            decoder_layers,\n",
    "            decoder_heads,\n",
    "            decoder_ff_size,\n",
    "            decoder_dropout,\n",
    "            self.device,\n",
    "        )\n",
    "\n",
    "        # Initialize the transformer model\n",
    "        self.transformer = Transformer(\n",
    "            encoder_part, decoder_part, self.device, PAD_TOKEN\n",
    "        ).to(self.device)\n",
    "        self.transformer.apply(self.initialize_weights)\n",
    "\n",
    "        # Define loss function and optimizer\n",
    "        self.loss_func = nn.CrossEntropyLoss(ignore_index=PAD_TOKEN)\n",
    "        self.optimizer = optim.Adam(self.transformer.parameters(), lr=lr)\n",
    "\n",
    "\n",
    "    def k_fold_cross_validation(self, k: int, epochs: int, saved_model_directory: str) -> None:\n",
    "        \"\"\"\n",
    "        Perform K-fold cross-validation.\n",
    "        \n",
    "        Args:\n",
    "            k (int): Number of folds.\n",
    "            epochs (int): Number of epochs for each fold.\n",
    "            saved_model_directory (str): Directory to save models.\n",
    "        \"\"\"\n",
    "        kf = KFold(n_splits=k, shuffle=True, random_state=42)\n",
    "        all_bleu_scores = []\n",
    "        all_losses = []\n",
    "        all_perplexities = []\n",
    "        \n",
    "        for fold, (train_index, val_index) in enumerate(kf.split(self.tokenized_input_lang)):\n",
    "            print(f\"Starting fold {fold + 1}/{k}\")\n",
    "            \n",
    "            # Create training and validation datasets\n",
    "            train_input = [self.tokenized_input_lang[i] for i in train_index]\n",
    "            train_output = [self.tokenized_output_lang[i] for i in train_index]\n",
    "            val_input = [self.tokenized_input_lang[i] for i in val_index]\n",
    "            val_output = [self.tokenized_output_lang[i] for i in val_index]\n",
    "            \n",
    "            # Create data loaders\n",
    "            self.data_loader = load_batches(train_input, train_output, self.batch_size, self.device)\n",
    "            val_loader = load_batches(val_input, val_output, self.batch_size, self.device)\n",
    "            \n",
    "            # Reinitialize the model for each fold\n",
    "            input_size = self.input_lang_dic.n_count\n",
    "            output_size = self.output_lang_dic.n_count\n",
    "            encoder_part = Encoder(\n",
    "                input_size, self.hidden_size, self.encoder_layers, self.encoder_heads,\n",
    "                self.encoder_ff_size, self.encoder_dropout, self.device\n",
    "            )\n",
    "            decoder_part = Decoder(\n",
    "                output_size, self.hidden_size, self.decoder_layers, self.decoder_heads,\n",
    "                self.decoder_ff_size, self.decoder_dropout, self.device\n",
    "            )\n",
    "            self.transformer = Transformer(\n",
    "                encoder_part, decoder_part, self.device, PAD_TOKEN\n",
    "            ).to(self.device)\n",
    "            self.transformer.apply(self.initialize_weights)\n",
    "            self.optimizer = optim.Adam(self.transformer.parameters(), lr=self.lr)\n",
    "            \n",
    "            # Train the model\n",
    "            for epoch in range(epochs):\n",
    "                start_time = time.time()\n",
    "                train_loss, perplexity = self.train_epoch()  # Train for one epoch\n",
    "                duration = time.time() - start_time  # Calculate epoch duration\n",
    "                estimated_remaining_time = (\n",
    "                    (k * epochs) - ((fold * epochs) + epoch + 1)\n",
    "                ) * duration  # Estimate remaining time\n",
    "\n",
    "                # Calculate BLEU score on validation dataset\n",
    "                bleu_score = self.calculate_bleu(val_loader)\n",
    "\n",
    "                # Log epoch statistics\n",
    "                logging.info(\n",
    "                    f\"Fold {fold + 1}/{k}, Epoch {epoch + 1}/{epochs}, Time: {duration:.1f}s, Estimated remaining time: {estimated_remaining_time:.1f}s\"\n",
    "                )\n",
    "                logging.info(\n",
    "                    f\"  Training Loss: {train_loss:.4f}, Perplexity: {perplexity:.4f}, BLEU Score: {bleu_score:.2f}\"\n",
    "                )\n",
    "\n",
    "                # Save model checkpoint every 5 epochs\n",
    "                if epoch % 5 == 0 or epoch == epochs - 1:\n",
    "                    self.save_model(epoch, saved_model_directory)\n",
    "\n",
    "            # Append results of the current fold\n",
    "            all_bleu_scores.append(bleu_score)\n",
    "            all_losses.append(train_loss)\n",
    "            all_perplexities.append(perplexity)\n",
    "\n",
    "        # Aggregate results\n",
    "        avg_bleu_score = np.mean(all_bleu_scores)\n",
    "        avg_loss = np.mean(all_losses)\n",
    "        avg_perplexity = np.mean(all_perplexities)\n",
    "        \n",
    "        logging.info(f\"K-Fold Cross-Validation Results:\")\n",
    "        logging.info(f\"  Average BLEU Score: {avg_bleu_score:.2f}\")\n",
    "        logging.info(f\"  Average Loss: {avg_loss:.4f}\")\n",
    "        logging.info(f\"  Average Perplexity: {avg_perplexity:.4f}\")\n",
    "\n",
    "        # Save final model\n",
    "        self.final_save_model(saved_model_directory)\n",
    "        \n",
    "\n",
    "    def train_epoch(self) -> Tuple[float, float]:\n",
    "        \"\"\"\n",
    "        Train the model for one epoch and calculate average training loss and perplexity.\n",
    "        \"\"\"\n",
    "        # Shuffle the data loader to prevent overfitting\n",
    "        shuffle(self.data_loader)\n",
    "        train_loss = 0\n",
    "\n",
    "        for input, target in self.data_loader:\n",
    "            # Skip empty batches\n",
    "            if input.size(0) == 0 or target.size(0) == 0:\n",
    "                logging.warning(\"Empty batch detected. Skipping...\")\n",
    "                continue\n",
    "\n",
    "            # Ensure tensors have at least 2 dimensions\n",
    "            if input.dim() == 1:\n",
    "                input = input.unsqueeze(0)\n",
    "            if target.dim() == 1:\n",
    "                target = target.unsqueeze(0)\n",
    "\n",
    "            # Debugging: Check if input values are within vocabulary size\n",
    "            if (input >= self.input_lang_dic.n_count).any() or (\n",
    "                target >= self.output_lang_dic.n_count\n",
    "            ).any():\n",
    "                print(f\"Input tensor has values outside the vocabulary size: {input}\")\n",
    "                print(f\"Target tensor has values outside the vocabulary size: {target}\")\n",
    "\n",
    "            # Zero gradients\n",
    "            self.optimizer.zero_grad()\n",
    "\n",
    "            # Forward pass through the transformer model\n",
    "            output, _ = self.transformer(input, target[:, :-1])\n",
    "\n",
    "            # Reshape output and target for loss calculation\n",
    "            output = output.reshape(-1, output.shape[-1])\n",
    "            target = target[:, 1:].reshape(-1)\n",
    "\n",
    "            # Calculate loss\n",
    "            loss = self.loss_func(output, target)\n",
    "            # Backpropagation\n",
    "            loss.backward()\n",
    "            # Update model parameters\n",
    "            self.optimizer.step()\n",
    "\n",
    "            train_loss += loss.item()\n",
    "\n",
    "        # Average training loss for the epoch\n",
    "        avg_loss = train_loss / len(self.data_loader)\n",
    "        # Perplexity is the exponential of the average loss\n",
    "        perplexity = np.exp(avg_loss)\n",
    "        return avg_loss, perplexity\n",
    "\n",
    "    def train(self, epochs: int, saved_model_directory: str) -> None:\n",
    "        \"\"\"\n",
    "        Train the model for a specified number of epochs.\n",
    "        \"\"\"\n",
    "        for epoch in range(epochs):\n",
    "            start_time = time.time()\n",
    "            train_loss, perplexity = self.train_epoch()  # Train for one epoch\n",
    "            duration = time.time() - start_time  # Calculate epoch duration\n",
    "            estimated_remaining_time = (\n",
    "                epochs - epoch - 1\n",
    "            ) * duration  # Estimate remaining time\n",
    "\n",
    "            # Log epoch statistics\n",
    "            logging.info(\n",
    "                f\"Epoch {epoch + 1}/{epochs}, Time: {duration:.1f}s, Estimated remaining time: {estimated_remaining_time:.1f}s\"\n",
    "            )\n",
    "            logging.info(\n",
    "                f\"  Training Loss: {train_loss:.4f}, Perplexity: {perplexity:.4f}\"\n",
    "            )\n",
    "\n",
    "            # Save model checkpoint every 5 epochs\n",
    "            if epoch % 5 == 0 or epoch == epochs - 1:\n",
    "                self.save_model(epoch, saved_model_directory)\n",
    "\n",
    "        logging.info(\"Training finished!\")  # Log completion message\n",
    "        self.final_save_model(saved_model_directory)  # Save final model\n",
    "\n",
    "    def save_model(self, epoch: int, saved_model_directory: str) -> None:\n",
    "        \"\"\"\n",
    "        Save the model checkpoint.\n",
    "\n",
    "        Args:\n",
    "            epoch (int): The current epoch number.\n",
    "            saved_model_directory (str): The directory where the model should be saved.\n",
    "        \"\"\"\n",
    "        # Create directory to save model checkpoints if it doesn't exist\n",
    "        directory = os.path.join(\n",
    "            saved_model_directory,\n",
    "            f\"{self.input_lang_dic.name}2{self.output_lang_dic.name}\",\n",
    "        )\n",
    "        if not os.path.exists(directory):\n",
    "            os.makedirs(directory)\n",
    "        # Define model path and save model state\n",
    "        model_path = os.path.join(directory, f\"transformer_epoch_{epoch}.pth\")\n",
    "        torch.save(self.transformer.state_dict(), model_path)\n",
    "        logging.info(f\"Model saved to {model_path}\")  # Log save message\n",
    "\n",
    "    def final_save_model(self, saved_model_directory: str) -> None:\n",
    "        \"\"\"\n",
    "        Save the final model checkpoint.\n",
    "        \"\"\"\n",
    "        # Create directory to save model checkpoints if it doesn't exist\n",
    "        directory = os.path.join(\n",
    "            saved_model_directory,\n",
    "            f\"{self.input_lang_dic.name}2{self.output_lang_dic.name}\",\n",
    "        )\n",
    "        if not os.path.exists(directory):\n",
    "            os.makedirs(directory)\n",
    "        # Define model path and save model state\n",
    "        model_path = os.path.join(directory, \"transformer_model.pt\")\n",
    "        torch.save(self.transformer.state_dict(), model_path)\n",
    "        logging.info(f\"Final model saved to {model_path}\")  # Log save message\n",
    "\n",
    "    def calculate_bleu(self, val_loader):\n",
    "        \"\"\"\n",
    "        Calculate BLEU score on the validation dataset.\n",
    "        \"\"\"\n",
    "        self.transformer.eval()\n",
    "        references = []\n",
    "        hypotheses = []\n",
    "        with torch.no_grad():\n",
    "            for input, target in val_loader:\n",
    "                input, target = input.to(self.device), target.to(self.device)\n",
    "                target_input = target[:, :-1]\n",
    "                target_output = target[:, 1:]\n",
    "                \n",
    "                output, _ = self.transformer(input, target_input)\n",
    "                predictions = output.argmax(dim=-1)\n",
    "                \n",
    "                for ref, hyp in zip(target_output, predictions):\n",
    "                    ref_text = ' '.join([self.output_lang_dic.index2word[idx.item()] for idx in ref if idx.item() not in [PAD_TOKEN, EOS_TOKEN]])\n",
    "                    hyp_text = ' '.join([self.output_lang_dic.index2word[idx.item()] for idx in hyp if idx.item() not in [PAD_TOKEN, EOS_TOKEN]])\n",
    "                    references.append(ref_text)\n",
    "                    hypotheses.append(hyp_text)\n",
    "        \n",
    "        bleu = sacrebleu.corpus_bleu(hypotheses, [references])\n",
    "        return bleu.score\n",
    "        \n",
    "\n",
    "def main() -> None:\n",
    "    # Example parameters\n",
    "    lang1 = 'english'\n",
    "    lang2 = 'french'\n",
    "    data_directory = 'data'\n",
    "    reverse = 0\n",
    "    MAX_LENGTH = 60\n",
    "    MAX_FILE_SIZE = 200000\n",
    "    batch_size = 128\n",
    "    lr = 0.0005\n",
    "    hidden_size = 256\n",
    "    encoder_layers = 3\n",
    "    decoder_layers = 3\n",
    "    encoder_heads = 8\n",
    "    decoder_heads = 8\n",
    "    encoder_ff_size = 512\n",
    "    decoder_ff_size = 512\n",
    "    encoder_dropout = 0.1\n",
    "    decoder_dropout = 0.1\n",
    "    device = 'cuda'\n",
    "    \n",
    "    # Initialize the trainer\n",
    "    trainer = Trainer(\n",
    "        lang1, lang2, data_directory, reverse, MAX_LENGTH, MAX_FILE_SIZE, batch_size, lr,\n",
    "        hidden_size, encoder_layers, decoder_layers, encoder_heads, decoder_heads,\n",
    "        encoder_ff_size, decoder_ff_size, encoder_dropout, decoder_dropout, device\n",
    "    )\n",
    "    \n",
    "    # Perform k-fold cross-validation\n",
    "    k = 5\n",
    "    epochs = 10\n",
    "    saved_model_directory = './saved_models'\n",
    "    trainer.k_fold_cross_validation(k, epochs, saved_model_directory)\n",
    "    # Start the training process\n",
    "    # trainer.train(epochs, saved_model_directory)\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()  # Run the main function if the script is executed\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "919f532e-c618-4021-957d-a4f5d472d7d4",
   "metadata": {},
   "source": [
    "# Translate and Summarize"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ee9f222",
   "metadata": {},
   "source": [
    "#### This cell contains the function to translate input text using the trained Transformer model and also summarize the text.\n",
    "#### It processes the input sequence, generates the translation, and returns the translated output.\n",
    "#### Before proceeding to translation, we produce a summarized version of the text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c0458761-567d-4a77-a762-386d5b78943c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_26167/2902028369.py:259: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  torch.load(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Skipping document due to insufficient structure: {'section_name': [], 'document': [], 'summary': []}\n",
      "Document 1 - Original: Let him know that you admire him and appreciate all of the good in him. Show him that you value him just as much, even if he is going through a tough time. Giving him unconditional positive regard will reassure him that he can count on you even when he’s not feeling his best. Take his mind off his s...\n",
      "Document 1 - Summary: Show him that you value him just as much, even if he is going through a tough time. If he is going t\n",
      "Document 1 - Translated Summary: montrez lui que de la valeur s il n en va pas beaucoup .\n",
      "\n",
      "Document 2 - Original: To recap our examples from the prior section:   Cash:  $20,000  Investments / Retirement:  $20,000  Home:  $150,000  Personal Property:  $25,000  Car:  $15,000  Life Insurance:  $10,000  Total:  $240,000 Again, from our sample estimates:   Home Loan:  $120,000  Auto Loan:  $15,000  Student Loans:  $...\n",
      "Document 2 - Summary: To recap our examples from the prior section:   Cash:  $20,000  Investments / Retirement:  $20,000  \n",
      "Document 2 - Translated Summary:  construis dans notre coin de nulle part en general en voiture . \n",
      "\n",
      "Document 3 - Original: Now that you have your music, revise your lyrics if any one word or phrase trips you up while singing them out loud. For instance, say that you use the word “particular” in one line, which you now find to be one syllable too many to enunciate clearly; try replacing it with a shorter synonym, like “c...\n",
      "Document 3 - Summary: Lyrically, this could be the “Yeah, yeah, yeah” in the chorus to the Beatles’ “She Loves You.” Music\n",
      "Document 3 - Translated Summary: fais dire que ce sera la folle\n",
      "\n",
      "Document 4 - Original: Choose a nice restaurant and eat dinner or go for drinks and dessert. Look for a place with a good menu, nice lighting, and enjoyable music that's not too loud. Your date will be impressed by your thoughtfulness. Avoid ordering certain foods on a first date: foods that are overly messy or difficult ...\n",
      "Document 4 - Summary: While you can look for a museum that neither of you have visited before, you can also go to a museum\n",
      "Document 4 - Translated Summary: en tant que vous ne pouvez pas vous rendre compte auparavant .\n",
      "\n",
      "ROUGE Scores: {'rouge1': 0.16702113307890853, 'rouge2': 0.044287801034038704, 'rougeL': 0.118403430183797}\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import pickle\n",
    "import random\n",
    "from typing import Tuple, Any, List\n",
    "from datasets import load_dataset\n",
    "from collections import Counter\n",
    "import re\n",
    "from rouge_score import rouge_scorer\n",
    "\n",
    "# Special tokens\n",
    "SOS_TOKEN = 1  # Start of Sentence token\n",
    "EOS_TOKEN = 2  # End of Sentence token\n",
    "\n",
    "def load_dictionary(directory: str) -> 'Dictionary':\n",
    "    \"\"\"\n",
    "    Load a language dictionary from a pickle file.\n",
    "\n",
    "    Args:\n",
    "        directory (str): Path to the pickle file.\n",
    "\n",
    "    Returns:\n",
    "        Dictionary: Loaded dictionary object.\n",
    "    \"\"\"\n",
    "    with open(directory, \"rb\") as f:\n",
    "        return pickle.load(f)\n",
    "\n",
    "def summarize_document(document: str, max_chars: int = 100) -> str:\n",
    "    \"\"\"\n",
    "    Summarize the document by extracting the most frequent sentences.\n",
    "\n",
    "    Args:\n",
    "        document (str): Document to be summarized.\n",
    "        max_chars (int): Maximum length of the summary in characters.\n",
    "\n",
    "    Returns:\n",
    "        str: Summary of the document.\n",
    "    \"\"\"\n",
    "    sentences = re.split(r'(?<!\\w\\.\\w.)(?<![A-Z][a-z]\\.)(?<=\\.|\\?)\\s', document)\n",
    "    if len(sentences) <= 1:\n",
    "        return document[:max_chars]\n",
    "\n",
    "    # Tokenize and count word frequencies\n",
    "    word_freq = Counter()\n",
    "    for sentence in sentences:\n",
    "        words = re.findall(r'\\w+', sentence.lower())\n",
    "        word_freq.update(words)\n",
    "\n",
    "    # Score sentences based on word frequencies\n",
    "    sentence_scores = {}\n",
    "    for sentence in sentences:\n",
    "        sentence_scores[sentence] = sum(word_freq.get(word.lower(), 0) for word in re.findall(r'\\w+', sentence))\n",
    "\n",
    "    # Sort sentences by score and extract the most frequent ones\n",
    "    sorted_sentences = sorted(sentence_scores, key=sentence_scores.get, reverse=True)\n",
    "    summary = ' '.join(sorted_sentences)\n",
    "\n",
    "    # Ensure summary does not exceed max_chars\n",
    "    return summary[:max_chars]\n",
    "\n",
    "def translate_sentence(\n",
    "    sentence: str,\n",
    "    input_dic: 'Dictionary',\n",
    "    output_dic: 'Dictionary',\n",
    "    model: 'Transformer',\n",
    "    device: torch.device,\n",
    "    max_len: int,\n",
    "    prob_threshold: float = 0.1,  # Probability threshold for stopping condition\n",
    ") -> Tuple[str, Any]:\n",
    "    \"\"\"\n",
    "    Translate a sentence from the input language to the output language using the transformer model.\n",
    "\n",
    "    Args:\n",
    "        sentence (str): Sentence to be translated.\n",
    "        input_dic (Dictionary): Input language dictionary.\n",
    "        output_dic (Dictionary): Output language dictionary.\n",
    "        model (Transformer): Transformer model for translation.\n",
    "        device (torch.device): Device to run the model on (CPU or GPU).\n",
    "        max_len (int): Maximum length of the output sentence.\n",
    "        prob_threshold (float): Probability threshold for stopping condition.\n",
    "\n",
    "    Returns:\n",
    "        Tuple[str, Any]: Translated sentence and attention weights.\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "    normalized_sentence = normalizeString(sentence)\n",
    "    tokens = tokenize(normalized_sentence, input_dic, max_len)\n",
    "\n",
    "    # Ensure tokens are within the valid range\n",
    "    tokens = [min(token, input_dic.n_count - 1) for token in tokens]\n",
    "\n",
    "    input_tensor = torch.LongTensor(tokens).unsqueeze(0).to(device)\n",
    "    input_mask = model.make_input_mask(input_tensor)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        encoded_input = model.encoder(input_tensor, input_mask)\n",
    "\n",
    "    target_tokens = [SOS_TOKEN]\n",
    "    generated_sentences = set()\n",
    "\n",
    "    for i in range(max_len):\n",
    "        target_tensor = torch.LongTensor(target_tokens).unsqueeze(0).to(device)\n",
    "        target_mask = model.make_target_mask(target_tensor)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            output, attention = model.decoder(\n",
    "                target_tensor, encoded_input, target_mask, input_mask\n",
    "            )\n",
    "\n",
    "        pred_token = output.argmax(2)[:, -1].item()\n",
    "        pred_prob = torch.softmax(output, dim=-1)[0, -1, pred_token].item()\n",
    "\n",
    "        # Penalize repetition of the same token\n",
    "        if len(target_tokens) > 1 and pred_token == target_tokens[-1]:\n",
    "            output[\n",
    "                0, -1, pred_token\n",
    "            ] -= 1.0  # Decrease the probability of the last token\n",
    "            pred_token = output.argmax(2)[:, -1].item()\n",
    "            pred_prob = torch.softmax(output, dim=-1)[0, -1, pred_token].item()\n",
    "\n",
    "        target_tokens.append(pred_token)\n",
    "\n",
    "        # Break if the probability of the predicted token is below the threshold\n",
    "        if pred_prob < prob_threshold:\n",
    "            break\n",
    "\n",
    "        # Check for sentence repetition and penalize\n",
    "        current_sentence = \" \".join(\n",
    "            [output_dic.index2word[t] for t in target_tokens[1:]]\n",
    "        )\n",
    "        if current_sentence in generated_sentences:\n",
    "            output[0, -1, pred_token] -= 1.0  # Penalize repetition\n",
    "            pred_token = output.argmax(2)[:, -1].item()\n",
    "            target_tokens[-1] = pred_token  # Update with new token\n",
    "        else:\n",
    "            generated_sentences.add(current_sentence)\n",
    "\n",
    "        # Break if the end of sentence token is predicted\n",
    "        if pred_token == EOS_TOKEN:\n",
    "            break\n",
    "\n",
    "    # Convert token IDs to words, ignoring the first token (SOS) and the last token (EOS)\n",
    "    target_results = [output_dic.index2word[i] for i in target_tokens if i != EOS_TOKEN]\n",
    "\n",
    "    return \" \".join(target_results[1:]), attention\n",
    "\n",
    "def translate_documents(\n",
    "    documents: List[str],\n",
    "    input_dic: 'Dictionary',\n",
    "    output_dic: 'Dictionary',\n",
    "    model: 'Transformer',\n",
    "    device: torch.device,\n",
    "    max_len: int,\n",
    "    prob_threshold: float = 0.1,\n",
    ") -> List[str]:\n",
    "    \"\"\"\n",
    "    Translate a list of documents from the input language to the output language using the transformer model.\n",
    "\n",
    "    Args:\n",
    "        documents (List[str]): List of documents to be translated.\n",
    "        input_dic (Dictionary): Input language dictionary.\n",
    "        output_dic (Dictionary): Output language dictionary.\n",
    "        model (Transformer): Transformer model for translation.\n",
    "        device (torch.device): Device to run the model on (CPU or GPU).\n",
    "        max_len (int): Maximum length of the output sentence.\n",
    "        prob_threshold (float): Probability threshold for stopping condition.\n",
    "\n",
    "    Returns:\n",
    "        List[str]: List of translated documents.\n",
    "    \"\"\"\n",
    "    translated_documents = []\n",
    "    for document in documents:\n",
    "        summary = summarize_document(document)\n",
    "        translation, _ = translate_sentence(\n",
    "            summary, input_dic, output_dic, model, device, max_len, prob_threshold\n",
    "        )\n",
    "        translated_documents.append(translation)\n",
    "    return translated_documents\n",
    "\n",
    "def calculate_rouge_scores(hypotheses: List[str], references: List[str]) -> dict:\n",
    "    \"\"\"\n",
    "    Calculate ROUGE scores for a list of hypotheses against reference summaries.\n",
    "\n",
    "    Args:\n",
    "        hypotheses (List[str]): List of generated summaries.\n",
    "        references (List[str]): List of reference summaries.\n",
    "\n",
    "    Returns:\n",
    "        dict: ROUGE scores dictionary.\n",
    "    \"\"\"\n",
    "    scorer = rouge_scorer.RougeScorer(['rouge1', 'rouge2', 'rougeL'], use_stemmer=True)\n",
    "    scores = {'rouge1': 0.0, 'rouge2': 0.0, 'rougeL': 0.0}\n",
    "    num_summaries = len(hypotheses)\n",
    "\n",
    "    for hypothesis, reference in zip(hypotheses, references):\n",
    "        if isinstance(reference, list):\n",
    "            reference = reference[0]  # Ensure reference is a single string\n",
    "        if isinstance(hypothesis, list):\n",
    "            hypothesis = hypothesis[0]  # Ensure hypothesis is a single string\n",
    "        score = scorer.score(reference, hypothesis)\n",
    "        for key in scores:\n",
    "            scores[key] += score[key].fmeasure\n",
    "\n",
    "    # Average scores\n",
    "    scores = {key: value / num_summaries for key, value in scores.items()}\n",
    "    return scores\n",
    "\n",
    "def main() -> None:\n",
    "    # Set parameters as variables\n",
    "    dataset_en = load_dataset(\"wiki_lingua\", \"english\")\n",
    "    input_lang = \"english\"\n",
    "    output_lang = \"french\"\n",
    "    models_dir = \"saved_models/\"\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "    # Hyperparameters\n",
    "    MAX_LENGTH = 60\n",
    "    hidden_size = 256\n",
    "    encoder_layers = 3\n",
    "    decoder_layers = 3\n",
    "    encoder_heads = 8\n",
    "    decoder_heads = 8\n",
    "    encoder_ff_size = 512\n",
    "    decoder_ff_size = 512\n",
    "    encoder_dropout = 0.1\n",
    "    decoder_dropout = 0.1\n",
    "\n",
    "    transformer_location = f\"{models_dir}{input_lang}2{output_lang}/\"\n",
    "\n",
    "    # Load dictionaries\n",
    "    input_lang_dic = load_dictionary(transformer_location + \"input_dic.pkl\")\n",
    "    output_lang_dic = load_dictionary(transformer_location + \"output_dic.pkl\")\n",
    "\n",
    "    input_size = input_lang_dic.n_count\n",
    "    output_size = output_lang_dic.n_count\n",
    "\n",
    "    # Define models\n",
    "    encoder_part = Encoder(\n",
    "        input_size,\n",
    "        hidden_size,\n",
    "        encoder_layers,\n",
    "        encoder_heads,\n",
    "        encoder_ff_size,\n",
    "        encoder_dropout,\n",
    "        device,\n",
    "    )\n",
    "    decoder_part = Decoder(\n",
    "        output_size,\n",
    "        hidden_size,\n",
    "        decoder_layers,\n",
    "        decoder_heads,\n",
    "        decoder_ff_size,\n",
    "        decoder_dropout,\n",
    "        device,\n",
    "    )\n",
    "\n",
    "    translator = Transformer(encoder_part, decoder_part, device).to(device)\n",
    "    # Correctly load model state dictionary\n",
    "    translator.load_state_dict(\n",
    "        torch.load(\n",
    "            transformer_location + \"transformer_model.pt\",\n",
    "            map_location=torch.device(\"cpu\"),  # remove this when cuda available NOTE: SAI\n",
    "        )\n",
    "    )\n",
    "\n",
    "    # Randomly select 5 documents\n",
    "    random_docs = random.sample(dataset_en[\"train\"][\"article\"], 5)\n",
    "    documents = []\n",
    "    reference_summaries = []\n",
    "    for doc in random_docs:\n",
    "        if len(doc[\"document\"]) > 1:\n",
    "            document_text = doc[\"document\"][1]\n",
    "            documents.append(document_text)\n",
    "            reference_summaries.append(doc[\"summary\"])\n",
    "        else:\n",
    "            print(f\"Skipping document due to insufficient structure: {doc}\")\n",
    "\n",
    "    # Translate the summarized documents\n",
    "    translations = translate_documents(\n",
    "        documents, input_lang_dic, output_lang_dic, translator, device, MAX_LENGTH\n",
    "    )\n",
    "\n",
    "    # Print results and calculate ROUGE scores\n",
    "    generated_summaries = [summarize_document(doc) for doc in documents]\n",
    "\n",
    "    for i, (original_doc, summary, translation, reference_summary) in enumerate(\n",
    "        zip(documents, generated_summaries, translations, reference_summaries)\n",
    "    ):\n",
    "        print(f\"Document {i+1} - Original: {original_doc[:300]}...\")  # Print a snippet of the original document\n",
    "        print(f\"Document {i+1} - Summary: {summary}\")\n",
    "        print(f\"Document {i+1} - Translated Summary: {translation}\\n\")\n",
    "\n",
    "    # Calculate ROUGE scores\n",
    "    rouge_scores = calculate_rouge_scores(generated_summaries, reference_summaries)\n",
    "    print(f\"ROUGE Scores: {rouge_scores}\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ece6c719-aa8a-4584-b3f2-7ec0b05fbfa0",
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}