{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":121150,"databundleVersionId":15833391,"isSourceIdPinned":false},{"sourceType":"datasetVersion","sourceId":14646609,"datasetId":9318078,"databundleVersionId":15486510},{"sourceType":"modelInstanceVersion","sourceId":789387,"databundleVersionId":16130305,"modelInstanceId":566051,"modelId":578456,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":789347,"databundleVersionId":16129846,"modelInstanceId":566051,"modelId":578456,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":800520,"databundleVersionId":16246597,"modelInstanceId":566051,"modelId":578456,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":787975,"databundleVersionId":16110960,"modelInstanceId":566051,"modelId":578456,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":795255,"databundleVersionId":16195580,"modelInstanceId":566051,"modelId":578456,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":745848,"databundleVersionId":15632297,"modelInstanceId":566051,"modelId":578456,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":799274,"databundleVersionId":16234058,"modelInstanceId":566051,"modelId":578456,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":799159,"databundleVersionId":16232781,"modelInstanceId":566051,"modelId":578456,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":798637,"databundleVersionId":16228970,"modelInstanceId":566051,"modelId":578456,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":746411,"databundleVersionId":15639258,"modelInstanceId":566051,"modelId":578456,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":764426,"databundleVersionId":15817052,"modelInstanceId":566051,"modelId":578456,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":795296,"databundleVersionId":16196217,"modelInstanceId":566051,"modelId":578456,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":795150,"databundleVersionId":16194417,"modelInstanceId":566051,"modelId":578456,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":753595,"databundleVersionId":15724484,"modelInstanceId":566051,"modelId":578456,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":786686,"databundleVersionId":16091123,"modelInstanceId":568583,"modelId":580914,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":769413,"databundleVersionId":15875819,"modelInstanceId":568583,"modelId":580914,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":797208,"databundleVersionId":16214010,"modelInstanceId":568583,"modelId":580914,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":794445,"databundleVersionId":16185886,"modelInstanceId":568583,"modelId":580914,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":783799,"databundleVersionId":16045617,"modelInstanceId":568583,"modelId":580914,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":770244,"databundleVersionId":15883122,"modelInstanceId":568583,"modelId":580914,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":769409,"databundleVersionId":15875786,"modelInstanceId":568583,"modelId":580914,"isSourceIdPinned":false}],"dockerImageVersionId":31287,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import gc\nimport math\nimport re\nimport warnings\nfrom contextlib import nullcontext\nfrom dataclasses import dataclass, field\nfrom pathlib import Path\nfrom typing import List\n\nimport numpy as np\nimport pandas as pd\nimport sacrebleu\nimport torch\nfrom torch.utils.data import DataLoader, Dataset, Sampler\nfrom transformers import AutoModelForSeq2SeqLM, AutoTokenizer\nfrom tqdm.auto import tqdm\n\nwarnings.filterwarnings(\"ignore\")\n\n# ---------------------------------------------------------------------------\n# 0.  Config\n# ---------------------------------------------------------------------------\n\n@dataclass\nclass Config:\n    data_path: str = \"/kaggle/input/competitions/deep-past-initiative-machine-translation/test.csv\"\n    output_dir: str = \"/kaggle/working/\"\n\n    model_paths: List[str] = field(default_factory=lambda: [\n        \"/kaggle/input/models/extergeist/byt5-training/transformers/default/15/folds/model_fold_1\",\n        \"/kaggle/input/models/extergeist/byt5-training/transformers/default/15/folds/model_fold_2\",\n        \"/kaggle/input/models/extergeist/byt5-groktuner/transformers/default/41/best_model_bidir\",\n    ])\n    model_weights: List[float] = field(default_factory=lambda: [0.988, 0.98, 0.396])\n\n    tokenizer_path: str = \"\"   # filled in __post_init__\n\n    # Inference\n    max_input_length: int = 512\n    max_new_tokens: int = 384\n    batch_size: int = 4\n    num_workers: int = 2\n    num_buckets: int = 6\n\n    # Beam search candidates\n    num_beams: int = 8\n    num_beam_cands: int = 4\n    length_penalty: float = 1.12\n    repetition_penalty: float = 1.2\n    early_stopping: bool = True\n\n    # Multi-temperature sampling candidates\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    top_p: float = 0.92\n\n    # MBR weights\n    mbr_pool_cap: int = 32\n    mbr_w_chrf:    float = 0.55   # character n-gram F-score — strongest signal\n    mbr_w_bleu:    float = 0.25   # word-level precision\n    mbr_w_jaccard: float = 0.20   # fast token overlap regulariser\n    mbr_w_length:  float = 0.10   # penalty for outlier lengths\n\n    # Hardware\n    use_bf16: bool = True\n    use_bucket_batching: bool = True\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        self.tokenizer_path = self.tokenizer_path or self.model_paths[1]\n        if self.device.type != \"cuda\":\n            self.use_bf16 = False\n        self.bf16_ok = (\n            self.use_bf16\n            and self.device.type == \"cuda\"\n            and getattr(torch.cuda, \"is_bf16_supported\", lambda: False)()\n        )\n\n    @property\n    def num_sample_cands(self) -> int:\n        return len(self.sample_temperatures) * self.num_sample_per_temp\n\n    @property\n    def total_cands_per_sample(self) -> int:\n        return self.num_beam_cands + self.num_sample_cands\n\n\n# ---------------------------------------------------------------------------\n# 1.  Preprocessing\n# ---------------------------------------------------------------------------\n\n_V2 = re.compile(r\"([aAeEiIuU])(?:2|₂)\")\n_V3 = re.compile(r\"([aAeEiIuU])(?:3|₃)\")\n_ACUTE = str.maketrans({\"a\": \"á\", \"e\": \"é\", \"i\": \"í\", \"u\": \"ú\",\n                         \"A\": \"Á\", \"E\": \"É\", \"I\": \"Í\", \"U\": \"Ú\"})\n_GRAVE = str.maketrans({\"a\": \"à\", \"e\": \"è\", \"i\": \"ì\", \"u\": \"ù\",\n                         \"A\": \"À\", \"E\": \"È\", \"I\": \"Ì\", \"U\": \"Ù\"})\n\n_CHAR_TRANS = str.maketrans({\n    \"ḫ\": \"h\", \"Ḫ\": \"H\", \"ʾ\": \"\",\n    \"₀\": \"0\", \"₁\": \"1\", \"₂\": \"2\", \"₃\": \"3\", \"₄\": \"4\",\n    \"₅\": \"5\", \"₆\": \"6\", \"₇\": \"7\", \"₈\": \"8\", \"₉\": \"9\",\n    \"—\": \"-\", \"–\": \"-\",\n})\n\n_GAP_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*\\)|\\(\\s*break\\s*\\)\"\n    r\"|\\(\\s*\\d+\\s+broken\\s+lines?\\s*\\)\",\n    re.I,\n)\n\n_DET_UPPER_RE = re.compile(r\"\\(([A-ZŠṬṢḪ\\u00C0-\\u00DE\\u0160\\u1E00-\\u1EFF0-9]{1,6})\\)\")\n_DET_LOWER_RE = re.compile(r\"\\(([a-zšṭṣḫ\\u00E0-\\u00FF\\u0161\\u1E01-\\u1EFF]{1,4})\\)\")\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_WS_RE = re.compile(r\"\\s+\")\n\n\ndef _ascii_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\ndef preprocess_batch(texts: List[str]) -> List[str]:\n    s = pd.Series(texts).fillna(\"\").astype(str)\n    s = s.apply(_ascii_diacritics)\n    s = s.str.replace(_DET_UPPER_RE, r\"\\1\", regex=True)\n    s = s.str.replace(_DET_LOWER_RE, r\"{\\1}\", regex=True)\n    s = s.str.replace(_GAP_RE, \"<gap>\", regex=True)\n    s = s.str.translate(_CHAR_TRANS)\n    s = s.str.replace(_EXACT_FRAC_RE, lambda m: _EXACT_FRAC_MAP[m.group(0)], regex=True)\n    s = s.str.replace(_WS_RE, \" \", regex=True).str.strip()\n    return s.tolist()\n\n\n# ---------------------------------------------------------------------------\n# 2.  Postprocessing\n# ---------------------------------------------------------------------------\n\n_SOFT_GRAM_RE = re.compile(\n    r\"\\(\\s*(?:fem|plur|pl|sing|singular|plural|\\?|!)\"\n    r\"(?:\\.\\s*(?:plur|plural|sing|singular))?\"\n    r\"\\.?\\s*[^)]*\\)\",\n    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_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 = {\n    \"I\": 1, \"II\": 2, \"III\": 3, \"IV\": 4, \"V\": 5, \"VI\": 6,\n    \"VII\": 7, \"VIII\": 8, \"IX\": 9, \"X\": 10, \"XI\": 11, \"XII\": 12,\n}\n_REPEAT_WORD_RE = re.compile(r\"\\b(\\w+)(?:\\s+\\1\\b)+\")\n_REPEAT_PUNCT_RE = re.compile(r\"([.,])\\1+\")\n_PUNCT_SPACE_RE = re.compile(r\"\\s+([.,:])\")\n_STRAY_MARKS_RE = re.compile(r\"<<[^>]*>>|<(?!gap\\b)[^>]*>\")\n_MULTI_GAP_RE = re.compile(r\"(?:<gap>\\s*){2,}\")\n_FORBIDDEN_TRANS = str.maketrans(\"\", \"\", '()——<>⌈⌋⌊[]+ʾ;')\n\n\ndef postprocess_batch(texts: List[str]) -> List[str]:\n    s = pd.Series(texts).fillna(\"\").astype(str)\n    s = s.str.replace(_GAP_RE, \"<gap>\", regex=True)\n    s = s.str.replace(_STRAY_MARKS_RE, \"\", 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(_CURLY_RE, \"\", regex=True)\n    s = s.str.replace(_MONTH_RE, lambda m: f\"Month {_ROMAN2INT.get(m.group(1).upper(), m.group(1))}\", regex=True)\n    s = s.str.replace(_MULTI_GAP_RE, \"<gap>\", regex=True)\n    s = s.str.replace(_EXACT_FRAC_RE, lambda m: _EXACT_FRAC_MAP[m.group(0)], 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()\n\n\n# ---------------------------------------------------------------------------\n# 3.  Dataset and bucket batching\n# ---------------------------------------------------------------------------\n\nclass AkkadianDataset(Dataset):\n    def __init__(self, df: pd.DataFrame):\n        self.ids = df[\"id\"].tolist()\n        proc = preprocess_batch(df[\"transliteration\"].tolist())\n        self.texts = [\"translate Akkadian to English: \" + t for t in proc]\n\n    def __len__(self):\n        return len(self.ids)\n\n    def __getitem__(self, idx):\n        return self.ids[idx], self.texts[idx]\n\n\nclass BucketBatchSampler(Sampler):\n    def __init__(self, dataset: AkkadianDataset, batch_size: int, num_buckets: int):\n        self.batch_size = batch_size\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    def __iter__(self):\n        for bucket in self.buckets:\n            for i in range(0, len(bucket), self.batch_size):\n                yield bucket[i: i + self.batch_size]\n\n    def __len__(self):\n        return sum(math.ceil(len(b) / self.batch_size) for b in self.buckets)\n\n\ndef make_collate(tokenizer, max_len: int):\n    def collate(batch):\n        ids = [x[0] for x in batch]\n        texts = [x[1] for x in batch]\n        enc = tokenizer(texts, max_length=max_len, padding=True,\n                        truncation=True, return_tensors=\"pt\")\n        return ids, enc\n    return collate\n\n\n# ---------------------------------------------------------------------------\n# 4.  Weight-blended model\n# ---------------------------------------------------------------------------\n\ndef load_blended_model(cfg: Config) -> torch.nn.Module:\n    total = sum(cfg.model_weights)\n    W = [w / total for w in cfg.model_weights]\n\n    print(\"Loading checkpoints for weight blending...\")\n    base = AutoModelForSeq2SeqLM.from_pretrained(cfg.model_paths[1])\n    sd = base.state_dict()\n\n    extra_sds = [\n        AutoModelForSeq2SeqLM.from_pretrained(p).state_dict()\n        for p in [cfg.model_paths[0], cfg.model_paths[2]]\n    ]\n\n    for k in sd:\n        val = W[1] * sd[k]\n        norm = W[1]\n        for w, esd in zip([W[0], W[2]], extra_sds):\n            if k in esd:\n                val = val + w * esd[k]\n                norm += w\n        sd[k] = val / norm\n\n    del extra_sds\n    gc.collect()\n\n    base.load_state_dict(sd)\n    model = base.to(cfg.device).eval().float()\n    print(f\"Blended model loaded ({sum(p.numel() for p in model.parameters()):,} params)\")\n    return model\n\n\n# ---------------------------------------------------------------------------\n# 5.  MBR selector  (chrF++ + BLEU + Jaccard + length prior)\n# ---------------------------------------------------------------------------\n\nclass MBRSelector:\n    \"\"\"\n    Weighted MBR using four complementary signals:\n      - chrF++ (w=0.55): character n-gram F-score, strongest signal for\n        morphologically rich languages like Akkadian.\n      - BLEU   (w=0.25): word-level precision, rewards fluent phrasing.\n      - Jaccard(w=0.20): fast token overlap, regularises the above two.\n      - Length (w=0.10): Gaussian prior that penalises outlier lengths.\n\n    The pairwise (chrF++ + BLEU + Jaccard) scores are normalised by their\n    combined weight before adding the length bonus, so the length term\n    stays consistently scaled regardless of the other weights.\n    \"\"\"\n\n    def __init__(self, cfg: Config):\n        self.pool_cap  = cfg.mbr_pool_cap\n        self.w_chrf    = cfg.mbr_w_chrf\n        self.w_bleu    = cfg.mbr_w_bleu\n        self.w_jaccard = cfg.mbr_w_jaccard\n        self.w_length  = cfg.mbr_w_length\n        self._pw_total = max(cfg.mbr_w_chrf + cfg.mbr_w_bleu + cfg.mbr_w_jaccard, 1e-9)\n\n        # Instantiate sacrebleu metrics once and reuse across all calls\n        self._chrf = sacrebleu.metrics.CHRF(word_order=2)\n        self._bleu = sacrebleu.metrics.BLEU(effective_order=True)\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.sentence_score(a, [b]).score)\n\n    def _bleu_score(self, a: str, b: str) -> float:\n        if not a or not b:\n            return 0.0\n        try:\n            return float(self._bleu.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(self, a: str, b: str) -> float:\n        score = (\n            self.w_chrf    * self._chrfpp(a, b)\n            + self.w_bleu    * self._bleu_score(a, b)\n            + self.w_jaccard * self._jaccard(a, b)\n        )\n        return score / self._pw_total\n\n    @staticmethod\n    def _length_bonus(lengths: List[int], idx: int) -> float:\n        if not lengths:\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)[: self.pool_cap]\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(cands[i], cands[j])\n                for j in range(n) if j != i\n            ) / max(1, n - 1)\n            lb = self._length_bonus(lengths, i)\n            scores.append(pw + self.w_length * lb)\n\n        return cands[int(np.argmax(scores))]\n\n\n# ---------------------------------------------------------------------------\n# 6.  Candidate generation\n# ---------------------------------------------------------------------------\n\ndef _bf16_ctx(cfg: Config):\n    if cfg.bf16_ok:\n        return torch.autocast(device_type=\"cuda\", dtype=torch.bfloat16)\n    return nullcontext()\n\n\ndef generate_candidates(\n    model,\n    tokenizer,\n    input_ids: torch.Tensor,\n    attention_mask: torch.Tensor,\n    cfg: Config,\n) -> List[List[str]]:\n    B   = input_ids.shape[0]\n    ctx = _bf16_ctx(cfg)\n\n    with ctx:\n        beam_out = model.generate(\n            input_ids=input_ids,\n            attention_mask=attention_mask,\n            do_sample=False,\n            num_beams=cfg.num_beams,\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 = tokenizer.batch_decode(beam_out, skip_special_tokens=True)\n\n        samp_by_temp: List[List[str]] = []\n        if cfg.use_sampling:\n            for temp in cfg.sample_temperatures:\n                try:\n                    samp_out = model.generate(\n                        input_ids=input_ids,\n                        attention_mask=attention_mask,\n                        do_sample=True,\n                        num_beams=1,\n                        top_p=cfg.top_p,\n                        temperature=temp,\n                        num_return_sequences=cfg.num_sample_per_temp,\n                        max_new_tokens=cfg.max_new_tokens,\n                        repetition_penalty=cfg.repetition_penalty,\n                        use_cache=True,\n                    )\n                    samp_by_temp.append(\n                        tokenizer.batch_decode(samp_out, skip_special_tokens=True)\n                    )\n                except Exception as e:\n                    print(f\"  Sampling @ temp={temp:.2f} failed: {e} — skipping\")\n                    samp_by_temp.append([\"\"] * (B * cfg.num_sample_per_temp))\n\n    Rb, Rs = cfg.num_beam_cands, cfg.num_sample_per_temp\n    pools  = []\n    for i in range(B):\n        pool = list(beam_texts[i * Rb: (i + 1) * Rb])\n        for temp_texts in samp_by_temp:\n            pool.extend(temp_texts[i * Rs: (i + 1) * Rs])\n        pools.append(pool)\n\n    return pools\n\n\n# ---------------------------------------------------------------------------\n# 7.  Main inference loop\n# ---------------------------------------------------------------------------\n\ndef run(cfg: Config):\n    print(f\"Device      : {cfg.device}\")\n    print(f\"BF16        : {cfg.bf16_ok}\")\n    print(f\"Cands/sample: {cfg.num_beam_cands} beam + \"\n          f\"{cfg.num_sample_cands} sampled = {cfg.total_cands_per_sample}\")\n    print(f\"MBR weights : chrF++={cfg.mbr_w_chrf}  BLEU={cfg.mbr_w_bleu}  \"\n          f\"Jaccard={cfg.mbr_w_jaccard}  length={cfg.mbr_w_length}\")\n\n    df = pd.read_csv(cfg.data_path, encoding=\"utf-8\")\n    print(f\"Test samples: {len(df)}\")\n\n    tokenizer = AutoTokenizer.from_pretrained(cfg.tokenizer_path)\n    model     = load_blended_model(cfg)\n    dataset   = AkkadianDataset(df)\n    collate   = make_collate(tokenizer, cfg.max_input_length)\n\n    if cfg.use_bucket_batching:\n        sampler = BucketBatchSampler(dataset, cfg.batch_size, cfg.num_buckets)\n        loader  = DataLoader(\n            dataset, batch_sampler=sampler,\n            num_workers=cfg.num_workers, collate_fn=collate,\n            pin_memory=(cfg.device.type == \"cuda\"),\n        )\n    else:\n        loader = DataLoader(\n            dataset, batch_size=cfg.batch_size, shuffle=False,\n            num_workers=cfg.num_workers, collate_fn=collate,\n            pin_memory=(cfg.device.type == \"cuda\"),\n        )\n\n    mbr       = MBRSelector(cfg)\n    raw_pools: dict = {}\n\n    with torch.inference_mode():\n        for batch_ids, enc in tqdm(loader, desc=\"Generating candidates\"):\n            input_ids = enc.input_ids.to(cfg.device, non_blocking=True)\n            attn      = enc.attention_mask.to(cfg.device, non_blocking=True)\n\n            try:\n                pools = generate_candidates(model, tokenizer, input_ids, attn, cfg)\n                for sid, pool in zip(batch_ids, pools):\n                    raw_pools[str(sid)] = pool\n            except RuntimeError as e:\n                if \"out of memory\" in str(e).lower():\n                    print(f\"  OOM — skipping batch, ids={batch_ids}\")\n                    torch.cuda.empty_cache()\n                    for sid in batch_ids:\n                        raw_pools.setdefault(str(sid), [])\n                else:\n                    raise\n\n            if torch.cuda.is_available():\n                torch.cuda.empty_cache()\n\n    results    = []\n    sample_ids = [str(s) for s in dataset.ids]\n\n    for sid in tqdm(sample_ids, desc=\"MBR selection\"):\n        raw     = raw_pools.get(sid, [])\n        cleaned = postprocess_batch(raw) if raw else []\n        chosen  = mbr.pick(cleaned)\n        if not chosen or not chosen.strip():\n            chosen = \"The tablet is too damaged to translate.\"\n        results.append((sid, chosen))\n\n    result_df = pd.DataFrame(results, columns=[\"id\", \"translation\"])\n\n    empty = result_df[\"translation\"].str.strip().eq(\"\").sum()\n    lens  = result_df[\"translation\"].str.len()\n    print(f\"\\nResults  : {len(result_df)} rows\")\n    print(f\"Empty    : {empty}\")\n    print(f\"Len mean : {lens.mean():.1f}  median: {lens.median():.1f}  \"\n          f\"min: {lens.min()}  max: {lens.max()}\")\n\n    out = Path(cfg.output_dir) / \"submission.csv\"\n    result_df.to_csv(out, index=False)\n    print(f\"\\nSaved → {out}\")\n    return result_df\n\n\n# ---------------------------------------------------------------------------\n# 8.  Entry point\n# ---------------------------------------------------------------------------\n\nif __name__ == \"__main__\":\n    cfg = Config()\n    run(cfg)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}