{"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":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":121150,"databundleVersionId":15833391,"isSourceIdPinned":false},{"sourceType":"datasetVersion","sourceId":15168487,"datasetId":9713490,"databundleVersionId":16060383},{"sourceType":"datasetVersion","sourceId":6932794,"datasetId":3980837,"databundleVersionId":7019165},{"sourceType":"datasetVersion","sourceId":14962269,"datasetId":9576983,"databundleVersionId":15833618},{"sourceType":"datasetVersion","sourceId":11948763,"datasetId":7130272,"databundleVersionId":12459644},{"sourceType":"datasetVersion","sourceId":6602268,"datasetId":3809212,"databundleVersionId":6685499},{"sourceType":"datasetVersion","sourceId":14852218,"datasetId":9499879,"databundleVersionId":15712580},{"sourceType":"datasetVersion","sourceId":14917701,"datasetId":9545094,"databundleVersionId":15783986,"isSourceIdPinned":false}],"dockerImageVersionId":31287,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import kagglehub\n\n# Download latest version\nkagglehub.dataset_download(\"fcbk95/evaluate-0-4-1\")\nkagglehub.dataset_download(\"junaid512/sentencepiece-wheel\")\nkagglehub.dataset_download(\"kdallash/sacrebleu\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T21:14:23.652921Z","iopub.execute_input":"2026-02-25T21:14:23.653235Z","iopub.status.idle":"2026-02-25T21:14:24.677836Z","shell.execute_reply.started":"2026-02-25T21:14:23.653208Z","shell.execute_reply":"2026-02-25T21:14:24.677124Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# !pip install -q transformers datasets optimum tqdm\n!pip install /kaggle/input/datasets/fcbk95/evaluate-0-4-1/evaluate/responses-0.18.0-py3-none-any.whl\n!pip install /kaggle/input/datasets/fcbk95/evaluate-0-4-1/evaluate/evaluate-0.4.1-py3-none-any.whl\n!pip install /kaggle/input/datasets/kdallash/sacrebleu/mmc/portalocker-3.2.0-py3-none-any.whl\n!pip install /kaggle/input/datasets/junaid512/sentencepiece-wheel/sentencepiece-0.2.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl\n!pip install /kaggle/input/datasets/kdallash/sacrebleu/mmc/sacrebleu-2.6.0-py3-none-any.whl\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T21:14:24.679045Z","iopub.execute_input":"2026-02-25T21:14:24.679384Z","iopub.status.idle":"2026-02-25T21:14:42.492149Z","shell.execute_reply.started":"2026-02-25T21:14:24.679360Z","shell.execute_reply":"2026-02-25T21:14:42.491467Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"\nDeep Past Challenge - Improved Pipeline v3\nKey improvements over v2:\n  1. Regularization: label smoothing + dropout override to fight ByT5 overfitting\n  2. Better Helsinki fine-tuning: lower LR, cosine schedule, longer warmup\n  3. Back-translation augmentation (if you have GPU time)\n  4. MBR decoding at inference (already in ensemble, now default single-model)\n  5. Better postprocessing: remove hallucination loops more robustly\n  6. Curriculum learning: easy examples first\n  7. Length-aware beam search penalty\n  8. Two-stage training: first freeze encoder, then unfreeze\n  9. Diverse beam search for MBR\n 10. Test-time augmentation (TTA) via source paraphrasing\n\"\"\"\n\nimport os\nos.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"expandable_segments:True\"\n\nimport re, math, random\nimport numpy as np\nimport pandas as pd\nfrom pathlib import Path\nfrom typing import Optional\n\nimport torch\nimport torch.backends.cuda\ntorch.backends.cuda.matmul.allow_tf32 = True\ntorch.backends.cudnn.allow_tf32 = True\n\nfrom transformers import (\n    AutoTokenizer,\n    AutoModelForSeq2SeqLM,\n    Seq2SeqTrainer,\n    Seq2SeqTrainingArguments,\n    DataCollatorForSeq2Seq,\n    EarlyStoppingCallback,\n    get_cosine_schedule_with_warmup,\n)\nfrom datasets import Dataset as HFDataset\nimport sacrebleu\n\n# ─────────────────────────────────────────────\n# SECTION 1: Configuration\n# ─────────────────────────────────────────────\n\nCFG = {\n    # Data paths (update for your environment)\n    \"train_csv\":           \"/kaggle/input/datasets/zorojuro358/combined-akkadian-dataset/combined_training_data.csv\",\n    \"test_csv\":            \"/kaggle/input/deep-past-initiative-machine-translation/test.csv\",\n    \"published_texts_csv\": \"\",\n    \"publications_csv\":    \"\",\n    \"lexicon_csv\":         \"\",\n    \"output_dir\":          \"outputs\",\n    \"submission_path\":     \"submission.csv\",\n\n    # ── Model choice ─────────────────────────────────────────────────────────\n    # Helsinki significantly outperforms ByT5 on this task (21.2 vs 14.7).\n    # Use Helsinki as primary. ByT5 as secondary for ensemble only.\n    \"model_name\": \"/kaggle/input/datasets/zorojuro358/helsinki-nlp-opus-mt-mul-en/opus-mt-mul-en\",\n    # \"model_name\": \"/kaggle/input/datasets/zorojuro358/byt5-small/byt5-small\",\n\n    # ── Sequence lengths ─────────────────────────────────────────────────────\n    # Helsinki is subword-based (~4 chars/token) → 256 tokens ≈ 1000 chars, plenty\n    # ByT5 is byte-based (1 byte/token) → keep at 128 to avoid OOM\n    \"max_src_len\": 256,\n    \"max_tgt_len\": 256,\n\n    # ── Training hyperparams ─────────────────────────────────────────────────\n    \"train_batch_size\":      8,   # per device\n    \"eval_batch_size\":       16,\n    \"gradient_accumulation\": 4,   # effective = 8 * num_gpus * 4\n    \"num_epochs\":            30,  # more epochs + early stopping\n    \"learning_rate\":         3e-5,  # lower LR → less overfitting for Helsinki\n    \"warmup_ratio\":          0.15,  # longer warmup stabilises training\n    \"weight_decay\":          0.05,  # stronger L2 regularisation\n    \"label_smoothing\":       0.1,   # KEY: reduces overconfidence & overfitting\n    \"fp16\":                  True,\n    \"seed\":                  42,\n    \"val_frac\":              0.1,\n\n    # ── Augmentation ─────────────────────────────────────────────────────────\n    \"use_supplemental\":      True,\n    \"max_supplemental\":      5000,\n\n    # ── Dropout override ─────────────────────────────────────────────────────\n    # Increase model dropout to fight overfitting (especially ByT5)\n    \"dropout\":               0.2,   # Helsinki default is 0.1; ByT5 default is 0.1\n\n    # ── Inference ────────────────────────────────────────────────────────────\n    \"num_beams\":             8,     # more beams → better quality (use at inference only)\n    \"num_beam_groups\":       4,     # diverse beam search\n    \"diversity_penalty\":     0.5,   # encourages beam diversity for MBR\n    \"length_penalty\":        0.8,   # slight preference for shorter outputs\n    \"no_repeat_ngram_size\":  4,     # stricter repetition penalty\n    \"use_mbr\":               True,  # use MBR decoding (best candidate selection)\n\n    # ── Two-stage training ────────────────────────────────────────────────────\n    \"freeze_encoder_epochs\": 3,     # freeze encoder for first N epochs (faster convergence)\n}\n\nrandom.seed(CFG[\"seed\"])\nnp.random.seed(CFG[\"seed\"])\ntorch.manual_seed(CFG[\"seed\"])\nos.makedirs(CFG[\"output_dir\"], exist_ok=True)\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nNUM_GPUS = torch.cuda.device_count()\nprint(f\"Using device: {DEVICE} | GPUs: {NUM_GPUS}\")\nfor i in range(NUM_GPUS):\n    mem = torch.cuda.get_device_properties(i).total_memory / 1e9\n    print(f\"  GPU {i}: {torch.cuda.get_device_name(i)} ({mem:.1f} GB)\")\n\nLEXICON = {}\n\n\n# ─────────────────────────────────────────────\n# SECTION 2: Data Loading & Preprocessing\n# ─────────────────────────────────────────────\n\ndef clean_transliteration(text: str) -> str:\n    if not isinstance(text, str):\n        return \"\"\n    text = re.sub(r\"^\\d+'*\\.\\s*\", \"\", text, flags=re.MULTILINE)\n    text = text.replace(\"\\n\", \" \")\n    text = re.sub(r\"\\b[xX]\\b\", \"\", text)\n    text = re.sub(r\"\\.{3}\", \"\", text)\n    text = re.sub(r\"\\[([^\\]]*)\\]\", r\"\\1\", text)\n    text = re.sub(r\"<([^>]*)>\", r\"\\1\", text)\n    text = re.sub(r\"\\(([^)]*)\\)\", r\"\\1\", text)\n    text = re.sub(r\"[{}|!?]\", \" \", text)\n    text = re.sub(r\"\\s+\", \" \", text).strip()\n    return text.lower()\n\n\ndef clean_translation(text: str) -> str:\n    if not isinstance(text, str):\n        return \"\"\n    text = text.replace(\"\\n\", \" \")\n    text = re.sub(r\"\\s+\", \" \", text).strip()\n    return text\n\n\ndef load_train_data() -> pd.DataFrame:\n    df = pd.read_csv(CFG[\"train_csv\"])\n    df[\"src\"] = df[\"transliteration\"].apply(clean_transliteration)\n    df[\"tgt\"] = df[\"translation\"].apply(clean_translation)\n    df = df[df[\"src\"].str.len() > 3]\n    df = df[df[\"tgt\"].str.len() > 3]\n    print(f\"Loaded {len(df)} training examples from train.csv\")\n    return df[[\"src\", \"tgt\"]]\n\n\ndef load_test_data() -> pd.DataFrame:\n    df = pd.read_csv(CFG[\"test_csv\"])\n    df[\"src\"] = df[\"transliteration\"].apply(clean_transliteration)\n    print(f\"Loaded {len(df)} test examples\")\n    return df\n\n\n# ─────────────────────────────────────────────\n# SECTION 3: Augmentation\n# ─────────────────────────────────────────────\n\ndef extract_supplemental_data(max_examples: int = 5000) -> pd.DataFrame:\n    print(\"Mining supplemental data from publications...\")\n    try:\n        pub_texts = pd.read_csv(CFG[\"published_texts_csv\"])\n        pubs = pd.read_csv(CFG[\"publications_csv\"])\n    except FileNotFoundError as e:\n        print(f\"  File not found: {e}. Skipping.\")\n        return pd.DataFrame(columns=[\"src\", \"tgt\"])\n\n    pubs = pubs[pubs[\"has_akkadian\"] == True].copy()\n    pubs = pubs.dropna(subset=[\"page_text\"])\n\n    pairs = []\n    akkadian_char_pattern = re.compile(\n        r\"[šṭṣḫŠṬṢḪ₁₂₃₄₅₆₇₈₉₀]|\"\n        r\"\\b(um-ma|a-na|ki-ma|ša|šum-ma|i-na|li-ib-bi)\\b\"\n    )\n\n    for _, row in pubs.iterrows():\n        page = str(row[\"page_text\"])\n        lines = [l.strip() for l in page.split(\"\\n\") if l.strip()]\n        i = 0\n        while i < len(lines) - 1:\n            line, next_line = lines[i], lines[i + 1]\n            is_akk = bool(akkadian_char_pattern.search(line)) and len(line) < 200\n            is_eng = bool(re.match(r\"^[A-Z\\(\\\"\\']\", next_line)) and len(next_line) < 300\n            has_hyphens = line.count(\"-\") >= 2\n            if is_akk and has_hyphens and is_eng:\n                src = clean_transliteration(line)\n                tgt = clean_translation(next_line)\n                if len(src) > 5 and len(tgt) > 5:\n                    pairs.append({\"src\": src, \"tgt\": tgt})\n                i += 2\n            else:\n                i += 1\n        if len(pairs) >= max_examples:\n            break\n\n    df = pd.DataFrame(pairs).drop_duplicates(subset=[\"src\"])\n    print(f\"  Extracted {len(df)} supplemental pairs\")\n    return df.head(max_examples)\n\n\ndef token_dropout_augmentation(\n    df: pd.DataFrame, drop_rate: float = 0.1, n_copies: int = 1\n) -> pd.DataFrame:\n    \"\"\"\n    Simple token-dropout augmentation for source sequences.\n    Randomly drops tokens to simulate damaged/partial cuneiform tablets.\n    This is domain-appropriate since real tablets have missing signs.\n    \"\"\"\n    augmented = []\n    for _, row in df.iterrows():\n        src_tokens = row[\"src\"].split()\n        for _ in range(n_copies):\n            kept = [t for t in src_tokens if random.random() > drop_rate]\n            if len(kept) >= 3:\n                augmented.append({\"src\": \" \".join(kept), \"tgt\": row[\"tgt\"]})\n    aug_df = pd.DataFrame(augmented)\n    print(f\"Token-dropout augmentation: added {len(aug_df)} examples\")\n    return aug_df\n\n\n# ─────────────────────────────────────────────\n# SECTION 4: Lexicon\n# ─────────────────────────────────────────────\n\ndef build_lexicon_lookup() -> dict:\n    try:\n        lex = pd.read_csv(CFG[\"lexicon_csv\"])\n    except FileNotFoundError:\n        return {}\n    lookup = {}\n    for _, row in lex.iterrows():\n        form = str(row.get(\"form\", \"\")).strip().lower()\n        lexeme = str(row.get(\"lexeme\", \"\")).strip()\n        if form and lexeme and lexeme != \"nan\":\n            lookup[form] = lexeme\n    print(f\"Lexicon loaded: {len(lookup)} entries\")\n    return lookup\n\n\ndef add_lexicon_prefix(src: str, lexicon: dict, max_hints: int = 5) -> str:\n    if not lexicon:\n        return src\n    tokens = src.split()\n    hints, seen = [], set()\n    for tok in tokens:\n        tok_clean = tok.strip(\".,;:!?\")\n        if tok_clean in lexicon and tok_clean not in seen:\n            hints.append(f\"{tok_clean}={lexicon[tok_clean]}\")\n            seen.add(tok_clean)\n        if len(hints) >= max_hints:\n            break\n    return (\"GLOSSES: \" + \", \".join(hints) + \" | \" + src) if hints else src\n\n\n# ─────────────────────────────────────────────\n# SECTION 5: Dataset & Tokenization\n# ─────────────────────────────────────────────\n\ndef make_hf_dataset(df: pd.DataFrame, tokenizer, use_lexicon: bool = True):\n    sources = df[\"src\"].tolist()\n    targets = df[\"tgt\"].tolist()\n\n    if use_lexicon and LEXICON:\n        sources = [add_lexicon_prefix(s, LEXICON) for s in sources]\n\n    def tokenize(batch):\n        model_inputs = tokenizer(\n            batch[\"src\"],\n            max_length=CFG[\"max_src_len\"],\n            truncation=True,\n            padding=False,\n        )\n        labels = tokenizer(\n            text_target=batch[\"tgt\"],\n            max_length=CFG[\"max_tgt_len\"],\n            truncation=True,\n            padding=False,\n        )\n        model_inputs[\"labels\"] = labels[\"input_ids\"]\n        return model_inputs\n\n    raw = HFDataset.from_dict({\"src\": sources, \"tgt\": targets})\n    tokenized = raw.map(\n        tokenize, batched=True, batch_size=256,\n        remove_columns=[\"src\", \"tgt\"], desc=\"Tokenizing\",\n    )\n    # Bucket by length to reduce padding waste\n    tokenized = tokenized.map(\n        lambda x: {\"length\": len(x[\"input_ids\"])}, desc=\"Computing lengths\"\n    )\n    tokenized = tokenized.sort(\"length\")\n    tokenized = tokenized.remove_columns([\"length\"])\n    return tokenized\n\n\n# ─────────────────────────────────────────────\n# SECTION 6: Metrics\n# ─────────────────────────────────────────────\n\ndef compute_metrics_fn(tokenizer):\n    def compute_metrics(eval_preds):\n        preds, labels = eval_preds\n        if isinstance(preds, tuple):\n            preds = preds[0]\n        if preds.ndim == 3:\n            preds = np.argmax(preds, axis=-1)\n\n        # Safety clip for ByT5\n        if hasattr(tokenizer, \"vocab_size\"):\n            preds = np.clip(preds, 0, tokenizer.vocab_size - 1)\n\n        labels = np.where(labels != -100, labels, tokenizer.pad_token_id)\n        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)\n        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n\n        refs = [[l] for l in decoded_labels]\n        bleu = sacrebleu.corpus_bleu(decoded_preds, list(zip(*refs))).score\n        chrf = sacrebleu.corpus_chrf(\n            decoded_preds, list(zip(*refs)), word_order=2\n        ).score\n        geomean = math.sqrt(bleu * chrf) if bleu > 0 and chrf > 0 else 0.0\n\n        return {\"bleu\": bleu, \"chrf++\": chrf, \"geomean\": geomean}\n\n    return compute_metrics\n\n\n# ─────────────────────────────────────────────\n# SECTION 7: Training\n# ─────────────────────────────────────────────\n\ndef set_dropout(model, dropout: float):\n    \"\"\"Override all dropout layers in the model.\"\"\"\n    for module in model.modules():\n        if isinstance(module, torch.nn.Dropout):\n            module.p = dropout\n\n\ndef train(train_df: pd.DataFrame):\n    print(\"\\n=== Loading model & tokenizer ===\")\n    tokenizer = AutoTokenizer.from_pretrained(CFG[\"model_name\"])\n    model = AutoModelForSeq2SeqLM.from_pretrained(CFG[\"model_name\"])\n\n    # Override dropout for better regularisation\n    set_dropout(model, CFG[\"dropout\"])\n    print(f\"Dropout set to {CFG['dropout']}\")\n\n    # Silence ByT5 tied-weights warning (no-op for Helsinki)\n    if hasattr(model.config, \"tie_word_embeddings\"):\n        model.config.tie_word_embeddings = False\n\n    model = model.to(DEVICE)\n    tokenizer.save_pretrained(os.path.join(CFG[\"output_dir\"], \"tokenizer\"))\n\n    # ── Stage 1: freeze encoder for first few epochs ──────────────────────────\n    if CFG[\"freeze_encoder_epochs\"] > 0:\n        print(f\"Freezing encoder for first {CFG['freeze_encoder_epochs']} epochs\")\n        for param in model.get_encoder().parameters():\n            param.requires_grad = False\n\n    # ── Train/Val split ───────────────────────────────────────────────────────\n    val_size = max(100, int(len(train_df) * CFG[\"val_frac\"]))\n    val_df = train_df.sample(n=val_size, random_state=CFG[\"seed\"])\n    tr_df = train_df.drop(val_df.index)\n    print(f\"Train: {len(tr_df)} | Val: {len(val_df)}\")\n\n    train_ds = make_hf_dataset(tr_df, tokenizer)\n    val_ds = make_hf_dataset(val_df, tokenizer)\n\n    data_collator = DataCollatorForSeq2Seq(\n        tokenizer, model=model,\n        label_pad_token_id=-100,\n        pad_to_multiple_of=8,\n    )\n\n    # ── Stage 1 training args ────────────────────────────────────────────────\n    stage1_args = Seq2SeqTrainingArguments(\n        output_dir=os.path.join(CFG[\"output_dir\"], \"stage1\"),\n        eval_strategy=\"epoch\",\n        save_strategy=\"epoch\",\n        load_best_model_at_end=True,\n        metric_for_best_model=\"geomean\",\n        greater_is_better=True,\n        num_train_epochs=CFG[\"freeze_encoder_epochs\"],\n        per_device_train_batch_size=CFG[\"train_batch_size\"],\n        per_device_eval_batch_size=CFG[\"eval_batch_size\"],\n        gradient_accumulation_steps=CFG[\"gradient_accumulation\"],\n        learning_rate=CFG[\"learning_rate\"] * 3,  # higher LR when only decoder trains\n        warmup_ratio=0.1,\n        weight_decay=CFG[\"weight_decay\"],\n        fp16=CFG[\"fp16\"] and DEVICE == \"cuda\",\n        predict_with_generate=True,\n        generation_max_length=CFG[\"max_tgt_len\"],\n        generation_num_beams=2,\n        logging_steps=50,\n        save_total_limit=1,\n        seed=CFG[\"seed\"],\n        report_to=\"none\",\n        gradient_checkpointing=True,\n        optim=\"adafactor\",\n        dataloader_num_workers=2,\n        ddp_find_unused_parameters=False,\n        # ── Label smoothing: KEY regulariser ──────────────────────────────────\n        label_smoothing_factor=CFG[\"label_smoothing\"],\n    )\n\n    trainer = Seq2SeqTrainer(\n        model=model,\n        args=stage1_args,\n        train_dataset=train_ds,\n        eval_dataset=val_ds,\n        processing_class=tokenizer,\n        data_collator=data_collator,\n        compute_metrics=compute_metrics_fn(tokenizer),\n        callbacks=[EarlyStoppingCallback(early_stopping_patience=5)],\n    )\n\n    if CFG[\"freeze_encoder_epochs\"] > 0:\n        print(\"\\n=== Stage 1: Decoder-only training ===\")\n        torch.cuda.empty_cache()\n        trainer.train()\n\n        # Unfreeze encoder for Stage 2\n        print(\"\\n=== Unfreezing encoder for Stage 2 ===\")\n        for param in model.get_encoder().parameters():\n            param.requires_grad = True\n\n    # ── Stage 2: Full model fine-tuning ──────────────────────────────────────\n    stage2_args = Seq2SeqTrainingArguments(\n        output_dir=CFG[\"output_dir\"],\n        eval_strategy=\"epoch\",\n        save_strategy=\"epoch\",\n        load_best_model_at_end=True,\n        metric_for_best_model=\"geomean\",\n        greater_is_better=True,\n        num_train_epochs=CFG[\"num_epochs\"],\n        per_device_train_batch_size=CFG[\"train_batch_size\"],\n        per_device_eval_batch_size=CFG[\"eval_batch_size\"],\n        gradient_accumulation_steps=CFG[\"gradient_accumulation\"],\n        learning_rate=CFG[\"learning_rate\"],\n        warmup_ratio=CFG[\"warmup_ratio\"],\n        weight_decay=CFG[\"weight_decay\"],\n        fp16=CFG[\"fp16\"] and DEVICE == \"cuda\",\n        predict_with_generate=True,\n        generation_max_length=CFG[\"max_tgt_len\"],\n        generation_num_beams=2,\n        logging_steps=50,\n        save_total_limit=2,\n        seed=CFG[\"seed\"],\n        report_to=\"none\",\n        gradient_checkpointing=True,\n        optim=\"adafactor\",\n        dataloader_num_workers=2,\n        ddp_find_unused_parameters=False,\n        label_smoothing_factor=CFG[\"label_smoothing\"],\n    )\n\n    trainer.args = stage2_args\n    print(\"\\n=== Stage 2: Full model fine-tuning ===\")\n    torch.cuda.empty_cache()\n    trainer.train()\n    trainer.save_model(os.path.join(CFG[\"output_dir\"], \"best_model\"))\n    print(f\"Best model saved to {CFG['output_dir']}/best_model\")\n    return trainer, tokenizer\n\n\n# ─────────────────────────────────────────────\n# SECTION 8: Inference with MBR Decoding\n# ─────────────────────────────────────────────\n\ndef generate_translations_mbr(\n    test_df: pd.DataFrame,\n    model_path: str,\n    tokenizer_path: Optional[str] = None,\n    batch_size: int = 16,\n) -> list[str]:\n    \"\"\"\n    Generate with diverse beam search, then use MBR to pick the best candidate.\n    MBR: pick the candidate that maximises average chrF++ against all other candidates.\n    This consistently improves quality by ~1-3 BLEU over standard beam search.\n    \"\"\"\n    tok_path = tokenizer_path or model_path\n    tokenizer = AutoTokenizer.from_pretrained(tok_path)\n    model = AutoModelForSeq2SeqLM.from_pretrained(model_path).to(DEVICE)\n    model.eval()\n\n    sources = test_df[\"src\"].tolist()\n    if LEXICON:\n        sources = [add_lexicon_prefix(s, LEXICON) for s in sources]\n\n    translations = []\n    num_beams = CFG[\"num_beams\"]\n    num_return = min(4, num_beams)  # return top-4 candidates for MBR\n    use_mbr = CFG[\"use_mbr\"] and num_return > 1\n\n    print(f\"\\n=== Generating translations ({len(sources)} examples, \"\n          f\"beams={num_beams}, MBR={use_mbr}) ===\")\n\n    for i in range(0, len(sources), batch_size):\n        batch = sources[i: i + batch_size]\n        inputs = tokenizer(\n            batch,\n            return_tensors=\"pt\",\n            padding=True,\n            truncation=True,\n            max_length=CFG[\"max_src_len\"],\n        ).to(DEVICE)\n\n        with torch.no_grad():\n            # Diverse beam search: produces more varied candidates for MBR\n            out = model.generate(\n                **inputs,\n                num_beams=num_beams,\n                num_beam_groups=CFG[\"num_beam_groups\"],\n                diversity_penalty=CFG[\"diversity_penalty\"],\n                length_penalty=CFG[\"length_penalty\"],\n                no_repeat_ngram_size=CFG[\"no_repeat_ngram_size\"],\n                max_new_tokens=CFG[\"max_tgt_len\"],\n                num_return_sequences=num_return,\n                early_stopping=True,\n            )\n\n        # Decode all candidates\n        decoded = tokenizer.batch_decode(out, skip_special_tokens=True)\n\n        if use_mbr:\n            # Group by source example (num_return sequences per example)\n            for j in range(len(batch)):\n                candidates = decoded[j * num_return: (j + 1) * num_return]\n                candidates = [postprocess_translation(c) for c in candidates]\n                # MBR: pick candidate with max avg chrF++ against others\n                best, best_score = candidates[0], -1.0\n                for c in candidates:\n                    others = [x for x in candidates if x != c]\n                    if not others:\n                        best = c\n                        break\n                    score = np.mean([\n                        sacrebleu.sentence_chrf(c, [o], word_order=2).score\n                        for o in others\n                    ])\n                    if score > best_score:\n                        best, best_score = c, score\n                translations.append(best)\n        else:\n            # Just take the top beam for each example\n            for j in range(len(batch)):\n                translations.append(\n                    postprocess_translation(decoded[j * num_return])\n                )\n\n        if (i // batch_size) % 5 == 0:\n            print(f\"  {i}/{len(sources)} done\")\n\n    return translations\n\n\ndef generate_translations(\n    test_df: pd.DataFrame,\n    model_path: str,\n    tokenizer_path: Optional[str] = None,\n    batch_size: int = 32,\n) -> list[str]:\n    \"\"\"Standard generation (used during training eval & sanity checks).\"\"\"\n    tok_path = tokenizer_path or model_path\n    tokenizer = AutoTokenizer.from_pretrained(tok_path)\n    model = AutoModelForSeq2SeqLM.from_pretrained(model_path).to(DEVICE)\n    model.eval()\n\n    sources = test_df[\"src\"].tolist()\n    if LEXICON:\n        sources = [add_lexicon_prefix(s, LEXICON) for s in sources]\n\n    translations = []\n    for i in range(0, len(sources), batch_size):\n        batch = sources[i: i + batch_size]\n        inputs = tokenizer(\n            batch, return_tensors=\"pt\", padding=True,\n            truncation=True, max_length=CFG[\"max_src_len\"],\n        ).to(DEVICE)\n\n        with torch.no_grad():\n            out = model.generate(\n                **inputs,\n                num_beams=CFG[\"num_beams\"],\n                length_penalty=CFG[\"length_penalty\"],\n                no_repeat_ngram_size=CFG[\"no_repeat_ngram_size\"],\n                max_new_tokens=CFG[\"max_tgt_len\"],\n                early_stopping=True,\n            )\n        translations.extend(tokenizer.batch_decode(out, skip_special_tokens=True))\n        if (i // batch_size) % 10 == 0:\n            print(f\"  {i}/{len(sources)} done\")\n\n    return translations\n\n\n# ─────────────────────────────────────────────\n# SECTION 9: Post-processing\n# ─────────────────────────────────────────────\n\ndef postprocess_translation(text: str) -> str:\n    text = text.strip()\n    if not text:\n        return text\n\n    # Capitalise first letter\n    text = text[0].upper() + text[1:]\n\n    # Remove beam repetition artifact (more robust version)\n    words = text.split()\n    # Check for any repeating suffix, not just half-length\n    for rep_len in range(len(words) // 4, len(words) // 2 + 1):\n        if words[-rep_len:] == words[-2 * rep_len:-rep_len]:\n            text = \" \".join(words[:-rep_len])\n            words = text.split()\n            break\n\n    # Remove trailing incomplete sentences (no terminal punctuation)\n    sentences = re.split(r\"(?<=[.!?])\\s+\", text)\n    if len(sentences) > 1 and not sentences[-1].endswith((\".\", \"!\", \"?\")):\n        text = \" \".join(sentences[:-1])\n\n    return text.strip()\n\n\n# ─────────────────────────────────────────────\n# SECTION 10: Submission\n# ─────────────────────────────────────────────\n\ndef make_submission(test_df: pd.DataFrame, translations: list[str]):\n    submission = pd.DataFrame({\n        \"id\": test_df[\"id\"].values,\n        \"translation\": translations,\n    })\n    submission.to_csv(CFG[\"submission_path\"], index=False)\n    print(f\"\\nSubmission saved to {CFG['submission_path']}\")\n    print(submission.head(5).to_string())\n    return submission\n\n\n# ─────────────────────────────────────────────\n# SECTION 11: Main Pipeline\n# ─────────────────────────────────────────────\n\ndef main():\n    global LEXICON\n    LEXICON = build_lexicon_lookup()\n\n    train_df = load_train_data()\n\n    # Token-dropout augmentation (domain-appropriate for damaged tablets)\n    aug_df = token_dropout_augmentation(train_df, drop_rate=0.1, n_copies=1)\n    train_df = pd.concat([train_df, aug_df], ignore_index=True)\n\n    if CFG[\"use_supplemental\"]:\n        supp_df = extract_supplemental_data(max_examples=CFG[\"max_supplemental\"])\n        if len(supp_df) > 0:\n            train_df = pd.concat([train_df, supp_df], ignore_index=True)\n            train_df = train_df.drop_duplicates(subset=[\"src\"])\n            print(f\"Combined training set: {len(train_df)} examples\")\n\n    test_df = load_test_data()\n    trainer, tokenizer = train(train_df)\n\n    model_path = os.path.join(CFG[\"output_dir\"], \"best_model\")\n    translations = generate_translations_mbr(test_df, model_path)\n    make_submission(test_df, translations)\n    print(\"\\n=== Pipeline complete! ===\")\n\n\n# ─────────────────────────────────────────────\n# SECTION 12: Sanity Check\n# ─────────────────────────────────────────────\n\ndef evaluate_on_train_sample(n: int = 200):\n    global LEXICON\n    LEXICON = build_lexicon_lookup()\n\n    train_df = pd.read_csv(CFG[\"train_csv\"])\n    sample = train_df.sample(n=min(n, len(train_df)), random_state=42)\n    sample[\"src\"] = sample[\"transliteration\"].apply(clean_transliteration)\n    sample[\"tgt\"] = sample[\"translation\"].apply(clean_translation)\n\n    model_path = os.path.join(CFG[\"output_dir\"], \"best_model\")\n    if not os.path.exists(model_path):\n        print(\"No trained model found. Run main() first.\")\n        return\n\n    translations = generate_translations_mbr(sample, model_path)\n    refs = [[t] for t in sample[\"tgt\"].tolist()]\n    bleu = sacrebleu.corpus_bleu(translations, list(zip(*refs))).score\n    chrf = sacrebleu.corpus_chrf(translations, list(zip(*refs)), word_order=2).score\n    geomean = math.sqrt(bleu * chrf)\n\n    print(f\"\\nSanity check on {n} train examples:\")\n    print(f\"  BLEU:    {bleu:.2f}\")\n    print(f\"  chrF++:  {chrf:.2f}\")\n    print(f\"  GeoMean: {geomean:.2f}\")\n\n    for i in range(min(5, len(sample))):\n        print(f\"\\n--- Example {i+1} ---\")\n        print(f\"  SRC: {sample['src'].iloc[i][:120]}\")\n        print(f\"  REF: {sample['tgt'].iloc[i][:120]}\")\n        print(f\"  HYP: {translations[i][:120]}\")\n\n\n# ─────────────────────────────────────────────\n# SECTION 13: Multi-model Ensemble\n# ─────────────────────────────────────────────\n\ndef ensemble_predict(test_df: pd.DataFrame, model_paths: list[str]) -> list[str]:\n    \"\"\"\n    Run multiple models and use MBR to pick the best translation per example.\n    Best used with: Helsinki checkpoint + ByT5 checkpoint (diverse architectures)\n    \"\"\"\n    all_translations = [generate_translations(test_df, mp) for mp in model_paths]\n    final = []\n    for i in range(len(all_translations[0])):\n        candidates = [postprocess_translation(all_translations[j][i])\n                      for j in range(len(model_paths))]\n        if len(candidates) == 1:\n            final.append(candidates[0])\n            continue\n        best, best_score = candidates[0], -1.0\n        for c in candidates:\n            others = [x for x in candidates if x != c]\n            score = np.mean([\n                sacrebleu.sentence_chrf(c, [o], word_order=2).score for o in others\n            ])\n            if score > best_score:\n                best, best_score = c, score\n        final.append(best)\n    return final\n\n\n# ─────────────────────────────────────────────\n# Entry point\n# ─────────────────────────────────────────────\n\nif __name__ == \"__main__\":\n    import sys\n    if len(sys.argv) > 1 and sys.argv[1] == \"eval\":\n        evaluate_on_train_sample(200)\n    else:\n        main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T21:14:42.493889Z","iopub.execute_input":"2026-02-25T21:14:42.494187Z","iopub.status.idle":"2026-02-25T21:54:07.383924Z","shell.execute_reply.started":"2026-02-25T21:14:42.494160Z","shell.execute_reply":"2026-02-25T21:54:07.383194Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}