{"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},"file_extension":".py","nbconvert_exporter":"python","pygments_lexer":"ipython3"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":121150,"databundleVersionId":15833391,"isSourceIdPinned":false},{"sourceType":"datasetVersion","sourceId":14899351,"datasetId":9366631,"databundleVersionId":15763950},{"sourceType":"modelInstanceVersion","sourceId":755680,"databundleVersionId":15750522,"modelInstanceId":577094,"modelId":589419,"isSourceIdPinned":false}],"dockerImageVersionId":31260,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"intro","cell_type":"markdown","source":"# Deep Past Challenge Public Baseline (LB 35.5)\n**Two-model candidate pooling + competition-aligned MBR reranking**\n\n\n## What we changed to improve the earlier starter\n- Kept the same two public ByT5 checkpoints, but pooled their candidates together before selection.\n- Switched the final reranker from `chrF++`-only MBR to a competition-aligned utility: `sqrt(sentence_BLEU * sentence_chrF++)`.\n- Increased stochastic candidates from `2` to `3` to widen the reranking pool.\n- Added a small agreement bonus so duplicate candidates from different decoding paths are preferred.\n\n## Provenance and thanks\n- The notebook is highly inspired by the one published by `@meenalsinha`\n- The ensemble uses public checkpoints credited in this repo to `@assiaben` (`byt5-akkadian-optimized-34x`) and `@mattiaangeli` (`byt5-akkadian-mbr-v2`).\n","metadata":{}},{"id":"imports-md","cell_type":"markdown","source":"## 1. Imports and runtime helpers\nThis section loads the dependencies, seeds the run, and defines the BF16/autocast helpers used later during generation.\n","metadata":{}},{"id":"imports-code","cell_type":"code","source":"from __future__ import annotations\n\nimport gc\nimport json\nimport logging\nimport math\nimport os\nimport random\nimport re\nfrom contextlib import nullcontext\nfrom dataclasses import dataclass\nfrom pathlib import Path\nfrom typing import Dict, List, Tuple\n\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom torch.utils.data import DataLoader, Dataset, Sampler\nfrom tqdm.auto import tqdm\nfrom transformers import AutoModelForSeq2SeqLM, AutoTokenizer\n\n\ndef set_seed(seed: int) -> None:\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed_all(seed)\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","metadata":{},"outputs":[],"execution_count":null},{"id":"config-md","cell_type":"markdown","source":"## 2. Configuration and logging\nAll Kaggle paths, decoding parameters, and reranking knobs live in one config object so people can tweak the baseline without hunting through the notebook.\n","metadata":{}},{"id":"config-code","cell_type":"code","source":"@dataclass\nclass EnsembleConfig:\n    test_data_path: str = \"/kaggle/input/deep-past-initiative-machine-translation/test.csv\"\n    output_dir: str = \"/kaggle/working/\"\n    # Public checkpoint credited in this repo to Kaggle user @assiaben.\n    model_a_path: str = \"/kaggle/input/final-byt5/byt5-akkadian-optimized-34x\"\n    # Public checkpoint by Kaggle user @mattiaangeli.\n    model_b_path: str = \"/kaggle/input/models/mattiaangeli/byt5-akkadian-mbr-v2/pytorch/default/1\"\n    # Optional additional model paths for larger ensembles (set manually in notebook/script).\n    extra_model_paths: Tuple[str, ...] = ()\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    # Main 35.5 gain: increase stochastic diversity from 2 to 3 candidates.\n    num_sample_cands: int = 3\n    mbr_top_p: float = 0.92\n    mbr_temperature: float = 0.75\n    mbr_pool_cap: int = 36\n\n    # Reward candidates that recur across beams, samples, or models.\n    agreement_bonus: float = 0.035\n    # Use the competition-style reranking utility instead of chrF-only MBR.\n    use_competition_utility: bool = True\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    checkpoint_freq: int = 200\n    seed: int = 42\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(parents=True, exist_ok=True)\n        if not torch.cuda.is_available():\n            self.use_mixed_precision = False\n            self.use_better_transformer = False\n        self.use_bf16_amp = bool(\n            self.use_mixed_precision\n            and self.device.type == \"cuda\"\n            and cuda_bf16_supported()\n        )\n        if self.num_beams < self.num_beam_cands:\n            raise ValueError(\"num_beams must be >= num_beam_cands\")\n\n\ndef resolve_test_data_path(path: str) -> str:\n    if os.path.exists(path):\n        return path\n    alternates = [\n        \"/kaggle/input/competitions/deep-past-initiative-machine-translation/test.csv\",\n        \"/kaggle/input/deep-past-initiative-machine-translation/test.csv\",\n    ]\n    for alt in alternates:\n        if os.path.exists(alt):\n            return alt\n    raise FileNotFoundError(\n        \"Could not find test.csv. Attach the competition dataset and set cfg.test_data_path accordingly.\"\n    )\n\n\ndef setup_logging(output_dir: str) -> logging.Logger:\n    Path(output_dir).mkdir(parents=True, exist_ok=True)\n    for h in logging.root.handlers[:]:\n        logging.root.removeHandler(h)\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_plus.log\"),\n        ],\n    )\n    return logging.getLogger(\"ensemble_mbr_plus\")\n","metadata":{},"outputs":[],"execution_count":null},{"id":"preprocess-md","cell_type":"markdown","source":"## 3. Preprocessing\nThe transliteration cleanup is intentionally light: normalize gap markers, clean determinatives, canonicalize decimals, and make the encoder input more consistent before tokenization.\n","metadata":{}},{"id":"preprocess-code","cell_type":"code","source":"# ---------------------------\n# Text preprocess / postprocess\n# ---------------------------\nV2_RE = re.compile(r\"([aAeEiIuU])(?:2|\\u2082)\")\nV3_RE = re.compile(r\"([aAeEiIuU])(?:3|\\u2083)\")\nACUTE = str.maketrans({\"a\": \"a\", \"e\": \"e\", \"i\": \"i\", \"u\": \"u\", \"A\": \"A\", \"E\": \"E\", \"I\": \"I\", \"U\": \"U\"})\nGRAVE = str.maketrans({\"a\": \"a\", \"e\": \"e\", \"i\": \"i\", \"u\": \"u\", \"A\": \"A\", \"E\": \"E\", \"I\": \"I\", \"U\": \"U\"})\n\n\ndef ascii_to_diacritics(text: str) -> str:\n    s = str(text or \"\")\n    s = s.replace(\"sz\", \"sh\").replace(\"SZ\", \"Sh\")\n    s = s.replace(\"s,\", \"s\").replace(\"S,\", \"S\")\n    s = s.replace(\"t,\", \"t\").replace(\"T,\", \"T\")\n    s = V2_RE.sub(lambda m: m.group(1).translate(ACUTE), s)\n    s = V3_RE.sub(lambda m: m.group(1).translate(GRAVE), s)\n    return s\n\n\nALLOWED_FRACS = [\n    (1 / 6, \"0.16666\"),\n    (1 / 4, \"0.25\"),\n    (1 / 3, \"0.33333\"),\n    (1 / 2, \"0.5\"),\n    (2 / 3, \"0.66666\"),\n    (3 / 4, \"0.75\"),\n    (5 / 6, \"0.83333\"),\n]\nFRAC_TOL = 2e-3\nFLOAT_RE = re.compile(r\"(?<![\\w/])(\\d+\\.\\d{4,})(?![\\w/])\")\nWS_RE = re.compile(r\"\\s+\")\n\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\nGAP_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,}|\\u2026+|\\[\\.+\\]\"\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\nCHAR_TRANS = str.maketrans({\n    \"\\u1e2b\": \"h\",\n    \"\\u1e2a\": \"H\",\n    \"\\u02be\": \"\",\n    \"\\u2080\": \"0\",\n    \"\\u2081\": \"1\",\n    \"\\u2082\": \"2\",\n    \"\\u2083\": \"3\",\n    \"\\u2084\": \"4\",\n    \"\\u2085\": \"5\",\n    \"\\u2086\": \"6\",\n    \"\\u2087\": \"7\",\n    \"\\u2088\": \"8\",\n    \"\\u2089\": \"9\",\n    \"\\u2014\": \"-\",\n    \"\\u2013\": \"-\",\n})\nSUB_X = \"\\u2093\"\n\nDET_UPPER_RE = re.compile(r\"\\(([A-Z0-9]{1,6})\\)\")\nDET_LOWER_RE = re.compile(r\"\\(([a-z]{1,4})\\)\")\nPN_RE = re.compile(r\"\\bPN\\b\")\nSOFT_GRAM_RE = re.compile(\n    r\"\\(\\s*(?:fem|plur|pl|sing|singular|plural|\\?|\\!)(?:\\.\\s*(?:plur|plural|sing|singular))?\\.?\\s*[^)]*\\)\",\n    re.I,\n)\nBARE_GRAM_RE = re.compile(r\"(?<!\\w)(?:fem|sing|pl|plural)\\.?(?!\\w)\\s*\", re.I)\nREPEAT_WORD_RE = re.compile(r\"\\b(\\w+)(?:\\s+\\1\\b)+\", re.I)\nREPEAT_PUNCT_RE = re.compile(r\"([.,])\\1+\")\nPUNCT_SPACE_RE = re.compile(r\"\\s+([.,:])\")\nFORBIDDEN_TRANS = str.maketrans(\"\", \"\", \"()<>[]+;\")\nMULTI_GAP_RE = re.compile(r\"(?:<gap>\\s*){2,}\")\n\n\nclass OptimizedPreprocessor:\n    def preprocess_batch(self, texts: List[str]) -> List[str]:\n        ser = pd.Series(texts, dtype=\"object\").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 = ser.str.replace(GAP_UNIFIED_RE, \"<gap>\", regex=True)\n        ser = ser.str.translate(CHAR_TRANS)\n        ser = ser.str.replace(SUB_X, \"\", regex=False)\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","metadata":{},"outputs":[],"execution_count":null},{"id":"postprocess-md","cell_type":"markdown","source":"## 4. Postprocessing\nAfter decoding, we standardize the candidate pool before reranking so equivalent outputs compare cleanly.\n","metadata":{}},{"id":"postprocess-code","cell_type":"code","source":"class VectorizedPostprocessor:\n    def postprocess_batch(self, translations: List[str]) -> List[str]:\n        ser = pd.Series(translations, dtype=\"object\").fillna(\"\").astype(str)\n        ser = ser.str.replace(GAP_UNIFIED_RE, \"<gap>\", regex=True)\n        ser = ser.str.replace(PN_RE, \"<gap>\", regex=True)\n        ser = ser.str.replace(SOFT_GRAM_RE, \" \", regex=True)\n        ser = ser.str.replace(BARE_GRAM_RE, \" \", regex=True)\n        ser = ser.str.replace(FLOAT_RE, lambda m: canon_decimal(float(m.group(1))), regex=True)\n        ser = ser.str.replace(MULTI_GAP_RE, \"<gap>\", regex=True)\n        ser = ser.str.replace(\"<gap>\", \"\\uFFF0\", regex=False)\n        ser = ser.str.translate(FORBIDDEN_TRANS)\n        ser = ser.str.replace(\"\\uFFF0\", \" <gap> \", regex=False)\n        ser = ser.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            ser = ser.str.replace(pat, r\"\\1\", regex=True)\n        ser = ser.str.replace(PUNCT_SPACE_RE, r\"\\1\", regex=True)\n        ser = ser.str.replace(REPEAT_PUNCT_RE, r\"\\1\", regex=True)\n        ser = ser.str.replace(WS_RE, \" \", regex=True).str.strip()\n        return ser.tolist()\n","metadata":{},"outputs":[],"execution_count":null},{"id":"dataset-md","cell_type":"markdown","source":"## 5. Dataset and bucketed batching\nWe preprocess once, prepend the translation instruction, and group similar lengths together to reduce padding waste on Kaggle GPUs.\n","metadata":{}},{"id":"dataset-code","cell_type":"code","source":"class AkkadianDataset(Dataset):\n    def __init__(self, df: pd.DataFrame, preprocessor: OptimizedPreprocessor):\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\n    def __len__(self) -> int:\n        return len(self.sample_ids)\n\n    def __getitem__(self, idx: int):\n        return self.sample_ids[idx], self.input_texts[idx]\n\n\nclass BucketBatchSampler(Sampler[List[int]]):\n    def __init__(self, dataset: AkkadianDataset, batch_size: int, num_buckets: int, shuffle: bool = False):\n        self.batch_size = batch_size\n        self.shuffle = shuffle\n\n        lengths = [len(t.split()) for t in dataset.input_texts]\n        sorted_idx = sorted(range(len(lengths)), key=lambda i: lengths[i])\n        bucket_size = max(1, len(sorted_idx) // max(1, num_buckets))\n        self.buckets = [\n            sorted_idx[i * bucket_size:None if i == num_buckets - 1 else (i + 1) * bucket_size]\n            for i in range(num_buckets)\n        ]\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","metadata":{},"outputs":[],"execution_count":null},{"id":"models-md","cell_type":"markdown","source":"## 6. Model loading and candidate generation\nEach model contributes beam-search candidates plus stochastic samples. Models are loaded one at a time to keep VRAM usage manageable.\n","metadata":{}},{"id":"models-code","cell_type":"code","source":"class ModelWrapper:\n    def __init__(self, model_path: str, cfg: EnsembleConfig, logger: logging.Logger, label: str):\n        self.cfg = cfg\n        self.logger = logger\n        self.label = label\n        self.tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)\n        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_path, local_files_only=True).to(cfg.device).eval()\n        self.logger.info(f\"[{label}] loaded: {model_path}\")\n\n        if cfg.use_better_transformer and cfg.device.type == \"cuda\":\n            try:\n                from optimum.bettertransformer import BetterTransformer\n\n                self.model = BetterTransformer.transform(self.model)\n                self.logger.info(f\"[{label}] BetterTransformer enabled\")\n            except Exception as e:\n                self.logger.info(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, enc\n\n    def generate_candidates(self, input_ids, attention_mask, beam_size: int) -> List[List[str]]:\n        cfg = self.cfg\n        bsz = int(input_ids.shape[0])\n        ctx = bf16_ctx(cfg.device, cfg.use_bf16_amp)\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            sample_texts: List[str] = []\n            if cfg.num_sample_cands > 0:\n                sample_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                sample_texts = self.tokenizer.batch_decode(sample_out, skip_special_tokens=True)\n\n        rb = int(cfg.num_beam_cands)\n        rs = int(cfg.num_sample_cands)\n        pools: List[List[str]] = []\n        for i in range(bsz):\n            p = list(beam_texts[i * rb:(i + 1) * rb])\n            if rs > 0:\n                p.extend(sample_texts[i * rs:(i + 1) * rs])\n            pools.append(p)\n        return pools\n\n    def unload(self):\n        label = self.label\n        try:\n            from optimum.bettertransformer import BetterTransformer\n\n            self.model = BetterTransformer.reverse(self.model)\n        except Exception:\n            pass\n\n        del self.model\n        del self.tokenizer\n        gc.collect()\n        if torch.cuda.is_available():\n            torch.cuda.empty_cache()\n            self.logger.info(f\"[{label}] unloaded\")\n","metadata":{},"outputs":[],"execution_count":null},{"id":"mbr-md","cell_type":"markdown","source":"## 7. Competition-aligned MBR selector\nThe key `35.5` change is here: we score pairwise agreement with a utility closer to the competition metric and add a small consensus bonus for repeated candidates.\n","metadata":{}},{"id":"mbr-code","cell_type":"code","source":"# ---------------------------\n# Competition-aware MBR selector\n# ---------------------------\ndef get_ngrams(tokens: List[str], n: int) -> Dict[Tuple[str, ...], int]:\n    counts: Dict[Tuple[str, ...], int] = {}\n    for i in range(len(tokens) - n + 1):\n        ng = tuple(tokens[i:i + n])\n        counts[ng] = counts.get(ng, 0) + 1\n    return counts\n\n\ndef sentence_bleu_smoothed(hypothesis: str, reference: str, max_n: int = 4) -> float:\n    hyp_tok = hypothesis.strip().split()\n    ref_tok = reference.strip().split()\n    if not hyp_tok or not ref_tok:\n        return 0.0\n\n    bp = min(1.0, math.exp(1.0 - len(ref_tok) / max(len(hyp_tok), 1)))\n    log_avg = 0.0\n    for n in range(1, max_n + 1):\n        hyp_ng = get_ngrams(hyp_tok, n)\n        ref_ng = get_ngrams(ref_tok, n)\n        clipped = 0\n        total = 0\n        for ng, cnt in hyp_ng.items():\n            clipped += min(cnt, ref_ng.get(ng, 0))\n            total += cnt\n        if total == 0:\n            return 0.0\n        p_n = (clipped / total) if n == 1 else ((clipped + 1.0) / (total + 1.0))\n        if p_n <= 0:\n            return 0.0\n        log_avg += math.log(p_n)\n    return float(bp * math.exp(log_avg / max_n))\n\n\ndef char_ngrams(text: str, n: int) -> Dict[str, int]:\n    s = text\n    if len(s) < n:\n        return {}\n    counts: Dict[str, int] = {}\n    for i in range(len(s) - n + 1):\n        ng = s[i:i + n]\n        counts[ng] = counts.get(ng, 0) + 1\n    return counts\n\n\ndef sentence_chrfpp(hypothesis: str, reference: str, char_order: int = 6, word_order: int = 2, beta: float = 2.0) -> float:\n    hyp = hypothesis.strip()\n    ref = reference.strip()\n    if not hyp or not ref:\n        return 0.0\n\n    precisions: List[float] = []\n    recalls: List[float] = []\n\n    for n in range(1, char_order + 1):\n        hyp_ng = char_ngrams(hyp, n)\n        ref_ng = char_ngrams(ref, n)\n        if not hyp_ng or not ref_ng:\n            continue\n        overlap = 0\n        hyp_total = 0\n        ref_total = 0\n        for ng, cnt in hyp_ng.items():\n            overlap += min(cnt, ref_ng.get(ng, 0))\n            hyp_total += cnt\n        for cnt in ref_ng.values():\n            ref_total += cnt\n        if hyp_total > 0:\n            precisions.append(overlap / hyp_total)\n        if ref_total > 0:\n            recalls.append(overlap / ref_total)\n\n    hyp_w = hyp.split()\n    ref_w = ref.split()\n    for n in range(1, word_order + 1):\n        hyp_ng = get_ngrams(hyp_w, n)\n        ref_ng = get_ngrams(ref_w, n)\n        if not hyp_ng or not ref_ng:\n            continue\n        overlap = 0\n        hyp_total = 0\n        ref_total = 0\n        for ng, cnt in hyp_ng.items():\n            overlap += min(cnt, ref_ng.get(ng, 0))\n            hyp_total += cnt\n        for cnt in ref_ng.values():\n            ref_total += cnt\n        if hyp_total > 0:\n            precisions.append(overlap / hyp_total)\n        if ref_total > 0:\n            recalls.append(overlap / ref_total)\n\n    if not precisions or not recalls:\n        return 0.0\n    p = sum(precisions) / len(precisions)\n    r = sum(recalls) / len(recalls)\n    if p <= 0 or r <= 0:\n        return 0.0\n    beta2 = beta * beta\n    return float((1 + beta2) * p * r / (beta2 * p + r))\n\n\nclass MBRSelector:\n    def __init__(self, pool_cap: int, use_competition_utility: bool, agreement_bonus: float):\n        self.pool_cap = int(pool_cap)\n        self.use_competition_utility = bool(use_competition_utility)\n        self.agreement_bonus = float(agreement_bonus)\n\n    @staticmethod\n    def dedup_with_counts(xs: List[str]) -> Tuple[List[str], Dict[str, int]]:\n        counts: Dict[str, int] = {}\n        order: List[str] = []\n        for x in xs:\n            t = str(x or \"\").strip()\n            if not t:\n                continue\n            if t not in counts:\n                order.append(t)\n                counts[t] = 0\n            counts[t] += 1\n        return order, counts\n\n    def utility(self, a: str, b: str) -> float:\n        if self.use_competition_utility:\n            bleu = sentence_bleu_smoothed(a, b)\n            chrf = sentence_chrfpp(a, b)\n            if bleu <= 0 or chrf <= 0:\n                return 0.0\n            return float(math.sqrt(bleu * chrf))\n        return sentence_chrfpp(a, b)\n\n    def pick(self, candidates: List[str]) -> str:\n        uniq, counts = self.dedup_with_counts(candidates)\n        if self.pool_cap > 0:\n            uniq = uniq[: self.pool_cap]\n        if not uniq:\n            return \"\"\n        if len(uniq) == 1:\n            return uniq[0]\n\n        total_count = sum(counts.get(c, 0) for c in uniq)\n        best_text = uniq[0]\n        best_score = -1e9\n        for cand in uniq:\n            support = 0.0\n            denom = 0\n            for other in uniq:\n                if other == cand:\n                    continue\n                w = counts.get(other, 1)\n                support += self.utility(cand, other) * w\n                denom += w\n            if denom > 0:\n                support /= denom\n            support += self.agreement_bonus * max(0, counts.get(cand, 1) - 1)\n            if support > best_score:\n                best_score = support\n                best_text = cand\n        return best_text\n","metadata":{},"outputs":[],"execution_count":null},{"id":"engine-md","cell_type":"markdown","source":"## 8. End-to-end ensemble engine\nThis engine runs the full pipeline: preprocess the test set, decode with each model, merge candidate pools, rerank with MBR, and write checkpoints during inference.\n","metadata":{}},{"id":"engine-code","cell_type":"code","source":"class EnsembleMBREngine:\n    def __init__(self, cfg: EnsembleConfig, 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_competition_utility=cfg.use_competition_utility,\n            agreement_bonus=cfg.agreement_bonus,\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(\n                dataset=dataset,\n                batch_size=self.cfg.batch_size,\n                num_buckets=self.cfg.num_buckets,\n                shuffle=False,\n            )\n            return DataLoader(\n                dataset,\n                batch_sampler=sampler,\n                num_workers=self.cfg.num_workers,\n                pin_memory=(self.cfg.device.type == \"cuda\"),\n                collate_fn=wrapper.collate,\n            )\n        return DataLoader(\n            dataset,\n            batch_size=self.cfg.batch_size,\n            shuffle=False,\n            num_workers=self.cfg.num_workers,\n            pin_memory=(self.cfg.device.type == \"cuda\"),\n            collate_fn=wrapper.collate,\n        )\n\n    def run_one_model(self, wrapper: ModelWrapper, dataset: AkkadianDataset) -> Dict[str, List[str]]:\n        dl = self.build_dataloader(dataset, wrapper)\n        pools_by_id: Dict[str, List[str]] = {}\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                beams = self.adaptive_beams(attn)\n                try:\n                    batch_pools = wrapper.generate_candidates(input_ids, attn, beams)\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\"[{wrapper.label}] OOM batch skipped\")\n                        if torch.cuda.is_available():\n                            torch.cuda.empty_cache()\n                    else:\n                        raise\n                if torch.cuda.is_available():\n                    torch.cuda.empty_cache()\n        return pools_by_id\n\n    def run(self, test_df: pd.DataFrame) -> pd.DataFrame:\n        cfg = self.cfg\n        model_paths = [cfg.model_a_path, cfg.model_b_path] + [p for p in cfg.extra_model_paths if str(p).strip()]\n        self.logger.info(\"=\" * 60)\n        self.logger.info(\"LB35+ | Multi-model ensemble + competition-MBR\")\n        for i, p in enumerate(model_paths, start=1):\n            self.logger.info(f\"Model {i}: {p}\")\n        self.logger.info(\n            f\"Pools/model: beam x {cfg.num_beam_cands} + sample x {cfg.num_sample_cands}\"\n        )\n        self.logger.info(f\"MBR utility: {'competition' if cfg.use_competition_utility else 'chrF-only'}\")\n        self.logger.info(f\"Agreement bonus: {cfg.agreement_bonus}\")\n        self.logger.info(\"=\" * 60)\n\n        dataset = AkkadianDataset(test_df, self.preprocessor)\n        sample_ids = [str(s) for s in dataset.sample_ids]\n\n        all_pools: Dict[str, List[str]] = {}\n        for i, model_path in enumerate(model_paths, start=1):\n            label = f\"Model-{i}\"\n            wrapper = ModelWrapper(model_path, cfg, self.logger, label)\n            pools = self.run_one_model(wrapper, dataset)\n            wrapper.unload()\n            del wrapper\n            for sid, cand_pool in pools.items():\n                all_pools.setdefault(sid, []).extend(cand_pool)\n\n        results: List[Tuple[str, str]] = []\n        for sid in tqdm(sample_ids, desc=\"MBR\"):\n            combined = all_pools.get(sid, [])\n            post = self.postprocessor.postprocess_batch(combined) if combined else []\n            chosen = self.mbr.pick(post)\n            if not chosen or not chosen.strip():\n                chosen = \"The tablet is too damaged to translate.\"\n            results.append((sid, chosen))\n            if cfg.checkpoint_freq > 0 and 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                self.logger.info(f\"checkpoint saved: {ckpt}\")\n\n        out = pd.DataFrame(results, columns=[\"id\", \"translation\"])\n        self.validate(out)\n        return out\n\n    def validate(self, df: pd.DataFrame):\n        empty = df[\"translation\"].astype(str).str.strip().eq(\"\").sum()\n        lens = df[\"translation\"].astype(str).str.len()\n        self.logger.info(f\"Empty translations: {empty}\")\n        self.logger.info(\n            f\"Length mean={lens.mean():.1f} median={lens.median():.1f} min={lens.min()} max={lens.max()}\"\n        )\n","metadata":{},"outputs":[],"execution_count":null},{"id":"run-md","cell_type":"markdown","source":"## 9. Run inference and export `submission.csv`\nBefore running this cell on Kaggle, attach:\n\n- the competition dataset containing `test.csv`\n- `final-byt5/byt5-akkadian-optimized-34x`\n- `mattiaangeli/byt5-akkadian-mbr-v2`\n","metadata":{}},{"id":"run-code","cell_type":"code","source":"def print_env(cfg: EnsembleConfig):\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        total = torch.cuda.get_device_properties(0).total_memory / 1e9\n        print(f\"GPU Mem : {total:.2f} GB\")\n        print(f\"BF16    : {cuda_bf16_supported()}\")\n    print(f\"BF16 AMP: {cfg.use_bf16_amp}\")\n\n\ndef main():\n    cfg = EnsembleConfig()\n    set_seed(cfg.seed)\n    logger = setup_logging(cfg.output_dir)\n    cfg.test_data_path = resolve_test_data_path(cfg.test_data_path)\n\n    print_env(cfg)\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 rows: {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 submission: {out_path}\")\n\n    cfg_snapshot = {\n        \"model_a\": cfg.model_a_path,\n        \"model_b\": cfg.model_b_path,\n        \"extra_model_paths\": list(cfg.extra_model_paths),\n        \"num_beam_cands\": cfg.num_beam_cands,\n        \"num_sample_cands\": cfg.num_sample_cands,\n        \"num_beams\": cfg.num_beams,\n        \"length_penalty\": cfg.length_penalty,\n        \"repetition_penalty\": cfg.repetition_penalty,\n        \"mbr_top_p\": cfg.mbr_top_p,\n        \"mbr_temperature\": cfg.mbr_temperature,\n        \"mbr_pool_cap\": cfg.mbr_pool_cap,\n        \"agreement_bonus\": cfg.agreement_bonus,\n        \"use_competition_utility\": cfg.use_competition_utility,\n        \"max_new_tokens\": cfg.max_new_tokens,\n        \"batch_size\": cfg.batch_size,\n        \"use_bf16_amp\": cfg.use_bf16_amp,\n    }\n    with open(Path(cfg.output_dir) / \"ensemble_mbr_plus_config.json\", \"w\", encoding=\"utf-8\") as f:\n        json.dump(cfg_snapshot, f, indent=2)\n\n    print(\"\\n\" + \"=\" * 60)\n    print(\"LB35+ ensemble complete\")\n    print(f\"Submission: {out_path}\")\n    print(f\"Rows      : {len(results_df)}\")\n    print(\"=\" * 60)\n\n\nif __name__ == \"__main__\":\n    main()\n","metadata":{},"outputs":[],"execution_count":null},{"id":"acknowledgements","cell_type":"markdown","source":"## Acknowledgements\nThanks to the Deep Past Challenge organizers and to the community work this baseline builds on.\n\n- Starter notebook lineage credited here to Kaggle user `@meenalsinha`\n- Public checkpoint `byt5-akkadian-optimized-34x`, credited in this repo to `@assiaben`\n- Public checkpoint `byt5-akkadian-mbr-v2` by `@mattiaangeli`\n\nThe notebook is reorganized for readability, but the inference logic follows the original `35.5` submission notebook.\n","metadata":{}}]}