{"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":"datasetVersion","sourceId":15297017,"datasetId":9750544,"databundleVersionId":16200425},{"sourceType":"modelInstanceVersion","sourceId":755680,"databundleVersionId":15750522,"modelInstanceId":577094,"modelId":589419}],"dockerImageVersionId":31286,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## 1. Imports and runtime setup\nThis section loads all required libraries and sets a few runtime defaults for Kaggle.","metadata":{}},{"cell_type":"code","source":"import torch\n\nif torch.cuda.is_available():\n    n = torch.cuda.device_count()\n    print(f\"CUDA devices: {n}\")\n    for i in range(n):\n        print(f\"GPU {i}: {torch.cuda.get_device_name(i)}\")\nelse:\n    print(\"No CUDA available\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-23T12:37:13.637257Z","iopub.execute_input":"2026-03-23T12:37:13.638075Z","iopub.status.idle":"2026-03-23T12:37:13.643777Z","shell.execute_reply.started":"2026-03-23T12:37:13.638040Z","shell.execute_reply":"2026-03-23T12:37:13.642878Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import sys\nimport subprocess\n\nsubprocess.run(\n    [sys.executable, \"-m\", \"pip\", \"install\", \"--no-index\",\n     \"--find-links\", \"/kaggle/input/datasets/cq9067/offline-whl\", \"sacrebleu\"],\n    check=True\n)\nsubprocess.run(\n    [sys.executable, \"-m\", \"pip\", \"install\", \"--no-index\",\n     \"--find-links\", \"/kaggle/input/datasets/cq9067/offline-whl\", \"optimum\"],\n    check=True\n)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-23T12:37:13.645094Z","iopub.execute_input":"2026-03-23T12:37:13.645468Z","iopub.status.idle":"2026-03-23T12:37:19.943938Z","shell.execute_reply.started":"2026-03-23T12:37:13.645445Z","shell.execute_reply":"2026-03-23T12:37:19.943245Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#!/usr/bin/env python\n# coding: utf-8\n\n# # Ensemble × MBR for Deep Past Challenge\n# **Cross-model candidate pooling + chrF++ consensus selection**\n# \n# This notebook builds a submission pipeline for the **Deep Past Challenge – Translate Akkadian to English** competition.\n# \n# ## Method summary\n# - **Model A**: `byt5-akkadian-optimized-34x`\n# - **Model B**: `byt5-akkadian-mbr-v2`\n# - Each model generates:\n#   - beam-search candidates\n#   - stochastic sampling candidates\n# - Candidate pools are merged across models\n# - Final output is selected by **Minimum Bayes Risk (MBR)** using **chrF++**\n# \n# ## Main design choices\n# - Cross-model diversity to reduce shared failure modes\n# - Preprocessing rules adapted from the stronger MBR-oriented pipeline\n# - BF16 autocast for safer mixed precision on ByT5\n# - Conservative postprocessing that preserves `/` when needed\n# \n# ### Huge thanks to @mattiaangeli and @assiaben for their pre-trained models\n\n# # **Author: Cho Royou**\n\n# ## 1. Imports and runtime setup\n# This section loads all required libraries and sets a few runtime defaults for Kaggle.\n\n# In[ ]:\n\n\nimport os\nimport gc\nimport re\nimport json\nimport math\nimport random\nimport logging\nimport warnings\nfrom concurrent.futures import ThreadPoolExecutor\n\nfrom pathlib import Path\nfrom dataclasses import dataclass\nfrom contextlib import nullcontext\nfrom typing import List, 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\n# ## 2. BF16 utility helpers\n# These helpers detect whether BF16 autocast is available and create the proper context manager.\n\n# In[ ]:\n\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\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\n# ## 3. Configuration\n# All main generation, batching, postprocessing, and checkpointing options are collected here.\n\n# In[ ]:\n\n\n@dataclass\nclass EnsembleMBRConfig:\n      test_data_path: str = \"/kaggle/input/competitions/deep-past-initiative-machine-translation/test.csv\"\n      output_dir: str     = \"/kaggle/working/\"\n      model_a_path: str   = \"/kaggle/input/models/mattiaangeli/byt5-akkadian-mbr-v2/pytorch/default/1\"\n      model_b_path: str   = \"/kaggle/input/datasets/assiaben/final-byt5/byt5-akkadian-optimized-34x\"\n    \n      max_input_length: int = 512\n      max_new_tokens: int   = 384\n      batch_size: int       = 8\n      num_workers: int      = 2\n      num_buckets: int      = 6\n\n      num_beam_cands: int = 6\n      num_beams: int      = 10\n      length_penalty: float = 1.2\n      early_stopping: bool  = True\n      repetition_penalty: float = 1.2\n\n      num_sample_cands: int = 10\n      mbr_top_p: float       = 0.96\n      mbr_temperature: float = 0.95\n      mbr_pool_cap: int      = 48\n\n      use_mixed_precision: bool    = False\n      use_better_transformer: bool = True\n      use_bucket_batching: bool    = True\n      use_adaptive_beams: bool     = False\n      aggressive_postprocessing: bool = True\n      use_dual_gpu: bool = True\n      parallel_model_inference: bool = True\n      gpu_ids: Tuple[int, int] = (0, 1)\n\n      use_light_postprocess_for_mbr: bool = True\n      apply_heavy_postprocess_after_mbr: bool = True\n      remove_gap_tokens_in_final_output: bool = False\n\n      use_mbr_length_calibration: bool = True\n      mbr_length_penalty_alpha: float = 0.12\n\n      use_src_length_penalty_in_mbr: bool = True\n      mbr_src_len_soft_ratio: float = 1.0\n      mbr_src_len_hard_ratio: float = 0.8\n      mbr_src_len_soft_penalty: float = 1.5\n      mbr_src_len_hard_penalty: float = 6.0\n\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\n\n\n# ## 4. Logging\n# Logs are written to both stdout and a file inside the working directory.\n\n# In[ ]:\n\n\ndef setup_logging(output_dir: str) -> logging.Logger:\n    Path(output_dir).mkdir(exist_ok=True, parents=True)\n\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\n# ## 5. Preprocessing rules\n# This block normalizes transliterations before feeding them into the models.\n# \n# ### Included transformations\n# - ASCII transliteration to diacritics\n# - Determinative handling\n# - Gap normalization\n# - Character cleanup\n# - Fraction and decimal normalization\n\n# In[ ]:\n\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\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\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\n    return f\"{x:.5f}\".rstrip(\"0\").rstrip(\".\")\n\n\n_TAG_BIGGAP_RE = re.compile(r\"<\\s*big[\\s_\\-]*gap\\s*>\", re.I)\n_TAG_GAP_RE    = re.compile(r\"<\\s*gap\\s*>\", re.I)\n_BARE_BIGGAP   = re.compile(r\"\\bbig[\\s_\\-]*gap\\b\", re.I)\n_ELLIPSIS_RE   = re.compile(r\"(?:\\.{3,}|…+|\\[\\.+\\])\")\n_BRACKET_X_RE  = re.compile(r\"(\\[\\s*x\\s*\\]|\\(\\s*x\\s*\\))\", re.I)\n_XTOKEN_RUN_RE = re.compile(r\"\\bx(?:\\s+x)+\\b\", re.I)\n_XRUN_RE       = re.compile(r\"(?<!\\w)x{2,}(?!\\w)\", re.I)\n_XTOK_RE       = re.compile(r\"(?<!\\w)x(?!\\w)\", re.I)\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\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\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\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\n        # Uppercase determinatives are unwrapped, lowercase ones are converted to braces.\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\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\n# ## 6. Postprocessing rules\n# This block standardizes model outputs before MBR selection.\n# \n# ### Included transformations\n# - Gap unification\n# - `PN` handling\n# - Commodity-specific normalization\n# - Shekel fraction normalization\n# - Grammar-marker removal\n# - Month normalization\n# - Cleanup of repeated fragments and whitespace\n\n# In[ ]:\n\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_CURLY_QUOTES_RE = re.compile(\"[\\u201c\\u201d\\u2018\\u2019]\")\n\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\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_FORBIDDEN_TRANS = str.maketrans(\"\", \"\", '()——<>⌈⌋⌊[]+ʾ;')\n\n_COMMODITY_RE = re.compile(r'-(gold|tax|textiles)\\b')\n_COMMODITY_REPL = {\n    \"gold\": \"pašallum gold\",\n    \"tax\": \"šadduātum tax\",\n    \"textiles\": \"kutānum textiles\",\n}\n\ndef _commodity_repl(m: re.Match) -> str:\n    return _COMMODITY_REPL[m.group(1)]\n\n\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_SLASH_ALT_RE = re.compile(r'(?<!\\d)\\s*/\\s*(?!\\d)\\S+')\n_STRAY_MARKS_RE = re.compile(r'<<[^>]*>>|<(?!gap\\b)[^>]*>')\n_MULTI_GAP_RE = re.compile(r'(?:<gap>\\s*){2,}')\n_FINAL_GAP_RE = re.compile(r\"\\s*<gap>\\s*\", re.I)\n\ndef _month_repl(m: re.Match) -> str:\n    return f\"Month {_ROMAN2INT.get(m.group(1).upper(), m.group(1))}\"\n\n\nclass VectorizedPostprocessor:\n    def light_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(_MULTI_GAP_RE, \"<gap>\", regex=True)\n        s = s.str.replace(_WS_RE, \" \", regex=True).str.strip()\n        return s.tolist()\n\n    def postprocess_batch(self, translations: List[str]) -> List[str]:\n        s = pd.Series(translations).fillna(\"\").astype(str)\n\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\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\n        s = s.str.replace(_STRAY_MARKS_RE, \"\", regex=True)\n        s = s.str.replace(_SLASH_ALT_RE, \"\", regex=True)\n\n        # Remove only curly quotes; keep straight quotes and apostrophes.\n        s = s.str.replace(_CURLY_QUOTES_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        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        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\n        return s.tolist()\n\n\n# ## 7. Dataset and bucketed batching\n# The dataset applies preprocessing once. The custom sampler groups similar sequence lengths together to reduce padding waste.\n\n# In[ ]:\n\n\nclass AkkadianDataset(Dataset):\n    def __init__(self, df: pd.DataFrame, preprocessor: OptimizedPreprocessor, logger: logging.Logger):\n        self.sample_ids = df[\"id\"].tolist()\n        proc = preprocessor.preprocess_batch(df[\"transliteration\"].tolist())\n        self.proc_texts = proc\n        self.src_char_lens = [len(t) for t in proc]\n        self.input_texts = [\"translate Akkadian to English: \" + 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\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\n        lengths = [len(t.split()) for _, t in dataset]\n        sorted_idx = sorted(range(len(lengths)), key=lambda i: lengths[i])\n\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\n# ## 8. Model wrapper\n# This wrapper handles:\n# - tokenizer/model loading\n# - optional BetterTransformer conversion\n# - candidate generation\n# - safe unload and VRAM cleanup\n\n# In[ ]:\n\n\nclass ModelWrapper:\n    def __init__(\n        self,\n        model_path: str,\n        cfg: EnsembleMBRConfig,\n        logger: logging.Logger,\n        label: str,\n        device: torch.device = None,\n    ):\n        self.cfg = cfg\n        self.logger = logger\n        self.label = label\n        self.device = device if device is not None else cfg.device\n        self.tokenizer = None\n        self.model = None\n\n        logger.info(f\"[{label}] Loading from {model_path} on {self.device}\")\n        self.tokenizer = AutoTokenizer.from_pretrained(model_path)\n        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_path).to(self.device).eval()\n\n        if self.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 self.device.type == \"cuda\":\n            used = torch.cuda.memory_allocated(self.device) / 1e9\n            logger.info(f\"[{label}] GPU mem used: {used:.2f} GB\")\n\n        if cfg.use_better_transformer and self.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\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, enc\n\n    def generate_candidates(self, input_ids, attention_mask, beam_size: int) -> List[List[str]]:\n        cfg = self.cfg\n        B = input_ids.shape[0]\n        ctx = _bf16_ctx(self.device, cfg.use_bf16_amp)\n\n        with ctx:\n            nb = max(beam_size, cfg.num_beam_cands)\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=cfg.num_beam_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            beam_texts = self.tokenizer.batch_decode(beam_out, skip_special_tokens=True)\n\n            samp_texts = []\n            if cfg.num_sample_cands > 0:\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=cfg.mbr_temperature,\n                    num_return_sequences=cfg.num_sample_cands,\n                    max_new_tokens=cfg.max_new_tokens,\n                    repetition_penalty=cfg.repetition_penalty,\n                    use_cache=True,\n                )\n                samp_texts = self.tokenizer.batch_decode(samp_out, skip_special_tokens=True)\n\n        Rb, Rs = cfg.num_beam_cands, cfg.num_sample_cands\n        pools = []\n        for i in range(B):\n            p = list(beam_texts[i*Rb:(i+1)*Rb])\n            if Rs > 0:\n                p += list(samp_texts[i*Rs:(i+1)*Rs])\n            pools.append(p)\n\n        return pools\n\n    def unload(self):\n        label = self.label\n\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 self.device.type == \"cuda\":\n            with torch.cuda.device(self.device):\n                torch.cuda.empty_cache()\n                torch.cuda.synchronize()\n                free = (\n                    torch.cuda.get_device_properties(self.device).total_memory\n                    - torch.cuda.memory_allocated(self.device)\n                ) / 1e9\n            self.logger.info(f\"[{label}] Unloaded. GPU free: {free:.2f} GB\")\n\n\n# ## 9. MBR selector\n# MBR chooses the candidate with the highest average chrF++ agreement against the other candidates in the pool.\n\n# In[ ]:\n\n\nclass MBRSelector:\n    def __init__(\n        self,\n        pool_cap: int = 32,\n        use_length_calibration: bool = False,\n        length_penalty_alpha: float = 0.12,\n        use_src_length_penalty: bool = False,\n        src_len_soft_ratio: float = 1.0,\n        src_len_hard_ratio: float = 0.8,\n        src_len_soft_penalty: float = 1.5,\n        src_len_hard_penalty: float = 6.0,\n    ):\n        self._metric = sacrebleu.metrics.CHRF(word_order=2)\n        self.pool_cap = pool_cap\n        self.use_length_calibration = use_length_calibration\n        self.length_penalty_alpha = max(0.0, float(length_penalty_alpha))\n        self.use_src_length_penalty = bool(use_src_length_penalty)\n        self.src_len_soft_ratio = max(0.0, float(src_len_soft_ratio))\n        self.src_len_hard_ratio = max(0.0, float(src_len_hard_ratio))\n        if self.src_len_hard_ratio > self.src_len_soft_ratio:\n            self.src_len_hard_ratio, self.src_len_soft_ratio = self.src_len_soft_ratio, self.src_len_hard_ratio\n        self.src_len_soft_penalty = max(0.0, float(src_len_soft_penalty))\n        self.src_len_hard_penalty = max(0.0, float(src_len_hard_penalty))\n\n    def _chrfpp(self, a: str, b: str) -> float:\n        a, b = (a or \"\").strip(), (b or \"\").strip()\n        if not a or not b:\n            return 0.0\n        return float(self._metric.sentence_score(a, [b]).score)\n\n    @staticmethod\n    def _dedup(xs: List[str]) -> List[str]:\n        seen, out = set(), []\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\n    def _src_length_penalty(self, cand_len: int, src_len: int) -> float:\n        if not self.use_src_length_penalty or src_len <= 0:\n            return 0.0\n        ratio = cand_len / max(1.0, float(src_len))\n        soft = self.src_len_soft_ratio\n        hard = self.src_len_hard_ratio\n\n        if ratio >= soft:\n            return 0.0\n        if ratio >= hard:\n            # Slightly short translations get a gentle linear penalty.\n            span = max(1e-6, soft - hard)\n            t = (soft - ratio) / span\n            return self.src_len_soft_penalty * t\n\n        # Clearly short translations get both soft and hard penalties.\n        severe = (hard - ratio) / max(1e-6, hard)\n        return self.src_len_soft_penalty + self.src_len_hard_penalty * severe\n\n    def pick(self, candidates: List[str], src_len: int = 0) -> str:\n        cands = self._dedup(candidates)\n\n        if self.pool_cap:\n            cands = cands[:self.pool_cap]\n\n        n = len(cands)\n        if n == 0:\n            return \"\"\n        if n == 1:\n            return cands[0]\n\n        target_len = float(np.median([len(c) for c in cands])) if self.use_length_calibration else 0.0\n        best_i, best_s = 0, -1e9\n        for i in range(n):\n            s = sum(self._chrfpp(cands[i], cands[j]) for j in range(n) if j != i)\n            s /= max(1, n - 1)\n            if self.use_length_calibration and target_len > 0:\n                len_ratio = abs(len(cands[i]) - target_len) / target_len\n                s -= self.length_penalty_alpha * (100.0 * len_ratio)\n            s -= self._src_length_penalty(len(cands[i]), src_len)\n            if s > best_s:\n                best_s, best_i = s, i\n\n        return cands[best_i]\n\n\n# ## 10. End-to-end ensemble engine\n# This class runs the full pipeline:\n# \n# 1. preprocess test inputs  \n# 2. generate candidates with Model A  \n# 3. generate candidates with Model B  \n# 4. merge candidate pools  \n# 5. postprocess outputs  \n# 6. select the final translation with MBR  \n# 7. write checkpoints and validate the submission\n\n# In[ ]:\n\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            use_length_calibration=cfg.use_mbr_length_calibration,\n            length_penalty_alpha=cfg.mbr_length_penalty_alpha,\n            use_src_length_penalty=cfg.use_src_length_penalty_in_mbr,\n            src_len_soft_ratio=cfg.mbr_src_len_soft_ratio,\n            src_len_hard_ratio=cfg.mbr_src_len_hard_ratio,\n            src_len_soft_penalty=cfg.mbr_src_len_soft_penalty,\n            src_len_hard_penalty=cfg.mbr_src_len_hard_penalty,\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\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(\n                dataset,\n                self.cfg.batch_size,\n                self.cfg.num_buckets,\n                self.logger,\n            )\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, enc in tqdm(dl, desc=f\"[{wrapper.label}]\"):\n                input_ids = enc.input_ids.to(wrapper.device, non_blocking=True)\n                attn = enc.attention_mask.to(wrapper.device, non_blocking=True)\n                beam_size = self._adaptive_beams(attn)\n\n                try:\n                    batch_pools = wrapper.generate_candidates(input_ids, attn, beam_size)\n                    for sid, pool in zip(batch_ids, batch_pools):\n                        pools_by_id[str(sid)] = pool\n\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                        if wrapper.device.type == \"cuda\":\n                            with torch.cuda.device(wrapper.device):\n                                torch.cuda.empty_cache()\n                        for sid in batch_ids:\n                            pools_by_id.setdefault(str(sid), [])\n                    else:\n                        raise\n\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 wrapper.device.type == \"cuda\":\n                    with torch.cuda.device(wrapper.device):\n                        torch.cuda.empty_cache()\n\n        return pools_by_id\n\n    def _resolve_model_devices(self) -> Tuple[torch.device, torch.device]:\n        cfg = self.cfg\n        if (\n            cfg.use_dual_gpu\n            and torch.cuda.is_available()\n            and torch.cuda.device_count() >= 2\n        ):\n            a_id, b_id = cfg.gpu_ids\n            if a_id == b_id:\n                self.logger.warning(\"Dual GPU enabled but gpu_ids are identical, fallback to single device.\")\n                d = torch.device(f\"cuda:{a_id}\")\n                return d, d\n            if a_id >= torch.cuda.device_count() or b_id >= torch.cuda.device_count():\n                self.logger.warning(\"gpu_ids exceed available GPU count, fallback to cuda:0.\")\n                d = torch.device(\"cuda:0\")\n                return d, d\n            return torch.device(f\"cuda:{a_id}\"), torch.device(f\"cuda:{b_id}\")\n\n        return self.cfg.device, self.cfg.device\n\n    def run(self, test_df: pd.DataFrame) -> pd.DataFrame:\n        cfg, logger = self.cfg, self.logger\n\n        logger.info(\"=\" * 60)\n        logger.info(\"Ensemble × MBR | Cross-model candidate pooling\")\n        logger.info(f\"  Model A : {cfg.model_a_path}\")\n        logger.info(f\"  Model B : {cfg.model_b_path}\")\n        logger.info(f\"  Pool per model: beam×{cfg.num_beam_cands} + sample×{cfg.num_sample_cands}\")\n        logger.info(f\"  MBR pool cap  : {cfg.mbr_pool_cap}\")\n        logger.info(f\"  Dual GPU      : {cfg.use_dual_gpu} (parallel={cfg.parallel_model_inference}, gpu_ids={cfg.gpu_ids})\")\n        logger.info(f\"  Light PP→MBR  : {cfg.use_light_postprocess_for_mbr}\")\n        logger.info(f\"  Heavy PP post : {cfg.apply_heavy_postprocess_after_mbr}\")\n        logger.info(f\"  Remove <gap>  : {cfg.remove_gap_tokens_in_final_output}\")\n        logger.info(\n            f\"  Len calib     : {cfg.use_mbr_length_calibration} (alpha={cfg.mbr_length_penalty_alpha})\"\n        )\n        logger.info(\n            \"  Src-len pen   : \"\n            f\"{cfg.use_src_length_penalty_in_mbr} \"\n            f\"(soft={cfg.mbr_src_len_soft_ratio}, hard={cfg.mbr_src_len_hard_ratio}, \"\n            f\"soft_pen={cfg.mbr_src_len_soft_penalty}, hard_pen={cfg.mbr_src_len_hard_penalty})\"\n        )\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)\n        sample_ids = [str(s) for s in dataset.sample_ids]\n        src_len_by_id = {str(sid): sl for sid, sl in zip(dataset.sample_ids, dataset.src_char_lens)}\n        dev_a, dev_b = self._resolve_model_devices()\n        logger.info(f\"  Model devices : A->{dev_a}, B->{dev_b}\")\n\n        logger.info(\"Phase 1/2 — Model inference\")\n        wrapper_a = ModelWrapper(cfg.model_a_path, cfg, logger, \"Model-A\", device=dev_a)\n        wrapper_b = ModelWrapper(cfg.model_b_path, cfg, logger, \"Model-B\", device=dev_b)\n\n        try:\n            if cfg.parallel_model_inference and dev_a != dev_b:\n                logger.info(\"Running Model-A and Model-B in parallel\")\n                with ThreadPoolExecutor(max_workers=2) as ex:\n                    fa = ex.submit(self._run_one_model, wrapper_a, dataset)\n                    fb = ex.submit(self._run_one_model, wrapper_b, dataset)\n                    pools_a = fa.result()\n                    pools_b = fb.result()\n            else:\n                logger.info(\"Running Model-A and Model-B sequentially\")\n                pools_a = self._run_one_model(wrapper_a, dataset)\n                pools_b = self._run_one_model(wrapper_b, dataset)\n        finally:\n            wrapper_a.unload()\n            wrapper_b.unload()\n            del wrapper_a\n            del wrapper_b\n\n        logger.info(\"Phase 3/3 — Pool merge + MBR selection\")\n        results: List[Tuple[str, str]] = []\n\n        for sid in tqdm(sample_ids, desc=\"MBR\"):\n            combined = pools_a.get(sid, []) + pools_b.get(sid, [])\n            if combined:\n                mbr_cands = (\n                    self.postprocessor.light_postprocess_batch(combined)\n                    if cfg.use_light_postprocess_for_mbr else combined\n                )\n            else:\n                mbr_cands = []\n\n            chosen = self.mbr.pick(mbr_cands, src_len=src_len_by_id.get(sid, 0))\n\n            if chosen and cfg.apply_heavy_postprocess_after_mbr:\n                chosen = self.postprocessor.postprocess_batch([chosen])[0]\n            if chosen and cfg.remove_gap_tokens_in_final_output:\n                chosen = _FINAL_GAP_RE.sub(\" \", chosen)\n                chosen = _WS_RE.sub(\" \", chosen).strip()\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\n        empty = df[\"translation\"].str.strip().eq(\"\").sum()\n        lens = df[\"translation\"].str.len()\n\n        logger.info(f\"Empty     : {empty} ({100 * empty / max(1, len(df)):.2f}%)\")\n        logger.info(\n            f\"Len mean  : {lens.mean():.1f}  median: {lens.median():.1f}  min: {lens.min()}  max: {lens.max()}\"\n        )\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\n# ## 11. Runtime summary helper\n# This small helper prints the environment before inference starts.\n\n# In[ ]:\n\n\ndef print_env(cfg: EnsembleMBRConfig):\n    print(f\"PyTorch  : {torch.__version__}\")\n    print(f\"CUDA     : {torch.cuda.is_available()}\")\n\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\n    print(f\"BF16 AMP : {cfg.use_bf16_amp}\")\n\n\n# ## 12. Run inference and create `submission.csv`\n# This final block executes the full pipeline and saves:\n# - `submission.csv`\n# - `ensemble_mbr.log`\n# - `ensemble_mbr_config.json`\n\n# In[ ]:\n\n\ncfg = EnsembleMBRConfig()\nlogger = setup_logging(cfg.output_dir)\n\nprint_env(cfg)\n\nlogger.info(f\"Loading test data: {cfg.test_data_path}\")\ntest_df = pd.read_csv(cfg.test_data_path, encoding=\"utf-8\")\nlogger.info(f\"Test samples: {len(test_df)}\")\n\nengine = EnsembleMBREngine(cfg, logger)\nresults_df = engine.run(test_df)\n\nout_path = Path(cfg.output_dir) / \"submission.csv\"\nresults_df.to_csv(out_path, index=False)\nlogger.info(f\"Saved → {out_path} ({len(results_df)} rows)\")\n\ncfg_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    \"num_sample_cands\": cfg.num_sample_cands,\n    \"length_penalty\": cfg.length_penalty,\n    \"mbr_top_p\": cfg.mbr_top_p,\n    \"mbr_temperature\": cfg.mbr_temperature,\n    \"mbr_pool_cap\": cfg.mbr_pool_cap,\n    \"use_dual_gpu\": cfg.use_dual_gpu,\n    \"parallel_model_inference\": cfg.parallel_model_inference,\n    \"gpu_ids\": list(cfg.gpu_ids),\n    \"use_light_postprocess_for_mbr\": cfg.use_light_postprocess_for_mbr,\n    \"apply_heavy_postprocess_after_mbr\": cfg.apply_heavy_postprocess_after_mbr,\n    \"remove_gap_tokens_in_final_output\": cfg.remove_gap_tokens_in_final_output,\n    \"use_mbr_length_calibration\": cfg.use_mbr_length_calibration,\n    \"mbr_length_penalty_alpha\": cfg.mbr_length_penalty_alpha,\n    \"use_src_length_penalty_in_mbr\": cfg.use_src_length_penalty_in_mbr,\n    \"mbr_src_len_soft_ratio\": cfg.mbr_src_len_soft_ratio,\n    \"mbr_src_len_hard_ratio\": cfg.mbr_src_len_hard_ratio,\n    \"mbr_src_len_soft_penalty\": cfg.mbr_src_len_soft_penalty,\n    \"mbr_src_len_hard_penalty\": cfg.mbr_src_len_hard_penalty,\n    \"max_new_tokens\": cfg.max_new_tokens,\n    \"repetition_penalty\": cfg.repetition_penalty,\n    \"use_bf16_amp\": cfg.use_bf16_amp,\n    \"batch_size\": cfg.batch_size,\n}\n\nwith open(Path(cfg.output_dir) / \"ensemble_mbr_config.json\", \"w\") as f:\n    json.dump(cfg_snap, f, indent=2)\n\nprint(\"\\n\" + \"=\" * 60)\nprint(\"Ensemble × MBR complete\")\nprint(f\"Submission : {out_path}\")\nprint(f\"Total rows : {len(results_df)}\")\nprint(\"=\" * 60)\n\n\n# ## Notes for Kaggle submission\n# - Make sure both model paths are attached as Kaggle datasets/models.\n# - If `optimum` is unavailable, the notebook will still run normally.\n# - For the full test set, you may want to tune:\n#   - `batch_size`\n#   - `num_sample_cands`\n#   - `max_new_tokens`\n#   - `mbr_pool_cap`\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-23T12:37:19.945479Z","iopub.execute_input":"2026-03-23T12:37:19.945741Z","iopub.status.idle":"2026-03-23T12:38:21.663281Z","shell.execute_reply.started":"2026-03-23T12:37:19.945721Z","shell.execute_reply":"2026-03-23T12:38:21.662672Z"}},"outputs":[],"execution_count":null}]}