{"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":14976537,"sourceType":"competition"},{"sourceId":14196765,"sourceType":"datasetVersion","datasetId":9053334}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os, re, math, random\nimport pandas as pd\nimport numpy as np\nimport torch\nimport pytorch_lightning as pl\nfrom torch.utils.data import DataLoader\nfrom transformers import AutoTokenizer, AutoModelForSeq2SeqLM\nfrom pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint, LearningRateMonitor\nfrom sklearn.model_selection import KFold\nimport sacrebleu\nimport matplotlib.pyplot as plt","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T08:14:24.362474Z","iopub.execute_input":"2025-12-18T08:14:24.362819Z","iopub.status.idle":"2025-12-18T08:14:24.436494Z","shell.execute_reply.started":"2025-12-18T08:14:24.362767Z","shell.execute_reply":"2025-12-18T08:14:24.435698Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"SEED = 2026\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\npl.seed_everything(SEED, workers=True)\nprint(\"Happy New Year - 2026\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T08:14:25.737282Z","iopub.execute_input":"2025-12-18T08:14:25.738200Z","iopub.status.idle":"2025-12-18T08:14:25.749925Z","shell.execute_reply.started":"2025-12-18T08:14:25.738166Z","shell.execute_reply":"2025-12-18T08:14:25.749193Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"Device:\", DEVICE)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T08:14:28.383880Z","iopub.execute_input":"2025-12-18T08:14:28.384602Z","iopub.status.idle":"2025-12-18T08:14:28.411748Z","shell.execute_reply.started":"2025-12-18T08:14:28.384570Z","shell.execute_reply":"2025-12-18T08:14:28.411034Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"TRAIN_PATH = \"/kaggle/input/deep-past-initiative-machine-translation/train.csv\"\nTEST_PATH = \"/kaggle/input/deep-past-initiative-machine-translation/test.csv\"\n\ntrain_raw = pd.read_csv(TRAIN_PATH)\ntest_df = pd.read_csv(TEST_PATH)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T08:14:43.340026Z","iopub.execute_input":"2025-12-18T08:14:43.340350Z","iopub.status.idle":"2025-12-18T08:14:43.405691Z","shell.execute_reply.started":"2025-12-18T08:14:43.340320Z","shell.execute_reply":"2025-12-18T08:14:43.405177Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def clean_ws(x):\n    return re.sub(r\"\\s+\", \" \", str(x)).strip()\n\ndef tok_len(s):\n    return max(1, len(s.split()))\n\ndef align_monotone(src_units, tgt_units, max_src_merge=2, max_tgt_merge=2):\n    i = j = 0\n    pairs = []\n\n    while i < len(src_units) and j < len(tgt_units):\n        best = None\n        for si in range(1, max_src_merge + 1):\n            for tj in range(1, max_tgt_merge + 1):\n                if i + si <= len(src_units) and j + tj <= len(tgt_units):\n                    s = \" \".join(src_units[i:i+si])\n                    t = \" \".join(tgt_units[j:j+tj])\n                    score = abs(math.log(tok_len(s) / tok_len(t)))\n                    if best is None or score < best[0]:\n                        best = (score, si, tj, s, t)\n        _, si, tj, s, t = best\n        pairs.append((s, t))\n        i += si\n        j += tj\n\n    if pairs:\n        if i < len(src_units):\n            pairs[-1] = (pairs[-1][0] + \" \" + \" \".join(src_units[i:]), pairs[-1][1])\n        if j < len(tgt_units):\n            pairs[-1] = (pairs[-1][0], pairs[-1][1] + \" \" + \" \".join(tgt_units[j:]))\n\n    return pairs\n\npairs = []\nfor _, r in train_raw.iterrows():\n    src = re.split(r\"\\n+\", str(r[\"transliteration\"]))\n    tgt = re.split(r\"(?<=[.!?])\\s+\", clean_ws(r[\"translation\"]))\n\n    src = [clean_ws(x) for x in src if clean_ws(x)]\n    tgt = [clean_ws(x) for x in tgt if clean_ws(x)]\n\n    if not src or not tgt:\n        continue\n\n    pairs.extend(align_monotone(src, tgt))\n\npairs_df = pd.DataFrame(pairs, columns=[\"src\", \"tgt\"]).drop_duplicates()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T08:14:46.596488Z","iopub.execute_input":"2025-12-18T08:14:46.596984Z","iopub.status.idle":"2025-12-18T08:14:47.028443Z","shell.execute_reply.started":"2025-12-18T08:14:46.596953Z","shell.execute_reply":"2025-12-18T08:14:47.027843Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pairs_df.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T08:14:55.613188Z","iopub.execute_input":"2025-12-18T08:14:55.613506Z","iopub.status.idle":"2025-12-18T08:14:55.620389Z","shell.execute_reply.started":"2025-12-18T08:14:55.613478Z","shell.execute_reply":"2025-12-18T08:14:55.619659Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"MODEL_NAME = \"/kaggle/input/akkadian-models\"\ntokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T08:14:57.137569Z","iopub.execute_input":"2025-12-18T08:14:57.138394Z","iopub.status.idle":"2025-12-18T08:14:57.314810Z","shell.execute_reply.started":"2025-12-18T08:14:57.138350Z","shell.execute_reply":"2025-12-18T08:14:57.314013Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class AkkadianT5DataModule(pl.LightningDataModule):\n    def __init__(self, tokenizer, train_df, val_df, batch_size=8):\n        super().__init__()\n        self.tokenizer = tokenizer\n        self.train_df = train_df\n        self.val_df = val_df\n        self.batch_size = batch_size\n\n    def collate(self, batch):\n        src, tgt = zip(*batch)\n        enc = self.tokenizer(list(src), padding=True, truncation=True, max_length=256, return_tensors=\"pt\")\n        dec = self.tokenizer(text_target=list(tgt), padding=True, truncation=True, max_length=128, return_tensors=\"pt\")\n        enc[\"labels\"] = dec[\"input_ids\"]\n        return enc\n\n    def train_dataloader(self):\n        data = list(zip(self.train_df.src.tolist(), self.train_df.tgt.tolist()))\n        return DataLoader(data, batch_size=self.batch_size, shuffle=True, collate_fn=self.collate)\n\n    def val_dataloader(self):\n        data = list(zip(self.val_df.src.tolist(), self.val_df.tgt.tolist()))\n        return DataLoader(data, batch_size=self.batch_size, shuffle=False, collate_fn=self.collate)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T08:14:59.633485Z","iopub.execute_input":"2025-12-18T08:14:59.634164Z","iopub.status.idle":"2025-12-18T08:14:59.640378Z","shell.execute_reply.started":"2025-12-18T08:14:59.634132Z","shell.execute_reply":"2025-12-18T08:14:59.639566Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class AkkadianT5AWP:\n    def __init__(self, model, adv_lr=1e-4, adv_eps=1e-2):\n        self.model = model\n        self.adv_lr = adv_lr\n        self.adv_eps = adv_eps\n        self.backup = {}\n\n    def attack(self):\n        for n, p in self.model.named_parameters():\n            if p.requires_grad and p.grad is not None:\n                self.backup[n] = p.data.clone()\n                r = self.adv_lr * p.grad / (torch.norm(p.grad) + 1e-6)\n                p.data.add_(r)\n                p.data = torch.clamp(p.data, self.backup[n] - self.adv_eps, self.backup[n] + self.adv_eps)\n\n    def restore(self):\n        for n, p in self.model.named_parameters():\n            if n in self.backup:\n                p.data = self.backup[n]\n        self.backup = {}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T08:15:01.048046Z","iopub.execute_input":"2025-12-18T08:15:01.048822Z","iopub.status.idle":"2025-12-18T08:15:01.054348Z","shell.execute_reply.started":"2025-12-18T08:15:01.048764Z","shell.execute_reply":"2025-12-18T08:15:01.053665Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class AkkadianT5Lightning(pl.LightningModule):\n    def __init__(self, model, tokenizer, lr=1e-4):\n        super().__init__()\n        self.model = model\n        self.tokenizer = tokenizer\n        self.lr = lr\n        self.automatic_optimization = False\n        self.awp = AkkadianT5AWP(model)\n        self.preds, self.refs = [], []\n        self.history = {\"loss\": [], \"geo\": []}\n\n    def training_step(self, batch, batch_idx):\n        opt = self.optimizers()\n\n        out = self.model(\n            input_ids=batch[\"input_ids\"],\n            attention_mask=batch[\"attention_mask\"],\n            labels=batch[\"labels\"],\n            label_smoothing_factor=0.1\n        )\n        self.manual_backward(out.loss)\n\n        if self.current_epoch > self.trainer.max_epochs * 0.4:\n            self.awp.attack()\n            out_adv = self.model(**batch)\n            self.manual_backward(out_adv.loss)\n            self.awp.restore()\n\n        if (batch_idx + 1) % self.trainer.accumulate_grad_batches == 0:\n            opt.step()\n            opt.zero_grad()\n\n        self.log(\"train_loss\", out.loss, prog_bar=True)\n        return out.loss\n\n    def validation_step(self, batch, batch_idx):\n        gen = self.model.generate(\n            batch[\"input_ids\"],\n            attention_mask=batch[\"attention_mask\"],\n            num_beams=4,\n            max_new_tokens=256,\n            length_penalty=0.8\n        )\n        self.preds.extend(self.tokenizer.batch_decode(gen, skip_special_tokens=True))\n        self.refs.extend(self.tokenizer.batch_decode(batch[\"labels\"], skip_special_tokens=True))\n\n    def on_validation_epoch_end(self):\n        bleu = sacrebleu.corpus_bleu(self.preds, [self.refs]).score\n        chrf = sacrebleu.corpus_chrf(self.preds, [self.refs]).score\n        geo = math.sqrt(bleu * chrf)\n        self.history[\"geo\"].append(geo)\n        self.log(\"geo_mean\", geo, prog_bar=True)\n        self.preds.clear()\n        self.refs.clear()\n\n    def configure_optimizers(self):\n        return torch.optim.AdamW(self.parameters(), lr=self.lr)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T08:15:02.395432Z","iopub.execute_input":"2025-12-18T08:15:02.395742Z","iopub.status.idle":"2025-12-18T08:15:02.404867Z","shell.execute_reply.started":"2025-12-18T08:15:02.395713Z","shell.execute_reply":"2025-12-18T08:15:02.404070Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class AkkadianT5Lightning(pl.LightningModule):\n    def __init__(self, model, tokenizer, lr=1e-4):\n        super().__init__()\n        self.model = model\n        self.tokenizer = tokenizer\n        self.lr = lr\n        self.automatic_optimization = False\n        self.awp = AkkadianT5AWP(model)\n        self.preds, self.refs = [], []\n        self.history = {\"loss\": [], \"geo\": []}\n\n    def training_step(self, batch, batch_idx):\n        opt = self.optimizers()\n\n        out = self.model(\n            input_ids=batch[\"input_ids\"],\n            attention_mask=batch[\"attention_mask\"],\n            labels=batch[\"labels\"]\n        )\n        self.manual_backward(out.loss)\n\n        if self.current_epoch > self.trainer.max_epochs * 0.4:\n            self.awp.attack()\n            out_adv = self.model(**batch)\n            self.manual_backward(out_adv.loss)\n            self.awp.restore()\n\n        if (batch_idx + 1) % self.trainer.accumulate_grad_batches == 0:\n            opt.step()\n            opt.zero_grad()\n\n        self.log(\"train_loss\", out.loss, prog_bar=True)\n        return out.loss\n\n    def validation_step(self, batch, batch_idx):\n        gen = self.model.generate(\n            batch[\"input_ids\"],\n            attention_mask=batch[\"attention_mask\"],\n            num_beams=4,\n            max_new_tokens=256,\n            length_penalty=0.8\n        )\n        self.preds.extend(self.tokenizer.batch_decode(gen, skip_special_tokens=True))\n        self.refs.extend(self.tokenizer.batch_decode(batch[\"labels\"], skip_special_tokens=True))\n\n    def on_validation_epoch_end(self):\n        bleu = sacrebleu.corpus_bleu(self.preds, [self.refs]).score\n        chrf = sacrebleu.corpus_chrf(self.preds, [self.refs]).score\n        geo = math.sqrt(bleu * chrf)\n        self.history[\"geo\"].append(geo)\n        self.log(\"geo_mean\", geo, prog_bar=True)\n        self.preds.clear()\n        self.refs.clear()\n\n    def configure_optimizers(self):\n        return torch.optim.AdamW(self.parameters(), lr=self.lr)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T08:15:04.185580Z","iopub.execute_input":"2025-12-18T08:15:04.186120Z","iopub.status.idle":"2025-12-18T08:15:04.194675Z","shell.execute_reply.started":"2025-12-18T08:15:04.186088Z","shell.execute_reply":"2025-12-18T08:15:04.194105Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"folds = 3\nkf = KFold(n_splits=folds, shuffle=True, random_state=SEED)\nhistories = []\n\nfor fold, (tr_idx, va_idx) in enumerate(kf.split(pairs_df)):\n    print(f\"\\nFOLD {fold+1}\")\n\n    tr = pairs_df.iloc[tr_idx]\n    va = pairs_df.iloc[va_idx]\n\n    tr = pd.concat([tr] * 2, ignore_index=True)\n\n    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)\n\n    dm = AkkadianT5DataModule(tokenizer, tr, va)\n\n    lit = AkkadianT5Lightning(model, tokenizer)\n\n    ckpt = ModelCheckpoint(monitor=\"geo_mean\", mode=\"max\", save_top_k=1)\n    early = EarlyStopping(monitor=\"geo_mean\", mode=\"max\", patience=3)\n    lr_mon = LearningRateMonitor(logging_interval=\"epoch\")\n\n    trainer = pl.Trainer(\n        max_epochs=3,\n        accelerator=\"gpu\" if torch.cuda.is_available() else \"cpu\",\n        precision=\"16-mixed\",\n        callbacks=[ckpt, early, lr_mon],\n    )                \n\n    trainer.fit(lit, dm)\n    histories.append(lit.history)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T08:15:24.797839Z","iopub.execute_input":"2025-12-18T08:15:24.798465Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"plt.figure(figsize=(12,3))\nfor h in histories:\n    plt.plot(h[\"geo\"])\nplt.title(\"BLEU × chrF++ per Epoch\")\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-17T17:18:55.686125Z","iopub.execute_input":"2025-12-17T17:18:55.686477Z","iopub.status.idle":"2025-12-17T17:18:55.857254Z","shell.execute_reply.started":"2025-12-17T17:18:55.686447Z","shell.execute_reply":"2025-12-17T17:18:55.856430Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from glob import glob\nbest_models = glob(\"/kaggle/working/lightning_logs/**/checkpoints/*.ckpt\")\ntotal_best_models = len(best_models)\nlist(best_models)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-17T17:21:30.952583Z","iopub.execute_input":"2025-12-17T17:21:30.953197Z","iopub.status.idle":"2025-12-17T17:21:30.958543Z","shell.execute_reply.started":"2025-12-17T17:21:30.953167Z","shell.execute_reply":"2025-12-17T17:21:30.957983Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\nimport torch\nimport pandas as pd\nfrom transformers import AutoTokenizer, AutoModelForSeq2SeqLM\n\ntokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\nhf_model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).to(\"cuda\")\n\nfold_ckpts = best_models\n\nall_preds = []\n\nfor ckpt_path in fold_ckpts:\n    lit_model = AkkadianT5Lightning(model=hf_model, tokenizer=tokenizer)\n    lit_model.load_state_dict(torch.load(ckpt_path)[\"state_dict\"])\n    lit_model.eval()\n    model = lit_model.model\n\n    preds = []\n    batch_size = 1\n    for i in range(0, len(test_df), batch_size):\n        batch_texts = test_df.transliteration.tolist()[i:i+batch_size]\n        enc = tokenizer(batch_texts, return_tensors=\"pt\", truncation=True, max_length=128, padding=True).to(\"cuda\")\n        with torch.no_grad():\n            gen = model.generate(**enc, num_beams=4, max_new_tokens=256)\n        for g in gen:\n            preds.append(clean_ws(tokenizer.decode(g, skip_special_tokens=True)))\n        torch.cuda.empty_cache()\n    \n    all_preds.append(preds)\n    \n    del lit_model, model, enc, gen\n    torch.cuda.empty_cache()\n\nfinal_preds = []\nfor i in range(len(test_df)):\n    fold_versions = [all_preds[f][i] for f in range(total_best_models)]\n    final_preds.append(max(set(fold_versions), key=fold_versions.count))\n\nsubmission = pd.DataFrame({\"id\": test_df.id, \"translation\": final_preds})\nsubmission.to_csv(\"submission.csv\", index=False)\nsubmission.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-17T17:21:37.170351Z","iopub.execute_input":"2025-12-17T17:21:37.170885Z","iopub.status.idle":"2025-12-17T17:21:47.676423Z","shell.execute_reply.started":"2025-12-17T17:21:37.170856Z","shell.execute_reply":"2025-12-17T17:21:47.675759Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}