{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"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":14911615,"datasetId":9528582,"databundleVersionId":15777356},{"sourceType":"datasetVersion","sourceId":14899351,"datasetId":9366631,"databundleVersionId":15763950},{"sourceType":"modelInstanceVersion","sourceId":800589,"databundleVersionId":16247520,"modelInstanceId":609559,"modelId":621394,"isSourceIdPinned":false}],"dockerImageVersionId":31287,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# ByT5 MBR Inference (Offline)\n\n**Purpose**: Generate test predictions using MBR (Minimum Bayes Risk) decoding.\n**Internet**: OFF (Kaggle code competition mode)\n\n## What is MBR?\nInstead of taking the single best beam search output, MBR:\n1. Generates **multiple candidates** (beam search + sampling)\n2. Scores each candidate by **average similarity** to all others (chrF++)\n3. Picks the **most \"central\"** translation → more robust output\n\n## Model Priority\n1. Fine-tuned model (from training notebook output)\n2. `mattiaangeli/byt5-akkadian-mbr` (public, scores 35.4)\n3. `assiaben/final-byt5` / `byt5-akkadian-optimized-34x` (public, scores 35.0)\n\n## Setup on Kaggle\n1. Add the fine-tuned model dataset (or public model) as input\n2. Add competition dataset\n3. Turn **Internet OFF**, enable **GPU**\n4. Run all cells → `submission.csv`","metadata":{}},{"cell_type":"code","source":"# sacrebleu is pre-installed on Kaggle, but install if running locally\ntry:\n    import sacrebleu as _test_sacrebleu\n    del _test_sacrebleu\n    print(\"sacrebleu already installed ✓\")\nexcept ImportError:\n    import subprocess\n    subprocess.check_call([\"pip\", \"install\", \"-q\", \"sacrebleu\"])\n    print(\"sacrebleu installed ✓\")\nprint(\"(needed for MBR decoding)\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-23T20:40:16.394411Z","iopub.execute_input":"2026-03-23T20:40:16.394890Z","iopub.status.idle":"2026-03-23T20:41:38.392652Z","shell.execute_reply.started":"2026-03-23T20:40:16.394856Z","shell.execute_reply":"2026-03-23T20:41:38.391966Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ═══════════════════════════════════════════════════════════════\n# IMPORTS & CONFIG\n# ═══════════════════════════════════════════════════════════════\n\nimport os\nimport re\nimport warnings\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom contextlib import nullcontext\nfrom torch.utils.data import Dataset, DataLoader, Sampler\nfrom transformers import AutoTokenizer, AutoModelForSeq2SeqLM\nimport sacrebleu\n\nwarnings.filterwarnings(\"ignore\")\n\n# ── Config ─────────────────────────────────────────────────────\nMAX_LENGTH = 256\nBATCH_SIZE = 2\nNUM_BUCKETS = 4\n\n# Generation params\nNUM_BEAMS = 6\nLENGTH_PENALTY = 1.0\nMAX_NEW_TOKENS = 256\nEARLY_STOPPING = True\n\n# MBR config (safer / lighter)\nUSE_MBR = True\nMBR_NUM_BEAM_CANDS = 3\nMBR_NUM_SAMPLE_CANDS = 2\nMBR_TOP_P = 0.9\nMBR_TEMPERATURES = [0.65]\nMBR_POOL_CAP = 12\nREPETITION_PENALTY = 1.03\n\n# BF16 optimization\nUSE_BF16 = True\n\nPREFIX = \"translate Akkadian to English: \"\n\n# ── BF16 helpers ───────────────────────────────────────────────\ndef _cuda_bf16_supported():\n    if not torch.cuda.is_available():\n        return False\n    try:\n        return bool(getattr(torch.cuda, \"is_bf16_supported\", lambda: False)())\n    except Exception:\n        return False\n\ndef _bf16_ctx(device, enabled):\n    if enabled and device.type == \"cuda\" and _cuda_bf16_supported():\n        return torch.autocast(device_type=\"cuda\", dtype=torch.bfloat16)\n    return nullcontext()\n\n# ── Path helpers ───────────────────────────────────────────────\ndef _list_subdirs(path):\n    if not os.path.exists(path):\n        return []\n    return [\n        os.path.join(path, d)\n        for d in os.listdir(path)\n        if os.path.isdir(os.path.join(path, d))\n    ]\n\ndef _looks_like_dataset_folder(path):\n    required_any = [\n        \"test.csv\",\n        \"train.csv\",\n        \"sample_submission.csv\",\n        \"published_texts.csv\",\n    ]\n    try:\n        files = set(os.listdir(path))\n        return any(f in files for f in required_any)\n    except Exception:\n        return False\n\ndef _looks_like_model_folder(path):\n    required_any = [\n        \"config.json\",\n        \"tokenizer_config.json\",\n        \"generation_config.json\",\n        \"spiece.model\",\n        \"special_tokens_map.json\",\n    ]\n    try:\n        files = set(os.listdir(path))\n        return \"config.json\" in files and any(f in files for f in required_any[1:])\n    except Exception:\n        return False\n\ndef _find_dataset_path():\n    candidates = []\n\n    local_candidates = [\n        \"./dataset\",\n        \"./data\",\n        \"./input\",\n        \".\",\n    ]\n    candidates.extend(local_candidates)\n\n    kaggle_root = \"/kaggle/input\"\n    if os.path.exists(kaggle_root):\n        candidates.extend(_list_subdirs(kaggle_root))\n        for sub in _list_subdirs(kaggle_root):\n            candidates.extend(_list_subdirs(sub))\n\n    seen = set()\n    filtered = []\n    for p in candidates:\n        p = os.path.abspath(p)\n        if p not in seen and os.path.isdir(p):\n            seen.add(p)\n            filtered.append(p)\n\n    for p in filtered:\n        if _looks_like_dataset_folder(p):\n            return p\n\n    if os.path.exists(\"/kaggle/input\"):\n        for root, dirs, files in os.walk(\"/kaggle/input\"):\n            if \"test.csv\" in files or \"train.csv\" in files:\n                return root\n\n    for root, dirs, files in os.walk(\".\"):\n        if \"test.csv\" in files or \"train.csv\" in files:\n            return root\n\n    return None\n\ndef _find_model_path():\n    candidates = []\n\n    local_candidates = [\n        \"./byt5-finetuned/best_model\",\n        \"./byt5-finetuned\",\n        \"./best_model\",\n        \"./model\",\n        \"./checkpoint\",\n        \".\",\n    ]\n    candidates.extend(local_candidates)\n\n    kaggle_root = \"/kaggle/input\"\n    if os.path.exists(kaggle_root):\n        candidates.extend(_list_subdirs(kaggle_root))\n        for sub in _list_subdirs(kaggle_root):\n            candidates.extend(_list_subdirs(sub))\n\n    seen = set()\n    filtered = []\n    for p in candidates:\n        p = os.path.abspath(p)\n        if p not in seen and os.path.isdir(p):\n            seen.add(p)\n            filtered.append(p)\n\n    for p in filtered:\n        if _looks_like_model_folder(p):\n            return p\n\n    for base in [\"/kaggle/input\", \".\"]:\n        if not os.path.exists(base):\n            continue\n        for root, dirs, files in os.walk(base):\n            if \"config.json\" in files and (\n                \"tokenizer_config.json\" in files or\n                \"spiece.model\" in files or\n                \"special_tokens_map.json\" in files\n            ):\n                return root\n\n    return None\n\nDATA_PATH = _find_dataset_path()\nMODEL_PATH = _find_model_path()\n\nprint(\"Detected DATA_PATH:\", DATA_PATH)\nprint(\"Detected MODEL_PATH:\", MODEL_PATH)\n\nif DATA_PATH is None:\n    print(\"\\nFolders under /kaggle/input:\")\n    if os.path.exists(\"/kaggle/input\"):\n        for x in sorted(os.listdir(\"/kaggle/input\")):\n            print(\" -\", x)\n    raise AssertionError(\"Could not find dataset path\")\n\nif MODEL_PATH is None:\n    print(\"\\nFolders under /kaggle/input:\")\n    if os.path.exists(\"/kaggle/input\"):\n        for x in sorted(os.listdir(\"/kaggle/input\")):\n            print(\" -\", x)\n    raise AssertionError(\"Could not find model path\")\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nprint(f\"\\nDATA_PATH:  {DATA_PATH}\")\nprint(f\"MODEL_PATH: {MODEL_PATH}\")\nprint(f\"DEVICE:     {DEVICE}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-23T20:41:38.393620Z","iopub.execute_input":"2026-03-23T20:41:38.393860Z","iopub.status.idle":"2026-03-23T20:41:55.262394Z","shell.execute_reply.started":"2026-03-23T20:41:38.393838Z","shell.execute_reply":"2026-03-23T20:41:55.261702Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\n\nprint(\"Contents of /kaggle/input:\")\nfor root, dirs, files in os.walk(\"/kaggle/input\"):\n    print(root)\n    for f in files[:10]:\n        print(\"   \", f)\n    print(\"---\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-23T20:41:55.264220Z","iopub.execute_input":"2026-03-23T20:41:55.264761Z","iopub.status.idle":"2026-03-23T20:41:55.274719Z","shell.execute_reply.started":"2026-03-23T20:41:55.264730Z","shell.execute_reply":"2026-03-23T20:41:55.273920Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Preprocessing — MINIMAL (Matches Training & 35.4 Notebook)\n\nOnly two transformations:\n- `...` / `…` → `<big_gap>`\n- `xx+` / ` x ` → `<gap>`\n\nPostprocessing cleans model output for submission (ḫ→h, forbidden chars, fractions, etc.).","metadata":{}},{"cell_type":"code","source":"# ═══════════════════════════════════════════════════════════════\n# PREPROCESSING — UPDATED FOR DATASET GAP REGULARIZATION\n# ═══════════════════════════════════════════════════════════════\n\nSUBSCRIPT_MAP = str.maketrans(\"₀₁₂₃₄₅₆₇₈₉\", \"0123456789\")\n\n# Gap patterns\n_PAT_LITERAL_SMALL_GAP = re.compile(r\"<\\s*gap\\s*>\", flags=re.I)\n_PAT_LITERAL_BIG_GAP = re.compile(r\"<\\s*big_gap\\s*>\", flags=re.I)\n\n_PAT_BIG_GAP = re.compile(\n    r\"(\\.{3,}|…+|\\[\\s*\\.{2,}\\s*\\]|\\(\\s*\\.{2,}\\s*\\))\",\n    flags=re.I\n)\n\n_PAT_SMALL_GAP = re.compile(\n    r\"(\\[\\s*x+\\s*\\]|\\(\\s*x+\\s*\\)|\\bxx+\\b|\\bx\\b)\",\n    flags=re.I\n)\n\ndef _normalize_gap_spacing(text: str) -> str:\n    text = re.sub(r\"\\s*<gap>\\s*\", \" <gap> \", text)\n    text = re.sub(r\"\\s*<big_gap>\\s*\", \" <big_gap> \", text)\n    text = re.sub(r\"\\s+\", \" \", text).strip()\n    return text\n\ndef _collapse_gap_sequences(text: str) -> str:\n    # many adjacent small gaps become one big gap\n    text = re.sub(r\"(<gap>\\s*){2,}\", \"<big_gap> \", text)\n    text = re.sub(r\"(<big_gap>\\s*){2,}\", \"<big_gap> \", text)\n    text = re.sub(r\"<big_gap>\\s*<gap>|<gap>\\s*<big_gap>\", \"<big_gap>\", text)\n    text = re.sub(r\"\\s+\", \" \", text).strip()\n    return text\n\ndef preprocess_transliteration(text):\n    \"\"\"\n    Safer preprocessing for new competition data.\n    Focus: normalize gap conventions only.\n    \"\"\"\n    if pd.isna(text) or not isinstance(text, str):\n        return \"\"\n\n    text = str(text).strip()\n\n    # protect already-correct tags\n    text = _PAT_LITERAL_BIG_GAP.sub(\" <big_gap> \", text)\n    text = _PAT_LITERAL_SMALL_GAP.sub(\" <gap> \", text)\n\n    # normalize raw gap variants\n    text = _PAT_BIG_GAP.sub(\" <big_gap> \", text)\n    text = _PAT_SMALL_GAP.sub(\" <gap> \", text)\n\n    # spacing + collapse\n    text = _normalize_gap_spacing(text)\n    text = _collapse_gap_sequences(text)\n\n    return text\n\n\ndef postprocess_output(text):\n    \"\"\"\n    Lighter postprocessing:\n    - normalize ḫ and subscripts\n    - normalize obvious gap variants\n    - clean spacing / punctuation lightly\n    - DO NOT aggressively deduplicate words/ngrams\n    \"\"\"\n    if not isinstance(text, str) or not text.strip():\n        return \"\"\n\n    text = str(text).strip()\n\n    # script normalization\n    text = text.replace(\"ḫ\", \"h\").replace(\"Ḫ\", \"H\")\n    text = text.translate(SUBSCRIPT_MAP)\n\n    # preserve already correct tags first\n    text = _PAT_LITERAL_BIG_GAP.sub(\" <big_gap> \", text)\n    text = _PAT_LITERAL_SMALL_GAP.sub(\" <gap> \", text)\n\n    # normalize common gap-like outputs\n    text = re.sub(r\"(\\[\\s*x+\\s*\\]|\\(\\s*x+\\s*\\)|\\bxx+\\b|\\bx\\b)\", \" <gap> \", text, flags=re.I)\n    text = re.sub(r\"(\\.{3,}|…+|\\[\\s*\\.{2,}\\s*\\])\", \" <big_gap> \", text)\n\n    # remove only truly noisy symbols, but keep normal punctuation\n    text = text.replace(\"⌈\", \" \").replace(\"⌉\", \" \")\n    text = text.replace(\"⌊\", \" \").replace(\"⌋\", \" \")\n    text = text.replace(\"ʾ\", \"'\")\n    text = text.replace(\"—\", \"-\").replace(\"–\", \"-\")\n\n    # light annotation cleanup only\n    text = re.sub(r\"\\((fem|plur|pl|sing|singular|plural)\\.?\\)\", \"\", text, flags=re.I)\n\n    # fraction normalization\n    text = re.sub(r'(\\d+)\\.5\\b', r'\\1½', text)\n    text = re.sub(r'\\b0\\.5\\b', '½', text)\n    text = re.sub(r'(\\d+)\\.25\\b', r'\\1¼', text)\n    text = re.sub(r'\\b0\\.25\\b', '¼', text)\n    text = re.sub(r'(\\d+)\\.75\\b', r'\\1¾', text)\n    text = re.sub(r'\\b0\\.75\\b', '¾', text)\n\n    # spacing + collapse gaps\n    text = _normalize_gap_spacing(text)\n    text = _collapse_gap_sequences(text)\n\n    # light punctuation spacing\n    text = re.sub(r\"\\s+([,.;:])\", r\"\\1\", text)\n    text = re.sub(r\"([,.;:])([^\\s])\", r\"\\1 \\2\", text)\n    text = re.sub(r\"\\s+\", \" \", text).strip()\n\n    return text","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-23T20:41:55.275729Z","iopub.execute_input":"2026-03-23T20:41:55.276121Z","iopub.status.idle":"2026-03-23T20:41:55.291693Z","shell.execute_reply.started":"2026-03-23T20:41:55.276075Z","shell.execute_reply":"2026-03-23T20:41:55.291066Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ═══════════════════════════════════════════════════════════════\n# LOAD MODEL & TOKENIZER\n# ═══════════════════════════════════════════════════════════════\n\nassert os.path.exists(MODEL_PATH), f\"❌ Model not found at: {MODEL_PATH}\"\nprint(f\"Loading model from: {MODEL_PATH}\")\n\ntokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, local_files_only=True)\nmodel = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH, local_files_only=True)\nmodel.to(DEVICE)\nmodel.eval()\n\n# Set matmul precision for better GPU perf\nif DEVICE.type == \"cuda\":\n    try:\n        torch.set_float32_matmul_precision(\"high\")\n    except Exception:\n        pass\n\nprint(f\"Model loaded ✓\")\nprint(f\"Parameters: {sum(p.numel() for p in model.parameters()):,}\")\nprint(f\"BF16 autocast: {USE_BF16 and _cuda_bf16_supported()}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-23T20:41:55.292340Z","iopub.execute_input":"2026-03-23T20:41:55.292510Z","iopub.status.idle":"2026-03-23T20:42:14.128445Z","shell.execute_reply.started":"2026-03-23T20:41:55.292493Z","shell.execute_reply":"2026-03-23T20:42:14.127692Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ═══════════════════════════════════════════════════════════════\n# LOAD & PREPROCESS TEST DATA\n# ═══════════════════════════════════════════════════════════════\n\ntest_df = pd.read_csv(os.path.join(DATA_PATH, \"test.csv\"))\n\n# keep original for diagnostics\ntest_df[\"transliteration_raw\"] = test_df[\"transliteration\"].astype(str)\ntest_df[\"src\"] = test_df[\"transliteration\"].apply(preprocess_transliteration)\n\n# sort in document order if fields exist\nsort_cols = [c for c in [\"text_id\", \"line_start\", \"line_end\", \"id\"] if c in test_df.columns]\nif sort_cols:\n    test_df = test_df.sort_values(sort_cols, kind=\"stable\").reset_index(drop=True)\n\nraw_has_gap = test_df[\"transliteration_raw\"].str.contains(r\"<\\s*gap\\s*>\", case=False, na=False).sum()\nraw_has_big_gap = test_df[\"transliteration_raw\"].str.contains(r\"<\\s*big_gap\\s*>\", case=False, na=False).sum()\nproc_has_gap = test_df[\"src\"].str.contains(r\"<gap>\", regex=False, na=False).sum()\nproc_has_big_gap = test_df[\"src\"].str.contains(r\"<big_gap>\", regex=False, na=False).sum()\nchanged = (test_df[\"transliteration_raw\"] != test_df[\"src\"]).sum()\n\nprint(f\"Test samples: {len(test_df)}\")\nprint(f\"Rows changed by preprocessing: {changed}/{len(test_df)} ({changed/len(test_df):.1%})\")\nprint(f\"Raw rows with <gap>: {raw_has_gap}\")\nprint(f\"Raw rows with <big_gap>: {raw_has_big_gap}\")\nprint(f\"Processed rows with <gap>: {proc_has_gap}\")\nprint(f\"Processed rows with <big_gap>: {proc_has_big_gap}\")\n\nprint(f\"\\nSample inputs:\")\nfor i in range(min(3, len(test_df))):\n    raw = test_df.iloc[i][\"transliteration_raw\"]\n    src = test_df.iloc[i][\"src\"]\n    print(f\"\\n[{i}] RAW: {raw[:140]}{'...' if len(raw) > 140 else ''}\")\n    print(f\"    SRC: {src[:140]}{'...' if len(src) > 140 else ''}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-23T20:42:14.129504Z","iopub.execute_input":"2026-03-23T20:42:14.129838Z","iopub.status.idle":"2026-03-23T20:42:14.162562Z","shell.execute_reply.started":"2026-03-23T20:42:14.129816Z","shell.execute_reply":"2026-03-23T20:42:14.161898Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## MBR Inference\n\n**MBR (Minimum Bayes Risk)** decoding generates multiple candidate translations and picks the most \"central\" one — the candidate with highest average chrF++ similarity to all others.\n\nCandidates: 4 from beam search + 4 from sampling = 8 per input.\nThis adds ~1-2x runtime but typically gains +0.5-1.0 points.","metadata":{}},{"cell_type":"code","source":"# ═══════════════════════════════════════════════════════════════\n# BUCKET BATCH SAMPLER + DATASET + MBR HELPERS\n# ═══════════════════════════════════════════════════════════════\n\n_CHRFPP = sacrebleu.metrics.CHRF(word_order=2)\n\ndef _dedup_keep_order(xs):\n    seen = set()\n    out = []\n    for x in xs:\n        x = str(x).strip()\n        if x and x not in seen:\n            out.append(x)\n            seen.add(x)\n    return out\n\ndef _sim_chrfpp(a, b):\n    a = (a or \"\").strip()\n    b = (b or \"\").strip()\n    if not a or not b:\n        return 0.0\n    return float(_CHRFPP.sentence_score(a, [b]).score)\n\ndef _token_len(text):\n    if not isinstance(text, str) or not text.strip():\n        return 0\n    return len(text.split())\n\ndef _has_gap_token(text):\n    if not isinstance(text, str):\n        return False\n    return (\"<gap>\" in text) or (\"<big_gap>\" in text)\n\ndef _source_gap_profile(src_text):\n    if not isinstance(src_text, str):\n        return {\"has_gap\": False, \"has_big_gap\": False, \"src_len\": 0}\n    return {\n        \"has_gap\": \"<gap>\" in src_text,\n        \"has_big_gap\": \"<big_gap>\" in src_text,\n        \"src_len\": len(src_text.split()),\n    }\n\ndef _candidate_quality_bonus(text, src_text=None, median_len=None):\n    if not text or not str(text).strip():\n        return -8.0\n\n    text = str(text).strip()\n    bonus = 0.0\n\n    bonus += 0.12 * len(re.findall(r\"<gap>\", text))\n    bonus += 0.18 * len(re.findall(r\"<big_gap>\", text))\n\n    bad_left = text.replace(\"<gap>\", \"\").replace(\"<big_gap>\", \"\")\n    if \"<\" in bad_left:\n        bonus -= 1.5\n    if \">\" in bad_left:\n        bonus -= 1.5\n\n    if re.search(r\"\\b(\\w+)(?:\\s+\\1\\b){2,}\", text, flags=re.I):\n        bonus -= 1.5\n\n    tok_len = _token_len(text)\n    if tok_len <= 1:\n        bonus -= 2.5\n    elif tok_len <= 2:\n        bonus -= 1.2\n\n    if src_text is not None:\n        prof = _source_gap_profile(src_text)\n\n        if (prof[\"has_gap\"] or prof[\"has_big_gap\"]) and not _has_gap_token(text):\n            bonus -= 0.6\n\n        if prof[\"src_len\"] >= 12 and tok_len <= 3:\n            bonus -= 1.2\n        elif prof[\"src_len\"] >= 8 and tok_len <= 2:\n            bonus -= 1.0\n\n    if median_len is not None and median_len > 0:\n        bonus -= 0.08 * abs(tok_len - median_len)\n\n    return bonus\n\ndef mbr_pick(candidates, src_text=None):\n    candidates = _dedup_keep_order(candidates)\n    if MBR_POOL_CAP:\n        candidates = candidates[:MBR_POOL_CAP]\n\n    n = len(candidates)\n    if n == 0:\n        return \"\"\n    if n == 1:\n        return candidates[0]\n\n    lens = [_token_len(c) for c in candidates]\n    median_len = int(np.median(lens)) if lens else None\n\n    best_i, best_s = 0, -1e18\n    for i in range(n):\n        sim = sum(_sim_chrfpp(candidates[i], candidates[j]) for j in range(n) if i != j)\n        sim /= max(1, n - 1)\n\n        score = sim + _candidate_quality_bonus(\n            candidates[i],\n            src_text=src_text,\n            median_len=median_len\n        )\n\n        if score > best_s:\n            best_s, best_i = score, i\n\n    return candidates[best_i]\n\ndef generate_sample_candidates(model, tokenizer, input_ids, attention_mask, num_return_sequences):\n    all_texts = []\n\n    temps = MBR_TEMPERATURES if isinstance(MBR_TEMPERATURES, (list, tuple)) else [MBR_TEMPERATURES]\n    per_temp = max(1, int(np.ceil(num_return_sequences / len(temps))))\n\n    for temp in temps:\n        with torch.inference_mode(), _bf16_ctx(DEVICE, USE_BF16):\n            outputs = model.generate(\n                input_ids=input_ids,\n                attention_mask=attention_mask,\n                do_sample=True,\n                top_p=MBR_TOP_P,\n                temperature=float(temp),\n                num_return_sequences=per_temp,\n                max_new_tokens=MAX_NEW_TOKENS,\n                repetition_penalty=REPETITION_PENALTY,\n                early_stopping=EARLY_STOPPING,\n                use_cache=True,\n            )\n\n        texts = tokenizer.batch_decode(outputs, skip_special_tokens=True)\n        all_texts.extend(texts)\n\n    return all_texts[:num_return_sequences]\n\nclass BucketBatchSampler(Sampler):\n    def __init__(self, lengths, batch_size, num_buckets=4):\n        self.batch_size = batch_size\n        sorted_indices = sorted(range(len(lengths)), key=lambda i: lengths[i])\n        bucket_size = max(1, len(sorted_indices) // num_buckets)\n        self.buckets = []\n        for i in range(num_buckets):\n            start = i * bucket_size\n            end = start + bucket_size if i < num_buckets - 1 else len(sorted_indices)\n            bucket = sorted_indices[start:end]\n            if bucket:\n                self.buckets.append(bucket)\n\n    def __iter__(self):\n        for bucket in self.buckets:\n            for i in range(0, len(bucket), self.batch_size):\n                yield bucket[i:i + self.batch_size]\n\n    def __len__(self):\n        return sum((len(b) + self.batch_size - 1) // self.batch_size for b in self.buckets)\n\nclass AkkadianDataset(Dataset):\n    def __init__(self, df):\n        self.df = df.reset_index(drop=True).copy()\n        self.ids = self.df[\"id\"].tolist()\n        self.src_texts = self.df[\"src\"].astype(str).tolist()\n        self.model_inputs = [PREFIX + x for x in self.src_texts]\n        self.lengths = [len(x.split()) for x in self.model_inputs]\n\n    def __len__(self):\n        return len(self.ids)\n\n    def __getitem__(self, idx):\n        return {\n            \"id\": self.ids[idx],\n            \"src_text\": self.src_texts[idx],\n            \"model_input\": self.model_inputs[idx],\n            \"length\": self.lengths[idx],\n        }\n\ndef collate_fn(batch):\n    ids = [x[\"id\"] for x in batch]\n    src_texts = [x[\"src_text\"] for x in batch]\n    model_inputs = [x[\"model_input\"] for x in batch]\n\n    encoded = tokenizer(\n        model_inputs,\n        max_length=MAX_LENGTH,\n        padding=True,\n        truncation=True,\n        return_tensors=\"pt\",\n    )\n\n    return {\n        \"id\": ids,\n        \"src_text\": src_texts,\n        \"input_ids\": encoded[\"input_ids\"],\n        \"attention_mask\": encoded[\"attention_mask\"],\n    }\n\ndef get_adaptive_beams(attention_mask):\n    lengths = attention_mask.sum(dim=1)\n    if lengths.max().item() < 100:\n        return max(4, NUM_BEAMS // 2)\n    return NUM_BEAMS\n\ndataset = AkkadianDataset(test_df)\nsampler = BucketBatchSampler(dataset.lengths, BATCH_SIZE, NUM_BUCKETS)\ndataloader = DataLoader(dataset, batch_sampler=sampler, collate_fn=collate_fn, num_workers=0)\n\nprint(f\"Dataset:    {len(dataset)} samples\")\nprint(f\"Batches:    {len(sampler)}\")\nprint(f\"Batch size: {BATCH_SIZE}\")\nprint(f\"MBR:        {USE_MBR} ({MBR_NUM_BEAM_CANDS} beam + {MBR_NUM_SAMPLE_CANDS} sample cands)\")\nprint(f\"Temps:      {MBR_TEMPERATURES}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-23T20:42:14.163519Z","iopub.execute_input":"2026-03-23T20:42:14.163959Z","iopub.status.idle":"2026-03-23T20:42:14.398181Z","shell.execute_reply.started":"2026-03-23T20:42:14.163938Z","shell.execute_reply":"2026-03-23T20:42:14.397296Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ═══════════════════════════════════════════════════════════════\n# INFERENCE — STABLE MBR DECODING\n# ═══════════════════════════════════════════════════════════════\n\nprint(\"Running stable inference...\")\nprint(f\"  Beams:              {NUM_BEAMS}\")\nprint(f\"  Length penalty:     {LENGTH_PENALTY}\")\nprint(f\"  Max new tokens:     {MAX_NEW_TOKENS}\")\nprint(f\"  MBR beam cands:     {MBR_NUM_BEAM_CANDS}\")\nprint(f\"  MBR sample cands:   {MBR_NUM_SAMPLE_CANDS}\")\nprint(f\"  MBR temperatures:   {MBR_TEMPERATURES}\")\nprint(f\"  MBR top_p:          {MBR_TOP_P}\")\nprint(f\"  Repetition penalty: {REPETITION_PENALTY}\")\nprint(f\"  BF16 autocast:      {USE_BF16 and _cuda_bf16_supported()}\")\nprint()\n\npredictions = []\nfailed_items = 0\nbatch_count = 0\n\nfor batch in dataloader:\n    batch_count += 1\n\n    try:\n        batch_ids = batch[\"id\"]\n        batch_src_texts = batch[\"src_text\"]\n        input_ids = batch[\"input_ids\"].to(DEVICE)\n        attention_mask = batch[\"attention_mask\"].to(DEVICE)\n\n        B = input_ids.shape[0]\n        beam_size = get_adaptive_beams(attention_mask)\n\n        final_texts = []\n\n        for i in range(B):\n            single_input_ids = input_ids[i:i+1]\n            single_attention_mask = attention_mask[i:i+1]\n            src_text = batch_src_texts[i]\n\n            candidate_pool = []\n\n            try:\n                # beam candidates\n                with torch.inference_mode(), _bf16_ctx(DEVICE, USE_BF16):\n                    nb = max(beam_size, MBR_NUM_BEAM_CANDS)\n                    beam_outputs = model.generate(\n                        input_ids=single_input_ids,\n                        attention_mask=single_attention_mask,\n                        do_sample=False,\n                        num_beams=nb,\n                        num_return_sequences=MBR_NUM_BEAM_CANDS,\n                        max_new_tokens=MAX_NEW_TOKENS,\n                        length_penalty=LENGTH_PENALTY,\n                        early_stopping=EARLY_STOPPING,\n                        repetition_penalty=REPETITION_PENALTY,\n                        use_cache=True,\n                    )\n\n                beam_texts = tokenizer.batch_decode(beam_outputs, skip_special_tokens=True)\n                candidate_pool.extend(beam_texts)\n\n                # sample candidates\n                if USE_MBR and MBR_NUM_SAMPLE_CANDS > 0:\n                    sample_texts = generate_sample_candidates(\n                        model,\n                        tokenizer,\n                        single_input_ids,\n                        single_attention_mask,\n                        MBR_NUM_SAMPLE_CANDS\n                    )\n                    candidate_pool.extend(sample_texts)\n\n                candidate_pool = [postprocess_output(x) for x in candidate_pool]\n                candidate_pool = [x for x in candidate_pool if isinstance(x, str) and x.strip()]\n\n                if USE_MBR:\n                    picked = mbr_pick(candidate_pool, src_text=src_text)\n                else:\n                    picked = candidate_pool[0] if candidate_pool else \"\"\n\n                if not isinstance(picked, str) or not picked.strip():\n                    picked = \"<gap>\"\n\n                final_texts.append(picked)\n\n            except RuntimeError as e:\n                if \"out of memory\" in str(e).lower():\n                    print(f\"OOM on item in batch {batch_count}; falling back to single greedy decode.\")\n                    if torch.cuda.is_available():\n                        torch.cuda.empty_cache()\n\n                    try:\n                        with torch.inference_mode(), _bf16_ctx(DEVICE, USE_BF16):\n                            fallback_outputs = model.generate(\n                                input_ids=single_input_ids,\n                                attention_mask=single_attention_mask,\n                                do_sample=False,\n                                num_beams=1,\n                                num_return_sequences=1,\n                                max_new_tokens=MAX_NEW_TOKENS,\n                                repetition_penalty=REPETITION_PENALTY,\n                                use_cache=True,\n                            )\n                        fallback_text = tokenizer.batch_decode(fallback_outputs, skip_special_tokens=True)[0]\n                        fallback_text = postprocess_output(fallback_text)\n                        final_texts.append(fallback_text if fallback_text.strip() else \"<gap>\")\n                    except Exception:\n                        final_texts.append(\"<gap>\")\n                        failed_items += 1\n                else:\n                    print(f\"Runtime error on item in batch {batch_count}: {e}\")\n                    final_texts.append(\"<gap>\")\n                    failed_items += 1\n\n            except Exception as e:\n                print(f\"Error on item in batch {batch_count}: {e}\")\n                final_texts.append(\"<gap>\")\n                failed_items += 1\n\n        predictions.extend(list(zip(batch_ids, final_texts)))\n\n        if batch_count % 20 == 0:\n            print(f\"Processed {len(predictions)} samples...\")\n\n        if torch.cuda.is_available():\n            torch.cuda.empty_cache()\n\n    except Exception as e:\n        print(f\"Batch-level error on batch {batch_count}: {e}\")\n        # fall back: add placeholder outputs so notebook never dies\n        for bad_id in batch[\"id\"]:\n            predictions.append((bad_id, \"<gap>\"))\n            failed_items += 1\n\nprint(f\"\\nFinished inference on {len(predictions)} samples.\")\nprint(f\"Failed items recovered with fallback: {failed_items}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-23T20:42:14.400428Z","iopub.execute_input":"2026-03-23T20:42:14.400848Z","iopub.status.idle":"2026-03-23T20:42:42.078667Z","shell.execute_reply.started":"2026-03-23T20:42:14.400826Z","shell.execute_reply":"2026-03-23T20:42:42.078054Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ═══════════════════════════════════════════════════════════════\n# CREATE SUBMISSION\n# ═══════════════════════════════════════════════════════════════\n\nif \"predictions\" not in globals():\n    raise RuntimeError(\"predictions was not created. Run Cell 10 first.\")\n\nif len(predictions) == 0:\n    raise RuntimeError(\"predictions is empty. Inference failed before producing outputs.\")\n\nsubmission = pd.DataFrame(predictions, columns=[\"id\", \"translation\"])\n\nsubmission[\"translation\"] = submission[\"translation\"].apply(\n    lambda t: t if isinstance(t, str) and t.strip() else \"<gap>\"\n)\n\nsubmission = submission.drop_duplicates(subset=[\"id\"], keep=\"first\").copy()\n\nif \"id\" in test_df.columns:\n    submission = test_df[[\"id\"]].merge(submission, on=\"id\", how=\"left\")\n    submission[\"translation\"] = submission[\"translation\"].fillna(\"<gap>\")\n\nempty = submission[\"translation\"].str.strip().eq(\"\").sum()\nlengths = submission[\"translation\"].str.len()\n\ngap_count = submission[\"translation\"].str.contains(r\"<gap>\", regex=False, na=False).sum()\nbig_gap_count = submission[\"translation\"].str.contains(r\"<big_gap>\", regex=False, na=False).sum()\nrepetition_count = submission[\"translation\"].str.contains(\n    r\"\\b(\\w+)(?:\\s+\\1\\b){2,}\", regex=True, na=False\n).sum()\n\nprint(f\"Submission shape: {submission.shape}\")\nprint(f\"Empty predictions: {empty}\")\nprint(\n    f\"Length stats: mean={lengths.mean():.0f}, \"\n    f\"median={lengths.median():.0f}, min={lengths.min()}, max={lengths.max()}\"\n)\nprint(f\"Rows with <gap>: {gap_count}\")\nprint(f\"Rows with <big_gap>: {big_gap_count}\")\nprint(f\"Rows with obvious triple repetition: {repetition_count}\")\n\nshort = ((lengths < 5) & (lengths > 0)).sum()\nif short > 0:\n    print(f\"{short} very short translations (<5 chars)\")\n\nprint(f\"\\n--- Sample predictions ---\")\nfor i in [0, len(submission)//2, len(submission)-1]:\n    if 0 <= i < len(submission):\n        row = submission.iloc[i]\n        test_row = test_df[test_df[\"id\"] == row[\"id\"]]\n        print(f\"ID {row['id']}:\")\n        if len(test_row) > 0:\n            print(f\"  SRC:  {str(test_row.iloc[0]['transliteration'])[:120]}\")\n        pred = row[\"translation\"][:120]\n        print(f\"  PRED: {pred}{'...' if len(row['translation']) > 120 else ''}\")\n        print()\n\nsubmission.to_csv(\"submission.csv\", index=False)\n\nprint(\"=\" * 60)\nprint(f\"submission.csv saved ({len(submission)} rows)\")\nprint(f\"MBR decoding: {'YES' if USE_MBR else 'NO'}\")\nprint(f\"Model: {os.path.basename(MODEL_PATH)}\")\nprint(\"=\" * 60)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-23T20:42:42.079568Z","iopub.execute_input":"2026-03-23T20:42:42.079867Z","iopub.status.idle":"2026-03-23T20:42:42.111884Z","shell.execute_reply.started":"2026-03-23T20:42:42.079844Z","shell.execute_reply":"2026-03-23T20:42:42.111049Z"}},"outputs":[],"execution_count":null}]}