{"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}],"dockerImageVersionId":31287,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Better Candidate Diversity and Stronger Consensus Selection\n\nThis notebook is **public progress update** for the **Deep Past Challenge - Translate Akkadian to English** competition.\n\nCompared with my previous public baseline, this version improved the public leaderboard by **+0.3** through a more deliberate inference-time ensemble strategy.\n\n## What changed\n\nThe core idea is still:\n\n- build a **candidate pool** from two ByT5 models,\n- normalize and clean the outputs,\n- use **Minimum Bayes Risk (MBR)** decoding to choose the most reliable translation.\n\nThe main improvements in this version are:\n\n1. **Richer candidate generation**\n   - Standard beam search remains the backbone.\n   - Multi-temperature sampling adds more lexical and structural diversity.\n   - The code is prepared for diverse beam search as an optional extension.\n\n2. **Stronger MBR scoring**\n   - The previous version used chrF++ only.\n   - This version uses a **weighted combination** of:\n     - **chrF++**\n     - **BLEU**\n     - **token-level Jaccard similarity**\n     - **length prior**\n   - This makes the selector more robust against both noisy long outputs and locally fluent but inconsistent candidates.\n\n3. **More careful text normalization**\n   - Better handling of gap markers.\n   - Improved preprocessing of diacritics, determinatives, fractions, and notation variants.\n   - Conservative postprocessing that keeps symbols that may carry meaning.\n\n4. **Cross-model pooling**\n   - Candidate sets from two different checkpoints are merged before MBR.\n   - This helps reduce model-specific bias and systematic decoding errors.\n\nThis is still an inference-time method, but it shows that **candidate diversity + better reranking** can produce a meaningful gain even without retraining.","metadata":{}},{"cell_type":"markdown","source":"## Notebook structure\n\nThis notebook is organized as follows:\n\n1. Environment and imports  \n2. Configuration  \n3. Preprocessing  \n4. Postprocessing  \n5. Dataset and batching  \n6. Model wrapper and candidate generation  \n7. Weighted MBR selector  \n8. End-to-end inference engine  \n9. Submission export\n\nThe code below is written to be **Kaggle submission-ready** while remaining readable enough for public sharing.","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport re\nfrom typing import Dict, List\n\ndef build_proper_noun_lexicon(csv_path: str) -> Dict[str, str]:\n    \"\"\"\n    读取 OA_Lexicon_eBL.csv，提取人名、地名、神名的映射表。\n    \"\"\"\n    df = pd.read_csv(csv_path, encoding='utf-8')\n    \n    # 只提取专有名词：PN(人名), GN(地名), DN(神名), RN(王名)\n    # 你可以根据实际数据中的 type 分类进行调整\n    target_types = ['PN', 'GN', 'DN', 'RN']\n    entity_df = df[df['type'].isin(target_types)].copy()\n    \n    lexicon_map = {}\n    \n    for _, row in entity_df.iterrows():\n        form = str(row['form']).strip()\n        norm = str(row['norm']).strip()\n        \n        if pd.isna(form) or pd.isna(norm) or form == 'nan' or norm == 'nan':\n            continue\n            \n        # 对 form 进行小写处理，并移除可能存在的方括号等噪声\n        clean_form = re.sub(r'[\\[\\]\\(\\)\\?\\!]', '', form).lower()\n        \n        # 保存映射：转写 form (小写) -> 标准化翻译 norm (原样，通常首字母大写)\n        if clean_form:\n            lexicon_map[clean_form] = norm\n            \n    print(f\"Loaded {len(lexicon_map)} proper nouns into lexicon map.\")\n    return lexicon_map\n\n# 在你的 Configuration 阶段加载它\nlexicon_csv_path = \"/kaggle/input/competitions/deep-past-initiative-machine-translation/OA_Lexicon_eBL.csv\"\nPROPER_NOUN_MAP = build_proper_noun_lexicon(lexicon_csv_path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T15:43:56.153356Z","iopub.execute_input":"2026-03-09T15:43:56.154095Z","iopub.status.idle":"2026-03-09T15:43:56.76501Z","shell.execute_reply.started":"2026-03-09T15:43:56.154065Z","shell.execute_reply":"2026-03-09T15:43:56.764294Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, gc, re, json, math, random, logging, warnings\nfrom pathlib import Path\nfrom dataclasses import dataclass, field\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-09T15:43:56.766232Z","iopub.execute_input":"2026-03-09T15:43:56.76645Z","iopub.status.idle":"2026-03-09T15:43:56.77346Z","shell.execute_reply.started":"2026-03-09T15:43:56.766432Z","shell.execute_reply":"2026-03-09T15:43:56.772847Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 1. Mixed precision helpers\n\nI use BF16 autocast whenever available.  \nFor ByT5 inference this is generally safer than FP16 while still reducing memory pressure.","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\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,"execution":{"iopub.status.busy":"2026-03-09T15:43:56.774208Z","iopub.execute_input":"2026-03-09T15:43:56.774447Z","iopub.status.idle":"2026-03-09T15:43:56.785431Z","shell.execute_reply.started":"2026-03-09T15:43:56.774405Z","shell.execute_reply":"2026-03-09T15:43:56.784709Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2. Configuration\n\nThis configuration keeps the public notebook transparent:\n\n- two different ByT5 checkpoints,\n- standard beam search,\n- optional diverse beam search,\n- multi-temperature sampling,\n- weighted MBR reranking.","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    \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    use_diverse_beam:    bool = False\n    num_diverse_cands:   int = 4\n    num_diverse_beams:   int = 8\n    num_beam_groups:     int = 4\n    diversity_penalty:   float = 0.8\n\n    use_sampling:        bool = True\n    sample_temperatures: List[float] = field(default_factory=lambda: [0.60, 0.80, 1.05])\n    num_sample_per_temp: int = 2\n    mbr_top_p:           float = 0.92\n\n    @property\n    def num_sample_cands(self) -> int:\n        return len(self.sample_temperatures) * self.num_sample_per_temp\n\n    mbr_pool_cap: int = 32\n\n    mbr_w_chrf:     float = 0.55\n    mbr_w_bleu:     float = 0.25\n    mbr_w_jaccard:  float = 0.20\n    mbr_w_length:   float = 0.10\n\n    use_mixed_precision:     bool = True\n    use_better_transformer:  bool = True\n    use_bucket_batching:     bool = True\n    use_adaptive_beams:      bool = True\n    aggressive_postprocessing: bool = True\n    checkpoint_freq: int = 200\n\n    def __post_init__(self):\n        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n        Path(self.output_dir).mkdir(exist_ok=True, parents=True)\n\n        if not torch.cuda.is_available():\n            self.use_mixed_precision = False\n            self.use_better_transformer = False\n\n        self.use_bf16_amp = bool(\n            self.use_mixed_precision\n            and self.device.type == \"cuda\"\n            and _cuda_bf16_supported()\n        )\n\n        assert self.num_beams >= self.num_beam_cands, \"num_beams must be >= num_beam_cands\"\n\n        if self.use_diverse_beam:\n            assert self.num_diverse_beams % self.num_beam_groups == 0\n            assert self.num_diverse_beams >= self.num_diverse_cands","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T15:43:56.786372Z","iopub.execute_input":"2026-03-09T15:43:56.787035Z","iopub.status.idle":"2026-03-09T15:43:56.804309Z","shell.execute_reply.started":"2026-03-09T15:43:56.787006Z","shell.execute_reply":"2026-03-09T15:43:56.803505Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 3. Logging utilities","metadata":{}},{"cell_type":"code","source":"def setup_logging(output_dir: str) -> logging.Logger:\n    Path(output_dir).mkdir(exist_ok=True, parents=True)\n    for h in logging.root.handlers[:]:\n        logging.root.removeHandler(h)\n\n    logging.basicConfig(\n        level=logging.INFO,\n        format=\"%(asctime)s | %(levelname)s | %(message)s\",\n        handlers=[\n            logging.StreamHandler(),\n            logging.FileHandler(Path(output_dir) / \"ensemble_mbr.log\"),\n        ],\n    )\n    return logging.getLogger(\"ensemble_mbr\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T15:43:56.806405Z","iopub.execute_input":"2026-03-09T15:43:56.806683Z","iopub.status.idle":"2026-03-09T15:43:56.819836Z","shell.execute_reply.started":"2026-03-09T15:43:56.806658Z","shell.execute_reply":"2026-03-09T15:43:56.819231Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 4. Preprocessing\n\nThis block normalizes Akkadian transliteration before tokenization.\n\nKey steps:\n\n- convert ASCII-style diacritic notation into normalized Unicode,\n- normalize gap markers,\n- standardize determinatives,\n- normalize decimal/fraction expressions,\n- clean whitespace and symbol variants.\n\nThis part is important because small notation mismatches can fragment the candidate space and hurt MBR consistency.","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_ALLOWED_FRACS = [\n    (1/6, \"0.16666\"), (1/4, \"0.25\"), (1/3, \"0.33333\"),\n    (1/2, \"0.5\"), (2/3, \"0.66666\"), (3/4, \"0.75\"), (5/6, \"0.83333\"),\n]\n_FRAC_TOL = 2e-3\n_FLOAT_RE = re.compile(r\"(?<![\\w/])(\\d+\\.\\d{4,})(?![\\w/])\")\n\ndef _canon_decimal(x: float) -> str:\n    ip = int(math.floor(x + 1e-12))\n    frac = x - ip\n    best = min(_ALLOWED_FRACS, key=lambda t: abs(frac - t[0]))\n    if abs(frac - best[0]) <= _FRAC_TOL:\n        dec = best[1]\n        if ip == 0:\n            return dec\n        return f\"{ip}{dec[1:]}\" if dec.startswith(\"0.\") else f\"{ip}+{dec}\"\n    return f\"{x:.5f}\".rstrip(\"0\").rstrip(\".\")\n\n_WS_RE = re.compile(r\"\\s+\")\n\n_GAP_UNIFIED_RE = re.compile(\n    r\"<\\s*big[\\s_\\-]*gap\\s*>\"\n    r\"|<\\s*gap\\s*>\"\n    r\"|\\bbig[\\s_\\-]*gap\\b\"\n    r\"|\\bx(?:\\s+x)+\\b\"\n    r\"|\\.{3,}|…+|\\[\\.+\\]\"\n    r\"|\\[\\s*x\\s*\\]|\\(\\s*x\\s*\\)\"\n    r\"|(?<!\\w)x{2,}(?!\\w)\"\n    r\"|(?<!\\w)x(?!\\w)\"\n    r\"|\\(\\s*large\\s+break\\s*\\)\"\n    r\"|\\(\\s*break\\s*\\)\"\n    r\"|\\(\\s*\\d+\\s+broken\\s+lines?\\s*\\)\",\n    re.I\n)\n\ndef _normalize_gaps_vec(ser: pd.Series) -> pd.Series:\n    return ser.str.replace(_GAP_UNIFIED_RE, \"<gap>\", regex=True)\n\n_CHAR_TRANS = str.maketrans({\n    \"ḫ\":\"h\",\"Ḫ\":\"H\",\"ʾ\":\"\",\n    \"₀\":\"0\",\"₁\":\"1\",\"₂\":\"2\",\"₃\":\"3\",\"₄\":\"4\",\n    \"₅\":\"5\",\"₆\":\"6\",\"₇\":\"7\",\"₈\":\"8\",\"₉\":\"9\",\n    \"—\":\"-\",\"–\":\"-\",\n})\n_SUB_X = \"ₓ\"\n\n_UNICODE_UPPER = r\"A-ZŠṬṢḪ\\u00C0-\\u00D6\\u00D8-\\u00DE\\u0160\\u1E00-\\u1EFF\"\n_UNICODE_LOWER = r\"a-zšṭṣḫ\\u00E0-\\u00F6\\u00F8-\\u00FF\\u0161\\u1E01-\\u1EFF\"\n_DET_UPPER_RE = re.compile(r\"\\(([\" + _UNICODE_UPPER + r\"0-9]{1,6})\\)\")\n_DET_LOWER_RE = re.compile(r\"\\(([\" + _UNICODE_LOWER + r\"]{1,4})\\)\")\n\n_PN_RE = re.compile(r\"\\bPN\\b\")\n_KUBABBAR_RE = re.compile(r\"KÙ\\.B\\.\")\n\n_EXACT_FRAC_RE = re.compile(r\"0\\.8333|0\\.6666|0\\.3333|0\\.1666|0\\.625|0\\.75|0\\.25|0\\.5\")\n_EXACT_FRAC_MAP = {\n    \"0.8333\": \"⅚\", \"0.6666\": \"⅔\", \"0.3333\": \"⅓\", \"0.1666\": \"⅙\",\n    \"0.625\": \"⅝\", \"0.75\": \"¾\", \"0.25\": \"¼\", \"0.5\": \"½\",\n}\n\ndef _frac_repl(m: re.Match) -> str:\n    return _EXACT_FRAC_MAP[m.group(0)]\n\nclass OptimizedPreprocessor:\n    def preprocess_batch(self, texts: List[str]) -> List[str]:\n        ser = pd.Series(texts).fillna(\"\").astype(str)\n        ser = ser.apply(_ascii_to_diacritics)\n        ser = ser.str.replace(_DET_UPPER_RE, r\"\\1\", regex=True)\n        ser = ser.str.replace(_DET_LOWER_RE, r\"{\\1}\", regex=True)\n        ser = _normalize_gaps_vec(ser)\n        ser = ser.str.translate(_CHAR_TRANS)\n        ser = ser.str.replace(_SUB_X, \"\", regex=False)\n        ser = ser.str.replace(_KUBABBAR_RE, \"KÙ.BABBAR\", regex=True)\n        ser = ser.str.replace(_EXACT_FRAC_RE, _frac_repl, regex=True)\n        ser = ser.str.replace(_FLOAT_RE, lambda m: _canon_decimal(float(m.group(1))), regex=True)\n        ser = ser.str.replace(_WS_RE, \" \", regex=True).str.strip()\n        return ser.tolist()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T15:43:56.820691Z","iopub.execute_input":"2026-03-09T15:43:56.820934Z","iopub.status.idle":"2026-03-09T15:43:56.83652Z","shell.execute_reply.started":"2026-03-09T15:43:56.820915Z","shell.execute_reply":"2026-03-09T15:43:56.835758Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 5. Postprocessing\n\nThis block is where a lot of the practical gain comes from.\n\nThe idea is not to aggressively rewrite the model output, but to remove avoidable formatting noise so that MBR compares candidates on their *meaningful content* rather than superficial inconsistencies.\n\nI keep `/`, straight quotes, and apostrophes when they may be meaningful.","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\"(?<!\\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_MONTH_RE = re.compile(r\"\\bMonth\\s+(XII|XI|X|IX|VIII|VII|VI|V|IV|III|II|I)\\b\", re.I)\n_ROMAN2INT = {\"I\":1,\"II\":2,\"III\":3,\"IV\":4,\"V\":5,\"VI\":6,\"VII\":7,\"VIII\":8,\"IX\":9,\"X\":10,\"XI\":11,\"XII\":12}\n_REPEAT_WORD_RE = re.compile(r\"\\b(\\w+)(?:\\s+\\1\\b)+\")\n_REPEAT_PUNCT_RE = re.compile(r\"([.,])\\1+\")\n_PUNCT_SPACE_RE = re.compile(r\"\\s+([.,:])\")\n_FORBIDDEN_TRANS = str.maketrans(\"\", \"\", '()——<>⌈⌋⌊[]+ʾ;')\n\n_COMMODITY_RE = re.compile(r'-(gold|tax|textiles)\\b')\n_COMMODITY_REPL = {\"gold\": \"pašallum gold\", \"tax\": \"šadduātum tax\", \"textiles\": \"kutānum textiles\"}\n\ndef _commodity_repl(m: re.Match) -> str:\n    return _COMMODITY_REPL[m.group(1)]\n\n_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\ndef _month_repl(m: re.Match) -> str:\n    return f\"Month {_ROMAN2INT.get(m.group(1).upper(), m.group(1))}\"\n\nclass VectorizedPostprocessor:\n    def postprocess_batch(self, translations: List[str]) -> List[str]:\n        s = pd.Series(translations).fillna(\"\").astype(str)\n        s = _normalize_gaps_vec(s)\n        s = s.str.replace(_PN_RE, \"<gap>\", regex=True)\n        s = s.str.replace(_COMMODITY_RE, _commodity_repl, regex=True)\n\n        for pat, repl in _SHEKEL_REPLS:\n            s = s.str.replace(pat, repl, regex=True)\n\n        s = s.str.replace(_EXACT_FRAC_RE, _frac_repl, regex=True)\n        s = s.str.replace(_FLOAT_RE, lambda m: _canon_decimal(float(m.group(1))), regex=True)\n        s = s.str.replace(_SOFT_GRAM_RE, \" \", regex=True)\n        s = s.str.replace(_BARE_GRAM_RE, \" \", regex=True)\n        s = s.str.replace(_UNCERTAIN_RE, \"\", regex=True)\n        s = s.str.replace(_STRAY_MARKS_RE, \"\", regex=True)\n        s = s.str.replace(_SLASH_ALT_RE, \"\", regex=True)\n        s = s.str.replace(_CURLY_QUOTES_RE, \"\", regex=True)\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        return s.tolist()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T15:43:56.837913Z","iopub.execute_input":"2026-03-09T15:43:56.838135Z","iopub.status.idle":"2026-03-09T15:43:56.855239Z","shell.execute_reply.started":"2026-03-09T15:43:56.838115Z","shell.execute_reply":"2026-03-09T15:43:56.854523Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 6. Dataset and bucket batching\n\nBucketed batching reduces padding waste and keeps throughput more stable for variable-length transliterations.","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\nclass BucketBatchSampler(Sampler):\n    def __init__(self, dataset, batch_size, num_buckets, logger, shuffle=False):\n        self.batch_size = batch_size\n        self.shuffle = shuffle\n        lengths = [len(t.split()) for _, t in dataset]\n        sorted_idx = sorted(range(len(lengths)), key=lambda i: lengths[i])\n        bsize = max(1, len(sorted_idx) // max(1, num_buckets))\n        self.buckets = [\n            sorted_idx[i*bsize : None if i == num_buckets-1 else (i+1)*bsize]\n            for i in range(num_buckets)\n        ]\n\n        for i, b in enumerate(self.buckets):\n            if b:\n                bl = [lengths[x] for x in b]\n                logger.info(f\"  Bucket {i}: {len(b)} samples, len [{min(bl)}, {max(bl)}]\")\n\n    def __iter__(self):\n        for bucket in self.buckets:\n            b = list(bucket)\n            if self.shuffle:\n                random.shuffle(b)\n            for i in range(0, len(b), self.batch_size):\n                yield b[i:i+self.batch_size]\n\n    def __len__(self):\n        return sum(math.ceil(len(b) / self.batch_size) for b in self.buckets)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T15:43:56.856179Z","iopub.execute_input":"2026-03-09T15:43:56.856404Z","iopub.status.idle":"2026-03-09T15:43:56.870286Z","shell.execute_reply.started":"2026-03-09T15:43:56.856385Z","shell.execute_reply":"2026-03-09T15:43:56.869695Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 7. Model wrapper and candidate generation\n\nThis is the main inference change compared with the earlier public version.\n\n### Candidate sources\n\nFor each model and each sample, the pool can include:\n\n- **standard beam candidates**\n- **diverse beam candidates** (optional)\n- **multi-temperature sampled candidates**\n\nIn my current public LB 34.9 configuration:\n\n- diverse beam is **prepared but disabled**\n- multi-temperature sampling is **enabled**\n- the effective per-model pool is:\n\n`4 beam + 6 sampled = 10 candidates`\n\nAcross two models, that gives about **20 candidates before deduplication**.\n\n### Why this helped\n\nThe previous version already benefited from cross-model pooling.  \nThis version improved further because the candidate pool became more heterogeneous, which gives MBR a better chance to select a consensus output that is both faithful and fluent.","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        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        Rb = cfg.num_beam_cands\n        Rd = cfg.num_diverse_cands if cfg.use_diverse_beam else 0\n        Rs = cfg.num_sample_per_temp\n\n        with ctx:\n            nb = max(beam_size, Rb)\n            beam_out = self.model.generate(\n                input_ids=input_ids,\n                attention_mask=attention_mask,\n                do_sample=False,\n                num_beams=nb,\n                num_return_sequences=Rb,\n                max_new_tokens=cfg.max_new_tokens,\n                length_penalty=cfg.length_penalty,\n                early_stopping=cfg.early_stopping,\n                repetition_penalty=cfg.repetition_penalty,\n                use_cache=True,\n            )\n            beam_texts = self.tokenizer.batch_decode(beam_out, skip_special_tokens=True)\n\n            diverse_texts = []\n            actual_Rd = 0\n            if cfg.use_diverse_beam:\n                try:\n                    div_out = self.model.generate(\n                        input_ids=input_ids,\n                        attention_mask=attention_mask,\n                        do_sample=False,\n                        num_beams=cfg.num_diverse_beams,\n                        num_beam_groups=cfg.num_beam_groups,\n                        diversity_penalty=cfg.diversity_penalty,\n                        num_return_sequences=cfg.num_diverse_cands,\n                        max_new_tokens=cfg.max_new_tokens,\n                        length_penalty=cfg.length_penalty,\n                        early_stopping=cfg.early_stopping,\n                        repetition_penalty=cfg.repetition_penalty,\n                        use_cache=True,\n                    )\n                    diverse_texts = self.tokenizer.batch_decode(div_out, skip_special_tokens=True)\n                    actual_Rd = cfg.num_diverse_cands\n                except Exception as e:\n                    self.logger.warning(f\"[{self.label}] Diverse beam failed ({e}), skipping\")\n\n            all_samp_texts = []\n            num_temps = 0\n            if cfg.use_sampling and cfg.sample_temperatures:\n                num_temps = len(cfg.sample_temperatures)\n                for temp in cfg.sample_temperatures:\n                    try:\n                        samp_out = self.model.generate(\n                            input_ids=input_ids,\n                            attention_mask=attention_mask,\n                            do_sample=True,\n                            num_beams=1,\n                            top_p=cfg.mbr_top_p,\n                            temperature=temp,\n                            num_return_sequences=Rs,\n                            max_new_tokens=cfg.max_new_tokens,\n                            repetition_penalty=cfg.repetition_penalty,\n                            use_cache=True,\n                        )\n                        temp_texts = self.tokenizer.batch_decode(samp_out, skip_special_tokens=True)\n                        all_samp_texts.extend(temp_texts)\n                    except Exception as e:\n                        self.logger.warning(\n                            f\"[{self.label}] Sampling @ temp={temp:.2f} failed ({e}), padding with empty strings\"\n                        )\n                        all_samp_texts.extend([\"\"] * (B * Rs))\n\n        pools = []\n        for i in range(B):\n            pool = []\n            pool.extend(beam_texts[i * Rb : (i + 1) * Rb])\n\n            if diverse_texts and actual_Rd > 0:\n                pool.extend(diverse_texts[i * actual_Rd : (i + 1) * actual_Rd])\n\n            if all_samp_texts and num_temps > 0:\n                for t_idx in range(num_temps):\n                    seg_start = t_idx * B * Rs + i * Rs\n                    pool.extend(all_samp_texts[seg_start : seg_start + Rs])\n\n            pools.append(pool)\n\n        if pools:\n            self.logger.info(\n                f\"[{self.label}] Pool per sample: \"\n                f\"beam={Rb} + diverse={actual_Rd} + sample={num_temps}x{Rs}={num_temps*Rs} \"\n                f\"= {len(pools[0])} total\"\n            )\n\n        return pools\n\n    def unload(self):\n        label = self.label\n        try:\n            from optimum.bettertransformer import BetterTransformer\n            self.model = BetterTransformer.reverse(self.model)\n        except Exception:\n            pass\n\n        del self.model\n        del self.tokenizer\n        self.model = None\n        self.tokenizer = None\n\n        gc.collect()\n        if torch.cuda.is_available():\n            torch.cuda.empty_cache()\n            torch.cuda.synchronize()\n            free = (torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_allocated()) / 1e9\n            self.logger.info(f\"[{label}] Unloaded. GPU free: {free:.2f} GB\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T15:43:56.871213Z","iopub.execute_input":"2026-03-09T15:43:56.871441Z","iopub.status.idle":"2026-03-09T15:43:56.889681Z","shell.execute_reply.started":"2026-03-09T15:43:56.871424Z","shell.execute_reply":"2026-03-09T15:43:56.888894Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 8. Weighted MBR selector\n\nThis is the second major improvement.\n\nEarlier, I used **chrF++ only**.  \nIn this version, I use a weighted combination of:\n\n- **chrF++** for character-level robustness,\n- **BLEU** for word-level precision,\n- **Jaccard** for fast token overlap regularization,\n- **length bonus** to reduce pathological outputs.\n\nThis is still a simple MBR implementation, but it is noticeably stronger in practice than relying on a single metric.","metadata":{}},{"cell_type":"code","source":"class MBRSelector:\n    # def __init__(\n    #     self,\n    #     pool_cap: int = 32,\n    #     w_chrf: float = 0.55,\n    #     w_bleu: float = 0.25,\n    #     w_jaccard: float = 0.20,\n    #     w_length: float = 0.10,\n    # ):\n    #     self._chrf_metric = sacrebleu.metrics.CHRF(word_order=2)\n    #     self._bleu_metric = sacrebleu.metrics.BLEU(effective_order=True)\n    #     self.pool_cap = pool_cap\n    #     self.w_chrf = w_chrf\n    #     self.w_bleu = w_bleu\n    #     self.w_jaccard = w_jaccard\n    #     self.w_length = w_length\n    #     self._pw_total = max(w_chrf + w_bleu + w_jaccard, 1e-9)\n\n    # def _chrfpp(self, a: str, b: str) -> float:\n    #     if not a or not b:\n    #         return 0.0\n    #     return float(self._chrf_metric.sentence_score(a, [b]).score)\n\n    # def _bleu(self, a: str, b: str) -> float:\n    #     if not a or not b:\n    #         return 0.0\n    #     try:\n    #         return float(self._bleu_metric.sentence_score(a, [b]).score)\n    #     except Exception:\n    #         return 0.0\n\n    # @staticmethod\n    # def _jaccard(a: str, b: str) -> float:\n    #     ta, tb = set(a.lower().split()), set(b.lower().split())\n    #     if not ta and not tb:\n    #         return 100.0\n    #     if not ta or not tb:\n    #         return 0.0\n    #     return 100.0 * len(ta & tb) / len(ta | tb)\n\n    # def _pairwise_score(self, a: str, b: str) -> float:\n    #     s = (\n    #         self.w_chrf * self._chrfpp(a, b)\n    #         + self.w_bleu * self._bleu(a, b)\n    #         + self.w_jaccard * self._jaccard(a, b)\n    #     )\n    #     return s / self._pw_total\n\n    # @staticmethod\n    # def _length_bonus(lengths: List[int], idx: int) -> float:\n    #     if len(lengths) == 0:\n    #         return 100.0\n    #     median = float(np.median(lengths))\n    #     sigma = max(median * 0.4, 5.0)\n    #     z = (lengths[idx] - median) / sigma\n    #     return 100.0 * math.exp(-0.5 * z * z)\n\n    # @staticmethod\n    # def _dedup(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    #     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    #     lengths = [len(c.split()) for c in cands]\n    #     scores = []\n\n    #     for i in range(n):\n    #         pw = sum(\n    #             self._pairwise_score(cands[i], cands[j])\n    #             for j in range(n) if j != i\n    #         ) / max(1, n - 1)\n\n    #         lb = self._length_bonus(lengths, i)\n    #         total = pw + self.w_length * lb\n    #         scores.append(total)\n\n    #     return cands[int(np.argmax(scores))]\n    def __init__(\n        self,\n        # lexicon_map: Dict[str, str],\n        pool_cap: int = 32,\n        w_chrf: float = 0.50,\n        w_bleu: float = 0.20,\n        w_jaccard: float = 0.15,\n        w_length: float = 0.05,\n        w_fidelity: float = 0.10, # 新增的专有名词忠实度权重\n    ):\n        self._chrf_metric = sacrebleu.metrics.CHRF(word_order=2)\n        self._bleu_metric = sacrebleu.metrics.BLEU(effective_order=True)\n        self.lexicon_map = PROPER_NOUN_MAP\n        self.pool_cap = pool_cap\n        \n        self.w_chrf = w_chrf\n        self.w_bleu = w_bleu\n        self.w_jaccard = w_jaccard\n        self.w_length = w_length\n        self.w_fidelity = w_fidelity\n        \n        self._pw_total = max(w_chrf + w_bleu + w_jaccard, 1e-9)\n\n    # ... 保持原有的 _chrfpp, _bleu, _jaccard, _pairwise_score, _length_bonus 不变 ...\n    def _chrfpp(self, a: str, b: str) -> float:\n        if not a or not b:\n            return 0.0\n        return float(self._chrf_metric.sentence_score(a, [b]).score)\n\n    def _bleu(self, a: str, b: str) -> float:\n        if not a or not b:\n            return 0.0\n        try:\n            return float(self._bleu_metric.sentence_score(a, [b]).score)\n        except Exception:\n            return 0.0\n\n    @staticmethod\n    def _jaccard(a: str, b: str) -> float:\n        ta, tb = set(a.lower().split()), set(b.lower().split())\n        if not ta and not tb:\n            return 100.0\n        if not ta or not tb:\n            return 0.0\n        return 100.0 * len(ta & tb) / len(ta | tb)\n\n    def _pairwise_score(self, a: str, b: str) -> float:\n        s = (\n            self.w_chrf * self._chrfpp(a, b)\n            + self.w_bleu * self._bleu(a, b)\n            + self.w_jaccard * self._jaccard(a, b)\n        )\n        return s / self._pw_total\n\n    @staticmethod\n    def _length_bonus(lengths: List[int], idx: int) -> float:\n        if len(lengths) == 0:\n            return 100.0\n        median = float(np.median(lengths))\n        sigma = max(median * 0.4, 5.0)\n        z = (lengths[idx] - median) / sigma\n        return 100.0 * math.exp(-0.5 * z * z)\n\n    def _lexical_fidelity_score(self, source_text: str, candidate: str) -> float:\n        \"\"\"\n        计算候选译文是否成功保留了源文本中的专有名词。\n        返回 0 到 100 的分数。\n        \"\"\"\n        # 1. 对原文进行极其简单的分词（保留连字符，因为阿卡德语转写如 a-šur 依赖连字符）\n        # 移除非字母、非数字、非连字符的符号\n        clean_source = re.sub(r'[^\\w\\-\\s]', '', source_text.lower())\n        source_tokens = clean_source.split()\n        \n        expected_entities = []\n        for token in source_tokens:\n            if token in self.lexicon_map:\n                expected_entities.append(self.lexicon_map[token].lower())\n                \n        # 如果源文本中没有识别出专有名词，则给所有候选词满分（不影响 MBR 相对排名）\n        if not expected_entities:\n            return 100.0\n            \n        cand_lower = candidate.lower()\n        match_count = 0\n        \n        # 2. 检查候选词中是否出现了预期的翻译\n        for entity in expected_entities:\n            # 放宽匹配条件：只要实体字符串包含在译文中即可\n            # 如果要更严谨，可以使用正则边界匹配：re.search(rf'\\b{re.escape(entity)}\\b', cand_lower)\n            if entity in cand_lower:\n                match_count += 1\n                \n        return (match_count / len(expected_entities)) * 100.0\n\n    def pick(self, source_text: str, candidates: List[str]) -> str:\n        cands = self._dedup(candidates)\n        if self.pool_cap:\n            cands = cands[:self.pool_cap]\n\n        n = len(cands)\n        if n == 0: return \"\"\n        if n == 1: return cands[0]\n\n        lengths = [len(c.split()) for c in cands]\n        scores = []\n\n        for i in range(n):\n            # 基础共识分 (Consensus)\n            pw = sum(\n                self._pairwise_score(cands[i], cands[j])\n                for j in range(n) if j != i\n            ) / max(1, n - 1)\n\n            # 长度先验分\n            lb = self._length_bonus(lengths, i)\n            \n            # 专有名词保真分 (Fidelity) - 传入源文本和当前候选词\n            fid = self._lexical_fidelity_score(source_text, cands[i])\n\n            # 加权求和\n            total = pw + (self.w_length * lb) + (self.w_fidelity * fid)\n            scores.append(total)\n\n        return cands[int(np.argmax(scores))]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T15:43:56.890617Z","iopub.execute_input":"2026-03-09T15:43:56.890824Z","iopub.status.idle":"2026-03-09T15:43:56.907844Z","shell.execute_reply.started":"2026-03-09T15:43:56.890804Z","shell.execute_reply":"2026-03-09T15:43:56.907275Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 9. End-to-end engine","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        #\n        self.mbr = MBRSelector(\n            pool_cap=cfg.mbr_pool_cap,\n            w_chrf=cfg.mbr_w_chrf,\n            w_bleu=cfg.mbr_w_bleu,\n            w_jaccard=cfg.mbr_w_jaccard,\n            w_length=cfg.mbr_w_length,\n        )\n\n    def _adaptive_beams(self, attn: torch.Tensor) -> int:\n        if not self.cfg.use_adaptive_beams:\n            return self.cfg.num_beams\n        med = float(attn.sum(dim=1).float().median().item())\n        short = max(self.cfg.num_beam_cands, self.cfg.num_beams // 2)\n        return short if med < 100 else self.cfg.num_beams\n\n    def _build_dataloader(self, dataset: AkkadianDataset, wrapper: ModelWrapper) -> DataLoader:\n        if self.cfg.use_bucket_batching:\n            sampler = BucketBatchSampler(dataset, self.cfg.batch_size, self.cfg.num_buckets, self.logger)\n            return DataLoader(\n                dataset,\n                batch_sampler=sampler,\n                num_workers=self.cfg.num_workers,\n                collate_fn=wrapper.collate,\n                pin_memory=(self.cfg.device.type == \"cuda\"),\n            )\n\n        return DataLoader(\n            dataset,\n            batch_size=self.cfg.batch_size,\n            shuffle=False,\n            num_workers=self.cfg.num_workers,\n            collate_fn=wrapper.collate,\n            pin_memory=(self.cfg.device.type == \"cuda\"),\n        )\n\n    def _run_one_model(self, wrapper: ModelWrapper, dataset: AkkadianDataset) -> dict:\n        dl = self._build_dataloader(dataset, wrapper)\n        pools_by_id = {}\n\n        with torch.inference_mode():\n            for batch_ids, 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                except RuntimeError as e:\n                    if \"out of memory\" in str(e).lower():\n                        self.logger.error(f\"OOM in [{wrapper.label}] — skipping batch\")\n                        torch.cuda.empty_cache()\n                        for sid in batch_ids:\n                            pools_by_id.setdefault(str(sid), [])\n                    else:\n                        raise\n                except Exception as e:\n                    self.logger.error(f\"[{wrapper.label}] batch error: {e}\")\n                    for sid in batch_ids:\n                        pools_by_id.setdefault(str(sid), [])\n\n                if torch.cuda.is_available():\n                    torch.cuda.empty_cache()\n\n        return pools_by_id\n\n    def run(self, test_df: pd.DataFrame) -> pd.DataFrame:\n        cfg, logger = self.cfg, self.logger\n\n        cands_per_sample = (\n            cfg.num_beam_cands\n            + (cfg.num_diverse_cands if cfg.use_diverse_beam else 0)\n            + cfg.num_sample_cands\n        )\n\n        logger.info(\"=\" * 60)\n        logger.info(\"Ensemble × MBR  |  Cross-model candidate pooling  v2\")\n        logger.info(f\"  Model A           : {cfg.model_a_path}\")\n        logger.info(f\"  Model B           : {cfg.model_b_path}\")\n        logger.info(f\"  Standard beam     : {cfg.num_beam_cands} (num_beams={cfg.num_beams})\")\n        logger.info(\n            f\"  Diverse beam      : {'ON' if cfg.use_diverse_beam else 'OFF'} \"\n            f\"cands={cfg.num_diverse_cands}, groups={cfg.num_beam_groups}, penalty={cfg.diversity_penalty}\"\n        )\n        logger.info(\n            f\"  Multi-temp sample : {'ON' if cfg.use_sampling else 'OFF'} \"\n            f\"temps={cfg.sample_temperatures}, {cfg.num_sample_per_temp}/temp \"\n            f\"→ {cfg.num_sample_cands} total\"\n        )\n        logger.info(\n            f\"  Cands/model/sample: {cands_per_sample} \"\n            f\"→ pool ≈ {cands_per_sample * 2} (2 models, pre-dedup)\"\n        )\n        logger.info(f\"  MBR pool cap      : {cfg.mbr_pool_cap}\")\n        logger.info(f\"  BF16 AMP          : {cfg.use_bf16_amp}\")\n        logger.info(f\"  batch_size        : {cfg.batch_size}\")\n        logger.info(\"=\" * 60)\n\n        dataset = AkkadianDataset(test_df, self.preprocessor, logger)\n        sample_ids = [str(s) for s in dataset.sample_ids]\n\n        logger.info(\"Phase 1/2 — Model A inference\")\n        wrapper_a = ModelWrapper(cfg.model_a_path, cfg, logger, \"Model-A\")\n        pools_a = self._run_one_model(wrapper_a, dataset)\n        wrapper_a.unload()\n        del wrapper_a\n\n        logger.info(\"Phase 2/2 — Model B inference\")\n        wrapper_b = ModelWrapper(cfg.model_b_path, cfg, logger, \"Model-B\")\n        pools_b = self._run_one_model(wrapper_b, dataset)\n        wrapper_b.unload()\n        del wrapper_b\n\n        logger.info(\"Phase 3/3 — Pool merge + MBR selection\")\n        results = []\n\n        # for sid in tqdm(sample_ids, desc=\"  MBR\"):\n        #     combined = pools_a.get(sid, []) + pools_b.get(sid, [])\n        #     pp = self.postprocessor.postprocess_batch(combined) if combined else []\n        #     chosen = self.mbr.pick(pp)\n\n        #     if not chosen or not chosen.strip():\n        #         chosen = \"The tablet is too damaged to translate.\"\n\n        #     results.append((sid, chosen))\n    \n        # test_df 需要被遍历，以便我们能拿到原始的 transliteration\n        for idx, row in tqdm(test_df.iterrows(), total=len(test_df), desc=\"  MBR\"):\n            sid = str(row['id'])\n            source_transliteration = str(row['transliteration']) # 获取原文\n            \n            combined = pools_a.get(sid, []) + pools_b.get(sid, [])\n            pp = self.postprocessor.postprocess_batch(combined) if combined else []\n            \n            # 将 source_transliteration 传入 pick 方法\n            chosen = self.mbr.pick(source_transliteration, 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            if len(results) % cfg.checkpoint_freq == 0:\n                ckpt = Path(cfg.output_dir) / f\"checkpoint_{len(results)}.csv\"\n                pd.DataFrame(results, columns=[\"id\", \"translation\"]).to_csv(ckpt, index=False)\n                logger.info(f\"  Checkpoint: {len(results)} rows → {ckpt}\")\n        \n        result_df = pd.DataFrame(results, columns=[\"id\", \"translation\"])\n        self._validate(result_df)\n        return result_df\n\n    def _validate(self, df: pd.DataFrame):\n        logger = self.logger\n        logger.info(\"=\" * 60)\n        empty = df[\"translation\"].str.strip().eq(\"\").sum()\n        lens = df[\"translation\"].str.len()\n        logger.info(f\"Empty     : {empty} ({100*empty/max(1,len(df)):.2f}%)\")\n        logger.info(f\"Len mean  : {lens.mean():.1f}  median: {lens.median():.1f}  min: {lens.min()}  max: {lens.max()}\")\n\n        for idx in [0, len(df)//4, len(df)//2, 3*len(df)//4, len(df)-1]:\n            row = df.iloc[idx]\n            logger.info(f\"  ID {row['id']}: {str(row['translation'])[:80]}\")\n\n        logger.info(\"=\" * 60)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T15:43:56.908781Z","iopub.execute_input":"2026-03-09T15:43:56.909059Z","iopub.status.idle":"2026-03-09T15:43:56.928624Z","shell.execute_reply.started":"2026-03-09T15:43:56.909032Z","shell.execute_reply":"2026-03-09T15:43:56.927958Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 10. Environment summary helper","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    if torch.cuda.is_available():\n        print(f\"GPU      : {torch.cuda.get_device_name(0)}\")\n        mem = torch.cuda.get_device_properties(0).total_memory / 1e9\n        print(f\"GPU Mem  : {mem:.2f} GB\")\n        print(f\"BF16     : {_cuda_bf16_supported()}\")\n    print(f\"BF16 AMP : {cfg.use_bf16_amp}\")\n    print()\n\n    cands_per_sample = (\n        cfg.num_beam_cands\n        + (cfg.num_diverse_cands if cfg.use_diverse_beam else 0)\n        + cfg.num_sample_cands\n    )\n    print(f\"Candidates/sample/model : {cands_per_sample}\")\n    print(f\"  ├─ standard beam      : {cfg.num_beam_cands} (num_beams={cfg.num_beams})\")\n    print(f\"  ├─ diverse beam       : {cfg.num_diverse_cands if cfg.use_diverse_beam else 0} (groups={cfg.num_beam_groups}, penalty={cfg.diversity_penalty})\")\n    print(f\"  └─ multi-temp sample  : {cfg.num_sample_cands} ({cfg.num_sample_per_temp}/temp × {cfg.sample_temperatures})\")\n    print(f\"Total pool (2 models)   : ~{cands_per_sample * 2} (before dedup)\")\n    print()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T15:43:56.929427Z","iopub.execute_input":"2026-03-09T15:43:56.929748Z","iopub.status.idle":"2026-03-09T15:43:56.944594Z","shell.execute_reply.started":"2026-03-09T15:43:56.929726Z","shell.execute_reply":"2026-03-09T15:43:56.9438Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 11. Run inference and export submission\n\nThis is the exact end-to-end submission path.","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)\")\nresults_df.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-09T15:43:56.945507Z","iopub.execute_input":"2026-03-09T15:43:56.945784Z","execution_failed":"2026-03-09T15:44:36.133Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"cfg_snap = {\n    \"model_a\": cfg.model_a_path,\n    \"model_b\": cfg.model_b_path,\n    \"num_beam_cands\": cfg.num_beam_cands,\n    \"num_beams\": cfg.num_beams,\n    \"length_penalty\": cfg.length_penalty,\n    \"repetition_penalty\": cfg.repetition_penalty,\n    \"use_diverse_beam\": cfg.use_diverse_beam,\n    \"num_diverse_cands\": cfg.num_diverse_cands,\n    \"num_diverse_beams\": cfg.num_diverse_beams,\n    \"num_beam_groups\": cfg.num_beam_groups,\n    \"diversity_penalty\": cfg.diversity_penalty,\n    \"use_sampling\": cfg.use_sampling,\n    \"sample_temperatures\": cfg.sample_temperatures,\n    \"num_sample_per_temp\": cfg.num_sample_per_temp,\n    \"num_sample_cands\": cfg.num_sample_cands,\n    \"mbr_top_p\": cfg.mbr_top_p,\n    \"mbr_w_chrf\": cfg.mbr_w_chrf,\n    \"mbr_w_bleu\": cfg.mbr_w_bleu,\n    \"mbr_w_jaccard\": cfg.mbr_w_jaccard,\n    \"mbr_w_length\": cfg.mbr_w_length,\n    \"mbr_pool_cap\": cfg.mbr_pool_cap,\n    \"max_new_tokens\": cfg.max_new_tokens,\n    \"use_bf16_amp\": cfg.use_bf16_amp,\n    \"batch_size\": cfg.batch_size,\n}\nwith open(Path(cfg.output_dir) / \"ensemble_mbr_config.json\", \"w\") as f:\n    json.dump(cfg_snap, f, indent=2)\n\nprint(\"Submission file:\", out_path)\nprint(\"Config saved to:\", Path(cfg.output_dir) / \"ensemble_mbr_config.json\")","metadata":{"trusted":true,"execution":{"execution_failed":"2026-03-09T15:44:36.134Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Final notes\n\n### What this notebook demonstrates\n\nThis notebook is not claiming a fundamentally new model architecture.  \nThe gain comes from **better inference design**:\n\n- more diverse candidate generation,\n- cleaner normalization,\n- stronger consensus reranking.\n\n### Most important practical lesson\n\nFor this competition, once the base model quality is already decent, a large part of the improvement can come from:\n\n- **candidate pool engineering**\n- **postprocessing discipline**\n- **MBR selection quality**\n","metadata":{}}]}