{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","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"},{"sourceId":14374989,"sourceType":"datasetVersion","datasetId":9147887},{"sourceId":14376272,"sourceType":"datasetVersion","datasetId":9181082},{"sourceId":14398637,"sourceType":"datasetVersion","datasetId":9195725},{"sourceId":14410665,"sourceType":"datasetVersion","datasetId":9203761}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import re\nimport pandas as pd\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nfrom transformers import AutoTokenizer, AutoModelForSeq2SeqLM\n\n# =========================\n# Config\n# =========================\nTEST_DATA_PATH = \"/kaggle/input/deep-past-initiative-machine-translation/test.csv\"\n\nMODEL1_PATH = \"/kaggle/input/byt5-base-big-data2\"\nMODEL2_PATH = \"/kaggle/input/byt5-akkadian-model\"\nMODEL3_PATH = \"/kaggle/input/train-gap-all-2/byt5-base-akkadian_gap_setence2\"\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nPREFIX = \"translate Akkadian to English: \"\nMAX_SOURCE_LEN = 512\n\nBATCH_SIZE = 8\n\n# Decoding defaults close to your baseline; tune later\nNUM_BEAMS = 6\nMAX_NEW_TOKENS = 384          # IMPORTANT: avoid truncation\nLENGTH_PENALTY = 1.05\nEARLY_STOPPING = True\n\n# IMPORTANT for score stability (ByT5 often dislikes fp16 decoding)\nUSE_AMP = False\n\n# =========================\n# Use YOUR original replace_gaps (recommended)\n# =========================\ndef replace_gaps(text):\n    if pd.isna(text):\n        return text\n    text = str(text)\n    text = re.sub(r'\\.3(?:\\s+\\.3)+\\.{3}(?:\\s+\\.{3})+\\s+\\.{3}(?:\\s+\\.{3})+', '<big_gap>', text)\n    text = re.sub(r'\\.3(?:\\s+\\.3)+\\.{3}(?:\\s+\\.{3})+', '<big_gap>', text)\n    text = re.sub(r'\\.{3}(?:\\s+\\.{3})+', '<big_gap>', text)\n    text = re.sub(r'xx', '<gap>', text)\n    text = re.sub(r' x ', ' <gap> ', text)\n    text = re.sub(r'……', '<big_gap>', text)\n    text = re.sub(r'\\.\\.\\.\\.\\.\\.', '<big_gap>', text)\n    text = re.sub(r'…', '<big_gap>', text)\n    text = re.sub(r'\\.\\.\\.', '<big_gap>', text)\n    return text\n\n# =========================\n# Weighted checkpoint averaging (your approach)\n# =========================\nperf1, perf2, perf3 = 0.98, 1.00, 0.40\ntotal = perf1 + perf2 + perf3\nw1, w2, w3 = perf1/total, perf2/total, perf3/total\nprint(f\"Weights: w1={w1:.4f}, w2={w2:.4f}, w3={w3:.4f}\")\n\nprint(\"Loading models...\")\nm1 = AutoModelForSeq2SeqLM.from_pretrained(MODEL1_PATH)\nm2 = AutoModelForSeq2SeqLM.from_pretrained(MODEL2_PATH)\nm3 = AutoModelForSeq2SeqLM.from_pretrained(MODEL3_PATH)\n\nsd1, sd2, sd3 = m1.state_dict(), m2.state_dict(), m3.state_dict()\n\nprint(\"Averaging weights...\")\nfinal_sd = sd2.copy()\nfor k in final_sd:\n    if k in sd1 and k in sd3:\n        final_sd[k] = w1 * sd1[k] + w2 * sd2[k] + w3 * sd3[k]\n    elif k in sd1:\n        final_sd[k] = w1 * sd1[k] + (w2 + w3) * sd2[k]\n    elif k in sd3:\n        final_sd[k] = w3 * sd3[k] + (w1 + w2) * sd2[k]\n\nmodel = AutoModelForSeq2SeqLM.from_pretrained(MODEL2_PATH)\nmodel.load_state_dict(final_sd)\nmodel.to(DEVICE).eval()\nmodel.float()   # keep fp32\n\ntokenizer = AutoTokenizer.from_pretrained(MODEL2_PATH)\n\n# =========================\n# Dataset + dynamic padding\n# =========================\ntest_df = pd.read_csv(TEST_DATA_PATH)\ntest_df[\"transliteration\"] = test_df[\"transliteration\"].apply(replace_gaps)\n\nclass InferenceDataset(Dataset):\n    def __init__(self, df):\n        self.ids = df[\"id\"].tolist()\n        self.texts = [PREFIX + t for t in df[\"transliteration\"].astype(str).tolist()]\n\n    def __len__(self):\n        return len(self.texts)\n\n    def __getitem__(self, idx):\n        return self.ids[idx], self.texts[idx]\n\ndef collate_fn(batch):\n    ids, texts = zip(*batch)\n    enc = tokenizer(\n        list(texts),\n        max_length=MAX_SOURCE_LEN,\n        truncation=True,\n        padding=True,\n        return_tensors=\"pt\"\n    )\n    return list(ids), enc[\"input_ids\"], enc[\"attention_mask\"]\n\nloader = DataLoader(\n    InferenceDataset(test_df),\n    batch_size=BATCH_SIZE,\n    shuffle=False,\n    num_workers=2,\n    pin_memory=(DEVICE.type == \"cuda\"),\n    collate_fn=collate_fn\n)\n\n# =========================\n# Inference\n# =========================\nall_ids, all_pred = [], []\ntorch.set_grad_enabled(False)\n\nwith torch.inference_mode():\n    for ids, input_ids, attention_mask in loader:\n        input_ids = input_ids.to(DEVICE)\n        attention_mask = attention_mask.to(DEVICE)\n\n        # Keep decoding in fp32 for score stability\n        if USE_AMP:\n            ctx = torch.autocast(device_type=\"cuda\", dtype=torch.float16)\n        else:\n            from contextlib import nullcontext\n            ctx = nullcontext()\n\n        with ctx:\n            outputs = model.generate(\n                input_ids=input_ids,\n                attention_mask=attention_mask,\n                num_beams=NUM_BEAMS,\n                max_new_tokens=MAX_NEW_TOKENS,\n                length_penalty=LENGTH_PENALTY,\n                early_stopping=EARLY_STOPPING,\n            )\n\n        decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)\n        decoded = [d.strip() if d.strip() else \"broken text\" for d in decoded]\n\n        all_ids.extend(ids)\n        all_pred.extend(decoded)\n\nsubmission = pd.DataFrame({\"id\": all_ids, \"translation\": all_pred})\nsubmission.to_csv(\"submission.csv\", index=False)\nprint(\"Saved submission.csv\")\nprint(submission.head())","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-13T13:32:22.18491Z","iopub.execute_input":"2026-01-13T13:32:22.185136Z","iopub.status.idle":"2026-01-13T13:34:32.198287Z","shell.execute_reply.started":"2026-01-13T13:32:22.185113Z","shell.execute_reply":"2026-01-13T13:34:32.197382Z"}},"outputs":[],"execution_count":null}]}