{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train Yorùbá Seq2Seq Diacritizer\n",
    "\n",
    "This notebook trains a **Seq2Seq (Encoder-Decoder)** model to restore diacritics to Yorùbá text.\n",
    "\n",
    "**Why Seq2Seq?**\n",
    "- Handles different input/output lengths (diacritics can change character count)\n",
    "- Uses **100% of the data** (676k pairs) instead of 16%\n",
    "- Expected accuracy: **93-95%**\n",
    "\n",
    "**Expected results:**\n",
    "- Training time: ~1-2 hours on Colab GPU\n",
    "- Word accuracy: 93-95%\n",
    "- Model size: 3-5MB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 1: Install dependencies\n",
    "!pip install -q datasets torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check GPU availability\n",
    "import torch\n",
    "print(f\"PyTorch version: {torch.__version__}\")\n",
    "print(f\"CUDA available: {torch.cuda.is_available()}\")\n",
    "if torch.cuda.is_available():\n",
    "    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n",
    "    device = torch.device(\"cuda\")\n",
    "else:\n",
    "    print(\"No GPU, using CPU (will be slower)\")\n",
    "    device = torch.device(\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 2: Load ALL data from HuggingFace (no length filtering!)\n",
    "from datasets import load_dataset\n",
    "import unicodedata\n",
    "\n",
    "print(\"Loading dataset...\")\n",
    "ds = load_dataset(\"bumie-e/Yoruba-diacritics-vs-non-diacritics\", split=\"train\")\n",
    "print(f\"Total samples: {len(ds)}\")\n",
    "\n",
    "# Extract ALL text pairs - Seq2Seq handles different lengths!\n",
    "undiacritized = []\n",
    "diacritized = []\n",
    "\n",
    "for item in ds:\n",
    "    undiac = item.get(\"no_diacritcs\", \"\")\n",
    "    diac = item.get(\"diacritcs\", \"\")\n",
    "    \n",
    "    # Normalize to NFC\n",
    "    undiac = unicodedata.normalize(\"NFC\", undiac)\n",
    "    diac = unicodedata.normalize(\"NFC\", diac)\n",
    "    \n",
    "    # Only skip empty or very long sentences\n",
    "    if undiac and diac and len(undiac) < 200 and len(diac) < 200:\n",
    "        undiacritized.append(undiac)\n",
    "        diacritized.append(diac)\n",
    "\n",
    "print(f\"Valid pairs: {len(undiacritized)} (using ALL data!)\")\n",
    "print(f\"\\nExample:\")\n",
    "print(f\"  Input:  {undiacritized[0]} (len={len(undiacritized[0])})\")\n",
    "print(f\"  Output: {diacritized[0]} (len={len(diacritized[0])})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 3: Build vocabulary with special tokens for Seq2Seq\n",
    "import json\n",
    "from typing import Dict, List, Tuple\n",
    "\n",
    "class CharVocab:\n",
    "    \"\"\"Character vocabulary with special tokens for Seq2Seq.\"\"\"\n",
    "    PAD = \"<PAD>\"\n",
    "    UNK = \"<UNK>\"\n",
    "    SOS = \"<SOS>\"  # Start of sequence\n",
    "    EOS = \"<EOS>\"  # End of sequence\n",
    "\n",
    "    def __init__(self):\n",
    "        self.char2idx: Dict[str, int] = {\n",
    "            self.PAD: 0, \n",
    "            self.UNK: 1,\n",
    "            self.SOS: 2,\n",
    "            self.EOS: 3,\n",
    "        }\n",
    "        self.idx2char: Dict[int, str] = {v: k for k, v in self.char2idx.items()}\n",
    "\n",
    "    def add_char(self, char: str) -> int:\n",
    "        if char not in self.char2idx:\n",
    "            idx = len(self.char2idx)\n",
    "            self.char2idx[char] = idx\n",
    "            self.idx2char[idx] = char\n",
    "        return self.char2idx[char]\n",
    "\n",
    "    def encode(self, text: str, add_special: bool = False) -> List[int]:\n",
    "        ids = [self.char2idx.get(c, 1) for c in text]\n",
    "        if add_special:\n",
    "            ids = [self.char2idx[self.SOS]] + ids + [self.char2idx[self.EOS]]\n",
    "        return ids\n",
    "\n",
    "    def decode(self, indices: List[int]) -> str:\n",
    "        chars = []\n",
    "        for i in indices:\n",
    "            if i == self.char2idx[self.EOS]:\n",
    "                break\n",
    "            if i not in (self.char2idx[self.PAD], self.char2idx[self.SOS], self.char2idx[self.UNK]):\n",
    "                chars.append(self.idx2char.get(i, \"\"))\n",
    "        return \"\".join(chars)\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self.char2idx)\n",
    "\n",
    "# Build vocabulary from all texts\n",
    "vocab = CharVocab()\n",
    "for text in undiacritized + diacritized:\n",
    "    for char in text:\n",
    "        vocab.add_char(char)\n",
    "\n",
    "print(f\"Vocabulary size: {len(vocab)} (including PAD, UNK, SOS, EOS)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 4: Create PyTorch datasets with SOS/EOS for decoder\n",
    "import random\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "\n",
    "class DiacritizationDataset(Dataset):\n",
    "    def __init__(self, undiac: List[str], diac: List[str], vocab: CharVocab):\n",
    "        self.undiac = undiac\n",
    "        self.diac = diac\n",
    "        self.vocab = vocab\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self.undiac)\n",
    "\n",
    "    def __getitem__(self, idx: int):\n",
    "        # Encoder input: just the characters\n",
    "        src = torch.tensor(self.vocab.encode(self.undiac[idx]), dtype=torch.long)\n",
    "        # Decoder input: SOS + characters (teacher forcing)\n",
    "        tgt_in = torch.tensor(\n",
    "            [self.vocab.char2idx[self.vocab.SOS]] + self.vocab.encode(self.diac[idx]), \n",
    "            dtype=torch.long\n",
    "        )\n",
    "        # Decoder target: characters + EOS\n",
    "        tgt_out = torch.tensor(\n",
    "            self.vocab.encode(self.diac[idx]) + [self.vocab.char2idx[self.vocab.EOS]], \n",
    "            dtype=torch.long\n",
    "        )\n",
    "        return src, tgt_in, tgt_out\n",
    "\n",
    "def collate_fn(batch):\n",
    "    srcs, tgt_ins, tgt_outs = zip(*batch)\n",
    "    src_lens = torch.tensor([len(s) for s in srcs])\n",
    "    srcs_padded = pad_sequence(srcs, batch_first=True, padding_value=0)\n",
    "    tgt_ins_padded = pad_sequence(tgt_ins, batch_first=True, padding_value=0)\n",
    "    tgt_outs_padded = pad_sequence(tgt_outs, batch_first=True, padding_value=0)\n",
    "    return srcs_padded, tgt_ins_padded, tgt_outs_padded, src_lens\n",
    "\n",
    "# Train/val split\n",
    "random.seed(42)\n",
    "indices = list(range(len(undiacritized)))\n",
    "random.shuffle(indices)\n",
    "\n",
    "split = int(0.95 * len(indices))\n",
    "train_idx = indices[:split]\n",
    "val_idx = indices[split:]\n",
    "\n",
    "train_dataset = DiacritizationDataset(\n",
    "    [undiacritized[i] for i in train_idx],\n",
    "    [diacritized[i] for i in train_idx],\n",
    "    vocab\n",
    ")\n",
    "val_dataset = DiacritizationDataset(\n",
    "    [undiacritized[i] for i in val_idx],\n",
    "    [diacritized[i] for i in val_idx],\n",
    "    vocab\n",
    ")\n",
    "\n",
    "BATCH_SIZE = 64\n",
    "train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)\n",
    "val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=collate_fn)\n",
    "\n",
    "print(f\"Train batches: {len(train_loader)}\")\n",
    "print(f\"Val batches: {len(val_loader)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 5: Define the Encoder-Decoder model with Attention\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class Encoder(nn.Module):\n",
    "    def __init__(self, vocab_size: int, embed_dim: int, hidden_dim: int, num_layers: int = 2):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)\n",
    "        self.lstm = nn.LSTM(\n",
    "            embed_dim, hidden_dim,\n",
    "            num_layers=num_layers,\n",
    "            batch_first=True,\n",
    "            bidirectional=True,\n",
    "            dropout=0.2 if num_layers > 1 else 0\n",
    "        )\n",
    "        # Project bidirectional hidden states to decoder size\n",
    "        self.fc_h = nn.Linear(hidden_dim * 2, hidden_dim)\n",
    "        self.fc_c = nn.Linear(hidden_dim * 2, hidden_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        emb = self.embedding(x)\n",
    "        outputs, (h, c) = self.lstm(emb)\n",
    "        # Combine bidirectional hidden states\n",
    "        h = torch.tanh(self.fc_h(torch.cat([h[-2], h[-1]], dim=1)))\n",
    "        c = torch.tanh(self.fc_c(torch.cat([c[-2], c[-1]], dim=1)))\n",
    "        return outputs, h.unsqueeze(0), c.unsqueeze(0)\n",
    "\n",
    "class Attention(nn.Module):\n",
    "    def __init__(self, enc_dim: int, dec_dim: int):\n",
    "        super().__init__()\n",
    "        self.attn = nn.Linear(enc_dim + dec_dim, dec_dim)\n",
    "        self.v = nn.Linear(dec_dim, 1, bias=False)\n",
    "\n",
    "    def forward(self, hidden, encoder_outputs, mask=None):\n",
    "        # hidden: [1, batch, dec_dim]\n",
    "        # encoder_outputs: [batch, src_len, enc_dim]\n",
    "        batch_size, src_len, _ = encoder_outputs.shape\n",
    "        \n",
    "        hidden = hidden.squeeze(0).unsqueeze(1).repeat(1, src_len, 1)\n",
    "        energy = torch.tanh(self.attn(torch.cat([hidden, encoder_outputs], dim=2)))\n",
    "        attention = self.v(energy).squeeze(2)\n",
    "        \n",
    "        if mask is not None:\n",
    "            attention = attention.masked_fill(mask == 0, -1e10)\n",
    "        \n",
    "        return F.softmax(attention, dim=1)\n",
    "\n",
    "class Decoder(nn.Module):\n",
    "    def __init__(self, vocab_size: int, embed_dim: int, hidden_dim: int, enc_dim: int):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)\n",
    "        self.attention = Attention(enc_dim, hidden_dim)\n",
    "        self.lstm = nn.LSTMCell(embed_dim + enc_dim, hidden_dim)\n",
    "        self.fc = nn.Linear(hidden_dim + enc_dim + embed_dim, vocab_size)\n",
    "        self.dropout = nn.Dropout(0.2)\n",
    "\n",
    "    def forward(self, input_char, hidden, cell, encoder_outputs):\n",
    "        # input_char: [batch]\n",
    "        emb = self.dropout(self.embedding(input_char))  # [batch, embed]\n",
    "        \n",
    "        attn_weights = self.attention(hidden.unsqueeze(0), encoder_outputs)  # [batch, src_len]\n",
    "        context = torch.bmm(attn_weights.unsqueeze(1), encoder_outputs).squeeze(1)  # [batch, enc_dim]\n",
    "        \n",
    "        lstm_input = torch.cat([emb, context], dim=1)\n",
    "        hidden, cell = self.lstm(lstm_input, (hidden, cell))\n",
    "        \n",
    "        output = self.fc(torch.cat([hidden, context, emb], dim=1))\n",
    "        return output, hidden, cell\n",
    "\n",
    "class Seq2SeqDiacritizer(nn.Module):\n",
    "    def __init__(self, vocab_size: int, embed_dim: int = 64, hidden_dim: int = 128, num_layers: int = 2):\n",
    "        super().__init__()\n",
    "        self.encoder = Encoder(vocab_size, embed_dim, hidden_dim, num_layers)\n",
    "        self.decoder = Decoder(vocab_size, embed_dim, hidden_dim, hidden_dim * 2)  # enc is bidirectional\n",
    "\n",
    "    def forward(self, src, tgt_in, teacher_forcing_ratio=0.5):\n",
    "        batch_size, tgt_len = tgt_in.shape\n",
    "        vocab_size = self.decoder.fc.out_features\n",
    "        \n",
    "        # Encode\n",
    "        encoder_outputs, h, c = self.encoder(src)\n",
    "        h = h.squeeze(0)\n",
    "        c = c.squeeze(0)\n",
    "        \n",
    "        # Decode\n",
    "        outputs = torch.zeros(batch_size, tgt_len, vocab_size, device=src.device)\n",
    "        input_char = tgt_in[:, 0]  # SOS token\n",
    "        \n",
    "        for t in range(tgt_len):\n",
    "            output, h, c = self.decoder(input_char, h, c, encoder_outputs)\n",
    "            outputs[:, t] = output\n",
    "            \n",
    "            if t < tgt_len - 1:\n",
    "                teacher_force = torch.rand(1).item() < teacher_forcing_ratio\n",
    "                input_char = tgt_in[:, t + 1] if teacher_force else output.argmax(1)\n",
    "        \n",
    "        return outputs\n",
    "\n",
    "# Create model\n",
    "EMBED_DIM = 64\n",
    "HIDDEN_DIM = 128\n",
    "NUM_LAYERS = 2\n",
    "\n",
    "model = Seq2SeqDiacritizer(\n",
    "    vocab_size=len(vocab),\n",
    "    embed_dim=EMBED_DIM,\n",
    "    hidden_dim=HIDDEN_DIM,\n",
    "    num_layers=NUM_LAYERS\n",
    ").to(device)\n",
    "\n",
    "total_params = sum(p.numel() for p in model.parameters())\n",
    "print(f\"Model parameters: {total_params:,}\")\n",
    "print(f\"Estimated model size: {total_params * 4 / 1024 / 1024:.1f} MB\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 6: Training loop with teacher forcing\n",
    "import time\n",
    "\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=0)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=2)\n",
    "\n",
    "NUM_EPOCHS = 10\n",
    "best_word_acc = 0\n",
    "best_model_state = None\n",
    "\n",
    "print(\"Starting training...\\n\")\n",
    "start_time = time.time()\n",
    "\n",
    "for epoch in range(1, NUM_EPOCHS + 1):\n",
    "    epoch_start = time.time()\n",
    "    \n",
    "    # Train\n",
    "    model.train()\n",
    "    train_loss = 0\n",
    "    teacher_forcing_ratio = max(0.5, 1.0 - epoch * 0.05)  # Decay from 1.0 to 0.5\n",
    "    \n",
    "    for batch_idx, (src, tgt_in, tgt_out, _) in enumerate(train_loader):\n",
    "        src, tgt_in, tgt_out = src.to(device), tgt_in.to(device), tgt_out.to(device)\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        logits = model(src, tgt_in, teacher_forcing_ratio)\n",
    "        \n",
    "        # Flatten for loss: [batch * seq, vocab] vs [batch * seq]\n",
    "        loss = criterion(logits.view(-1, logits.size(-1)), tgt_out.view(-1))\n",
    "        loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "        optimizer.step()\n",
    "        train_loss += loss.item()\n",
    "        \n",
    "        if (batch_idx + 1) % 500 == 0:\n",
    "            print(f\"  Batch {batch_idx + 1}/{len(train_loader)}, Loss: {loss.item():.4f}\")\n",
    "    \n",
    "    train_loss /= len(train_loader)\n",
    "    \n",
    "    # Evaluate (with no teacher forcing)\n",
    "    model.eval()\n",
    "    val_loss = 0\n",
    "    total_words = correct_words = 0\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for src, tgt_in, tgt_out, src_lens in val_loader:\n",
    "            src, tgt_in, tgt_out = src.to(device), tgt_in.to(device), tgt_out.to(device)\n",
    "            \n",
    "            logits = model(src, tgt_in, teacher_forcing_ratio=0.0)  # No teacher forcing\n",
    "            loss = criterion(logits.view(-1, logits.size(-1)), tgt_out.view(-1))\n",
    "            val_loss += loss.item()\n",
    "            \n",
    "            preds = logits.argmax(dim=-1)\n",
    "            \n",
    "            # Calculate word accuracy\n",
    "            for i in range(len(src)):\n",
    "                pred_seq = preds[i].cpu().tolist()\n",
    "                tgt_seq = tgt_out[i].cpu().tolist()\n",
    "                \n",
    "                pred_text = vocab.decode(pred_seq)\n",
    "                tgt_text = vocab.decode(tgt_seq)\n",
    "                \n",
    "                for pw, tw in zip(pred_text.split(), tgt_text.split()):\n",
    "                    total_words += 1\n",
    "                    if pw == tw:\n",
    "                        correct_words += 1\n",
    "    \n",
    "    val_loss /= len(val_loader)\n",
    "    word_acc = correct_words / total_words if total_words > 0 else 0\n",
    "    \n",
    "    scheduler.step(val_loss)\n",
    "    epoch_time = time.time() - epoch_start\n",
    "    \n",
    "    print(f\"Epoch {epoch}/{NUM_EPOCHS} ({epoch_time:.0f}s) | \"\n",
    "          f\"Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f} | \"\n",
    "          f\"Word Acc: {word_acc:.1%} | TF: {teacher_forcing_ratio:.0%}\")\n",
    "    \n",
    "    if word_acc > best_word_acc:\n",
    "        best_word_acc = word_acc\n",
    "        best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}\n",
    "        print(f\"  ↑ New best model!\")\n",
    "\n",
    "total_time = time.time() - start_time\n",
    "print(f\"\\nTraining complete in {total_time/60:.1f} minutes\")\n",
    "print(f\"Best word accuracy: {best_word_acc:.1%}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 7: Save the model\n",
    "checkpoint = {\n",
    "    \"model_type\": \"seq2seq\",\n",
    "    \"model_state_dict\": best_model_state,\n",
    "    \"vocab\": vocab.char2idx,\n",
    "    \"config\": {\n",
    "        \"vocab_size\": len(vocab),\n",
    "        \"embed_dim\": EMBED_DIM,\n",
    "        \"hidden_dim\": HIDDEN_DIM,\n",
    "        \"num_layers\": NUM_LAYERS,\n",
    "    },\n",
    "    \"metrics\": {\n",
    "        \"word_acc\": best_word_acc,\n",
    "    },\n",
    "}\n",
    "\n",
    "torch.save(checkpoint, \"yoruba_seq2seq_diacritizer.pt\")\n",
    "\n",
    "import os\n",
    "model_size = os.path.getsize(\"yoruba_seq2seq_diacritizer.pt\") / (1024 * 1024)\n",
    "print(f\"Model saved: yoruba_seq2seq_diacritizer.pt ({model_size:.1f} MB)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 8: Test the model with autoregressive decoding\n",
    "model.load_state_dict({k: v.to(device) for k, v in best_model_state.items()})\n",
    "model.eval()\n",
    "\n",
    "def diacritize_seq2seq(text: str, max_len: int = 300) -> str:\n",
    "    \"\"\"Diacritize text using autoregressive decoding.\"\"\"\n",
    "    with torch.no_grad():\n",
    "        # Encode\n",
    "        src = torch.tensor([vocab.encode(text)], dtype=torch.long).to(device)\n",
    "        encoder_outputs, h, c = model.encoder(src)\n",
    "        h = h.squeeze(0)\n",
    "        c = c.squeeze(0)\n",
    "        \n",
    "        # Decode autoregressively\n",
    "        output_chars = []\n",
    "        input_char = torch.tensor([vocab.char2idx[vocab.SOS]], device=device)\n",
    "        \n",
    "        for _ in range(max_len):\n",
    "            output, h, c = model.decoder(input_char, h, c, encoder_outputs)\n",
    "            pred_idx = output.argmax(1).item()\n",
    "            \n",
    "            if pred_idx == vocab.char2idx[vocab.EOS]:\n",
    "                break\n",
    "            if pred_idx not in (vocab.char2idx[vocab.PAD], vocab.char2idx[vocab.SOS]):\n",
    "                output_chars.append(vocab.idx2char.get(pred_idx, \"\"))\n",
    "            \n",
    "            input_char = torch.tensor([pred_idx], device=device)\n",
    "        \n",
    "        return \"\".join(output_chars)\n",
    "\n",
    "# Test examples\n",
    "test_sentences = [\n",
    "    \"Ojo dara pupo\",\n",
    "    \"E ku ishe o\",\n",
    "    \"Mo fe ran re\",\n",
    "    \"Bawo ni o se wa\",\n",
    "    \"Olorun a bukun fun e\",\n",
    "]\n",
    "\n",
    "print(\"\\nTest Results:\")\n",
    "print(\"=\" * 60)\n",
    "for sent in test_sentences:\n",
    "    result = diacritize_seq2seq(sent)\n",
    "    print(f\"Input:  {sent}\")\n",
    "    print(f\"Output: {result}\")\n",
    "    print(\"-\" * 40)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 9: Download the model (Colab)\n",
    "try:\n",
    "    from google.colab import files\n",
    "    files.download(\"yoruba_seq2seq_diacritizer.pt\")\n",
    "    print(\"Download started! Check your downloads folder.\")\n",
    "except:\n",
    "    print(\"Not running on Colab. Model saved locally.\")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
