{"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},{"sourceType":"datasetVersion","sourceId":14962950,"datasetId":9577356,"databundleVersionId":15834364},{"sourceType":"datasetVersion","sourceId":6932794,"datasetId":3980837,"databundleVersionId":7019165},{"sourceType":"datasetVersion","sourceId":14899351,"datasetId":9366631,"databundleVersionId":15763950},{"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}],"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":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-03-07T09:21:12.207125Z","iopub.execute_input":"2026-03-07T09:21:12.207356Z","iopub.status.idle":"2026-03-07T09:21:13.499358Z","shell.execute_reply.started":"2026-03-07T09:21:12.207335Z","shell.execute_reply":"2026-03-07T09:21:13.498570Z"}},"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-03-07T09:21:24.089375Z","iopub.execute_input":"2026-03-07T09:21:24.090105Z","iopub.status.idle":"2026-03-07T09:21:43.334585Z","shell.execute_reply.started":"2026-03-07T09:21:24.090076Z","shell.execute_reply":"2026-03-07T09:21:43.333629Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"\nDeep Past Challenge - Train ALL Models + Ensemble\n==================================================\nRun this single notebook on Kaggle. It will:\n  1. Train Helsinki  (opus-mt-mul-en)   → saves to outputs/helsinki/\n  3. Train ByT5-small                   → saves to outputs/byt5/\n  4. Ensemble all three with MBR        → submission.csv\n\nOn T4 x2 (32GB total):\n  - Helsinki  : ~2-3 hrs\n  - NLLB      : ~3-4 hrs\n  - ByT5      : ~2-3 hrs\n  - Inference : ~20 min\n  Total       : fits in a 12hr Kaggle session\n\nIf you only have one GPU session, train one model per session,\nsave the checkpoint, then run Section 4 (ensemble) in a final session.\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\nimport torch\nimport torch.backends.cuda\nimport sacrebleu\n\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)\nfrom datasets import Dataset as HFDataset\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nNUM_GPUS = torch.cuda.device_count()\nprint(f\"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\n# ═════════════════════════════════════════════════════════════════════════════\n# SECTION 0 — PATHS  (edit these to match your Kaggle dataset locations)\n# ═════════════════════════════════════════════════════════════════════════════\n\nDATA = {\n    \"train_csv\": \"/kaggle/input/competitions/deep-past-initiative-machine-translation/train.csv\",\n    \"test_csv\":  \"/kaggle/input/competitions/deep-past-initiative-machine-translation/test.csv\",\n    \"lexicon_csv\": \"\",   # leave empty if you don't have it\n}\n\nMODEL_PATHS = {\n    # Paste the exact paths from your Kaggle dataset cards\n    \"helsinki\": \"/kaggle/input/datasets/zorojuro358/helsinki-nlp-opus-mt-mul-en/opus-mt-mul-en\",\n    # \"nllb\":     \"facebook/nllb-200-distilled-600M\",   # downloads from HF hub\n    \"byt5\":     \"/kaggle/input/datasets/assiaben/final-byt5/byt5-akkadian-optimized-34x\",\n}\n\nOUTPUT_BASE = \"/kaggle/working\"\nos.makedirs(OUTPUT_BASE, exist_ok=True)\n\n# ═════════════════════════════════════════════════════════════════════════════\n# SECTION 1 — SHARED HYPERPARAMETERS\n# ═════════════════════════════════════════════════════════════════════════════\n\nSHARED = {\n    \"val_frac\":              0.1,\n    \"seed\":                  42,\n    \"label_smoothing\":       0.1,    # key regulariser — prevents overfitting\n    \"weight_decay\":          0.05,\n    \"warmup_ratio\":          0.15,\n    \"fp16\":                  True,\n    \"gradient_accumulation\": 4,\n    \"num_epochs\":            25,\n    \"early_stopping_patience\": 4,\n    \"train_batch_size\":      8,\n    \"eval_batch_size\":       16,\n    \"save_total_limit\":      3,      # keep 3 checkpoints for checkpoint ensembling\n}\n\n# Per-model overrides\nMODEL_CFG = {\n    \"helsinki\": {\n        \"max_src_len\":  256,\n        \"max_tgt_len\":  256,\n        \"learning_rate\": 3e-5,\n        \"dropout\":       0.15,\n        \"num_beams\":     5,\n    },\n    \"nllb\": {\n        \"max_src_len\":  256,\n        \"max_tgt_len\":  256,\n        \"learning_rate\": 2e-5,     # NLLB is large, needs lower LR\n        \"dropout\":       0.1,\n        \"num_beams\":     5,\n        \"src_lang\":      \"tur_Latn\",   # closest NLLB code to Akkadian\n        \"tgt_lang\":      \"eng_Latn\",\n    },\n    \"byt5\": {\n        \"max_src_len\":  128,\n        \"max_tgt_len\":  128,\n        \"learning_rate\": 5e-4,     # ByT5 uses adafactor → higher LR is fine\n        \"dropout\":       0.2,      # higher dropout to fight ByT5 overfitting\n        \"num_beams\":     5,\n    },\n}\n\nrandom.seed(SHARED[\"seed\"])\nnp.random.seed(SHARED[\"seed\"])\ntorch.manual_seed(SHARED[\"seed\"])\n\n# ═════════════════════════════════════════════════════════════════════════════\n# SECTION 2 — DATA UTILITIES\n# ═════════════════════════════════════════════════════════════════════════════\n\ndef clean_src(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_tgt(text: str) -> str:\n    if not isinstance(text, str):\n        return \"\"\n    return re.sub(r\"\\s+\", \" \", text.replace(\"\\n\", \" \")).strip()\n\n\ndef load_data():\n    train = pd.read_csv(DATA[\"train_csv\"])\n    train[\"src\"] = train[\"transliteration\"].apply(clean_src)\n    train[\"tgt\"] = train[\"translation\"].apply(clean_tgt)\n    train = train[(train[\"src\"].str.len() > 3) & (train[\"tgt\"].str.len() > 3)]\n\n    test = pd.read_csv(DATA[\"test_csv\"])\n    test[\"src\"] = test[\"transliteration\"].apply(clean_src)\n\n    print(f\"Train: {len(train)}  |  Test: {len(test)}\")\n    return train[[\"src\", \"tgt\"]], test\n\n\ndef token_dropout_augment(df: pd.DataFrame, rate=0.1) -> pd.DataFrame:\n    \"\"\"Drop random tokens to simulate damaged tablet text.\"\"\"\n    rows = []\n    for _, r in df.iterrows():\n        tokens = r[\"src\"].split()\n        kept = [t for t in tokens if random.random() > rate]\n        if len(kept) >= 3:\n            rows.append({\"src\": \" \".join(kept), \"tgt\": r[\"tgt\"]})\n    aug = pd.DataFrame(rows)\n    print(f\"Augmentation: +{len(aug)} token-dropout examples\")\n    return aug\n\n\ndef make_hf_dataset(df, tokenizer, max_src, max_tgt):\n    def tok(batch):\n        enc = tokenizer(\n            batch[\"src\"], max_length=max_src, truncation=True, padding=False\n        )\n        lab = tokenizer(\n            text_target=batch[\"tgt\"], max_length=max_tgt,\n            truncation=True, padding=False\n        )\n        enc[\"labels\"] = lab[\"input_ids\"]\n        return enc\n\n    ds = HFDataset.from_dict({\"src\": df[\"src\"].tolist(), \"tgt\": df[\"tgt\"].tolist()})\n    ds = ds.map(tok, batched=True, batch_size=256, remove_columns=[\"src\", \"tgt\"])\n    ds = ds.map(lambda x: {\"length\": len(x[\"input_ids\"])})\n    ds = ds.sort(\"length\").remove_columns([\"length\"])\n    return ds\n\n\n# ═════════════════════════════════════════════════════════════════════════════\n# SECTION 3 — METRIC\n# ═════════════════════════════════════════════════════════════════════════════\n\ndef make_compute_metrics(tokenizer, vocab_size=None):\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        if vocab_size:\n            preds = np.clip(preds, 0, vocab_size - 1)\n        labels = np.where(labels != -100, labels, tokenizer.pad_token_id)\n        hyps = tokenizer.batch_decode(preds, skip_special_tokens=True)\n        refs = tokenizer.batch_decode(labels, skip_special_tokens=True)\n        ref_wrapped = [[r] for r in refs]\n        bleu = sacrebleu.corpus_bleu(hyps, list(zip(*ref_wrapped))).score\n        chrf = sacrebleu.corpus_chrf(hyps, list(zip(*ref_wrapped)), word_order=2).score\n        geo  = math.sqrt(bleu * chrf) if bleu > 0 and chrf > 0 else 0.0\n        return {\"bleu\": bleu, \"chrf++\": chrf, \"geomean\": geo}\n    return compute_metrics\n\n\n# ═════════════════════════════════════════════════════════════════════════════\n# SECTION 4 — TRAINING FUNCTION  (works for all three models)\n# ═════════════════════════════════════════════════════════════════════════════\n\ndef set_dropout(model, p):\n    for m in model.modules():\n        if isinstance(m, torch.nn.Dropout):\n            m.p = p\n\n\ndef train_model(model_key: str, train_df: pd.DataFrame):\n    \"\"\"\n    Train one model and save the best checkpoint.\n    Returns path to best checkpoint.\n    \"\"\"\n    mcfg = {**SHARED, **MODEL_CFG[model_key]}\n    model_path = MODEL_PATHS[model_key]\n    out_dir = os.path.join(OUTPUT_BASE, model_key)\n    os.makedirs(out_dir, exist_ok=True)\n\n    print(f\"\\n{'='*60}\")\n    print(f\"  TRAINING: {model_key.upper()}\")\n    print(f\"{'='*60}\")\n    print(f\"  Model path : {model_path}\")\n    print(f\"  Output dir : {out_dir}\")\n    print(f\"  LR={mcfg['learning_rate']}  dropout={mcfg['dropout']}  \"\n          f\"epochs={mcfg['num_epochs']}  label_smooth={mcfg['label_smoothing']}\")\n\n    # ── Load tokenizer ────────────────────────────────────────────────────────\n    if model_key == \"nllb\":\n        from transformers import NllbTokenizer\n        tokenizer = NllbTokenizer.from_pretrained(\n            model_path, src_lang=mcfg[\"src_lang\"]\n        )\n    else:\n        tokenizer = AutoTokenizer.from_pretrained(model_path)\n\n    # ── Load model ────────────────────────────────────────────────────────────\n    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)\n    set_dropout(model, mcfg[\"dropout\"])\n\n    if model_key == \"byt5\":\n        model.config.tie_word_embeddings = False   # silence ByT5 warning\n\n    vocab_size = tokenizer.vocab_size if hasattr(tokenizer, \"vocab_size\") else None\n    model = model.to(DEVICE)\n    tokenizer.save_pretrained(os.path.join(out_dir, \"tokenizer\"))\n\n    # ── Split ─────────────────────────────────────────────────────────────────\n    val_n = max(100, int(len(train_df) * mcfg[\"val_frac\"]))\n    val_df = train_df.sample(n=val_n, random_state=mcfg[\"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, mcfg[\"max_src_len\"], mcfg[\"max_tgt_len\"])\n    val_ds   = make_hf_dataset(val_df, tokenizer, mcfg[\"max_src_len\"], mcfg[\"max_tgt_len\"])\n\n    collator = DataCollatorForSeq2Seq(\n        tokenizer, model=model, label_pad_token_id=-100, pad_to_multiple_of=8\n    )\n\n    # ── Training args ─────────────────────────────────────────────────────────\n    # adafactor for ByT5 (memory efficient), adamw for Helsinki/NLLB\n    optimizer = \"adafactor\" if model_key == \"byt5\" else \"adamw_torch\"\n\n    args = Seq2SeqTrainingArguments(\n        output_dir=out_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=mcfg[\"num_epochs\"],\n        per_device_train_batch_size=mcfg[\"train_batch_size\"],\n        per_device_eval_batch_size=mcfg[\"eval_batch_size\"],\n        gradient_accumulation_steps=mcfg[\"gradient_accumulation\"],\n        learning_rate=mcfg[\"learning_rate\"],\n        warmup_ratio=mcfg[\"warmup_ratio\"],\n        weight_decay=mcfg[\"weight_decay\"],\n        label_smoothing_factor=mcfg[\"label_smoothing\"],\n        fp16=mcfg[\"fp16\"] and DEVICE == \"cuda\",\n        predict_with_generate=True,\n        generation_max_length=mcfg[\"max_tgt_len\"],\n        generation_num_beams=2,          # fast during training eval\n        logging_steps=50,\n        save_total_limit=mcfg[\"save_total_limit\"],\n        seed=mcfg[\"seed\"],\n        report_to=\"none\",\n        gradient_checkpointing=True,\n        optim=optimizer,\n        dataloader_num_workers=2,\n        dataloader_pin_memory=False,\n        ddp_find_unused_parameters=False,\n    )\n\n    trainer = Seq2SeqTrainer(\n        model=model,\n        args=args,\n        train_dataset=train_ds,\n        eval_dataset=val_ds,\n        processing_class=tokenizer,\n        data_collator=collator,\n        compute_metrics=make_compute_metrics(tokenizer, vocab_size),\n        callbacks=[EarlyStoppingCallback(\n            early_stopping_patience=mcfg[\"early_stopping_patience\"]\n        )],\n    )\n\n    torch.cuda.empty_cache()\n    trainer.train()\n\n    best_dir = os.path.join(out_dir, \"best_model\")\n    trainer.save_model(best_dir)\n    tokenizer.save_pretrained(best_dir)\n    print(f\"\\n  ✓ Best model saved → {best_dir}\")\n\n    # Free VRAM before next model\n    del model, trainer\n    torch.cuda.empty_cache()\n\n    return best_dir\n\n\n# ═════════════════════════════════════════════════════════════════════════════\n# SECTION 5 — INFERENCE\n# ═════════════════════════════════════════════════════════════════════════════\n\ndef postprocess(text: str) -> str:\n    text = text.strip()\n    if not text:\n        return text\n    text = text[0].upper() + text[1:]\n    words = text.split()\n    for rep_len in range(max(1, 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    return text.strip()\n\n\ndef generate(sources: list[str], model_key: str, model_dir: str,\n             batch_size: int = 32) -> list[str]:\n    \"\"\"Load model, generate all translations, unload to free VRAM.\"\"\"\n    mcfg = MODEL_CFG[model_key]\n    print(f\"\\n── Inference: {model_key} from {model_dir} ──\")\n\n    if model_key == \"nllb\":\n        from transformers import NllbTokenizer\n        tokenizer = NllbTokenizer.from_pretrained(\n            model_dir, src_lang=mcfg[\"src_lang\"]\n        )\n        forced_bos = tokenizer.lang_code_to_id[mcfg[\"tgt_lang\"]]\n    else:\n        tokenizer = AutoTokenizer.from_pretrained(model_dir)\n        forced_bos = None\n\n    model = AutoModelForSeq2SeqLM.from_pretrained(model_dir).to(DEVICE)\n    model.eval()\n\n    results = []\n    for i in range(0, len(sources), batch_size):\n        batch = sources[i: i + batch_size]\n        enc = tokenizer(\n            batch, return_tensors=\"pt\", padding=True,\n            truncation=True, max_length=mcfg[\"max_src_len\"],\n        ).to(DEVICE)\n\n        gen_kwargs = dict(\n            num_beams=mcfg[\"num_beams\"],\n            length_penalty=0.8,\n            no_repeat_ngram_size=4,\n            max_new_tokens=mcfg[\"max_tgt_len\"],\n            early_stopping=True,\n        )\n        if forced_bos:\n            gen_kwargs[\"forced_bos_token_id\"] = forced_bos\n\n        with torch.no_grad():\n            out = model.generate(**enc, **gen_kwargs)\n\n        results.extend([\n            postprocess(t)\n            for t in tokenizer.batch_decode(out, skip_special_tokens=True)\n        ])\n\n        if (i // batch_size) % 5 == 0:\n            print(f\"  {i}/{len(sources)}\")\n\n    del model\n    torch.cuda.empty_cache()\n    print(f\"  ✓ {model_key} done, VRAM freed\")\n    return results\n\n\n# ═════════════════════════════════════════════════════════════════════════════\n# SECTION 6 — MBR ENSEMBLE\n# ═════════════════════════════════════════════════════════════════════════════\n\n# Weights reflect expected quality (adjust after you see val scores)\nMODEL_WEIGHTS = {\n    \"helsinki\": 0.5,\n    \"byt5\":     1.0,\n}\n\n\ndef mbr_select(candidates: list[str], weights: list[float]) -> str:\n    \"\"\"Pick candidate with highest weighted avg chrF++ against all others.\"\"\"\n    if len(candidates) == 1:\n        return candidates[0]\n    total_w = sum(weights)\n    norm_w = [w / total_w for w in weights]\n    best, best_score = candidates[0], -1.0\n    for i, cand in enumerate(candidates):\n        score = sum(\n            norm_w[j] * sacrebleu.sentence_chrf(cand, [other], word_order=2).score\n            for j, other in enumerate(candidates) if j != i\n        )\n        if score > best_score:\n            best, best_score = cand, score\n    return best\n\n\ndef ensemble(all_preds: dict[str, list[str]]) -> list[str]:\n    \"\"\"\n    all_preds: {\"helsinki\": [...], \"nllb\": [...], \"byt5\": [...]}\n    Returns one best translation per example.\n    \"\"\"\n    keys = list(all_preds.keys())\n    weights = [MODEL_WEIGHTS.get(k, 1.0) for k in keys]\n    n = len(all_preds[keys[0]])\n    results = []\n    for i in range(n):\n        candidates = [all_preds[k][i] for k in keys]\n        results.append(mbr_select(candidates, weights))\n    return results\n\n\n# ═════════════════════════════════════════════════════════════════════════════\n# SECTION 7 — MAIN PIPELINE\n# ═════════════════════════════════════════════════════════════════════════════\n\ndef main(\n    models_to_train: list[str] = [\"helsinki\", \"nllb\", \"byt5\"],\n    models_to_infer: list[str] = None,\n):\n    \"\"\"\n    models_to_train : which models to fine-tune now\n    models_to_infer : which model checkpoints to use for ensemble\n                      (defaults to same as models_to_train)\n\n    ── Common usage patterns ────────────────────────────────────────────────\n\n    Session 1 — train Helsinki only:\n        main(models_to_train=[\"helsinki\"])\n\n    Session 2 — train NLLB only:\n        main(models_to_train=[\"nllb\"])\n\n    Session 3 — train ByT5 only:\n        main(models_to_train=[\"byt5\"])\n\n    Session 4 — ensemble all (all checkpoints already saved):\n        main(models_to_train=[], models_to_infer=[\"helsinki\", \"nllb\", \"byt5\"])\n\n    All in one shot (if you have 12hr session):\n        main()\n    \"\"\"\n    if models_to_infer is None:\n        models_to_infer = models_to_train\n\n    # ── Load & augment data ───────────────────────────────────────────────────\n    train_df, test_df = load_data()\n    aug_df = token_dropout_augment(train_df, rate=0.1)\n    train_df = pd.concat([train_df, aug_df], ignore_index=True).drop_duplicates(\"src\")\n    print(f\"Final training set: {len(train_df)} examples\")\n\n    sources = test_df[\"src\"].tolist()\n\n    # ── Train ─────────────────────────────────────────────────────────────────\n    trained_dirs = {}\n    for key in models_to_train:\n        best_dir = train_model(key, train_df)\n        trained_dirs[key] = best_dir\n\n    # ── Infer ─────────────────────────────────────────────────────────────────\n    all_preds = {}\n    for key in models_to_infer:\n        # Use just-trained dir, or fall back to pre-saved checkpoint\n        model_dir = trained_dirs.get(\n            key, os.path.join(OUTPUT_BASE, key, \"best_model\")\n        )\n        if not os.path.exists(model_dir):\n            print(f\"⚠️  No checkpoint found for {key} at {model_dir}. Skipping.\")\n            continue\n        all_preds[key] = generate(sources, key, model_dir)\n\n    if not all_preds:\n        print(\"No models available for inference. Train at least one model first.\")\n        return\n\n    # ── Ensemble or single-model ──────────────────────────────────────────────\n    if len(all_preds) == 1:\n        key = list(all_preds.keys())[0]\n        print(f\"\\nOnly one model available ({key}), using directly.\")\n        translations = all_preds[key]\n    else:\n        print(f\"\\nEnsembling {list(all_preds.keys())} with MBR...\")\n        translations = ensemble(all_preds)\n\n    # ── Save submission ───────────────────────────────────────────────────────\n    sub = pd.DataFrame({\n        \"id\": test_df[\"id\"].values,\n        \"translation\": translations,\n    })\n    sub_path = os.path.join(OUTPUT_BASE, \"submission.csv\")\n    sub.to_csv(sub_path, index=False)\n    print(f\"\\n✓ Submission saved → {sub_path}\")\n    print(sub.head(5).to_string())\n\n\n# ═════════════════════════════════════════════════════════════════════════════\n# ENTRY POINT\n# Change this call to match your session:\n# ═════════════════════════════════════════════════════════════════════════════\n\nif __name__ == \"__main__\":\n\n    # ── OPTION 1: Train everything in one 12hr session ────────────────────────\n    # main()\n\n    # ── OPTION 2: Train one model per session ────────────────────────────────\n    # Session 1:  main(models_to_train=[\"helsinki\"])\n    # Session 2:  main(models_to_train=[\"nllb\"])\n    # Session 3:  main(models_to_train=[\"byt5\"])\n    # Session 4 (ensemble only, all models already trained):\n    #   main(models_to_train=[], models_to_infer=[\"helsinki\", \"nllb\", \"byt5\"])\n\n    # ── OPTION 3: Skip NLLB, just Helsinki + ByT5 ────────────────────────────\n    main(models_to_train=[\"helsinki\", \"byt5\"])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-07T10:00:06.064452Z","iopub.execute_input":"2026-03-07T10:00:06.065311Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}