examples/seq2seq.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "if 'has_changed_dir' not in globals():\n",
    "    repo_path = os.path.abspath(os.path.join('..'))\n",
    "    \n",
    "    if repo_path not in sys.path:\n",
    "        sys.path.append(repo_path)\n",
    "    \n",
    "    os.chdir(repo_path)\n",
    "    \n",
    "    globals()['has_changed_dir'] = True\n",
    "\n",
    "import math\n",
    "from pathlib import Path\n",
    "from typing import Optional\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from tokenizers.implementations import ByteLevelBPETokenizer\n",
    "from tokenizers.processors import TemplateProcessing\n",
    "from tqdm import tqdm\n",
    "\n",
    "import neunet\n",
    "import neunet.nn as nn\n",
    "from datasets import load_dataset  # type: ignore\n",
    "from neunet import Tensor\n",
    "from neunet.optim import Adam\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Seq2Seq Transformer for language translation from English to German\n",
    "\"\"\"\n",
    "\n",
    "# [Model implementation]\n",
    "\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, d_model, n_heads, dropout=0.1):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        self.d_model = d_model\n",
    "        self.n_heads = n_heads\n",
    "        self.scale = math.sqrt(d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "        if d_model % n_heads != 0:\n",
    "            raise ValueError(\"d_model must be divisible by n_heads\")\n",
    "\n",
    "        self.depth = d_model // n_heads\n",
    "\n",
    "        self.wq = nn.Linear(d_model, d_model, bias = False)\n",
    "        self.wk = nn.Linear(d_model, d_model, bias = False)\n",
    "        self.wv = nn.Linear(d_model, d_model, bias = False)\n",
    "\n",
    "        self.fc = nn.Linear(d_model, d_model)\n",
    "\n",
    "    def forward(self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor]=None):\n",
    "        batch_size = q.shape[0]\n",
    "        q = self.wq(q).contiguous().reshape(batch_size, -1, self.n_heads, self.depth).transpose(0, 2, 1, 3)\n",
    "        k = self.wk(k).contiguous().reshape(batch_size, -1, self.n_heads, self.depth).transpose(0, 2, 1, 3)\n",
    "        v = self.wv(v).contiguous().reshape(batch_size, -1, self.n_heads, self.depth).transpose(0, 2, 1, 3)\n",
    "\n",
    "        # q = self.wq(q).contiguous().reshape(batch_size, self.n_heads, -1, self.depth)\n",
    "        # k = self.wk(k).contiguous().reshape(batch_size, self.n_heads, -1, self.depth)\n",
    "        # v = self.wv(v).contiguous().reshape(batch_size, self.n_heads, -1, self.depth)  \n",
    "\n",
    "        scores = neunet.matmul(q, k.transpose(0, 1, 3, 2)) / self.scale\n",
    "        if mask is not None:\n",
    "            # mask = mask.unsqueeze(1)#.repeat(1, self.n_heads, 1, 1)\n",
    "            mask = mask[:, None, ...]\n",
    "            # scores = scores.masked_fill(mask == 0, -1e9)\n",
    "            scores = neunet.where(mask == 0, -1e9, scores)\n",
    "\n",
    "        attn = self.dropout(nn.Softmax(axis = -1)(scores))\n",
    "\n",
    "        x = neunet.matmul(attn, v)\n",
    "        x = x.contiguous().transpose(0, 2, 1, 3).reshape(batch_size, -1, self.n_heads * self.depth)\n",
    "        # x = x.contiguous().reshape(batch_size, -1, self.n_heads * self.depth)\n",
    "        x = self.fc(x)\n",
    "\n",
    "        return x, attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionwiseFeedForward(nn.Module):\n",
    "    def __init__(self, d_model, d_ff, dropout=0.1):\n",
    "        super(PositionwiseFeedForward, self).__init__()\n",
    "        self.fc_1 = nn.Linear(d_model, d_ff)\n",
    "        self.fc_2 = nn.Linear(d_ff, d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "        self.activation = nn.ReLU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fc_1(x)\n",
    "        x = self.activation(x)\n",
    "        x = self.dropout(x)\n",
    "        x = self.fc_2(x)\n",
    "\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderLayer(nn.Module):\n",
    "    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)\n",
    "        self.ffn = PositionwiseFeedForward(d_model, d_ff, dropout)\n",
    "        self.norm1 = nn.LayerNorm(d_model, eps = 0.001)\n",
    "        self.norm2 = nn.LayerNorm(d_model, eps = 0.001)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, src, src_mask=None):\n",
    "        # Self-attention\n",
    "        _src, _ = self.self_attn(src, src, src, src_mask)\n",
    "        src = src + self.dropout(_src)\n",
    "        src = self.norm1(src)\n",
    "\n",
    "        # Feed-forward network\n",
    "        _src = self.ffn(src)\n",
    "        src = src + self.dropout(_src)\n",
    "        src = self.norm2(src)\n",
    "\n",
    "        return src"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderLayer(nn.Module):\n",
    "    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):\n",
    "        super(DecoderLayer, self).__init__()\n",
    "        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)\n",
    "        self.cross_attn = MultiHeadAttention(d_model, n_heads, dropout)\n",
    "        self.ffn = PositionwiseFeedForward(d_model, d_ff, dropout)\n",
    "        self.norm1 = nn.LayerNorm(d_model, eps = 0.001)\n",
    "        self.norm2 = nn.LayerNorm(d_model, eps = 0.001)\n",
    "        self.norm3 = nn.LayerNorm(d_model, eps = 0.001)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, tgt, tgt_mask, src, src_mask):\n",
    "        # Masked self-attention (for the target sequence)\n",
    "        _tgt, _ = self.self_attn(tgt, tgt, tgt, tgt_mask)\n",
    "        tgt = tgt + self.dropout(_tgt)\n",
    "        tgt = self.norm1(tgt)\n",
    "\n",
    "        # Cross-attention (attending to the encoder output)\n",
    "        _tgt, attn = self.cross_attn(tgt, src, src, src_mask)\n",
    "        tgt = tgt + self.dropout(_tgt)\n",
    "        tgt = self.norm2(tgt)\n",
    "\n",
    "        # Feed-forward network\n",
    "        _tgt = self.ffn(tgt)\n",
    "        tgt = tgt + self.dropout(_tgt)\n",
    "        tgt = self.norm3(tgt)\n",
    "\n",
    "        return tgt, attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "    def __init__(self, d_model, max_len=5000):\n",
    "        \"\"\"\n",
    "        Inputs\n",
    "            d_model - Hidden dimensionality of the input.\n",
    "            max_len - Maximum length of a sequence to expect.\n",
    "            https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/tutorial6/Transformers_and_MHAttention.html\n",
    "        \"\"\"\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "        pe = neunet.zeros(max_len, d_model, requires_grad=False)\n",
    "        position = neunet.arange(0, max_len, dtype=neunet.float32)[:, None, ...]\n",
    "        div_term = neunet.exp(neunet.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))\n",
    "        pe[:, 0::2] = neunet.sin(position * div_term)\n",
    "        pe[:, 1::2] = neunet.cos(position * div_term)\n",
    "        # print(pe.shape, pe[None, ...].shape,  pe[None, ...].data.shape)\n",
    "        self.pe = pe[None, ...]\n",
    "\n",
    "    def forward(self, x):\n",
    "        # print(x.shape, self.pe.shape, self.pe[:, :x.shape[1]].shape)\n",
    "        x = x + self.pe[:, :x.shape[1]] # (batch_size, seq_len, d_model)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, src_vocab_size, d_model, n_heads, d_ff, n_layers, dropout=0.1, max_len=5000):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.token_embedding = nn.Embedding(src_vocab_size, d_model)\n",
    "        self.position_embedding = PositionalEncoding(d_model, max_len)\n",
    "        self.layers = nn.ModuleList([EncoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)])\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "        self.scale = math.sqrt(d_model)\n",
    "\n",
    "    def forward(self, src, mask=None):\n",
    "        src = self.token_embedding(src)  * self.scale\n",
    "        src = self.position_embedding(src)\n",
    "        src = self.dropout(src)\n",
    "\n",
    "        for layer in self.layers:\n",
    "            src = layer(src, mask)\n",
    "        return src"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, tgt_vocab_size, d_model, n_heads, d_ff, n_layers, dropout=0.1, max_len=5000):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.token_embedding = nn.Embedding(tgt_vocab_size, d_model)\n",
    "        self.position_embedding = PositionalEncoding(d_model, max_len)\n",
    "        self.layers = nn.ModuleList([DecoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)])\n",
    "\n",
    "        self.fc_out = nn.Linear(d_model, tgt_vocab_size)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "        self.scale = math.sqrt(d_model)\n",
    "\n",
    "\n",
    "    def forward(self, tgt, tgt_mask, src, src_mask):\n",
    "        tgt = self.token_embedding(tgt) * self.scale\n",
    "        tgt = self.position_embedding(tgt)\n",
    "        tgt = self.dropout(tgt)\n",
    "        for layer in self.layers:\n",
    "            tgt, attn = layer(tgt, tgt_mask, src, src_mask)\n",
    "\n",
    "        out = self.fc_out(tgt)\n",
    "        return out, attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Seq2SeqTransformer(nn.Module):\n",
    "    def __init__(self, encoder, decoder, pad_idx) -> None:\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "\n",
    "        self.pad_idx = pad_idx\n",
    "\n",
    "    def get_pad_mask(self,  x: np.ndarray) -> np.ndarray:\n",
    "        #x: (batch_size, seq_len)\n",
    "        return (x != self.pad_idx).astype(int)[:, np.newaxis, :]\n",
    "\n",
    "    def get_sub_mask(self, x: np.ndarray) -> np.ndarray:\n",
    "        #x: (batch_size, seq_len)\n",
    "        seq_len = x.shape[1]\n",
    "        subsequent_mask = np.triu(np.ones((seq_len, seq_len)), k = 1).astype(int)\n",
    "        subsequent_mask = np.logical_not(subsequent_mask)\n",
    "        return subsequent_mask\n",
    "\n",
    "    def forward(self, src: np.ndarray, tgt: np.ndarray) -> tuple[Tensor, Tensor]:\n",
    "        #src: (batch_size, source_seq_len)\n",
    "        #tgt: (batch_size, target_seq_len)\n",
    "\n",
    "        # src_mask: (batch_size, 1, seq_len)\n",
    "        # tgt_mask: (batch_size, seq_len, seq_len)\n",
    "        src_mask = self.get_pad_mask(src)\n",
    "\n",
    "        tgt_mask = self.get_pad_mask(tgt) & self.get_sub_mask(tgt)\n",
    "\n",
    "        src, src_mask = neunet.tensor(src, dtype=neunet.int32, device=device), neunet.tensor(src_mask, dtype=neunet.int32, device=device)\n",
    "        tgt, tgt_mask = neunet.tensor(tgt, dtype=neunet.int32, device=device), neunet.tensor(tgt_mask, dtype=neunet.int32, device=device)\n",
    "        \n",
    "        enc_src = self.encoder(src, src_mask)\n",
    "\n",
    "        out, attention = self.decoder(tgt, tgt_mask, enc_src, src_mask)\n",
    "        # out: (batch_size, target_seq_len, vocab_size)\n",
    "        # attention: (batch_size, heads_num, target_seq_len, source_seq_len)\n",
    "        return out, attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# [Data preprocessing]\n",
    "\n",
    "BATCH_SIZE = 32\n",
    "\n",
    "PAD_TOKEN = '<pad>' # noqa: S105\n",
    "SOS_TOKEN = '<sos>' # noqa: S105\n",
    "EOS_TOKEN = '<eos>' # noqa: S105\n",
    "# UNK_TOKEN = '<unk>'\n",
    "\n",
    "DATASET_PATH = Path(\"./datasets/multi30k/\")\n",
    "SAVE_PATH = Path(\"./saved models/seq2seq/\")\n",
    "\n",
    "if not DATASET_PATH.exists():\n",
    "    data = load_dataset(\"bentrevett/multi30k\", cache_dir=\"datasets/multi30k\")\n",
    "\n",
    "    for split, split_dataset in data.items():\n",
    "        with Path(f\"./datasets/multi30k/{split}.en\").open('w', encoding='utf-8') as f:\n",
    "            for item in split_dataset:\n",
    "                f.write(item['en'] + '\\n')\n",
    "\n",
    "        with Path(f\"./datasets/multi30k/{split}.de\").open('w', encoding='utf-8') as f:\n",
    "            for item in split_dataset:\n",
    "                f.write(item['de'] + '\\n')\n",
    "\n",
    "FILE_PATHS = [DATASET_PATH / \"train.en\", DATASET_PATH / \"train.de\", DATASET_PATH / \"val.en\", DATASET_PATH / \"val.de\", DATASET_PATH / \"test.en\", DATASET_PATH / \"test.de\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# [Train and load Tokenizer]\n",
    "if not (SAVE_PATH / \"vocab\").exists():\n",
    "    tokenizer = ByteLevelBPETokenizer()\n",
    "\n",
    "    tokenizer.train(files=[str(path) for path in FILE_PATHS], vocab_size=15000, min_frequency=1, special_tokens=[\n",
    "        PAD_TOKEN,\n",
    "        SOS_TOKEN,\n",
    "        EOS_TOKEN,\n",
    "        # UNK_TOKEN\n",
    "    ])\n",
    "\n",
    "    (SAVE_PATH / \"vocab\").mkdir(parents=True)\n",
    "    tokenizer.save_model(str(SAVE_PATH / \"vocab\"), \"multi30k-tokenizer\")\n",
    "\n",
    "tokenizer = ByteLevelBPETokenizer(\n",
    "    str(SAVE_PATH / \"vocab/multi30k-tokenizer-vocab.json\"),\n",
    "    str(SAVE_PATH / \"vocab/multi30k-tokenizer-merges.txt\"),\n",
    ")\n",
    "\n",
    "\n",
    "\n",
    "PAD_INDEX = tokenizer.token_to_id(PAD_TOKEN)\n",
    "SOS_INDEX = tokenizer.token_to_id(SOS_TOKEN)\n",
    "EOS_INDEX = tokenizer.token_to_id(EOS_TOKEN)\n",
    "# UNK_INDEX = tokenizer.token_to_id(UNK_TOKEN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DataPreprocessor():\n",
    "    def __init__(self, tokenizer: ByteLevelBPETokenizer):\n",
    "        self.tokenizer = tokenizer\n",
    "\n",
    "        self.tokenizer._tokenizer.post_processor  = TemplateProcessing( # noqa SLF001\n",
    "            single=f\"{SOS_TOKEN} $A {EOS_TOKEN}\",\n",
    "            special_tokens=[\n",
    "                (f\"{SOS_TOKEN}\", tokenizer.token_to_id(f\"{SOS_TOKEN}\")),\n",
    "                (f\"{EOS_TOKEN}\", tokenizer.token_to_id(f\"{EOS_TOKEN}\")),\n",
    "            ],\n",
    "        )\n",
    "\n",
    "        # self.tokenizer.enable_truncation(max_length=128)\n",
    "        self.tokenizer.enable_padding(pad_token = PAD_TOKEN)\n",
    "        \n",
    "    def tokenize(self, paths: list[str], batch_size: int, lines_limit: Optional[int] = None) -> list[np.ndarray]:\n",
    "        examples = []\n",
    "\n",
    "        for src_file in paths:\n",
    "            print(f\"Processing {src_file}\")\n",
    "            path_src_file = Path(src_file)\n",
    "            lines = path_src_file.read_text(encoding=\"utf-8\").splitlines()\n",
    "\n",
    "            if lines_limit:\n",
    "                lines = lines[:lines_limit]\n",
    "\n",
    "            for i in range(0, len(lines), batch_size):\n",
    "                examples.append(np.array([x.ids for x in self.tokenizer.encode_batch(lines[i:i+batch_size])]))\n",
    "\n",
    "\n",
    "        return examples\n",
    "\n",
    "    def __call__(self, paths: list[str], batch_size: int, lines_limit: Optional[int] = None) -> list[np.ndarray]:\n",
    "        return self.tokenize(paths, batch_size, lines_limit)\n",
    "\n",
    "\n",
    "data_post_processor = DataPreprocessor(tokenizer)\n",
    "\n",
    "train_src = data_post_processor([str(DATASET_PATH / \"train.en\")], batch_size = BATCH_SIZE)\n",
    "train_tgt = data_post_processor([str(DATASET_PATH / \"train.de\")], batch_size = BATCH_SIZE)\n",
    "\n",
    "val_src = data_post_processor([str(DATASET_PATH / \"val.en\")], batch_size = BATCH_SIZE)\n",
    "val_tgt = data_post_processor([str(DATASET_PATH / \"val.de\")], batch_size = BATCH_SIZE)\n",
    "\n",
    "test_src = data_post_processor([str(DATASET_PATH / \"test.en\")], batch_size = BATCH_SIZE)\n",
    "test_tgt = data_post_processor([str(DATASET_PATH / \"test.de\")], batch_size = BATCH_SIZE)\n",
    "\n",
    "\n",
    "train_data = train_src, train_tgt\n",
    "val_data = val_src, val_tgt\n",
    "test_data = test_src, test_tgt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# [Model intialization]\n",
    "\n",
    "encoder = Encoder(\n",
    "    src_vocab_size = tokenizer.get_vocab_size(),\n",
    "    d_model = 256, # 512\n",
    "    n_heads = 8,\n",
    "    d_ff = 512, # 2048\n",
    "    n_layers = 3, # 6\n",
    "    dropout = 0.1\n",
    ")\n",
    "\n",
    "decoder = Decoder(\n",
    "    tgt_vocab_size = tokenizer.get_vocab_size(),\n",
    "    d_model = 256, # 512\n",
    "    n_heads = 8,\n",
    "    d_ff = 512, # 2048\n",
    "    n_layers = 3, # 6\n",
    "    dropout = 0.1\n",
    ")\n",
    "\n",
    "\n",
    "model = Seq2SeqTransformer(\n",
    "    encoder = encoder,\n",
    "    decoder = decoder,\n",
    "    pad_idx = PAD_INDEX\n",
    ")\n",
    "\n",
    "model = model.to(device)\n",
    "\n",
    "optimizer = Adam(model.parameters(), lr=3e-4, betas=(0.9, 0.98), eps = 1e-9)\n",
    "\n",
    "loss_function = nn.CrossEntropyLoss(ignore_index = PAD_INDEX)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# [train, eval, predict methods definition]\n",
    "\n",
    "def train_step(source: list[np.ndarray], target: list[np.ndarray], epoch: int, epochs: int) -> float:\n",
    "    total_loss = 0.0\n",
    "    model.train()\n",
    "\n",
    "    tqdm_range = tqdm(enumerate(zip(source, target, strict=False)), total = len(source))\n",
    "    for batch_num, (source_batch, target_batch) in tqdm_range:\n",
    "\n",
    "        output, _ = model.forward(source_batch, target_batch[:,:-1])\n",
    "        \n",
    "        output = output.reshape(output.shape[0] * output.shape[1], output.shape[2])\n",
    "\n",
    "        loss = loss_function(output, neunet.tensor(target_batch[:, 1:].flatten(), device=device, dtype=neunet.int32))\n",
    "        loss.backward()\n",
    "\n",
    "        optimizer.step()\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        total_loss += loss.detach().item()\n",
    "\n",
    "\n",
    "        tqdm_range.set_description(\n",
    "                f\"training | loss: {loss.detach().item():.7f} | perplexity: {np.exp(loss.detach().item()):.7f} | epoch {epoch + 1}/{epochs}\" #loss: {loss:.4f}\n",
    "            )\n",
    "\n",
    "        if batch_num == (len(source) - 1):\n",
    "            epoch_loss = total_loss / len(source)\n",
    "\n",
    "            tqdm_range.set_description(\n",
    "                    f\"training | avg loss: {epoch_loss:.7f} | avg perplexity: {np.exp(epoch_loss):.7f} | epoch {epoch + 1}/{epochs}\"\n",
    "            )\n",
    "\n",
    "    return epoch_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval(source: list[np.ndarray], target: list[np.ndarray]) -> float:\n",
    "    total_loss = 0.0\n",
    "    model.eval()\n",
    "\n",
    "    tqdm_range = tqdm(enumerate(zip(source, target, strict=False)), total = len(source))\n",
    "    for batch_num, (source_batch, target_batch) in tqdm_range:\n",
    "        \n",
    "        output, _ = model.forward(source_batch, target_batch[:,:-1])\n",
    "        \n",
    "        output = output.reshape(output.shape[0] * output.shape[1], output.shape[2])\n",
    "        \n",
    "        loss = loss_function(output, neunet.tensor(target_batch[:, 1:].flatten(), device=device, dtype=neunet.int32))\n",
    "        total_loss += loss.detach().item()\n",
    "        \n",
    "        tqdm_range.set_description(\n",
    "                f\"testing  | loss: {loss.detach().item():.7f} | perplexity: {np.exp(loss.detach().item()):.7f}\"\n",
    "            )\n",
    "\n",
    "        if batch_num == (len(source) - 1):\n",
    "            epoch_loss = total_loss / len(source)\n",
    "\n",
    "            tqdm_range.set_description(\n",
    "                    f\"testing  | avg loss: {epoch_loss:.7f} | avg perplexity: {np.exp(epoch_loss):.7f}\"\n",
    "            )\n",
    "\n",
    "    return epoch_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(train_data: tuple[list[np.ndarray], list[np.ndarray]], val_data: tuple[list[np.ndarray], list[np.ndarray]], epochs: int, save_every_epochs: int, save_path: Optional[str] = None, validation_check: bool = False):\n",
    "    best_val_loss = float('inf')\n",
    "    \n",
    "    train_loss_history = []\n",
    "    val_loss_history = []\n",
    "\n",
    "    train_source, train_target = train_data\n",
    "    val_source, val_target = val_data\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "\n",
    "        train_loss_history.append(train_step(train_source, train_target, epoch, epochs))\n",
    "        val_loss_history.append(eval(val_source, val_target))\n",
    "\n",
    "\n",
    "        if (save_path is not None) and ((epoch + 1) % save_every_epochs == 0):\n",
    "            if not Path(save_path).exists():\n",
    "                Path(save_path).mkdir(parents=True, exist_ok=True)\n",
    "            if validation_check == False:\n",
    "\n",
    "                neunet.save(model.state_dict(), f\"{save_path}/seq2seq_{epoch + 1}.nt\")\n",
    "            else:\n",
    "                if val_loss_history[-1] < best_val_loss:\n",
    "                    best_val_loss = val_loss_history[-1]\n",
    "                    \n",
    "                    neunet.save(model.state_dict(), f\"{save_path}/seq2seq_{epoch + 1}.nt\")\n",
    "                else:\n",
    "                    print('Current validation loss is higher than previous. Not saved.')\n",
    "                    break\n",
    "            \n",
    "    return train_loss_history, val_loss_history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(sentence: str, max_length: int = 50) -> tuple[str, Tensor]:\n",
    "    model.eval()\n",
    "\n",
    "    tokens = tokenizer.encode(sentence)\n",
    "    src_ids = tokens.ids\n",
    "    # src_ids = [SOS_INDEX] + src_ids + [EOS_INDEX] # Special tokens already here\n",
    "    \n",
    "    src = np.asarray(src_ids).reshape(1, -1)\n",
    "    src_mask =  model.get_pad_mask(src)\n",
    "\n",
    "    src, src_mask = neunet.tensor(src, dtype=neunet.int32, device=device), neunet.tensor(src_mask, dtype=neunet.int32, device=device)\n",
    "\n",
    "    enc_src = model.encoder.forward(src, src_mask)\n",
    "\n",
    "    tgt_tokens = [SOS_INDEX]\n",
    "\n",
    "    for _ in range(max_length):\n",
    "        tgt = np.asarray(tgt_tokens).reshape(1, -1)\n",
    "        tgt_mask = model.get_pad_mask(tgt) & model.get_sub_mask(tgt)\n",
    "\n",
    "        tgt, tgt_mask = neunet.tensor(tgt, dtype=neunet.int32, device=device), neunet.tensor(tgt_mask, dtype=neunet.int32, device=device)\n",
    "\n",
    "        outputs, attention = model.decoder.forward(tgt, tgt_mask, enc_src, src_mask)\n",
    "        \n",
    "        tgt_next_token = outputs.detach().cpu().numpy().argmax(axis=-1)[:, -1].item()\n",
    "        tgt_tokens.append(tgt_next_token)\n",
    "\n",
    "        if tgt_next_token == EOS_INDEX or len(tgt_tokens) >= max_length:\n",
    "            break\n",
    "    \n",
    "    \n",
    "    # Remove special tokens\n",
    "    if SOS_INDEX in tgt_tokens:\n",
    "        tgt_tokens.remove(SOS_INDEX)\n",
    "    if EOS_INDEX in tgt_tokens:\n",
    "        tgt_tokens.remove(EOS_INDEX)\n",
    "\n",
    "    decoded_sentence = tokenizer.decode(tgt_tokens)\n",
    "\n",
    "    return decoded_sentence, attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# [Train the Model]\n",
    "\n",
    "# model.load_state_dict(neunet.load(\"./saved models/seq2seq/seq2seq_8.nt\"))\n",
    "\n",
    "train_loss_history, val_loss_history = None, None\n",
    "train_loss_history, val_loss_history = train(train_data, val_data, epochs=30, save_every_epochs=1, save_path = str(SAVE_PATH), validation_check=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# [Model inferecnce and Plot]\n",
    "\n",
    "def plot_loss_history(train_loss_history, val_loss_history):\n",
    "    plt.plot(train_loss_history)\n",
    "    plt.plot(val_loss_history)\n",
    "    plt.title('Loss history')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.legend(['train', 'val'], loc='upper left')\n",
    "    plt.show()\n",
    "        \n",
    "if train_loss_history is not None and val_loss_history is not None:\n",
    "    plot_loss_history(train_loss_history, val_loss_history)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_test_data: list[dict[str, str]] = []\n",
    "\n",
    "en_file = [l.strip() for l in Path(DATASET_PATH / \"test.en\").open('r', encoding='utf-8')]\n",
    "de_file = [l.strip() for l in Path(DATASET_PATH / \"test.de\").open('r', encoding='utf-8')]\n",
    "\n",
    "for i in range(len(en_file)):\n",
    "    if en_file[i] == '' or de_file[i] == '':\n",
    "        continue\n",
    "    en_seq, de_seq = en_file[i], de_file[i]\n",
    "\n",
    "    raw_test_data.append({'en': en_seq, 'de': de_seq})\n",
    "    \n",
    "sentences_num = 10\n",
    "\n",
    "random_indices = np.random.randint(0, len(raw_test_data), sentences_num)\n",
    "sentences_selection = [raw_test_data[i] for i in random_indices]\n",
    "\n",
    "# [Translate sentences from validation set]\n",
    "for i, example in enumerate(sentences_selection):\n",
    "    print(f\"\\nExample №{i + 1}\")\n",
    "    print(f\"Input sentence: {example['en']}\")\n",
    "    print(f\"Decoded sentence: {predict(example['en'])[0]}\")\n",
    "    print(f\"Target sentence: {example['de']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_attention(sentence: str, translation: str, attention: Tensor, heads_num: int = 8, rows_num: int = 2, cols_num: int = 4):\n",
    "    if rows_num * cols_num != heads_num:\n",
    "        raise ValueError(\"heads_num must be equal to rows_num * cols_num\")\n",
    "    attention = attention.detach().cpu().numpy().squeeze()\n",
    "\n",
    "    sentence = tokenizer.encode(sentence, add_special_tokens=False).tokens\n",
    "    translation = tokenizer.encode(translation, add_special_tokens=False).tokens\n",
    "\n",
    "\n",
    "    fig = plt.figure(figsize = (15, 25))\n",
    "    \n",
    "    for h in range(heads_num):\n",
    "        \n",
    "        ax = fig.add_subplot(rows_num, cols_num, h + 1)\n",
    "        ax.set_xlabel(f'Head {h + 1}')\n",
    "        \n",
    "        ax.matshow(attention[h], cmap = 'inferno')\n",
    "\n",
    "        ax.tick_params(labelsize = 7)\n",
    "\n",
    "        ax.set_xticks(range(len(sentence)))\n",
    "        ax.set_yticks(range(len(translation)))\n",
    "\n",
    "        ax.set_xticklabels(sentence, rotation=90)\n",
    "        ax.set_yticklabels(translation)\n",
    "\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "# [Plot Attention]\n",
    "sentence = sentences_selection[0]['en']\n",
    "print(f\"\\nInput sentence: {sentence}\")\n",
    "decoded_sentence, attention =  predict(sentence)\n",
    "print(f\"Decoded sentence: {decoded_sentence}\")\n",
    "\n",
    "plot_attention(sentence, decoded_sentence, attention)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}