{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":121150,"databundleVersionId":15061024,"sourceType":"competition"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Ignore warnings for cleaner logs\nimport warnings\nwarnings.simplefilter(\"ignore\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install -q sacrebleu wandb","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Imports\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader, Subset\nfrom transformers import AutoTokenizer, T5ForConditionalGeneration\nimport pandas as pd\nimport numpy as np\nfrom tqdm.notebook import tqdm\nimport os\nfrom kaggle_secrets import UserSecretsClient\nfrom sklearn.model_selection import train_test_split\nimport sacrebleu\nimport wandb\nimport re\nimport math\nfrom sacrebleu.metrics import CHRF\nfrom pprint import pprint","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Comment out if you dont want WANDB logging\nuser_secrets = UserSecretsClient()\nwandb_key = user_secrets.get_secret(\"wandb\")\nwandb.login(key=wandb_key)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Training Configuration\nclass CONFIG():\n    def __init__(self):\n        self.MODEL_NAME  = \"google/byt5-base\"\n        self.MAX_SOURCE_LEN = 512\n        self.MAX_TARGET_LEN = 512\n        self.SEED = 42\n        self.TEST_SIZE = 0.1\n        self.BIDIRECTIONAL_SPLIT = False\n        self.SENTENCE_ALIGNMENT = False\n        self.ASCII_NORM = True\n        self.GAP_NORM = True\n        self.BATCH_SIZE = 2\n        self.EPOCHS = 10\n        self.LR = 2e-4\n        self.WEIGHT_DECAY = 0.01\n        self.SAVE_PATH = \"./byt5_base_dpc_v11\"\n        self.ES_PATIENCE = 3\n        self.DECODER_BEAM_WIDTH = 4\n        self.WANDB_EXPT_NAME = \"byt5-base-v11\"\n        self.GRAD_ACC_STEPS = 2\n        self.LABEL_SMOOTHING = 0.01\n        self.REPETITION_PENALTY = 1.2\n        self.DESC = \"Low Epochs with LS=0.01\"\n\ncfg = CONFIG()\n\ntorch.manual_seed(cfg.SEED)\nif torch.cuda.is_available():\n    torch.cuda.manual_seed(cfg.SEED)\n    torch.cuda.manual_seed_all(cfg.SEED)\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nPREFIX = \"translate Akkadian to English: \"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Comment out if you dont want WANDB logging\nwandb.init(\n    project=\"assyrian-mt\",\n    name=cfg.WANDB_EXPT_NAME,\n    config=cfg\n)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Training Details: \")\npprint(cfg.__dict__)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Data\ntrain = pd.read_csv('/kaggle/input/deep-past-initiative-machine-translation/train.csv')\nsentences_oare = pd.read_csv('/kaggle/input/deep-past-initiative-machine-translation/Sentences_Oare_FirstWord_LinNum.csv')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Serntence Alignment\nif cfg.SENTENCE_ALIGNMENT:\n    df = train.merge(\n        sentences_oare,\n        left_on=\"oare_id\",\n        right_on=\"text_uuid\",\n        how=\"left\",\n        suffixes=(\"_train\", \"_sentcs\")\n    )\n    \n    # -----------------\n    # Helpers\n    # -------------------------\n    chrf = CHRF(word_order=2)\n    \n    def split_translation_clauses(text):\n        return [\n            t.strip()\n            for t in text.replace(\";\", \".\").split(\".\")\n            if t.strip()\n        ]\n    \n    # -------------------------\n    # Alignment logic\n    # -------------------------\n    aligned_pairs = []\n    \n    for oare_id, g in df.groupby(\"oare_id\"):\n    \n        # tablet-level data (always available)\n        full_translit = g.iloc[0][\"transliteration\"]\n        full_translation = g.iloc[0][\"translation_train\"]\n    \n        # check if sentence-level alignment exists\n        has_sentence_alignment = (\n            g[\"first_word_number\"].notna().any()\n            and g[\"translation_sentcs\"].notna().any()\n        )\n    \n        # --------------------------------------------------\n        # CASE 1: sentence-level alignment EXISTS → expand\n        # --------------------------------------------------\n        if has_sentence_alignment:\n    \n            g = g.dropna(subset=[\"first_word_number\"]).sort_values(\"first_word_number\")\n    \n            full_tokens = full_translit.split()\n            translation_clauses = split_translation_clauses(full_translation)\n    \n            # sentence boundaries\n            word_starts = g[\"first_word_number\"].astype(int).tolist()\n            word_starts.append(len(full_tokens) + 1)\n    \n            for idx, row in enumerate(g.itertuples()):\n    \n                start = int(row.first_word_number - 1)\n                end = int(word_starts[idx + 1] - 1)\n                translit_sentence = \" \".join(full_tokens[start:end]).strip()\n                anchor = str(row.translation_sentcs).strip()\n    \n                if not translit_sentence or not anchor:\n                    continue\n    \n                # fuzzy match translation clauses\n                scores = []\n                for j, clause in enumerate(translation_clauses):\n                    score = chrf.sentence_score(anchor, [clause]).score\n                    scores.append((score, j))\n    \n                if not scores:\n                    continue\n    \n                best_idx = max(scores)[1]\n                candidate = translation_clauses[best_idx]\n    \n                # expansion heuristic\n                if best_idx + 1 < len(translation_clauses):\n                    expanded = candidate + \", \" + translation_clauses[best_idx + 1]\n                    if chrf.sentence_score(anchor, [expanded]).score > \\\n                       chrf.sentence_score(anchor, [candidate]).score:\n                        candidate = expanded\n    \n                aligned_pairs.append({\n                    \"oare_id\": oare_id,\n                    \"translit\": translit_sentence,\n                    \"translation\": candidate,\n                    \"alignment_type\": \"sentence\",\n                    \"anchor\": anchor\n                })\n    \n        # --------------------------------------------------\n        # CASE 2: NO sentence alignment → keep tablet-level\n        # --------------------------------------------------\n        else:\n            aligned_pairs.append({\n                \"oare_id\": oare_id,\n                \"translit\": full_translit,\n                \"translation\": full_translation,\n                \"alignment_type\": \"tablet\",\n                \"anchor\": None\n            })\n    # -------------------------\n    # Save result\n    # -------------------------\n    train = pd.DataFrame(aligned_pairs)\n    train['transliteration'] = train['translit']\n    print(\"Total aligned pairs:\", len(train))\n    print(train[\"alignment_type\"].value_counts())","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ASCII Normalization and Gap Normalization\nASCII_TO_DIACRITIC = {\n    # Sibilants\n    \"sz\": \"š\",\n    \"SZ\": \"Š\",\n    \"Sz\": \"Š\",\n    \"sh\": \"š\",\n    \"SH\": \"Š\",\n    \"Sh\": \"Š\",\n    # Emphatics (comma notation)\n    \"s,\": \"ṣ\",\n    \"S,\": \"Ṣ\",\n    \"t,\": \"ṭ\",\n    \"T,\": \"Ṭ\",\n    \"z,\": \"ẓ\",\n    \"Z,\": \"Ẓ\",\n    # Emphatics (dot-under notation)\n    \".s\": \"ṣ\",\n    \".S\": \"Ṣ\",\n    \".t\": \"ṭ\",\n    \".T\": \"Ṭ\",\n    \".z\": \"ẓ\",\n    \".Z\": \"Ẓ\",\n    # H variants\n    \"h,\": \"ḫ\",\n    \"H,\": \"Ḫ\",\n    \".h\": \"ḫ\",\n    \".H\": \"Ḫ\",\n    \"hh\": \"ḫ\",\n    \"HH\": \"Ḫ\",\n    # Aleph\n    \"'\": \"ʾ\",\n    # Subscript numbers (common ASCII alternatives)\n    \"s2\": \"š\",\n    \"S2\": \"Š\",\n    \"s3\": \"ś\",\n    \"S3\": \"Ś\",\n}\n\nVOWEL_SUBSCRIPTS = {\n    \"a2\": \"á\",\n    \"a3\": \"à\",\n    \"e2\": \"é\",\n    \"e3\": \"è\",\n    \"i2\": \"í\",\n    \"i3\": \"ì\",\n    \"u2\": \"ú\",\n    \"u3\": \"ù\",\n    \"A2\": \"Á\",\n    \"A3\": \"À\",\n    \"E2\": \"É\",\n    \"E3\": \"È\",\n    \"I2\": \"Í\",\n    \"I3\": \"Ì\",\n    \"U2\": \"Ú\",\n    \"U3\": \"Ù\",\n}\n\n\ndef normalize_ascii_to_diacritics(text: str) -> str:\n    \"\"\"Convert ASCII transliteration conventions to proper diacritics.\"\"\"\n    result = text\n    for ascii_seq, diacritic in sorted(\n        ASCII_TO_DIACRITIC.items(), key=lambda x: -len(x[0])\n    ):\n        result = result.replace(ascii_seq, diacritic)\n    for ascii_seq, diacritic in VOWEL_SUBSCRIPTS.items():\n        result = result.replace(ascii_seq, diacritic)\n    return result\n\n\ndef normalize_gaps(text: str) -> str:\n    \"\"\"Normalize gap/damage markers: x→<gap>, x x x→<big_gap>.\"\"\"\n    tokens = text.split()\n    result = []\n    i = 0\n\n    while i < len(tokens):\n        token = tokens[i]\n        if token.lower() == \"x\":\n            x_count = 1\n            while i + x_count < len(tokens) and tokens[i + x_count].lower() == \"x\":\n                x_count += 1\n            result.append(\"<gap>\" if x_count == 1 else \"<big_gap>\")\n            i += x_count\n        else:\n            if token.lower().startswith(\"x-\"):\n                result.append(\"<gap>\" + token[1:])\n            elif token.lower().endswith(\"-x\"):\n                result.append(token[:-1] + \"-<gap>\")\n            else:\n                result.append(token)\n            i += 1\n\n    text_result = \" \".join(result)\n    text_result = re.sub(r\"<gap>\\s+<big_gap>\", \"<big_gap>\", text_result)\n    text_result = re.sub(r\"<big_gap>\\s+<gap>\", \"<big_gap>\", text_result)\n    text_result = re.sub(r\"(<gap>\\s*){2,}\", \"<big_gap> \", text_result)\n    return text_result.strip()\n\n\ndef apply_host_normalizations(text, ascii_norm_flag, gap_norm_flag):\n    \"\"\"Apply all host-recommended normalizations.\"\"\"\n    \n    if ascii_norm_flag:\n        text = normalize_ascii_to_diacritics(text)\n    if gap_norm_flag:\n        text = normalize_gaps(text)\n    return text","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Pre-processing\ntrain['transliteration'] = train['transliteration'].apply(lambda x: apply_host_normalizations(\n    x, cfg.ASCII_NORM, cfg.GAP_NORM))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Model and Tokenizer\ntokenizer = AutoTokenizer.from_pretrained(cfg.MODEL_NAME)\nmodel = T5ForConditionalGeneration.from_pretrained(cfg.MODEL_NAME).to(DEVICE)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# NMT Dataset Class\nclass NMTDataset(Dataset):\n    def __init__(self, df, tokenizer, max_source_len=512, max_target_len=256):\n        self.df = df\n        self.tokenizer = tokenizer\n        self.max_source_len = max_source_len\n        self.max_target_len = max_target_len\n\n        self.sources = self.df[\"input_text\"].tolist()\n        self.targets = self.df[\"target_text\"].tolist()\n\n    def __len__(self):\n        return len(self.sources)\n\n    def __getitem__(self, idx):\n        src = self.sources[idx]\n        tgt = self.targets[idx]\n\n        src_enc = self.tokenizer(\n            src,\n            truncation=True,\n            padding=\"max_length\",\n            max_length=self.max_source_len,\n            return_tensors=\"pt\",\n        )\n\n        tgt_enc = self.tokenizer(\n            tgt,\n            truncation=True,\n            padding=\"max_length\",\n            max_length=self.max_target_len,\n            return_tensors=\"pt\",\n        )\n\n        labels = tgt_enc[\"input_ids\"]\n        labels[labels == self.tokenizer.pad_token_id] = -100\n\n        return {\n            \"input_ids\": src_enc[\"input_ids\"].squeeze(0),\n            \"attention_mask\": src_enc[\"attention_mask\"].squeeze(0),\n            \"labels\": labels.squeeze(0),\n        }","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Train-Val Split and Dataloaders\ntrain_idx, valid_idx = train_test_split(\n    np.arange(len(train)),\n    test_size=cfg.TEST_SIZE,            \n    shuffle=True,                  \n    random_state=42\n)\n\ntrain_df = train.iloc[train_idx].reset_index(drop=True)\nvalid_df = train.iloc[valid_idx].reset_index(drop=True)\n\ndef create_bidirectional_data(df):\n    # Akkadian -> English\n    df_fwd = df.copy()\n    df_fwd['input_text'] = \"translate Akkadian to English: \" + df_fwd['transliteration'].astype(str)\n    df_fwd['target_text'] = df_fwd['translation'].astype(str)\n    \n    # English -> Akkadian\n    df_bwd = df.copy()\n    df_bwd['input_text'] = \"translate English to Akkadian: \" + df_bwd['translation'].astype(str)\n    df_bwd['target_text'] = df_bwd['transliteration'].astype(str)\n    \n    # Combined Data\n    df_combined = pd.concat([df_fwd, df_bwd], ignore_index=True)\n    df_combined = df_combined.sample(frac=1, random_state=42).reset_index(drop=True)\n    \n    return df_combined\n\ndef create_unidirectional_data(df):\n    df['input_text'] = \"translate Akkadian to English: \" + df['transliteration'].astype(str)\n    df['target_text'] = df['translation'].astype(str)\n    return df\n\n# Train Split\nif cfg.BIDIRECTIONAL_SPLIT:\n    train_df = create_bidirectional_data(train_df)\nelse:\n    train_df = create_unidirectional_data(train_df)\n\n# Validation Split\nvalid_df = create_unidirectional_data(valid_df)\n\nprint(f\"Train samples: {len(train_df)}\")\nprint(f\"Val samples:   {len(valid_df)}\")\n\ntrain_ds = NMTDataset(train_df, tokenizer, cfg.MAX_SOURCE_LEN, cfg.MAX_TARGET_LEN)\nvalid_ds = NMTDataset(valid_df, tokenizer, cfg.MAX_SOURCE_LEN, cfg.MAX_TARGET_LEN)\n\ntrain_loader = DataLoader(\n    train_ds,\n    batch_size=cfg.BATCH_SIZE,\n    shuffle=True\n)\nvalid_loader = DataLoader(\n    valid_ds,\n    batch_size=cfg.BATCH_SIZE,\n    shuffle=False\n)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Loss, Optimizer, Scheduler\nloss_fn = nn.CrossEntropyLoss(\n    ignore_index=-100,\n    label_smoothing=cfg.LABEL_SMOOTHING,\n)\n\noptimizer = torch.optim.AdamW(model.parameters(), lr=cfg.LR, weight_decay=cfg.WEIGHT_DECAY, fused=True)\n\nnum_update_steps_per_epoch = math.ceil(len(train_loader) / cfg.GRAD_ACC_STEPS)\nnum_training_steps = num_update_steps_per_epoch * cfg.EPOCHS\n\nscheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n    optimizer,\n    T_max=num_training_steps,\n    eta_min=0.0,\n)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Computing Grad Norm\ndef compute_grad_norm(model):\n    total_norm = 0.0\n    for p in model.parameters():\n        if p.grad is not None:\n            param_norm = p.grad.data.norm(2)\n            total_norm += param_norm.item() ** 2\n    return total_norm ** 0.5","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Train one epoch\ndef train_epoch(model, loader, optimizer, scheduler, loss_fn, grad_accum_steps=4):\n    model.train()\n    total_loss = 0.0\n    optimizer.zero_grad()\n\n    last_grad_norm = 0.0\n\n    for step, batch in enumerate(tqdm(loader)):\n        batch = {k: v.to(DEVICE) for k, v in batch.items()}\n\n        outputs = model(\n            input_ids=batch[\"input_ids\"],\n            attention_mask=batch[\"attention_mask\"],\n            labels=batch[\"labels\"],\n            use_cache=False,\n        )\n\n        logits = outputs.logits\n        labels = batch[\"labels\"]\n\n        loss = loss_fn(\n            logits.view(-1, logits.size(-1)),\n            labels.view(-1),\n        )\n\n        # scale loss for gradient accumulation\n        loss = loss / grad_accum_steps\n        loss.backward()\n\n        # optimizer + scheduler step ONLY after accumulation\n        if (step + 1) % grad_accum_steps == 0 or (step + 1) == len(loader):\n\n            # Gradient clipping (HF default)\n            last_grad_norm = torch.nn.utils.clip_grad_norm_(\n                model.parameters(), 1.0\n            )\n\n            optimizer.step()\n            scheduler.step()\n            optimizer.zero_grad()\n\n        total_loss += loss.item()\n\n    return total_loss / len(loader), last_grad_norm","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Compute Val Loss\n@torch.no_grad()\ndef validate(model, loader):\n    model.eval()\n    total_loss = 0.0\n\n    for batch in tqdm(loader, desc='Validation: ', total=len(loader)):\n        batch = {k: v.to(DEVICE) for k, v in batch.items()}\n\n        outputs = model(\n            input_ids=batch[\"input_ids\"],\n            attention_mask=batch[\"attention_mask\"],\n            labels=batch[\"labels\"], \n            use_cache=False,\n        )\n\n        logits = outputs.logits\n        labels = batch[\"labels\"]\n\n        loss = loss_fn(\n            logits.view(-1, logits.size(-1)),\n            labels.view(-1),\n        )\n\n        total_loss += loss.item()\n\n    return total_loss / len(loader)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Compute Val Metrics with Inference Settings\n@torch.no_grad()\ndef validate_generation(model, loader, tokenizer, max_len, num_beams):\n    model.eval()\n\n    preds, refs = [], []\n\n    for batch in tqdm(loader, desc='Metrics: ', total=len(loader)):\n        input_ids = batch[\"input_ids\"].to(DEVICE)\n        attention_mask = batch[\"attention_mask\"].to(DEVICE)\n        labels = batch[\"labels\"]\n\n        generated_ids = model.generate(\n            input_ids=input_ids,\n            attention_mask=attention_mask,\n            max_length=max_len,\n            num_beams=num_beams,\n            length_penalty=1.0,\n            early_stopping=True,\n            repetition_penalty=cfg.REPETITION_PENALTY\n        )\n\n        pred_texts = tokenizer.batch_decode(\n            generated_ids, skip_special_tokens=True\n        )\n\n        labels = labels.clone()\n        labels[labels == -100] = tokenizer.pad_token_id\n\n        ref_texts = tokenizer.batch_decode(\n            labels,\n            skip_special_tokens=True,\n        )\n\n        preds.extend(pred_texts)\n        refs.extend(ref_texts)\n\n    bleu = sacrebleu.corpus_bleu(\n        preds,\n        [refs],\n        tokenize=\"13a\",\n        lowercase=False,\n    ).score\n\n    chrf = sacrebleu.corpus_chrf(\n        preds,\n        [refs],\n        word_order=2   # chrF++\n    ).score\n\n    geo_mean = (bleu * chrf) ** 0.5\n\n    return bleu, chrf, geo_mean\n    {\n        \"BLEU\": bleu,\n        \"chrF++\": chrf,\n        \"GEO_MEAN\": geo_mean,\n    }","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Training Loop\nbest_loss = float(\"inf\")\nearly_stop_counter = 0\n\nprint(\"Training started\")\nfor epoch in range(cfg.EPOCHS):\n    current_lr = optimizer.param_groups[0][\"lr\"]\n    train_loss, grad_norm = train_epoch(model, train_loader, optimizer, \n                                        scheduler, loss_fn, cfg.GRAD_ACC_STEPS)\n    \n    val_loss = validate(model, valid_loader)\n    # scheduler.step(val_loss)\n    \n    print(\n        f\"Epoch {epoch:02d} | \"\n        f\"Training Loss: {train_loss:.4f} | \"\n        f\"Grad Norm: {grad_norm:.4f} | \"\n        f\"Val Loss: {val_loss:.3f} | \"\n        f\"LR: {current_lr:.2e}\"\n    )\n\n    # Comment out if you dont want WANDB logging\n    wandb.log(\n        {\n            \"Train Loss\": train_loss,\n            \"Grad Norm\": grad_norm,\n            \"Val Loss\": val_loss,\n            \"learning_rate\": current_lr\n        }\n    )\n    # Till here\n\n    # ===== Early stopping =====\n    if val_loss < best_loss:\n        best_loss = val_loss\n        early_stop_counter = 0\n        model.save_pretrained(cfg.SAVE_PATH)\n        tokenizer.save_pretrained(cfg.SAVE_PATH)\n        print(\"Best Model Saved!\")\n    else:\n        early_stop_counter += 1\n\n    if early_stop_counter >= cfg.ES_PATIENCE:\n        print(\"Early stopping triggered!\")\n        break","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Loading best model\nprint(\"Loading best model\")\ntokenizer = AutoTokenizer.from_pretrained(cfg.SAVE_PATH)\nmodel = T5ForConditionalGeneration.from_pretrained(cfg.SAVE_PATH).to(DEVICE);","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Generation Metrics on Val\nprint(\"Generation Metrics\")\nbleu, chrf, geo_mean = validate_generation(model, valid_loader, tokenizer, \n                                               cfg.MAX_TARGET_LEN, cfg.DECODER_BEAM_WIDTH)\n\nprint(\n    f\"BLEU: {bleu:.4f} | \"\n    f\"CHRF++: {chrf:.4f} | \"\n    f\"Metric: {geo_mean:.4f} | \"\n)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}