{"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},{"sourceType":"datasetVersion","sourceId":14899351,"datasetId":9366631,"databundleVersionId":15763950},{"sourceType":"modelInstanceVersion","sourceId":754941,"databundleVersionId":15740514,"modelInstanceId":573226,"modelId":585575},{"sourceType":"modelInstanceVersion","sourceId":755680,"databundleVersionId":15750522,"modelInstanceId":577094,"modelId":589419}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# %% cell 0\n#!/usr/bin/env python3\n\"\"\"\nDeep Past Challenge — Akkadian-to-English Translation\nEnsemble MBR inference with corrected postprocessing (FIXED version)\n\nCorrections applied based on host v3 update + community discussion:\n  1. 5/12 shekel → ⅓ shekel 15 grains (was ⅔)\n  2. Commodity regex respects word boundaries (no import-tax corruption)\n  3. Parentheses preserved in translations (appear in test set)\n  4. Curly quotes → straight quotes (not deleted)\n  5. ḫ/Ḫ → h/H in postprocessing (not in test set)\n  6. Safer slash-alternative regex (protects fractions)\n  7. Stray mark removal (.., xx, <<>>, <>)\n\"\"\"\n\n\n# ======================================================================\n# Better Candidate Diversity and Stronger Consensus Selection\n# ======================================================================\n\n# ======================================================================\n# Notebook structure\n# ======================================================================\n\n# ======================================================================\n# Corrections applied (based on host recommendations + community discussion)\n# ======================================================================\nimport os, gc, re, json, math, random, logging, warnings\nfrom pathlib import Path\nfrom dataclasses import dataclass, field\nfrom contextlib import nullcontext\nfrom typing import Dict, List, Optional, Tuple\n\nimport pandas as pd\nimport numpy as np\nimport torch\nfrom torch.utils.data import Dataset, DataLoader, Sampler\nfrom transformers import AutoTokenizer, AutoModelForSeq2SeqLM\nfrom tqdm.auto import tqdm\nimport sacrebleu\n\nwarnings.filterwarnings(\"ignore\")\nos.environ.setdefault(\"OMP_NUM_THREADS\", \"4\")\nos.environ.setdefault(\"MKL_NUM_THREADS\", \"4\")\nos.environ.setdefault(\"TOKENIZERS_PARALLELISM\", \"true\")\n\n# 1. Mixed precision helpers\n#\n# I use BF16 autocast whenever available.\n# For ByT5 inference this is generally safer than FP16 while still reducing memory pressure.\n\ndef _cuda_bf16_supported() -> bool:\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: torch.device, enabled: bool):\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\ndef _first_existing(paths: List[str], required: bool = True) -> Optional[str]:\n    for p in paths:\n        if p and Path(p).exists():\n            return p\n    if required:\n        raise FileNotFoundError(f\"No valid path found in candidates: {paths}\")\n    return None\n\n# 2. Configuration\n#\n# This configuration keeps the public notebook transparent:\n#\n# - two different ByT5 checkpoints,\n# - standard beam search,\n# - optional diverse beam search,\n# - multi-temperature sampling,\n# - weighted MBR reranking.\n\n@dataclass\nclass EnsembleMBRConfig:\n    test_data_path_candidates: List[str] = field(default_factory=lambda: [\n        \"/kaggle/input/competitions/deep-past-initiative-machine-translation/test.csv\",\n        \"/kaggle/input/deep-past-initiative-machine-translation/test.csv\",\n        \"datasets/test.csv\",\n    ])\n    model_a_path_candidates: List[str] = field(default_factory=lambda: [\n        \"/kaggle/input/datasets/assiaben/final-byt5/byt5-akkadian-optimized-34x\",\n        \"/kaggle/input/final-byt5/byt5-akkadian-optimized-34x\",\n    ])\n    model_b_path_candidates: List[str] = field(default_factory=lambda: [\n        \"/kaggle/input/models/mattiaangeli/byt5-akkadian-mbr-v2/pytorch/default/1\",\n        \"/kaggle/input/models/mattiaangeli/byt5-akkadian-mbr/pytorch/default/6\",\n    ])\n\n    test_data_path: str = \"\"\n    output_dir: str = \"/kaggle/working/\"\n    model_a_path: str = \"\"\n    model_b_path: str = \"\"\n\n    max_input_length: int = 512\n    max_new_tokens:   int = 768\n    batch_size:       int = 2\n    num_workers:      int = 2\n    num_buckets:      int = 6\n\n    num_beam_cands:      int = 10\n    num_beams:           int = 20\n    length_penalty:      float = 1.3\n    early_stopping:      bool = True\n    repetition_penalty:  float = 1.2\n\n    use_diverse_beam:    bool = False\n    num_diverse_cands:   int = 2\n    num_diverse_beams:   int = 8\n    num_beam_groups:     int = 4\n    diversity_penalty:   float = 0.8\n\n    use_sampling:        bool = True\n    sample_temperatures: List[float] = field(default_factory=lambda: [0.60, 0.80, 1.05])\n    num_sample_per_temp: int = 2\n    mbr_top_p:           float = 0.95\n\n    # Multi-prompt diversity (inspired by EMNLP 2024 multi-prompt MBR)\n    use_prompt_variants: bool = False # changed this from true to false\n    base_prompt: str = \"translate Akkadian to English: \"\n    prompt_variants: List[str] = field(default_factory=lambda: [\n        \"translate Old Assyrian to English: \",\n    ])\n    num_prompt_cands: int = 1\n\n    @property\n    def num_sample_cands(self) -> int:\n        return len(self.sample_temperatures) * self.num_sample_per_temp\n\n    mbr_pool_cap: int = 32\n\n    mbr_w_chrf:     float = 0.55\n    mbr_w_bleu:     float = 0.25\n    mbr_w_jaccard:  float = 0.20\n    mbr_w_length:   float = 0.10\n    mbr_w_support:  float = 0.08\n    mbr_support_alpha: float = 0.65\n\n    # CBP-style bootstrap pre-pruning (inspired by Cheng & Vlachos, 2023)\n    use_cbp_preprune: bool = True # changed this from true to false to test.\n    cbp_trigger_n: int = 20\n    cbp_bootstrap_rounds: int = 48\n    cbp_keep_n: int = 16\n\n    use_mixed_precision:     bool = True\n    use_better_transformer:  bool = True\n    use_bucket_batching:     bool = True\n    use_adaptive_beams:      bool = True\n    aggressive_postprocessing: bool = True\n    checkpoint_freq: int = 200\n\n    def __post_init__(self):\n        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n        Path(self.output_dir).mkdir(exist_ok=True, parents=True)\n\n        if not torch.cuda.is_available():\n            self.use_mixed_precision = False\n            self.use_better_transformer = False\n\n        self.use_bf16_amp = bool(\n            self.use_mixed_precision\n            and self.device.type == \"cuda\"\n            and _cuda_bf16_supported()\n        )\n\n        assert self.num_beams >= self.num_beam_cands, \"num_beams must be >= num_beam_cands\"\n\n        if self.use_diverse_beam:\n            assert self.num_diverse_beams % self.num_beam_groups == 0\n            assert self.num_diverse_beams >= self.num_diverse_cands\n\n    def resolve_paths(self):\n        self.test_data_path = _first_existing(self.test_data_path_candidates, required=True)\n        self.model_a_path = _first_existing(self.model_a_path_candidates, required=True)\n        self.model_b_path = _first_existing(self.model_b_path_candidates, required=True)\n\n# 3. Logging utilities\n\ndef setup_logging(output_dir: str) -> logging.Logger:\n    Path(output_dir).mkdir(exist_ok=True, parents=True)\n    for h in logging.root.handlers[:]:\n        logging.root.removeHandler(h)\n\n    logging.basicConfig(\n        level=logging.INFO,\n        format=\"%(asctime)s | %(levelname)s | %(message)s\",\n        handlers=[\n            logging.StreamHandler(),\n            logging.FileHandler(Path(output_dir) / \"ensemble_mbr.log\"),\n        ],\n    )\n    return logging.getLogger(\"ensemble_mbr\")\n\n# 4. Preprocessing\n#\n# This block normalizes Akkadian transliteration before tokenization.\n#\n# Key steps:\n#\n# - convert ASCII-style diacritic notation into normalized Unicode,\n# - normalize gap markers,\n# - standardize determinatives,\n# - normalize decimal/fraction expressions,\n# - clean whitespace and symbol variants.\n#\n# This part is important because small notation mismatches can fragment the candidate space and hurt MBR consistency.\n\n_V2 = re.compile(r\"([aAeEiIuU])(?:2|₂)\")\n_V3 = re.compile(r\"([aAeEiIuU])(?:3|₃)\")\n_ACUTE = str.maketrans({\"a\":\"á\",\"e\":\"é\",\"i\":\"í\",\"u\":\"ú\",\"A\":\"Á\",\"E\":\"É\",\"I\":\"Í\",\"U\":\"Ú\"})\n_GRAVE = str.maketrans({\"a\":\"à\",\"e\":\"è\",\"i\":\"ì\",\"u\":\"ù\",\"A\":\"À\",\"E\":\"È\",\"I\":\"Ì\",\"U\":\"Ù\"})\n\ndef _ascii_to_diacritics(s: str) -> str:\n    s = s.replace(\"sz\", \"š\").replace(\"SZ\", \"Š\")\n    s = s.replace(\"s,\", \"ṣ\").replace(\"S,\", \"Ṣ\")\n    s = s.replace(\"t,\", \"ṭ\").replace(\"T,\", \"Ṭ\")\n    s = _V2.sub(lambda m: m.group(1).translate(_ACUTE), s)\n    s = _V3.sub(lambda m: m.group(1).translate(_GRAVE), s)\n    return s\n\n_ALLOWED_FRACS = [\n    (1/6, \"0.16666\"), (1/4, \"0.25\"), (1/3, \"0.33333\"),\n    (1/2, \"0.5\"), (2/3, \"0.66666\"), (3/4, \"0.75\"), (5/6, \"0.83333\"),\n]\n_FRAC_TOL = 2e-3\n_FLOAT_RE = re.compile(r\"(?<![\\w/])(\\d+\\.\\d{4,})(?![\\w/])\")\n\ndef _canon_decimal(x: float) -> str:\n    ip = int(math.floor(x + 1e-12))\n    frac = x - ip\n    best = min(_ALLOWED_FRACS, key=lambda t: abs(frac - t[0]))\n    if abs(frac - best[0]) <= _FRAC_TOL:\n        dec = best[1]\n        if ip == 0:\n            return dec\n        return f\"{ip}{dec[1:]}\" if dec.startswith(\"0.\") else f\"{ip}+{dec}\"\n    return f\"{x:.5f}\".rstrip(\"0\").rstrip(\".\")\n\n_WS_RE = re.compile(r\"\\s+\")\n\n_GAP_UNIFIED_RE = re.compile(\n    r\"<\\s*big[\\s_\\-]*gap\\s*>\"\n    r\"|<\\s*gap\\s*>\"\n    r\"|\\bbig[\\s_\\-]*gap\\b\"\n    r\"|\\bx(?:\\s+x)+\\b\"\n    r\"|\\.{3,}|…+|\\[\\.+\\]\"\n    r\"|\\[\\s*x\\s*\\]|\\(\\s*x\\s*\\)\"\n    r\"|(?<!\\w)x{2,}(?!\\w)\"\n    r\"|(?<!\\w)x(?!\\w)\"\n    r\"|\\(\\s*large\\s+break\\s*\\)\"\n    r\"|\\(\\s*break\\s*\\)\"\n    r\"|\\(\\s*\\d+\\s+broken\\s+lines?\\s*\\)\",\n    re.I\n)\n\ndef _normalize_gaps_vec(ser: pd.Series) -> pd.Series:\n    return ser.str.replace(_GAP_UNIFIED_RE, \"<gap>\", regex=True)\n\n_CHAR_TRANS = str.maketrans({\n    \"ḫ\":\"h\",\"Ḫ\":\"H\",\"ʾ\":\"\",\n    \"₀\":\"0\",\"₁\":\"1\",\"₂\":\"2\",\"₃\":\"3\",\"₄\":\"4\",\n    \"₅\":\"5\",\"₆\":\"6\",\"₇\":\"7\",\"₈\":\"8\",\"₉\":\"9\",\n    \"—\":\"-\",\"–\":\"-\",\n})\n_SUB_X = \"ₓ\"\n\n_UNICODE_UPPER = r\"A-ZŠṬṢḪ\\u00C0-\\u00D6\\u00D8-\\u00DE\\u0160\\u1E00-\\u1EFF\"\n_UNICODE_LOWER = r\"a-zšṭṣḫ\\u00E0-\\u00F6\\u00F8-\\u00FF\\u0161\\u1E01-\\u1EFF\"\n_DET_UPPER_RE = re.compile(r\"\\(([\" + _UNICODE_UPPER + r\"0-9]{1,6})\\)\")\n_DET_LOWER_RE = re.compile(r\"\\(([\" + _UNICODE_LOWER + r\"]{1,4})\\)\")\n\n_PN_RE = re.compile(r\"\\bPN\\b\")\n_KUBABBAR_RE = re.compile(r\"KÙ\\.B\\.\")\n\n_EXACT_FRAC_RE = re.compile(r\"0\\.8333|0\\.6666|0\\.3333|0\\.1666|0\\.625|0\\.75|0\\.25|0\\.5\")\n_EXACT_FRAC_MAP = {\n    \"0.8333\": \"⅚\", \"0.6666\": \"⅔\", \"0.3333\": \"⅓\", \"0.1666\": \"⅙\",\n    \"0.625\": \"⅝\", \"0.75\": \"¾\", \"0.25\": \"¼\", \"0.5\": \"½\",\n}\n\ndef _frac_repl(m: re.Match) -> str:\n    return _EXACT_FRAC_MAP[m.group(0)]\n\nclass OptimizedPreprocessor:\n    def preprocess_batch(self, texts: List[str]) -> List[str]:\n        ser = pd.Series(texts).fillna(\"\").astype(str)\n        ser = ser.apply(_ascii_to_diacritics)\n        ser = ser.str.replace(_DET_UPPER_RE, r\"\\1\", regex=True)\n        ser = ser.str.replace(_DET_LOWER_RE, r\"{\\1}\", regex=True)\n        ser = _normalize_gaps_vec(ser)\n        ser = ser.str.translate(_CHAR_TRANS)\n        ser = ser.str.replace(_SUB_X, \"\", regex=False)\n        ser = ser.str.replace(_KUBABBAR_RE, \"KÙ.BABBAR\", regex=True)\n        ser = ser.str.replace(_EXACT_FRAC_RE, _frac_repl, regex=True)\n        ser = ser.str.replace(_FLOAT_RE, lambda m: _canon_decimal(float(m.group(1))), regex=True)\n        ser = ser.str.replace(_WS_RE, \" \", regex=True).str.strip()\n        return ser.tolist()\n\n# 5. Postprocessing\n#\n# This block is where a lot of the practical gain comes from.\n#\n# The idea is not to aggressively rewrite the model output, but to remove avoidable formatting noise so that MBR compares candidates on their *meaningful content* rather than superficial inconsistencies.\n#\n# I keep `/`, straight quotes, parentheses, and apostrophes when they may be meaningful.\n#\n# **Key corrections in this version:**\n# - Parentheses `()` are **preserved** (they appear in the test set)\n# - Curly quotes are **converted to straight quotes** (not deleted)\n# - Commodity replacements only match space-prefixed hyphens\n# - 5/12 shekel correctly maps to ⅓ shekel 15 grains\n# - ḫ/Ḫ characters normalized to h/H\n\n_SOFT_GRAM_RE = re.compile(\n    r\"\\(\\s*(?:fem|plur|pl|sing|singular|plural|\\?|\\!)\"\n    r\"(?:\\.\\s*(?:plur|plural|sing|singular))?\"\n    r\"\\.?\\s*[^)]*\\)\", re.I\n)\n_BARE_GRAM_RE = re.compile(r\"(?<!\\w)(?:fem|sing|pl|plural)\\.?(?!\\w)\\s*\", re.I)\n_UNCERTAIN_RE = re.compile(r\"\\(\\?\\)\")\n# FIX #4: curly quotes → straight quotes (not deletion)\n_CURLY_DQ_RE = re.compile(\"[\\u201c\\u201d]\")   # \" \" → \"\n_CURLY_SQ_RE = re.compile(\"[\\u2018\\u2019]\")   # ' ' → '\n_MONTH_RE = re.compile(r\"\\bMonth\\s+(XII|XI|X|IX|VIII|VII|VI|V|IV|III|II|I)\\b\", re.I)\n_ROMAN2INT = {\"I\":1,\"II\":2,\"III\":3,\"IV\":4,\"V\":5,\"VI\":6,\"VII\":7,\"VIII\":8,\"IX\":9,\"X\":10,\"XI\":11,\"XII\":12}\n_REPEAT_WORD_RE = re.compile(r\"\\b(\\w+)(?:\\s+\\1\\b)+\")\n_REPEAT_PUNCT_RE = re.compile(r\"([.,])\\1+\")\n_PUNCT_SPACE_RE = re.compile(r\"\\s+([.,:])\") \n\n# FIX #3: Removed '(' and ')' from forbidden chars — host confirmed\n# parentheses appear in the test set, e.g. \"(silver)\", \"(and)\"\n_FORBIDDEN_TRANS = str.maketrans(\"\", \"\", '——<>⌈⌋⌊[]+ʾ;')\n\n# FIX #2: Commodity regex now requires a SPACE before the hyphen,\n# so \"import-tax\" and \"kutānu-textiles\" are NOT corrupted.\n# Only standalone \" -gold\", \" -tax\", \" -textiles\" are replaced.\n_COMMODITY_RE = re.compile(r'(?<=\\s)-(gold|tax|textiles)\\b')\n_COMMODITY_REPL = {\"gold\": \"pašallum gold\", \"tax\": \"šadduātum tax\", \"textiles\": \"kutānum textiles\"}\n\ndef _commodity_repl(m: re.Match) -> str:\n    return _COMMODITY_REPL[m.group(1)]\n\n# FIX #1: 5/12 shekel → ⅓ shekel 15 grains (was ⅔, host confirmed)\n_SHEKEL_REPLS = [\n    (re.compile(r'5\\s+11\\s*/\\s*12\\s+shekels?', re.I), '6 shekels less 15 grains'),\n    (re.compile(r'5\\s*/\\s*12\\s+shekels?', re.I), '⅓ shekel 15 grains'),\n    (re.compile(r'7\\s*/\\s*12\\s+shekels?', re.I), '½ shekel 15 grains'),\n    (re.compile(r'1\\s*/\\s*12\\s*(?:\\(shekel\\)|\\bshekel)?', re.I), '15 grains'),\n]\n\n# FIX #6: Safer slash-alternative regex — require word chars around the slash\n# and explicitly exclude digit/digit patterns (fractions like 1/3)\n_SLASH_ALT_RE = re.compile(r'(?<![0-9/])\\s+/\\s+(?![0-9])\\S+')\n\n_STRAY_MARKS_RE = re.compile(r'<<[^>]*>>|<(?!gap\\b)[^>]*>')\n_MULTI_GAP_RE = re.compile(r'(?:<gap>\\s*){2,}')\n\n# FIX #7: Additional stray marks the host recommends removing\n_EXTRA_STRAY_RE = re.compile(r'(?<!\\w)(?:\\.\\.+|xx+)(?!\\w)')\n\n# FIX #5: ḫ/Ḫ normalization for translations (host confirmed not in test set)\n_HACEK_TRANS = str.maketrans({\"ḫ\": \"h\", \"Ḫ\": \"H\"})\n\ndef _month_repl(m: re.Match) -> str:\n    return f\"Month {_ROMAN2INT.get(m.group(1).upper(), m.group(1))}\"\n\nclass VectorizedPostprocessor:\n    def postprocess_batch(self, translations: List[str]) -> List[str]:\n        s = pd.Series(translations).fillna(\"\").astype(str)\n        s = _normalize_gaps_vec(s)\n        s = s.str.replace(_PN_RE, \"<gap>\", regex=True)\n        s = s.str.replace(_COMMODITY_RE, _commodity_repl, regex=True)\n\n        for pat, repl in _SHEKEL_REPLS:\n            s = s.str.replace(pat, repl, regex=True)\n\n        s = s.str.replace(_EXACT_FRAC_RE, _frac_repl, regex=True)\n        s = s.str.replace(_FLOAT_RE, lambda m: _canon_decimal(float(m.group(1))), regex=True)\n        s = s.str.replace(_SOFT_GRAM_RE, \" \", regex=True)\n        s = s.str.replace(_BARE_GRAM_RE, \" \", regex=True)\n        s = s.str.replace(_UNCERTAIN_RE, \"\", regex=True)\n        s = s.str.replace(_STRAY_MARKS_RE, \"\", regex=True)\n        s = s.str.replace(_EXTRA_STRAY_RE, \"\", regex=True)       # FIX #7\n        s = s.str.replace(_SLASH_ALT_RE, \"\", regex=True)\n\n        # FIX #4: curly → straight quotes (not deletion)\n        s = s.str.replace(_CURLY_DQ_RE, '\"', regex=True)\n        s = s.str.replace(_CURLY_SQ_RE, \"'\", regex=True)\n\n        s = s.str.replace(_MONTH_RE, _month_repl, regex=True)\n        s = s.str.replace(_MULTI_GAP_RE, \"<gap>\", regex=True)\n\n        # Protect <gap> markers before stripping forbidden chars\n        s = s.str.replace(\"<gap>\", \"\\x00GAP\\x00\", regex=False)\n        s = s.str.translate(_FORBIDDEN_TRANS)\n        s = s.str.replace(\"\\x00GAP\\x00\", \" <gap> \", regex=False)\n\n        # FIX #5: ḫ/Ḫ → h/H in translations\n        s = s.str.translate(_HACEK_TRANS)\n\n        s = s.str.replace(_REPEAT_WORD_RE, r\"\\1\", regex=True)\n        for n in range(4, 1, -1):\n            pat = r\"\\b((?:\\w+\\s+){\" + str(n-1) + r\"}\\w+)(?:\\s+\\1\\b)+\"\n            s = s.str.replace(pat, r\"\\1\", regex=True)\n\n        s = s.str.replace(_PUNCT_SPACE_RE, r\"\\1\", regex=True)\n        s = s.str.replace(_REPEAT_PUNCT_RE, r\"\\1\", regex=True)\n        s = s.str.replace(_WS_RE, \" \", regex=True).str.strip()\n        return s.tolist()\n\n# 6. Dataset and bucket batching\n#\n# Bucketed batching reduces padding waste and keeps throughput more stable for variable-length transliterations.\n\nclass AkkadianDataset(Dataset):\n    def __init__(self, df: pd.DataFrame, preprocessor: OptimizedPreprocessor, logger: logging.Logger, prompt: str):\n        self.sample_ids = df[\"id\"].tolist()\n        proc = preprocessor.preprocess_batch(df[\"transliteration\"].tolist())\n        self.input_texts = [prompt + t for t in proc]\n        logger.info(f\"Dataset: {len(self.sample_ids)} samples\")\n\n    def __len__(self):\n        return len(self.sample_ids)\n\n    def __getitem__(self, idx):\n        return self.sample_ids[idx], self.input_texts[idx]\n\nclass BucketBatchSampler(Sampler):\n    def __init__(self, dataset, batch_size, num_buckets, logger, shuffle=False):\n        self.batch_size = batch_size\n        self.shuffle = shuffle\n        lengths = [len(t.split()) for _, t in dataset]\n        sorted_idx = sorted(range(len(lengths)), key=lambda i: lengths[i])\n        bsize = max(1, len(sorted_idx) // max(1, num_buckets))\n        self.buckets = [\n            sorted_idx[i*bsize : None if i == num_buckets-1 else (i+1)*bsize]\n            for i in range(num_buckets)\n        ]\n\n        for i, b in enumerate(self.buckets):\n            if b:\n                bl = [lengths[x] for x in b]\n                logger.info(f\"  Bucket {i}: {len(b)} samples, len [{min(bl)}, {max(bl)}]\")\n\n    def __iter__(self):\n        for bucket in self.buckets:\n            b = list(bucket)\n            if self.shuffle:\n                random.shuffle(b)\n            for i in range(0, len(b), self.batch_size):\n                yield b[i:i+self.batch_size]\n\n    def __len__(self):\n        return sum(math.ceil(len(b) / self.batch_size) for b in self.buckets)\n\n# 7. Model wrapper and candidate generation\n#\n# This is the main inference change compared with the earlier public version.\n#\n# Candidate sources\n#\n# For each model and each sample, the pool can include:\n#\n# - **standard beam candidates**\n# - **diverse beam candidates** (optional)\n# - **multi-temperature sampled candidates**\n#\n# In my current public LB 34.9 configuration:\n#\n# - diverse beam is **prepared but disabled**\n# - multi-temperature sampling is **enabled**\n# - the effective per-model pool is:\n#\n# `4 beam + 6 sampled = 10 candidates`\n#\n# Across two models, that gives about **20 candidates before deduplication**.\n#\n# Why this helped\n#\n# The previous version already benefited from cross-model pooling.\n# This version improved further because the candidate pool became more heterogeneous, which gives MBR a better chance to select a consensus output that is both faithful and fluent.\n\nclass ModelWrapper:\n    def __init__(self, model_path: str, cfg: EnsembleMBRConfig, logger: logging.Logger, label: str):\n        self.cfg = cfg\n        self.logger = logger\n        self.label = label\n        self.tokenizer = None\n        self.model = None\n\n        logger.info(f\"[{label}] Loading from {model_path}\")\n        self.tokenizer = AutoTokenizer.from_pretrained(model_path)\n        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_path).to(cfg.device).eval()\n\n        if cfg.device.type == \"cuda\":\n            try:\n                torch.set_float32_matmul_precision(\"high\")\n            except Exception:\n                pass\n\n        n = sum(p.numel() for p in self.model.parameters())\n        logger.info(f\"[{label}] {n:,} parameters\")\n\n        if cfg.device.type == \"cuda\":\n            used = torch.cuda.memory_allocated() / 1e9\n            logger.info(f\"[{label}] GPU mem used: {used:.2f} GB\")\n\n        if cfg.use_better_transformer and cfg.device.type == \"cuda\":\n            try:\n                from optimum.bettertransformer import BetterTransformer\n                self.model = BetterTransformer.transform(self.model)\n                logger.info(f\"[{label}] BetterTransformer applied\")\n            except Exception as e:\n                logger.warning(f\"[{label}] BetterTransformer skipped: {e}\")\n\n    def collate(self, batch_samples):\n        ids = [s[0] for s in batch_samples]\n        texts = [s[1] for s in batch_samples]\n        enc = self.tokenizer(\n            texts,\n            max_length=self.cfg.max_input_length,\n            padding=True,\n            truncation=True,\n            return_tensors=\"pt\",\n        )\n        return ids, texts, enc\n\n    def generate_candidates(self, input_ids, attention_mask, input_texts: List[str], beam_size: int) -> List[List[str]]:\n        cfg = self.cfg\n        B = input_ids.shape[0]\n        ctx = _bf16_ctx(cfg.device, cfg.use_bf16_amp)\n\n        Rb = cfg.num_beam_cands\n        Rd = cfg.num_diverse_cands if cfg.use_diverse_beam else 0\n        Rs = cfg.num_sample_per_temp\n\n        with ctx:\n            nb = max(beam_size, Rb)\n            beam_out = self.model.generate(\n                input_ids=input_ids,\n                attention_mask=attention_mask,\n                do_sample=False,\n                num_beams=nb,\n                num_return_sequences=Rb,\n                max_new_tokens=cfg.max_new_tokens,\n                length_penalty=cfg.length_penalty,\n                early_stopping=cfg.early_stopping,\n                repetition_penalty=cfg.repetition_penalty,\n                use_cache=True,\n            )\n            beam_texts = self.tokenizer.batch_decode(beam_out, skip_special_tokens=True)\n\n            diverse_texts = []\n            actual_Rd = 0\n            if cfg.use_diverse_beam:\n                try:\n                    div_out = self.model.generate(\n                        input_ids=input_ids,\n                        attention_mask=attention_mask,\n                        do_sample=False,\n                        num_beams=cfg.num_diverse_beams,\n                        num_beam_groups=cfg.num_beam_groups,\n                        diversity_penalty=cfg.diversity_penalty,\n                        num_return_sequences=cfg.num_diverse_cands,\n                        max_new_tokens=cfg.max_new_tokens,\n                        length_penalty=cfg.length_penalty,\n                        early_stopping=cfg.early_stopping,\n                        repetition_penalty=cfg.repetition_penalty,\n                        use_cache=True,\n                    )\n                    diverse_texts = self.tokenizer.batch_decode(div_out, skip_special_tokens=True)\n                    actual_Rd = cfg.num_diverse_cands\n                except Exception as e:\n                    self.logger.warning(f\"[{self.label}] Diverse beam failed ({e}), skipping\")\n\n            all_samp_texts = []\n            num_temps = 0\n            if cfg.use_sampling and cfg.sample_temperatures:\n                num_temps = len(cfg.sample_temperatures)\n                for temp in cfg.sample_temperatures:\n                    try:\n                        samp_out = self.model.generate(\n                            input_ids=input_ids,\n                            attention_mask=attention_mask,\n                            do_sample=True,\n                            num_beams=1,\n                            top_p=cfg.mbr_top_p,\n                            temperature=temp,\n                            num_return_sequences=Rs,\n                            max_new_tokens=cfg.max_new_tokens,\n                            repetition_penalty=cfg.repetition_penalty,\n                            use_cache=True,\n                        )\n                        temp_texts = self.tokenizer.batch_decode(samp_out, skip_special_tokens=True)\n                        all_samp_texts.extend(temp_texts)\n                    except Exception as e:\n                        self.logger.warning(\n                            f\"[{self.label}] Sampling @ temp={temp:.2f} failed ({e}), padding with empty strings\"\n                        )\n                        all_samp_texts.extend([\"\"] * (B * Rs))\n\n            prompt_variant_texts = []\n            num_prompt_variants = 0\n            Rp = int(max(0, cfg.num_prompt_cands))\n            if cfg.use_prompt_variants and cfg.prompt_variants and Rp > 0:\n                # Generate low-count deterministic candidates from alternate prompts.\n                for prompt in cfg.prompt_variants:\n                    try:\n                        alt_texts = []\n                        base = cfg.base_prompt\n                        for t in input_texts:\n                            if isinstance(t, str) and t.startswith(base):\n                                alt_texts.append(prompt + t[len(base):])\n                            else:\n                                alt_texts.append(prompt + str(t))\n\n                        alt_enc = self.tokenizer(\n                            alt_texts,\n                            max_length=cfg.max_input_length,\n                            padding=True,\n                            truncation=True,\n                            return_tensors=\"pt\",\n                        )\n                        alt_input_ids = alt_enc.input_ids.to(cfg.device, non_blocking=True)\n                        alt_attn = alt_enc.attention_mask.to(cfg.device, non_blocking=True)\n\n                        alt_out = self.model.generate(\n                            input_ids=alt_input_ids,\n                            attention_mask=alt_attn,\n                            do_sample=False,\n                            num_beams=max(beam_size, Rp),\n                            num_return_sequences=Rp,\n                            max_new_tokens=cfg.max_new_tokens,\n                            length_penalty=cfg.length_penalty,\n                            early_stopping=cfg.early_stopping,\n                            repetition_penalty=cfg.repetition_penalty,\n                            use_cache=True,\n                        )\n                        prompt_variant_texts.extend(self.tokenizer.batch_decode(alt_out, skip_special_tokens=True))\n                        num_prompt_variants += 1\n                    except Exception as e:\n                        self.logger.warning(f\"[{self.label}] Prompt variant failed ({e}), skipping\")\n\n        pools = []\n        for i in range(B):\n            pool = []\n            pool.extend(beam_texts[i * Rb : (i + 1) * Rb])\n\n            if diverse_texts and actual_Rd > 0:\n                pool.extend(diverse_texts[i * actual_Rd : (i + 1) * actual_Rd])\n\n            if all_samp_texts and num_temps > 0:\n                for t_idx in range(num_temps):\n                    seg_start = t_idx * B * Rs + i * Rs\n                    pool.extend(all_samp_texts[seg_start : seg_start + Rs])\n\n            if prompt_variant_texts and num_prompt_variants > 0:\n                for p_idx in range(num_prompt_variants):\n                    seg_start = p_idx * B * Rp + i * Rp\n                    pool.extend(prompt_variant_texts[seg_start : seg_start + Rp])\n\n            pools.append(pool)\n\n        if pools:\n            prompt_total = num_prompt_variants * Rp\n            self.logger.info(\n                f\"[{self.label}] Pool per sample: \"\n                f\"beam={Rb} + diverse={actual_Rd} + sample={num_temps}x{Rs}={num_temps*Rs} + \"\n                f\"prompt={num_prompt_variants}x{Rp}={prompt_total} \"\n                f\"= {len(pools[0])} total\"\n            )\n\n        return pools\n\n    def unload(self):\n        label = self.label\n        try:\n            from optimum.bettertransformer import BetterTransformer\n            self.model = BetterTransformer.reverse(self.model)\n        except Exception:\n            pass\n\n        del self.model\n        del self.tokenizer\n        self.model = None\n        self.tokenizer = None\n\n        gc.collect()\n        if torch.cuda.is_available():\n            torch.cuda.empty_cache()\n            torch.cuda.synchronize()\n            free = (torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_allocated()) / 1e9\n            self.logger.info(f\"[{label}] Unloaded. GPU free: {free:.2f} GB\")\n\n# 8. Weighted MBR selector\n#\n# This is the second major improvement.\n#\n# Earlier, I used **chrF++ only**.\n# In this version, I use a weighted combination of:\n#\n# - **chrF++** for character-level robustness,\n# - **BLEU** for word-level precision,\n# - **Jaccard** for fast token overlap regularization,\n# - **length bonus** to reduce pathological outputs.\n#\n# This is still a simple MBR implementation, but it is noticeably stronger in practice than relying on a single metric.\n\nclass MBRSelector:\n    def __init__(\n        self,\n        pool_cap: int = 32,\n        w_chrf: float = 0.55,\n        w_bleu: float = 0.25,\n        w_jaccard: float = 0.20,\n        w_length: float = 0.10,\n        w_support: float = 0.08,\n        support_alpha: float = 0.65,\n        use_cbp_preprune: bool = True,\n        cbp_trigger_n: int = 20,\n        cbp_bootstrap_rounds: int = 48,\n        cbp_keep_n: int = 16,\n    ):\n        self._chrf_metric = sacrebleu.metrics.CHRF(word_order=2)\n        self._bleu_metric = sacrebleu.metrics.BLEU(effective_order=True)\n        self.pool_cap = pool_cap\n        self.w_chrf = w_chrf\n        self.w_bleu = w_bleu\n        self.w_jaccard = w_jaccard\n        self.w_length = w_length\n        self.w_support = w_support\n        self.support_alpha = support_alpha\n        self.use_cbp_preprune = use_cbp_preprune\n        self.cbp_trigger_n = cbp_trigger_n\n        self.cbp_bootstrap_rounds = cbp_bootstrap_rounds\n        self.cbp_keep_n = cbp_keep_n\n        self._pw_total = max(w_chrf + w_bleu + w_jaccard, 1e-9)\n\n    def _chrfpp(self, a: str, b: str) -> float:\n        if not a or not b:\n            return 0.0\n        return float(self._chrf_metric.sentence_score(a, [b]).score)\n\n    def _bleu(self, a: str, b: str) -> float:\n        if not a or not b:\n            return 0.0\n        try:\n            return float(self._bleu_metric.sentence_score(a, [b]).score)\n        except Exception:\n            return 0.0\n\n    @staticmethod\n    def _jaccard(a: str, b: str) -> float:\n        ta, tb = set(a.lower().split()), set(b.lower().split())\n        if not ta and not tb:\n            return 100.0\n        if not ta or not tb:\n            return 0.0\n        return 100.0 * len(ta & tb) / len(ta | tb)\n\n    def _pairwise_score(self, a: str, b: str) -> float:\n        s = (\n            self.w_chrf * self._chrfpp(a, b)\n            + self.w_bleu * self._bleu(a, b)\n            + self.w_jaccard * self._jaccard(a, b)\n        )\n        return s / self._pw_total\n\n    @staticmethod\n    def _length_bonus(lengths: List[int], idx: int) -> float:\n        if len(lengths) == 0:\n            return 100.0\n        median = float(np.median(lengths))\n        sigma = max(median * 0.4, 5.0)\n        z = (lengths[idx] - median) / sigma\n        return 100.0 * math.exp(-0.5 * z * z)\n\n    @staticmethod\n    def _dedup_with_counts(xs: List[str]) -> Tuple[List[str], List[int]]:\n        freq: Dict[str, int] = {}\n        order: List[str] = []\n        for x in xs:\n            x = str(x).strip()\n            if not x:\n                continue\n            if x not in freq:\n                freq[x] = 0\n                order.append(x)\n            freq[x] += 1\n        counts = [freq[x] for x in order]\n        return order, counts\n\n    def _cbp_preprune(self, cands: List[str], counts: List[int]) -> Tuple[List[str], List[int]]:\n        n = len(cands)\n        if (\n            not self.use_cbp_preprune\n            or n <= max(2, self.cbp_trigger_n)\n            or self.cbp_keep_n <= 0\n            or self.cbp_keep_n >= n\n        ):\n            return cands, counts\n\n        jmat = np.zeros((n, n), dtype=np.float32)\n        for i in range(n):\n            for j in range(i + 1, n):\n                s = float(self._jaccard(cands[i], cands[j]))\n                jmat[i, j] = s\n                jmat[j, i] = s\n\n        rng = np.random.default_rng(42)\n        wins = np.zeros(n, dtype=np.float64)\n        rounds = max(8, int(self.cbp_bootstrap_rounds))\n        for _ in range(rounds):\n            refs = rng.integers(0, n, size=n)\n            util = jmat[:, refs].mean(axis=1)\n            wins[int(np.argmax(util))] += 1.0\n\n        conf = wins / max(float(rounds), 1.0)\n        cnt = np.asarray(counts, dtype=np.float64)\n        sup = np.power(np.maximum(cnt, 1.0), self.support_alpha)\n        sup = sup / max(float(np.max(sup)), 1e-9)\n        final = 0.75 * conf + 0.25 * sup\n\n        keep = min(n, max(2, int(self.cbp_keep_n)))\n        idx = np.argsort(-final)[:keep]\n        idx = sorted(idx.tolist())\n        return [cands[i] for i in idx], [counts[i] for i in idx]\n\n    def pick(self, candidates: List[str]) -> str:\n        cands, counts = self._dedup_with_counts(candidates)\n        if self.pool_cap:\n            cands = cands[:self.pool_cap]\n            counts = counts[:self.pool_cap]\n\n        cands, counts = self._cbp_preprune(cands, counts)\n\n        n = len(cands)\n        if n == 0:\n            return \"\"\n        if n == 1:\n            return cands[0]\n\n        lengths = [len(c.split()) for c in cands]\n        cnt = np.asarray(counts, dtype=np.float64)\n        support = np.power(np.maximum(cnt, 1.0), self.support_alpha)\n        scores = []\n\n        for i in range(n):\n            num = 0.0\n            den = 0.0\n            for j in range(n):\n                if j == i:\n                    continue\n                w = float(support[j])\n                num += w * self._pairwise_score(cands[i], cands[j])\n                den += w\n            pw = num / max(den, 1e-9)\n\n            lb = self._length_bonus(lengths, i)\n            sb = 100.0 * float(support[i] / max(np.max(support), 1e-9))\n            total = pw + self.w_length * lb + self.w_support * sb\n            scores.append(total)\n\n        return cands[int(np.argmax(scores))]\n\n# 9. End-to-end engine\n\nclass EnsembleMBREngine:\n    def __init__(self, cfg: EnsembleMBRConfig, logger: logging.Logger):\n        self.cfg = cfg\n        self.logger = logger\n        self.preprocessor = OptimizedPreprocessor()\n        self.postprocessor = VectorizedPostprocessor()\n        self.mbr = MBRSelector(\n            pool_cap=cfg.mbr_pool_cap,\n            w_chrf=cfg.mbr_w_chrf,\n            w_bleu=cfg.mbr_w_bleu,\n            w_jaccard=cfg.mbr_w_jaccard,\n            w_length=cfg.mbr_w_length,\n            w_support=cfg.mbr_w_support,\n            support_alpha=cfg.mbr_support_alpha,\n            use_cbp_preprune=cfg.use_cbp_preprune,\n            cbp_trigger_n=cfg.cbp_trigger_n,\n            cbp_bootstrap_rounds=cfg.cbp_bootstrap_rounds,\n            cbp_keep_n=cfg.cbp_keep_n,\n        )\n\n    def _adaptive_beams(self, attn: torch.Tensor) -> int:\n        if not self.cfg.use_adaptive_beams:\n            return self.cfg.num_beams\n        med = float(attn.sum(dim=1).float().median().item())\n        short = max(self.cfg.num_beam_cands, self.cfg.num_beams // 2)\n        return short if med < 100 else self.cfg.num_beams\n\n    def _build_dataloader(self, dataset: AkkadianDataset, wrapper: ModelWrapper) -> DataLoader:\n        if self.cfg.use_bucket_batching:\n            sampler = BucketBatchSampler(dataset, self.cfg.batch_size, self.cfg.num_buckets, self.logger)\n            return DataLoader(\n                dataset,\n                batch_sampler=sampler,\n                num_workers=self.cfg.num_workers,\n                collate_fn=wrapper.collate,\n                pin_memory=(self.cfg.device.type == \"cuda\"),\n            )\n\n        return DataLoader(\n            dataset,\n            batch_size=self.cfg.batch_size,\n            shuffle=False,\n            num_workers=self.cfg.num_workers,\n            collate_fn=wrapper.collate,\n            pin_memory=(self.cfg.device.type == \"cuda\"),\n        )\n\n    def _run_one_model(self, wrapper: ModelWrapper, dataset: AkkadianDataset) -> dict:\n        dl = self._build_dataloader(dataset, wrapper)\n        pools_by_id = {}\n\n        with torch.inference_mode():\n            for batch_ids, batch_texts, enc in tqdm(dl, desc=f\"  [{wrapper.label}]\"):\n                input_ids = enc.input_ids.to(self.cfg.device, non_blocking=True)\n                attn = enc.attention_mask.to(self.cfg.device, non_blocking=True)\n                beam_size = self._adaptive_beams(attn)\n\n                try:\n                    batch_pools = wrapper.generate_candidates(input_ids, attn, batch_texts, beam_size)\n                    for sid, pool in zip(batch_ids, batch_pools):\n                        pools_by_id[str(sid)] = pool\n                except RuntimeError as e:\n                    if \"out of memory\" in str(e).lower():\n                        self.logger.error(f\"OOM in [{wrapper.label}] — skipping batch\")\n                        torch.cuda.empty_cache()\n                        for sid in batch_ids:\n                            pools_by_id.setdefault(str(sid), [])\n                    else:\n                        raise\n                except Exception as e:\n                    self.logger.error(f\"[{wrapper.label}] batch error: {e}\")\n                    for sid in batch_ids:\n                        pools_by_id.setdefault(str(sid), [])\n\n                if torch.cuda.is_available():\n                    torch.cuda.empty_cache()\n\n        return pools_by_id\n\n    def run(self, test_df: pd.DataFrame) -> pd.DataFrame:\n        cfg, logger = self.cfg, self.logger\n\n        cands_per_sample = (\n            cfg.num_beam_cands\n            + (cfg.num_diverse_cands if cfg.use_diverse_beam else 0)\n            + cfg.num_sample_cands\n            + (len(cfg.prompt_variants) * cfg.num_prompt_cands if cfg.use_prompt_variants else 0)\n        )\n\n        logger.info(\"=\" * 60)\n        logger.info(\"Ensemble × MBR  |  Cross-model candidate pooling  v2\")\n        logger.info(f\"  Model A           : {cfg.model_a_path}\")\n        logger.info(f\"  Model B           : {cfg.model_b_path}\")\n        logger.info(f\"  Standard beam     : {cfg.num_beam_cands} (num_beams={cfg.num_beams})\")\n        logger.info(\n            f\"  Diverse beam      : {'ON' if cfg.use_diverse_beam else 'OFF'} \"\n            f\"cands={cfg.num_diverse_cands}, groups={cfg.num_beam_groups}, penalty={cfg.diversity_penalty}\"\n        )\n        logger.info(\n            f\"  Multi-temp sample : {'ON' if cfg.use_sampling else 'OFF'} \"\n            f\"temps={cfg.sample_temperatures}, {cfg.num_sample_per_temp}/temp \"\n            f\"→ {cfg.num_sample_cands} total\"\n        )\n        logger.info(\n            f\"  Multi-prompt      : {'ON' if cfg.use_prompt_variants else 'OFF'} \"\n            f\"variants={len(cfg.prompt_variants)}, {cfg.num_prompt_cands}/variant \"\n            f\"→ {len(cfg.prompt_variants) * cfg.num_prompt_cands if cfg.use_prompt_variants else 0} total\"\n        )\n        logger.info(\n            f\"  CBP pre-prune     : {'ON' if cfg.use_cbp_preprune else 'OFF'} \"\n            f\"(trigger={cfg.cbp_trigger_n}, keep={cfg.cbp_keep_n}, rounds={cfg.cbp_bootstrap_rounds})\"\n        )\n        logger.info(\n            f\"  Cands/model/sample: {cands_per_sample} \"\n            f\"→ pool ≈ {cands_per_sample * 2} (2 models, pre-dedup)\"\n        )\n        logger.info(f\"  MBR pool cap      : {cfg.mbr_pool_cap}\")\n        logger.info(f\"  BF16 AMP          : {cfg.use_bf16_amp}\")\n        logger.info(f\"  batch_size        : {cfg.batch_size}\")\n        logger.info(\"=\" * 60)\n\n        dataset = AkkadianDataset(test_df, self.preprocessor, logger, prompt=cfg.base_prompt)\n        sample_ids = [str(s) for s in dataset.sample_ids]\n\n        logger.info(\"Phase 1/2 — Model A inference\")\n        wrapper_a = ModelWrapper(cfg.model_a_path, cfg, logger, \"Model-A\")\n        pools_a = self._run_one_model(wrapper_a, dataset)\n        wrapper_a.unload()\n        del wrapper_a\n\n        logger.info(\"Phase 2/2 — Model B inference\")\n        wrapper_b = ModelWrapper(cfg.model_b_path, cfg, logger, \"Model-B\")\n        pools_b = self._run_one_model(wrapper_b, dataset)\n        wrapper_b.unload()\n        del wrapper_b\n\n        logger.info(\"Phase 3/3 — Pool merge + MBR selection\")\n        results = []\n\n        for sid in tqdm(sample_ids, desc=\"  MBR\"):\n            combined = pools_a.get(sid, []) + pools_b.get(sid, [])\n            pp = self.postprocessor.postprocess_batch(combined) if combined else []\n            chosen = self.mbr.pick(pp)\n\n            if not chosen or not chosen.strip():\n                chosen = \"The tablet is too damaged to translate.\"\n\n            results.append((sid, chosen))\n\n            if len(results) % cfg.checkpoint_freq == 0:\n                ckpt = Path(cfg.output_dir) / f\"checkpoint_{len(results)}.csv\"\n                pd.DataFrame(results, columns=[\"id\", \"translation\"]).to_csv(ckpt, index=False)\n                logger.info(f\"  Checkpoint: {len(results)} rows → {ckpt}\")\n\n        result_df = pd.DataFrame(results, columns=[\"id\", \"translation\"])\n        self._validate(result_df)\n        return result_df\n\n    def _validate(self, df: pd.DataFrame):\n        logger = self.logger\n        logger.info(\"=\" * 60)\n        empty = df[\"translation\"].str.strip().eq(\"\").sum()\n        lens = df[\"translation\"].str.len()\n        logger.info(f\"Empty     : {empty} ({100*empty/max(1,len(df)):.2f}%)\")\n        logger.info(f\"Len mean  : {lens.mean():.1f}  median: {lens.median():.1f}  min: {lens.min()}  max: {lens.max()}\")\n\n        for idx in [0, len(df)//4, len(df)//2, 3*len(df)//4, len(df)-1]:\n            row = df.iloc[idx]\n            logger.info(f\"  ID {row['id']}: {str(row['translation'])[:80]}\")\n\n        logger.info(\"=\" * 60)\n\n# 10. Environment summary helper\n\ndef print_env(cfg: EnsembleMBRConfig):\n    print(f\"PyTorch  : {torch.__version__}\")\n    print(f\"CUDA     : {torch.cuda.is_available()}\")\n    if torch.cuda.is_available():\n        print(f\"GPU      : {torch.cuda.get_device_name(0)}\")\n        mem = torch.cuda.get_device_properties(0).total_memory / 1e9\n        print(f\"GPU Mem  : {mem:.2f} GB\")\n        print(f\"BF16     : {_cuda_bf16_supported()}\")\n    print(f\"BF16 AMP : {cfg.use_bf16_amp}\")\n    print()\n\n    cands_per_sample = (\n        cfg.num_beam_cands\n        + (cfg.num_diverse_cands if cfg.use_diverse_beam else 0)\n        + cfg.num_sample_cands\n        + (len(cfg.prompt_variants) * cfg.num_prompt_cands if cfg.use_prompt_variants else 0)\n    )\n    print(f\"Candidates/sample/model : {cands_per_sample}\")\n    print(f\"  ├─ standard beam      : {cfg.num_beam_cands} (num_beams={cfg.num_beams})\")\n    print(f\"  ├─ diverse beam       : {cfg.num_diverse_cands if cfg.use_diverse_beam else 0} (groups={cfg.num_beam_groups}, penalty={cfg.diversity_penalty})\")\n    print(f\"  ├─ multi-temp sample  : {cfg.num_sample_cands} ({cfg.num_sample_per_temp}/temp × {cfg.sample_temperatures})\")\n    print(f\"  └─ multi-prompt       : {len(cfg.prompt_variants) * cfg.num_prompt_cands if cfg.use_prompt_variants else 0} ({cfg.num_prompt_cands}/variant × {cfg.prompt_variants if cfg.use_prompt_variants else []})\")\n    print(f\"Total pool (2 models)   : ~{cands_per_sample * 2} (before dedup)\")\n    print(f\"CBP pre-prune           : {cfg.use_cbp_preprune} (trigger={cfg.cbp_trigger_n}, keep={cfg.cbp_keep_n}, rounds={cfg.cbp_bootstrap_rounds})\")\n    print()\n\n# 11. Run inference and export submission\n#\n# This is the exact end-to-end submission path.\n\n\n# ======================================================================\n# Final notes\n# ======================================================================\n\n\n# ======================================================================\n# Main execution\n# ======================================================================\n\nif __name__ == \"__main__\":\n    cfg = EnsembleMBRConfig()\n    cfg.resolve_paths()\n    logger = setup_logging(cfg.output_dir)\n\n    print_env(cfg)\n\n    logger.info(f\"Loading test data: {cfg.test_data_path}\")\n    test_df = pd.read_csv(cfg.test_data_path, encoding=\"utf-8\")\n    logger.info(f\"Test samples: {len(test_df)}\")\n\n    engine = EnsembleMBREngine(cfg, logger)\n    results_df = engine.run(test_df)\n\n    out_path = Path(cfg.output_dir) / \"submission.csv\"\n    results_df.to_csv(out_path, index=False)\n    logger.info(f\"Saved → {out_path}  ({len(results_df)} rows)\")\n    results_df.head()\n    cfg_snap = {\n        \"model_a\": cfg.model_a_path,\n        \"model_b\": cfg.model_b_path,\n        \"num_beam_cands\": cfg.num_beam_cands,\n        \"num_beams\": cfg.num_beams,\n        \"length_penalty\": cfg.length_penalty,\n        \"repetition_penalty\": cfg.repetition_penalty,\n        \"use_diverse_beam\": cfg.use_diverse_beam,\n        \"num_diverse_cands\": cfg.num_diverse_cands,\n        \"num_diverse_beams\": cfg.num_diverse_beams,\n        \"num_beam_groups\": cfg.num_beam_groups,\n        \"diversity_penalty\": cfg.diversity_penalty,\n        \"use_sampling\": cfg.use_sampling,\n        \"sample_temperatures\": cfg.sample_temperatures,\n        \"num_sample_per_temp\": cfg.num_sample_per_temp,\n        \"num_sample_cands\": cfg.num_sample_cands,\n        \"use_prompt_variants\": cfg.use_prompt_variants,\n        \"base_prompt\": cfg.base_prompt,\n        \"prompt_variants\": cfg.prompt_variants,\n        \"num_prompt_cands\": cfg.num_prompt_cands,\n        \"mbr_top_p\": cfg.mbr_top_p,\n        \"mbr_w_chrf\": cfg.mbr_w_chrf,\n        \"mbr_w_bleu\": cfg.mbr_w_bleu,\n        \"mbr_w_jaccard\": cfg.mbr_w_jaccard,\n        \"mbr_w_length\": cfg.mbr_w_length,\n        \"mbr_w_support\": cfg.mbr_w_support,\n        \"mbr_support_alpha\": cfg.mbr_support_alpha,\n        \"use_cbp_preprune\": cfg.use_cbp_preprune,\n        \"cbp_trigger_n\": cfg.cbp_trigger_n,\n        \"cbp_bootstrap_rounds\": cfg.cbp_bootstrap_rounds,\n        \"cbp_keep_n\": cfg.cbp_keep_n,\n        \"mbr_pool_cap\": cfg.mbr_pool_cap,\n        \"max_new_tokens\": cfg.max_new_tokens,\n        \"use_bf16_amp\": cfg.use_bf16_amp,\n        \"batch_size\": cfg.batch_size,\n    }\n    with open(Path(cfg.output_dir) / \"ensemble_mbr_config.json\", \"w\") as f:\n        json.dump(cfg_snap, f, indent=2)\n\n    print(\"Submission file:\", out_path)\n    print(\"Config saved to:\", Path(cfg.output_dir) / \"ensemble_mbr_config.json\")\n\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}