{"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":"gpu","dataSources":[{"sourceType":"competition","sourceId":121150,"databundleVersionId":15833391},{"sourceType":"datasetVersion","sourceId":14899351,"datasetId":9366631,"databundleVersionId":15763950},{"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":"markdown","source":"# Ensemble × MBR for Deep Past Challenge\n**Cross-model candidate pooling + chrF++ consensus selection**\n\nThis 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","metadata":{}},{"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 os\nimport gc\nimport re\nimport json\nimport math\nimport random\nimport logging\nimport warnings\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\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-06T06:46:05.443861Z","iopub.execute_input":"2026-03-06T06:46:05.444148Z","iopub.status.idle":"2026-03-06T06:46:24.428193Z","shell.execute_reply.started":"2026-03-06T06:46:05.444119Z","shell.execute_reply":"2026-03-06T06:46:24.427172Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2. BF16 utility helpers\nThese helpers detect whether BF16 autocast is available and create the proper context manager.","metadata":{}},{"cell_type":"code","source":"def _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()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 3. Configuration\nAll main generation, batching, postprocessing, and checkpointing options are collected here.","metadata":{}},{"cell_type":"code","source":"@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/datasets/assiaben/final-byt5/byt5-akkadian-optimized-34x\"\n model_b_path: str = \"/kaggle/input/models/mattiaangeli/byt5-akkadian-mbr-v2/pytorch/default/1\"\n\n max_input_length: int = 512\n max_new_tokens: int = 384\n batch_size: int = 2\n num_workers: int = 2\n num_buckets: int = 6\n\n num_beam_cands: int = 4\n num_beams: int = 8\n length_penalty: float = 1.3\n early_stopping: bool = True\n repetition_penalty: float = 1.2\n\n num_sample_cands: int = 2\n mbr_top_p: float = 0.92\n mbr_temperature: float = 0.75\n mbr_pool_cap: int = 32\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","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 4. Logging\nLogs are written to both stdout and a file inside the working directory.","metadata":{}},{"cell_type":"code","source":"def 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\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 5. Preprocessing rules\nThis 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","metadata":{}},{"cell_type":"code","source":"_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\"(? 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\"(?\"\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\"|(? pd.Series:\n return ser.str.replace(_GAP_UNIFIED_RE, \"\", 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()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 6. Postprocessing rules\nThis 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","metadata":{}},{"cell_type":"code","source":"_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\"(?⌈⌋⌊[]+ʾ;')\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'(?]*>>|<(?!gap\\b)[^>]*>')\n_MULTI_GAP_RE = re.compile(r'(?:\\s*){2,}')\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 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, \"\", 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, \"\", regex=True)\n\n s = s.str.replace(\"\", \"\\x00GAP\\x00\", regex=False)\n s = s.str.translate(_FORBIDDEN_TRANS)\n s = s.str.replace(\"\\x00GAP\\x00\", \" \", 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()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 7. Dataset and bucketed batching\nThe dataset applies preprocessing once. The custom sampler groups similar sequence lengths together to reduce padding waste.","metadata":{}},{"cell_type":"code","source":"class 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.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)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 8. Model wrapper\nThis wrapper handles:\n- tokenizer/model loading\n- optional BetterTransformer conversion\n- candidate generation\n- safe unload and VRAM cleanup","metadata":{}},{"cell_type":"code","source":"class 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\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(cfg.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 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\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 9. MBR selector\nMBR chooses the candidate with the highest average chrF++ agreement against the other candidates in the pool.","metadata":{}},{"cell_type":"code","source":"class MBRSelector:\n def __init__(self, pool_cap: int = 32):\n self._metric = sacrebleu.metrics.CHRF(word_order=2)\n self.pool_cap = pool_cap\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 pick(self, candidates: List[str]) -> 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 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 s > best_s:\n best_s, best_i = s, i\n\n return cands[best_i]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 10. End-to-end ensemble engine\nThis class runs the full pipeline:\n\n1. preprocess test inputs \n2. generate candidates with Model A \n3. generate candidates with Model B \n4. merge candidate pools \n5. postprocess outputs \n6. select the final translation with MBR \n7. write checkpoints and validate the submission","metadata":{}},{"cell_type":"code","source":"class 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(pool_cap=cfg.mbr_pool_cap)\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(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, 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 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 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 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\" 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\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: 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 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\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)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 11. Runtime summary helper\nThis small helper prints the environment before inference starts.","metadata":{}},{"cell_type":"code","source":"def 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}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 12. Run inference and create `submission.csv`\nThis final block executes the full pipeline and saves:\n- `submission.csv`\n- `ensemble_mbr.log`\n- `ensemble_mbr_config.json`","metadata":{}},{"cell_type":"code","source":"cfg = 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 \"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)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 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`","metadata":{}}]}