{"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},{"sourceId":286997242,"sourceType":"kernelVersion"}],"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-18T09:31:18.692135Z","iopub.execute_input":"2025-12-18T09:31:18.692979Z","iopub.status.idle":"2025-12-18T09:31:18.769940Z","shell.execute_reply.started":"2025-12-18T09:31:18.692942Z","shell.execute_reply":"2025-12-18T09:31:18.769367Z"}},"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-18T09:31:20.165554Z","iopub.execute_input":"2025-12-18T09:31:20.166419Z","iopub.status.idle":"2025-12-18T09:31:20.178411Z","shell.execute_reply.started":"2025-12-18T09:31:20.166384Z","shell.execute_reply":"2025-12-18T09:31:20.177735Z"}},"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-18T09:31:21.747741Z","iopub.execute_input":"2025-12-18T09:31:21.748046Z","iopub.status.idle":"2025-12-18T09:31:21.803384Z","shell.execute_reply.started":"2025-12-18T09:31:21.748015Z","shell.execute_reply":"2025-12-18T09:31:21.802791Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def akk_clean(x):\n    return re.sub(r\"\\s+\", \" \", str(x)).strip()\n\npairs = []\nfor _, r in train_raw.iterrows():\n    src = re.split(r\"\\n+\", str(r[\"transliteration\"]))\n    tgt = re.split(r\"(?<=[.!?])\\s+\", akk_clean(r[\"translation\"]))\n    src = [akk_clean(x) for x in src if akk_clean(x)]\n    tgt = [akk_clean(x) for x in tgt if akk_clean(x)]\n    for i in range(min(len(src), len(tgt))):\n        pairs.append((src[i], tgt[i]))\n\npairs_df = pd.DataFrame(pairs, columns=[\"src\", \"tgt\"]).drop_duplicates()\npairs_df = pairs_df.sample(frac=1.0, random_state=SEED).reset_index(drop=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T09:31:22.484342Z","iopub.execute_input":"2025-12-18T09:31:22.485052Z","iopub.status.idle":"2025-12-18T09:31:22.859394Z","shell.execute_reply.started":"2025-12-18T09:31:22.485020Z","shell.execute_reply":"2025-12-18T09:31:22.858647Z"}},"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-18T09:31:37.322839Z","iopub.execute_input":"2025-12-18T09:31:37.323167Z","iopub.status.idle":"2025-12-18T09:31:37.506073Z","shell.execute_reply.started":"2025-12-18T09:31:37.323126Z","shell.execute_reply":"2025-12-18T09:31:37.505308Z"}},"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-18T09:31:39.889131Z","iopub.execute_input":"2025-12-18T09:31:39.889850Z","iopub.status.idle":"2025-12-18T09:31:39.898474Z","shell.execute_reply.started":"2025-12-18T09:31:39.889808Z","shell.execute_reply":"2025-12-18T09:31:39.897786Z"}},"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-18T09:31:41.059124Z","iopub.execute_input":"2025-12-18T09:31:41.059867Z","iopub.status.idle":"2025-12-18T09:31:41.066077Z","shell.execute_reply.started":"2025-12-18T09:31:41.059836Z","shell.execute_reply":"2025-12-18T09:31:41.065213Z"}},"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=128,\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-18T09:31:41.686735Z","iopub.execute_input":"2025-12-18T09:31:41.687053Z","iopub.status.idle":"2025-12-18T09:31:41.696357Z","shell.execute_reply.started":"2025-12-18T09:31:41.687026Z","shell.execute_reply":"2025-12-18T09:31:41.695539Z"}},"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=128,\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-18T09:31:42.624268Z","iopub.execute_input":"2025-12-18T09:31:42.624565Z","iopub.status.idle":"2025-12-18T09:31:42.633739Z","shell.execute_reply.started":"2025-12-18T09:31:42.624539Z","shell.execute_reply":"2025-12-18T09:31:42.632971Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# folds = 4\n# kf = KFold(n_splits=folds, shuffle=True, random_state=SEED)\n# histories = []\n\n# for 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] * 6, 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=1,\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-18T09:31:43.795888Z","iopub.execute_input":"2025-12-18T09:31:43.796696Z","iopub.status.idle":"2025-12-18T09:31:43.800106Z","shell.execute_reply.started":"2025-12-18T09:31:43.796663Z","shell.execute_reply":"2025-12-18T09:31:43.799385Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# plt.figure(figsize=(12,4))\n# for h in histories:\n#     plt.plot(h[\"geo\"])\n# plt.title(\"BLEU × chrF++ per Epoch\")\n# plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T09:31:45.265788Z","iopub.execute_input":"2025-12-18T09:31:45.266083Z","iopub.status.idle":"2025-12-18T09:31:45.269434Z","shell.execute_reply.started":"2025-12-18T09:31:45.266054Z","shell.execute_reply":"2025-12-18T09:31:45.268720Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from glob import glob\nbest_models = glob(\"/kaggle/input/akkadian-translate-lightning-cv-36-00-train/lightning_logs/version_*/checkpoints/epoch=*.ckpt\")\ntotal_best_models = len(best_models)\nlist(best_models)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T09:31:53.078082Z","iopub.execute_input":"2025-12-18T09:31:53.078896Z","iopub.status.idle":"2025-12-18T09:31:53.115667Z","shell.execute_reply.started":"2025-12-18T09:31:53.078860Z","shell.execute_reply":"2025-12-18T09:31:53.115067Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from collections import Counter\n\ndef akk_norm(x):\n    return re.sub(r\"\\s+\", \" \", str(x)).strip()\n    \n\ntokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n\nfold_ckpts = best_models\ntotal_best_models = len(fold_ckpts)\n\nall_preds = []\n\nfor ckpt_path in fold_ckpts:\n    hf_model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).to(\"cuda\")\n\n    lit_model = AkkadianT5Lightning(\n        model=hf_model,\n        tokenizer=tokenizer\n    )\n\n    state = torch.load(ckpt_path, map_location=\"cuda\")\n    lit_model.load_state_dict(state[\"state_dict\"])\n    lit_model = lit_model.to(\"cuda\")\n    lit_model.eval()\n\n    model = lit_model.model\n\n    preds = []\n    batch_size = 1\n\n    for i in range(0, len(test_df), batch_size):\n        batch_texts = test_df.transliteration.iloc[i:i+batch_size].tolist()\n        enc = tokenizer(\n            batch_texts,\n            return_tensors=\"pt\",\n            truncation=True,\n            max_length=128,\n            padding=True\n        ).to(\"cuda\")\n\n        with torch.no_grad():\n            gen = model.generate(\n                **enc,\n                num_beams=4,\n                max_new_tokens=256\n            )\n\n        for g in gen:\n            preds.append(akk_norm(tokenizer.decode(g, skip_special_tokens=True)))\n\n    all_preds.append(preds)\n\n    del lit_model, model, hf_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(Counter(fold_versions).most_common(1)[0][0])\n\nsubmission = pd.DataFrame({\n    \"id\": test_df.id,\n    \"translation\": final_preds\n})\n\nsubmission.to_csv(\"submission.csv\", index=False)\nsubmission.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T09:32:13.172942Z","iopub.execute_input":"2025-12-18T09:32:13.173580Z","iopub.status.idle":"2025-12-18T09:33:42.650530Z","shell.execute_reply.started":"2025-12-18T09:32:13.173552Z","shell.execute_reply":"2025-12-18T09:33:42.649743Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for _, row in submission.iterrows():\n    print(f\"ID:-  {row['id']} translation:-  {row['translation']}\\n\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T09:46:17.733885Z","iopub.execute_input":"2025-12-18T09:46:17.734211Z","iopub.status.idle":"2025-12-18T09:46:17.739217Z","shell.execute_reply.started":"2025-12-18T09:46:17.734181Z","shell.execute_reply":"2025-12-18T09:46:17.738574Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import re\nfrom spellchecker import SpellChecker\n\nspell = SpellChecker()\n\ndef fix_english(text):\n    if not isinstance(text, str):\n        return text\n\n    words = text.split()\n    fixed_words = []\n\n    for word in words:\n        if word == '...':\n            continue\n        clean = re.sub(r\"[^\\w']\", \"\", word)\n        if clean.lower() in spell:\n            fixed_words.append(word)\n        else:\n            corrected = spell.correction(clean)\n            if corrected:\n                fixed_words.append(word.replace(clean, corrected))\n            else:\n                fixed_words.append(word)\n\n    # Fix capitalization after periods\n    text = \" \".join(fixed_words)\n    text = re.sub(r\"([.!?])\\s*(\\w)\", lambda m: m.group(1) + \" \" + m.group(2).upper(), text)\n\n    return text","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T09:45:59.583845Z","iopub.execute_input":"2025-12-18T09:45:59.584477Z","iopub.status.idle":"2025-12-18T09:45:59.711128Z","shell.execute_reply.started":"2025-12-18T09:45:59.584443Z","shell.execute_reply":"2025-12-18T09:45:59.710368Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission[\"translation\"] = submission[\"translation\"].apply(fix_english)\nfor _, row in submission.iterrows():\n    print(f\"ID:-  {row['id']} translation:-  {row['translation']}\\n\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T09:46:01.198621Z","iopub.execute_input":"2025-12-18T09:46:01.198918Z","iopub.status.idle":"2025-12-18T09:46:01.205627Z","shell.execute_reply.started":"2025-12-18T09:46:01.198891Z","shell.execute_reply":"2025-12-18T09:46:01.204829Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission.to_csv(\"submission.csv\", index=False)\nsubmission.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T09:46:41.148138Z","iopub.execute_input":"2025-12-18T09:46:41.148819Z","iopub.status.idle":"2025-12-18T09:46:41.161071Z","shell.execute_reply.started":"2025-12-18T09:46:41.148784Z","shell.execute_reply":"2025-12-18T09:46:41.160216Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}