{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","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":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":121150,"databundleVersionId":15833391,"isSourceIdPinned":false}],"dockerImageVersionId":31287,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"178b8634-7bc9-41f7-8d88-97010bf978b3","cell_type":"markdown","source":"# Deep Past Initiative – Best Combined Notebook\n\nCombines the strongest elements from all three starter notebooks:\n\n| Feature | Source |\n|---|---|\n| ByT5 character-level model | Train NB |\n| Sentence alignment (data augmentation) | Train NB |\n| **Gap token normalisation** applied to **training data too** | Ensemble NB |\n| chrF++ evaluation metric | Train NB |\n| TR-TRY rule-based retrieval fallback (active, not commented out) | Infer NB |\n| Weight-averaged ensemble of best checkpoints | Ensemble NB |\n| Period-termination & gap restoration post-processing | Infer NB |\n\n### Key fixes over the originals\n- Gap normalisation is now applied **before training** (consistent train/test representation)\n- TR-TRY retrieval is **active** and used as a hybrid fallback (not commented out)\n- Weight averaging uses **normalised weights** derived from validation chrF++ scores\n- Warmup scheduler + label smoothing added to training\n- Improved sentence aligner: also splits on Akkadian semicolons and pipe characters\n- Post-processing enforces sentence-final punctuation and restores gap tokens","metadata":{}},{"id":"07a53f59-15cc-4e2c-9e92-d12ff3838ee1","cell_type":"markdown","source":"## 1. Install & Imports","metadata":{}},{"id":"6406ec44-cf48-4405-8809-5d01fd40a1f6","cell_type":"code","source":"!pip install -q evaluate sacrebleu scikit-learn sentence-transformers","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T09:46:45.251888Z","iopub.execute_input":"2026-03-09T09:46:45.252701Z","iopub.status.idle":"2026-03-09T09:46:50.503280Z","shell.execute_reply.started":"2026-03-09T09:46:45.252663Z","shell.execute_reply":"2026-03-09T09:46:50.502541Z"}},"outputs":[],"execution_count":null},{"id":"27749bfa-053a-4dcb-ae5f-004d607a298b","cell_type":"code","source":"import os\nimport gc\nimport re\nimport unicodedata\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom collections import Counter\nfrom sklearn.feature_extraction.text import TfidfVectorizer\nfrom sklearn.metrics.pairwise import cosine_similarity\nfrom torch.utils.data import DataLoader, Dataset as TorchDataset\nfrom datasets import Dataset\nfrom transformers import (\n    AutoTokenizer,\n    AutoModelForSeq2SeqLM,\n    DataCollatorForSeq2Seq,\n    Seq2SeqTrainingArguments,\n    Seq2SeqTrainer,\n    get_cosine_schedule_with_warmup,\n)\nimport evaluate\nfrom tqdm.auto import tqdm\n\nprint(\"PyTorch:\", torch.__version__)\nprint(\"CUDA available:\", torch.cuda.is_available())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T09:46:50.504911Z","iopub.execute_input":"2026-03-09T09:46:50.505285Z","iopub.status.idle":"2026-03-09T09:47:11.453932Z","shell.execute_reply.started":"2026-03-09T09:46:50.505213Z","shell.execute_reply":"2026-03-09T09:47:11.453166Z"}},"outputs":[],"execution_count":null},{"id":"2fd263bd-12fa-446e-bb74-6db16e2d3798","cell_type":"markdown","source":"## 2. Configuration","metadata":{}},{"id":"1a2091ac-8ad8-4e6b-ba55-8f41a09a1790","cell_type":"code","source":"# ============================================================\n# 🔧 DEBUG SWITCH\n#   True  → 1 epoch, 128 train rows, 32 test rows — fast end-to-end check\n#   False → full run (35 epochs, all data)\n# ============================================================\nDEBUG = False\n\nclass Config:\n    # ByT5 operates at byte level — ideal for noisy Akkadian transliterations\n    # with diacritics, subscript numbers, editorial symbols, and gap markers.\n    # Use byt5-base if your GPU has >= 16 GB VRAM; otherwise byt5-small.\n    MODEL_NAME = \"google/byt5-small\"   # swap to \"google/byt5-base\" if VRAM allows\n\n    MAX_LENGTH = 512\n    PREFIX = \"translate Akkadian to English: \"\n\n    # Effective batch size = PER_DEVICE_BATCH * GRAD_ACCUM = 4 * 4 = 16\n    PER_DEVICE_BATCH = 4\n    GRAD_ACCUM = 4\n\n    EPOCHS = 35                 # more epochs than the baseline (was 20)\n    LEARNING_RATE = 2e-4\n    WARMUP_RATIO = 0.05         # 5% warmup steps\n    LABEL_SMOOTHING = 0.1       # smoothing helps generalisation on small data\n    WEIGHT_DECAY = 0.01\n\n    # TR-TRY: similarity threshold above which retrieval beats the neural model\n    # (only used when the retrieved score is very high — near-exact match)\n    RETRIEVAL_THRESHOLD = 0.82\n\n    OUTPUT_DIR = \"/kaggle/working/byt5-akkadian-model\"\n    INPUT_DIR = \"/kaggle/input/competitions/deep-past-initiative-machine-translation\"\n\n\ndef seed_everything(seed=42):\n    torch.manual_seed(seed)\n    np.random.seed(seed)\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed_all(seed)\n\nseed_everything()\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"Device:\", DEVICE)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T09:47:11.454904Z","iopub.execute_input":"2026-03-09T09:47:11.455484Z","iopub.status.idle":"2026-03-09T09:47:11.467795Z","shell.execute_reply.started":"2026-03-09T09:47:11.455458Z","shell.execute_reply":"2026-03-09T09:47:11.467069Z"}},"outputs":[],"execution_count":null},{"id":"17295e8d-fed4-460a-a76d-db51df32289e","cell_type":"markdown","source":"## 3. Text Preprocessing\n\n### 3a. Gap Token Normalisation\nApplied to **both training and test** data so the model sees a consistent representation\nof damaged/illegible text. (The original ensemble notebook only applied this at inference time.)","metadata":{}},{"id":"4311a6f4-1de5-47e0-811a-6b3125d0c615","cell_type":"code","source":"def normalise_gaps(text: str) -> str:\n    \"\"\"Replace gap/damage markers in Akkadian transliterations with canonical tokens.\n    \n    Akkadian cuneiform tablets often have broken/illegible sections marked with:\n      x, xx, ..., …, ……, etc.\n    We collapse all of these into two clean tokens that ByT5 can learn to pass through.\n    \"\"\"\n    if pd.isna(text):\n        return \"\"\n    text = str(text)\n\n    # Multi-ellipsis patterns → big gap\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'……', ' <big_gap> ', text)\n    text = re.sub(r'\\.{6}', ' <big_gap> ', text)\n    text = re.sub(r'…', ' <big_gap> ', text)\n    text = re.sub(r'\\.{3}', ' <big_gap> ', text)\n\n    # Single gap markers → small gap\n    text = re.sub(r'\\bxx\\b', ' <gap> ', text)\n    text = re.sub(r' x ', ' <gap> ', text)\n    text = re.sub(r'^x ', '<gap> ', text)\n    text = re.sub(r' x$', ' <gap>', text)\n\n    # Collapse multiple spaces\n    text = re.sub(r'\\s+', ' ', text).strip()\n    return text\n\n\ndef restore_gaps(text: str) -> str:\n    \"\"\"Reverse normalise_gaps: convert canonical tokens back to readable notation.\"\"\"\n    if pd.isna(text):\n        return \"\"\n    text = re.sub(r'<big_gap>', '...', str(text))\n    text = re.sub(r'<gap>', 'x', text)\n    return text.strip()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T09:47:11.469449Z","iopub.execute_input":"2026-03-09T09:47:11.469791Z","iopub.status.idle":"2026-03-09T09:47:11.483198Z","shell.execute_reply.started":"2026-03-09T09:47:11.469757Z","shell.execute_reply":"2026-03-09T09:47:11.482542Z"}},"outputs":[],"execution_count":null},{"id":"b49411f8-17fe-4eab-8eb7-7aa8f0500f07","cell_type":"markdown","source":"### 3b. Akkadian Normalisation for Retrieval Index\nUsed only by the TR-TRY retrieval system — strips diacritics and noise for better TF-IDF matching.","metadata":{}},{"id":"0b0391ca-edcb-4f44-a585-6a4f2b07b2dd","cell_type":"code","source":"def normalise_akkadian_for_retrieval(text: str) -> str:\n    \"\"\"Normalise Akkadian source for TF-IDF indexing.\n    \n    Strips subscript numbers, diacritics, editorial brackets, and\n    determinative markers so that phonetically similar words match.\n    \"\"\"\n    if pd.isna(text):\n        return \"\"\n    text = str(text).lower().strip()\n\n    # Expand determinatives: {d}, {m}, {ki} → DET_d etc.\n    text = re.sub(r'\\{([^}]+)\\}', r' DET_\\1 ', text)\n    text = re.sub(r'\\(([a-z]{1,6})\\)', r' DET_\\1 ', text)\n\n    # Remove subscript/superscript numbers (homophone disambiguation)\n    text = re.sub(r'[₀-₉0-9]', '', text)\n\n    # Unify gap markers\n    text = re.sub(r'\\.{2,}|…|\\bx{1,3}\\b', ' GAP ', text)\n\n    # Remove editorial brackets\n    text = re.sub(r'[\\[\\]<>]', ' ', text)\n\n    # Strip diacritics (š→s, ṣ→s, ṭ→t, etc.)\n    text = unicodedata.normalize('NFKD', text)\n    text = ''.join(ch for ch in text if not unicodedata.combining(ch))\n\n    # Remove apostrophe-like signs (ayin, aleph)\n    text = re.sub(r\"[ʾʿˀˁ']\", '', text)\n\n    # Keep only ASCII alphanum, hyphen, underscore\n    text = re.sub(r'[^a-z0-9\\-_ ]+', ' ', text)\n    text = re.sub(r'\\s+', ' ', text).strip()\n    return text","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T09:47:11.484421Z","iopub.execute_input":"2026-03-09T09:47:11.484790Z","iopub.status.idle":"2026-03-09T09:47:11.501993Z","shell.execute_reply.started":"2026-03-09T09:47:11.484767Z","shell.execute_reply":"2026-03-09T09:47:11.501166Z"}},"outputs":[],"execution_count":null},{"id":"b4b8ef00-802c-453d-8eb1-b2b48f69a947","cell_type":"markdown","source":"### 3c. Sentence Aligner (Data Augmentation)\n\nThe training set is document-level; the test set is sentence-level.\nThis aligner heuristically splits training documents into sentence pairs,\npotentially 2–3× expanding the training set.\n\nImprovements over the baseline:\n- Also splits Akkadian on `;` and `|` (common line delimiters in cuneiform)\n- Minimum token-length filter raised to 5 characters to reduce noise\n- Keeps original document if any split fragment is suspiciously short","metadata":{}},{"id":"73fb4c17-bff1-482c-8bfd-04233ed02716","cell_type":"code","source":"def improved_sentence_aligner(df: pd.DataFrame) -> pd.DataFrame:\n    \"\"\"Split document-level pairs into sentence-level pairs where the split is reliable.\"\"\"\n    aligned = []\n\n    for _, row in df.iterrows():\n        src = str(row['transliteration'])\n        tgt = str(row['translation'])\n\n        # --- English: split on sentence-ending punctuation ---\n        tgt_sents = [t.strip() for t in re.split(r'(?<=[.!?])\\s+', tgt) if t.strip()]\n\n        # --- Akkadian: split on newlines, semicolons, or pipe characters ---\n        src_lines = [s.strip() for s in re.split(r'[\\n;|]', src) if s.strip()]\n\n        # Only split when counts match and every fragment is long enough\n        counts_match = len(tgt_sents) > 1 and len(tgt_sents) == len(src_lines)\n        fragments_ok = all(len(s) >= 5 and len(t) >= 5\n                           for s, t in zip(src_lines, tgt_sents))\n\n        if counts_match and fragments_ok:\n            for s, t in zip(src_lines, tgt_sents):\n                aligned.append({'transliteration': s, 'translation': t})\n        else:\n            # Safe fallback: keep the original document pair\n            aligned.append({'transliteration': src, 'translation': tgt})\n\n    return pd.DataFrame(aligned)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T09:47:11.503042Z","iopub.execute_input":"2026-03-09T09:47:11.503385Z","iopub.status.idle":"2026-03-09T09:47:11.518902Z","shell.execute_reply.started":"2026-03-09T09:47:11.503350Z","shell.execute_reply":"2026-03-09T09:47:11.518296Z"}},"outputs":[],"execution_count":null},{"id":"16880390-191e-486b-82d0-6abbaa8733f0","cell_type":"markdown","source":"## 4. Load & Prepare Training Data","metadata":{}},{"id":"855230a0-36c3-437d-9a96-77129e2a7f9d","cell_type":"code","source":"train_df = pd.read_csv(f\"{Config.INPUT_DIR}/train.csv\")\ntest_df  = pd.read_csv(f\"{Config.INPUT_DIR}/test.csv\")\nprint(f\"Raw train docs  : {len(train_df)}\")\nprint(f\"Test examples   : {len(test_df)}\")\n\n# Step 1: Apply gap normalisation to the source side of training data.\n# This is the key fix vs. the originals — the model now sees the same\n# <gap>/<big_gap> tokens at training time as it will at inference time.\ntrain_df['transliteration'] = train_df['transliteration'].apply(normalise_gaps)\n\n# Step 2: Expand documents into sentence pairs where reliable.\ntrain_expanded = improved_sentence_aligner(train_df)\nprint(f\"After alignment : {len(train_expanded)} sentence pairs\")\n\n# Step 3: Build HuggingFace Dataset and create 90/10 train/val split.\nhf_dataset = Dataset.from_pandas(train_expanded.reset_index(drop=True))\nsplits = hf_dataset.train_test_split(test_size=0.1, seed=42)\nprint(f\"Train split     : {len(splits['train'])}\")\nprint(f\"Val split       : {len(splits['test'])}\")\n# ── Debug slice ────────────────────────────────────────────\nif DEBUG:\n    train_expanded = train_expanded.head(128).reset_index(drop=True)\n    test_df        = test_df.head(32).reset_index(drop=True)\n    print(f\"[DEBUG] Sliced to {len(train_expanded)} train rows, {len(test_df)} test rows\")\n# ───────────────────────────────────────────────────────────\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T09:47:11.519740Z","iopub.execute_input":"2026-03-09T09:47:11.520068Z","iopub.status.idle":"2026-03-09T09:47:11.841181Z","shell.execute_reply.started":"2026-03-09T09:47:11.520046Z","shell.execute_reply":"2026-03-09T09:47:11.840573Z"}},"outputs":[],"execution_count":null},{"id":"0aa718ff-0d93-46a1-b81a-b2dc33817721","cell_type":"markdown","source":"## 5. Tokenisation","metadata":{}},{"id":"d2e68332-2aa5-4753-81ee-a3056a42c398","cell_type":"code","source":"tokenizer = AutoTokenizer.from_pretrained(Config.MODEL_NAME)\n\ndef preprocess(examples):\n    inputs  = [Config.PREFIX + str(x) for x in examples['transliteration']]\n    targets = [str(x) for x in examples['translation']]\n\n    model_inputs = tokenizer(\n        inputs, max_length=Config.MAX_LENGTH, truncation=True\n    )\n    labels = tokenizer(\n        targets, max_length=Config.MAX_LENGTH, truncation=True\n    )\n    model_inputs['labels'] = labels['input_ids']\n    return model_inputs\n\ntokenised_train = splits['train'].map(preprocess, batched=True, remove_columns=splits['train'].column_names)\ntokenised_val   = splits['test'].map(preprocess,  batched=True, remove_columns=splits['test'].column_names)\nprint(\"Tokenisation complete.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T09:47:11.841985Z","iopub.execute_input":"2026-03-09T09:47:11.842256Z","iopub.status.idle":"2026-03-09T09:48:00.641025Z","shell.execute_reply.started":"2026-03-09T09:47:11.842204Z","shell.execute_reply":"2026-03-09T09:48:00.640093Z"}},"outputs":[],"execution_count":null},{"id":"d0e4be15-7113-4b69-bdfe-518e594503e2","cell_type":"markdown","source":"## 6. Model Training\n\nKey changes vs baseline:\n- `label_smoothing_factor=0.1` — regularises on small/noisy data\n- `warmup_ratio=0.05` — prevents early LR spike\n- `load_best_model_at_end=True` — saves the best chrF++ checkpoint automatically\n- `fp16=False` — required; ByT5 produces NaN losses in mixed precision","metadata":{}},{"id":"02fef158-fb2e-4fdb-b76a-5a797818b14f","cell_type":"code","source":"# ── Debug overrides ────────────────────────────────────────\nif DEBUG:\n    print(\"⚠️  DEBUG MODE — 1 epoch, reduced dataset, fast validation\")\n    Config.EPOCHS        = 1\n    Config.GRAD_ACCUM    = 1   # no need to accumulate over tiny dataset\n_epochs = Config.EPOCHS\n# ───────────────────────────────────────────────────────────\n\n# Compute warmup_steps explicitly (warmup_ratio is deprecated in transformers v5)\nsteps_per_epoch = len(tokenised_train) // (Config.PER_DEVICE_BATCH * Config.GRAD_ACCUM)\ntotal_steps = steps_per_epoch * Config.EPOCHS\nwarmup_steps = int(total_steps * Config.WARMUP_RATIO)\nprint(f\"Total steps: {total_steps}, Warmup steps: {warmup_steps}\")\n\ngc.collect()\ntorch.cuda.empty_cache() if torch.cuda.is_available() else None\n\nmodel = AutoModelForSeq2SeqLM.from_pretrained(Config.MODEL_NAME, tie_word_embeddings=False)\ndata_collator = DataCollatorForSeq2Seq(tokenizer, model=model, pad_to_multiple_of=8)\n\nchrf_metric = evaluate.load('chrf')\n\ndef compute_metrics(eval_preds):\n    preds, labels = eval_preds\n    if isinstance(preds, tuple):\n        preds = preds[0]\n    decoded_preds  = tokenizer.batch_decode(preds, skip_special_tokens=True)\n    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)\n    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n    decoded_preds  = [p.strip() for p in decoded_preds]\n    decoded_labels = [[l.strip()] for l in decoded_labels]\n    result = chrf_metric.compute(predictions=decoded_preds, references=decoded_labels)\n    return {'chrf': result['score']}\n\ntraining_args = Seq2SeqTrainingArguments(\n    output_dir=Config.OUTPUT_DIR,\n\n    # Evaluation & checkpointing\n    eval_strategy='epoch',\n    save_strategy='epoch',\n    load_best_model_at_end=True,   # keeps the best chrF++ checkpoint\n    metric_for_best_model='chrf',\n    greater_is_better=True,\n    save_total_limit=3,\n\n    # Optimisation\n    learning_rate=Config.LEARNING_RATE,\n    warmup_steps=warmup_steps,\n    weight_decay=Config.WEIGHT_DECAY,\n    label_smoothing_factor=Config.LABEL_SMOOTHING,\n    num_train_epochs=_epochs,\n\n    # Batch sizes — fp32 requires smaller per-device batch, compensated with accumulation\n    per_device_train_batch_size=Config.PER_DEVICE_BATCH,\n    per_device_eval_batch_size=Config.PER_DEVICE_BATCH,\n    gradient_accumulation_steps=Config.GRAD_ACCUM,\n    fp16=False,   # MUST be False for ByT5 — mixed precision causes NaN losses\n\n    predict_with_generate=True,\n    logging_steps=20,\n    report_to='none',\n)\n\ntrainer = Seq2SeqTrainer(\n    model=model,\n    args=training_args,\n    train_dataset=tokenised_train,\n    eval_dataset=tokenised_val,\n    data_collator=data_collator,\n    processing_class=tokenizer,\n    compute_metrics=compute_metrics,\n)\n\nprint(f\"Training {Config.MODEL_NAME} for {Config.EPOCHS} epochs (fp32)...\")\ntrainer.train()\n\n# Save the best checkpoint\ntrainer.save_model(Config.OUTPUT_DIR)\ntokenizer.save_pretrained(Config.OUTPUT_DIR)\nprint(f\"Best model saved to {Config.OUTPUT_DIR}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T09:48:00.642088Z","iopub.execute_input":"2026-03-09T09:48:00.642421Z","iopub.status.idle":"2026-03-09T10:03:54.505066Z","shell.execute_reply.started":"2026-03-09T09:48:00.642396Z","shell.execute_reply":"2026-03-09T10:03:54.504447Z"}},"outputs":[],"execution_count":null},{"id":"7b860606-ab08-4811-9f11-e4bd0d13f135","cell_type":"markdown","source":"## 7. Build TR-TRY Translation Memory\n\nA TF-IDF retrieval index built over the full training set.\nAt inference time, if a test source has cosine similarity > `RETRIEVAL_THRESHOLD`\nagainst a training source, we use the retrieved translation directly —\nthis is essentially free and very accurate for near-duplicate inputs.\n\nFor everything else, the neural ByT5 model is used.","metadata":{}},{"id":"cd996be9-f10a-4cae-a335-58e8a95e4d4d","cell_type":"code","source":"# Build on the full (unexpanded) training set to avoid duplicate retrievals\ntrain_raw = pd.read_csv(f\"{Config.INPUT_DIR}/train.csv\")\ntrain_raw['transliteration'] = train_raw['transliteration'].apply(normalise_gaps)\ntrain_raw['src_norm'] = train_raw['transliteration'].apply(normalise_akkadian_for_retrieval)\ntrain_raw['tgt_text'] = train_raw['translation'].astype(str)\n\nprint(\"Building TF-IDF retrieval index...\")\ntfidf = TfidfVectorizer(\n    analyzer='char_wb',\n    ngram_range=(2, 4),   # character n-grams — robust to morphological variation\n    min_df=1,\n    sublinear_tf=True,\n)\ntrain_vectors = tfidf.fit_transform(train_raw['src_norm'])\nprint(f\"Index built: {train_vectors.shape[0]} documents, {train_vectors.shape[1]} features\")\n\n\ndef retrieve_translation(query_src: str, top_k: int = 1):\n    \"\"\"Return (best_score, best_translation) for a source string.\"\"\"\n    q_norm = normalise_akkadian_for_retrieval(query_src)\n    q_vec  = tfidf.transform([q_norm])\n    sims   = cosine_similarity(q_vec, train_vectors)[0]\n    best_idx   = sims.argmax()\n    best_score = sims[best_idx]\n    best_tgt   = train_raw.iloc[best_idx]['tgt_text']\n    return float(best_score), best_tgt","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T10:03:54.507204Z","iopub.execute_input":"2026-03-09T10:03:54.507479Z","iopub.status.idle":"2026-03-09T10:03:55.533674Z","shell.execute_reply.started":"2026-03-09T10:03:54.507457Z","shell.execute_reply":"2026-03-09T10:03:55.532905Z"}},"outputs":[],"execution_count":null},{"id":"0f9e6558-efb4-4629-8fc2-2959810ce011","cell_type":"markdown","source":"## 8. Optional: Weight-Average a Second Checkpoint\n\nIf you have trained a second model (e.g., with a different seed or more epochs),\nweight-averaging it with the primary model often improves performance at zero inference cost.\n\nSet `SECOND_MODEL_PATH` and adjust `w1`/`w2` to reflect relative validation scores.\nSkip this cell if you only have one checkpoint.","metadata":{}},{"id":"72c78b70-cad5-4657-83df-0bde84312019","cell_type":"code","source":"SECOND_MODEL_PATH = None   # e.g. \"/kaggle/input/byt5-extra/byt5-akkadian-model-seed99\"\n\nif SECOND_MODEL_PATH and os.path.exists(SECOND_MODEL_PATH):\n    print(f\"Loading second checkpoint from {SECOND_MODEL_PATH} for weight averaging...\")\n\n    model2 = AutoModelForSeq2SeqLM.from_pretrained(SECOND_MODEL_PATH)\n    sd_primary = model.state_dict()\n    sd_second  = model2.state_dict()\n\n    # Set w1, w2 proportional to validation chrF++ scores.\n    # Example: primary=42.1, secondary=40.0  →  w1=42.1/82.1, w2=40.0/82.1\n    chrf_primary   = 1.00   # replace with actual val score\n    chrf_secondary = 0.94   # replace with actual val score\n    total = chrf_primary + chrf_secondary\n    w1 = chrf_primary   / total\n    w2 = chrf_secondary / total\n    print(f\"Ensemble weights — primary: {w1:.3f}, secondary: {w2:.3f}\")\n\n    for key in sd_primary:\n        sd_primary[key] = w1 * sd_primary[key] + w2 * sd_second[key].to(sd_primary[key].device)\n\n    model.load_state_dict(sd_primary)\n    del model2, sd_second\n    gc.collect()\n    print(\"Weight averaging complete.\")\nelse:\n    print(\"Single model mode (no second checkpoint provided).\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T10:03:55.534635Z","iopub.execute_input":"2026-03-09T10:03:55.535023Z","iopub.status.idle":"2026-03-09T10:03:57.987063Z","shell.execute_reply.started":"2026-03-09T10:03:55.534996Z","shell.execute_reply":"2026-03-09T10:03:57.986058Z"}},"outputs":[],"execution_count":null},{"id":"b86c8b3c-9d60-483c-bec7-b3a3e704f906","cell_type":"markdown","source":"## 9. Inference\n\nHybrid strategy:\n1. For each test example, check retrieval similarity against training data\n2. If similarity ≥ `RETRIEVAL_THRESHOLD` → use retrieved translation (very reliable)\n3. Otherwise → use ByT5 neural translation\n\nAfter generation, apply post-processing.","metadata":{}},{"id":"8a15d428-6fde-4f85-9f06-48dcf29bbaa7","cell_type":"code","source":"model = model.to(DEVICE)\nmodel.eval()\nmodel.float()   # ensure fp32 throughout\n\n# Prepare test data with gap normalisation\ntest_df['transliteration_norm'] = test_df['transliteration'].apply(normalise_gaps)\n\n# --- Custom Dataset for batched inference ---\nclass InferenceDataset(TorchDataset):\n    def __init__(self, texts, tokenizer, prefix, max_length):\n        self.texts = [prefix + str(t) for t in texts]\n        self.tokenizer = tokenizer\n        self.max_length = max_length\n\n    def __len__(self):\n        return len(self.texts)\n\n    def __getitem__(self, idx):\n        enc = self.tokenizer(\n            self.texts[idx],\n            max_length=self.max_length,\n            padding='max_length',\n            truncation=True,\n            return_tensors='pt',\n        )\n        return {\n            'input_ids':      enc['input_ids'].squeeze(0),\n            'attention_mask': enc['attention_mask'].squeeze(0),\n        }\n\n\n# Build a list of which indices need neural inference vs. retrieval\nretrieval_scores  = []\nretrieval_results = []\nneural_needed     = []   # indices that will go to the neural model\n\nprint(f\"Running retrieval pre-filter{'  [DEBUG: {len(test_df)} rows]' if DEBUG else ''}...\")\nfor i, src in enumerate(tqdm(test_df['transliteration_norm'])):\n    score, retrieved_tgt = retrieve_translation(src)\n    retrieval_scores.append(score)\n    retrieval_results.append(retrieved_tgt)\n    if score < Config.RETRIEVAL_THRESHOLD:\n        neural_needed.append(i)\n\nretrieval_hits = len(test_df) - len(neural_needed)\nprint(f\"Retrieval hits  : {retrieval_hits}/{len(test_df)} ({100*retrieval_hits/len(test_df):.1f}%)\")\nprint(f\"Neural inference: {len(neural_needed)} examples\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T10:03:57.988151Z","iopub.execute_input":"2026-03-09T10:03:57.988648Z","iopub.status.idle":"2026-03-09T10:03:59.798423Z","shell.execute_reply.started":"2026-03-09T10:03:57.988621Z","shell.execute_reply":"2026-03-09T10:03:59.797440Z"}},"outputs":[],"execution_count":null},{"id":"20c5f2e8-810a-4435-a30c-cd209c9c7aef","cell_type":"code","source":"# Run ByT5 inference only on examples that retrieval couldn't handle\nneural_texts = test_df.iloc[neural_needed]['transliteration_norm'].tolist()\nneural_preds = {}\n\nif neural_texts:\n    neural_dataset = InferenceDataset(\n        neural_texts, tokenizer, Config.PREFIX, Config.MAX_LENGTH\n    )\n    neural_loader = DataLoader(neural_dataset, batch_size=16, shuffle=False)\n\n    raw_preds = []\n    with torch.inference_mode():\n        for batch in tqdm(neural_loader, desc='ByT5 inference'):\n            ids  = batch['input_ids'].to(DEVICE)\n            mask = batch['attention_mask'].to(DEVICE)\n            out  = model.generate(\n                input_ids=ids,\n                attention_mask=mask,\n                max_length=Config.MAX_LENGTH,\n                num_beams=4,\n                early_stopping=True,\n            )\n            decoded = tokenizer.batch_decode(out, skip_special_tokens=True)\n            raw_preds.extend([d.strip() for d in decoded])\n\n    for orig_idx, pred in zip(neural_needed, raw_preds):\n        neural_preds[orig_idx] = pred\n\nprint(\"Inference complete.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T10:03:59.799573Z","iopub.execute_input":"2026-03-09T10:03:59.799809Z","iopub.status.idle":"2026-03-09T10:04:25.021883Z","shell.execute_reply.started":"2026-03-09T10:03:59.799786Z","shell.execute_reply":"2026-03-09T10:04:25.021205Z"}},"outputs":[],"execution_count":null},{"id":"1a127b19-ca20-48e2-8806-1b3e626dead7","cell_type":"markdown","source":"## 10. Post-Processing\n\n- Restore gap tokens (`<gap>` → `x`, `<big_gap>` → `...`)\n- Enforce sentence-final period\n- Fallback `\"broken text\"` for empty predictions","metadata":{}},{"id":"eab6b3ca-5dae-450d-8d3e-0fd76acec79b","cell_type":"code","source":"def post_process(text: str) -> str:\n    \"\"\"Clean and finalise a translation output.\"\"\"\n    if not text or not text.strip():\n        return 'broken text'\n\n    # Restore gap tokens to readable notation\n    text = restore_gaps(text)\n\n    # Strip stray whitespace\n    text = text.strip()\n\n    # Ensure sentence ends with punctuation (common in formal translations)\n    if text and text[-1] not in '.!?':\n        text += '.'\n\n    return text\n\n\n# Assemble final predictions: retrieval where confident, neural elsewhere\nfinal_translations = []\nfor i in range(len(test_df)):\n    if retrieval_scores[i] >= Config.RETRIEVAL_THRESHOLD:\n        raw = retrieval_results[i]\n    else:\n        raw = neural_preds.get(i, '')\n    final_translations.append(post_process(raw))\n\nprint(\"Post-processing complete.\")\nprint(\"Sample outputs:\")\nfor i in range(min(5, len(final_translations))):\n    print(f\"  [{i}] {final_translations[i][:120]}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T10:04:25.023283Z","iopub.execute_input":"2026-03-09T10:04:25.023981Z","iopub.status.idle":"2026-03-09T10:04:25.030412Z","shell.execute_reply.started":"2026-03-09T10:04:25.023945Z","shell.execute_reply":"2026-03-09T10:04:25.029753Z"}},"outputs":[],"execution_count":null},{"id":"63002624-0fb3-47f7-acaf-2a74bc301f91","cell_type":"markdown","source":"## 11. Submission","metadata":{}},{"id":"fc52f1ad-7946-4a81-bbb0-aa255f79ce3a","cell_type":"code","source":"submission = pd.DataFrame({\n    'id':          test_df['id'],\n    'translation': final_translations,\n})\n\nsubmission.to_csv('submission.csv', index=False)\nprint(f\"Saved submission.csv with {len(submission)} rows.\")\nsubmission.head(10)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T10:04:25.031571Z","iopub.execute_input":"2026-03-09T10:04:25.032132Z","iopub.status.idle":"2026-03-09T10:04:25.065659Z","shell.execute_reply.started":"2026-03-09T10:04:25.032109Z","shell.execute_reply":"2026-03-09T10:04:25.065045Z"}},"outputs":[],"execution_count":null}]}