{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":121150,"databundleVersionId":15833391},{"sourceType":"datasetVersion","sourceId":15105457,"datasetId":9671515,"databundleVersionId":15991409},{"sourceType":"datasetVersion","sourceId":14899351,"datasetId":9366631,"databundleVersionId":15763950},{"sourceType":"modelInstanceVersion","sourceId":755680,"databundleVersionId":15750522,"modelInstanceId":577094,"modelId":589419}],"dockerImageVersionId":31329,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"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\")\n\n\n# 1. Mixed precision helpers\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\ndef _bf16_ctx(device: torch.device, enabled: bool):\n if enabled and device.type == \"cuda\" and _cuda_bf16_supported():\n return torch.autocast(device_type=\"cuda\", dtype=torch.bfloat16)\n return nullcontext()\n\n# 2. Configuration\n\n@dataclass\nclass EnsembleMBRConfig:\n test_data_path: str = \"/kaggle/input/competitions/deep-past-initiative-machine-translation/test.csv\"\n output_dir: str = \"/kaggle/working/\"\n # 1. The 3-Model Powerhouse\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 = 128\n batch_size: int = 16\n num_workers: int = 2\n num_buckets: int = 6\n num_beams: int = 2\n num_beam_cands: int = 2\n repetition_penalty: float = 1.0 # <--- CRITICAL: Turns off Alphabet Exhaustion\n sample_temperatures: List[float] = field(default_factory=list)\n num_sample_per_temp: int = 3\n mbr_pool_cap: int = 48\n length_penalty: float = 1.3\n early_stopping: bool = True\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 use_sampling: bool = False\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 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 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\n\n# 3. Logging utilities\n\ndef 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\")\n\n# 4. Preprocessing\n\n_V2 = re.compile(r\"([aAeEiIuU])(?:2|₂)\")\n_V3 = re.compile(r\"([aAeEiIuU])(?:3|₃)\")\n_ACUTE = str.maketrans({\"a\":\"á\",\"e\":\"é\",\"i\":\"í\",\"u\":\"ú\",\"A\":\"Á\",\"E\":\"É\",\"I\":\"Í\",\"U\":\"Ú\"})\n_GRAVE = str.maketrans({\"a\":\"à\",\"e\":\"è\",\"i\":\"ì\",\"u\":\"ù\",\"A\":\"À\",\"E\":\"È\",\"I\":\"Ì\",\"U\":\"Ù\"})\n\ndef _ascii_to_diacritics(s: str) -> str:\n s = s.replace(\"sz\", \"š\").replace(\"SZ\", \"Š\")\n s = s.replace(\"s,\", \"ṣ\").replace(\"S,\", \"Ṣ\")\n s = s.replace(\"t,\", \"ṭ\").replace(\"T,\", \"Ṭ\")\n s = _V2.sub(lambda m: m.group(1).translate(_ACUTE), s)\n s = _V3.sub(lambda m: m.group(1).translate(_GRAVE), s)\n return s\n\n_ALLOWED_FRACS = [\n (1/6, \"0.16666\"), (1/4, \"0.25\"), (1/3, \"0.33333\"),\n (1/2, \"0.5\"), (2/3, \"0.66666\"), (3/4, \"0.75\"), (5/6, \"0.83333\"),\n]\n_FRAC_TOL = 2e-3\n_FLOAT_RE = re.compile(r\"(? str:\n ip = int(math.floor(x + 1e-12))\n frac = x - ip\n best = min(_ALLOWED_FRACS, key=lambda t: abs(frac - t[0]))\n 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\"|(? pd.Series:\n return ser.str.replace(_GAP_UNIFIED_RE, \"\", 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()\n\n# 5. Postprocessing\n\n_SOFT_GRAM_RE = re.compile(\n r\"\\(\\s*(?:fem|plur|pl|sing|singular|plural|\\?|\\!)\"\n r\"(?:\\.\\s*(?:plur|plural|sing|singular))?\"\n r\"\\.?\\s*[^)]*\\)\", re.I\n)\n_BARE_GRAM_RE = re.compile(r\"(?⌈⌋⌊[]+ʾ;')\n_COMMODITY_RE = re.compile(r'(?<=\\s)-(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_SLASH_ALT_RE = re.compile(r'(?]*>>|<(?!gap\\b)[^>]*>')\n_MULTI_GAP_RE = re.compile(r'(?:\\s*){2,}')\n_EXTRA_STRAY_RE = re.compile(r'(? 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, \"\", 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(_EXTRA_STRAY_RE, \"\", regex=True)\n s = s.str.replace(_SLASH_ALT_RE, \"\", regex=True)\n s = s.str.replace(_CURLY_DQ_RE, '\"', regex=True)\n s = s.str.replace(_CURLY_SQ_RE, \"'\", regex=True)\n\n s = s.str.replace(_MONTH_RE, _month_repl, regex=True)\n s = s.str.replace(_MULTI_GAP_RE, \"\", regex=True)\n s = s.str.replace(\"\", \"\\x00GAP\\x00\", regex=False)\n s = s.str.translate(_FORBIDDEN_TRANS)\n s = s.str.replace(\"\\x00GAP\\x00\", \" \", regex=False)\n s = s.str.translate(_HACEK_TRANS)\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# 6. Dataset and bucket batching\n\nclass AkkadianDataset(Dataset):\n def __init__(self, df: pd.DataFrame, preprocessor: OptimizedPreprocessor, logger: logging.Logger):\n self.sample_ids = df[\"id\"].tolist()\n proc = preprocessor.preprocess_batch(df[\"transliteration\"].tolist())\n self.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)\n\n# 7. Model wrapper and candidate generation\n\nclass 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 self.model.config.use_cache = True\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 # 🚀 HYPER-BEAM EXPANSION\n # Override the config to force the 15-beam expansion\n hyper_beams = 15 \n\n with ctx:\n beam_out = self.model.generate(\n input_ids=input_ids,\n attention_mask=attention_mask,\n do_sample=False, # STRICTLY FALSE: No hallucinations allowed\n num_beams=hyper_beams, # Massive search tree\n num_return_sequences=hyper_beams, # Return all 15 top pure candidates\n max_new_tokens=cfg.max_new_tokens,\n length_penalty=1.2, # 1.2 encourages complete sentences\n early_stopping=True,\n repetition_penalty=cfg.repetition_penalty,\n use_cache=True,\n )\n \n # Decode the 15 candidates\n beam_texts = self.tokenizer.batch_decode(beam_out, skip_special_tokens=True)\n\n # We completely bypass the div_out and samp_out blocks to save compute time\n # and prevent the MBR selector from picking a hallucinated temperature sample.\n \n # Reshape the flat list back into a list of candidate lists per batch item\n candidates = []\n for i in range(B):\n start_idx = i * hyper_beams\n end_idx = start_idx + hyper_beams\n candidates.append(beam_texts[start_idx:end_idx])\n \n return candidates\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\")\n\n# 8. Weighted MBR selector\n\nclass 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\n# 9. End-to-end engine\n\nclass EnsembleMBREngine:\n def __init__(self, cfg: EnsembleMBRConfig, logger: logging.Logger):\n self.cfg = cfg\n self.logger = logger\n self.preprocessor = OptimizedPreprocessor()\n self.postprocessor = VectorizedPostprocessor()\n self.mbr = MBRSelector(\n pool_cap=cfg.mbr_pool_cap,\n 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 logger.info(\"=\" * 60)\n logger.info(\"Titan-Prime 2-Model Ensemble + Aggressive RAG Override\")\n logger.info(f\" Model A : {cfg.model_a_path}\")\n logger.info(f\" Model B : {cfg.model_b_path}\")\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 # Phase 1/2 — Model A\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 # Phase 2/2 — Model B\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 \n\n # --- THE EXPANDED LEGAL OVERRIDE INITIALIZATION ---\n logger.info(\"Loading Legal Train Override Dictionary...\")\n from rapidfuzz import process, fuzz\n \n train_dict = {}\n \n # Load the official training data\n train_df = pd.read_csv(\"/kaggle/input/competitions/deep-past-initiative-machine-translation/train.csv\")\n for _, row in train_df.dropna(subset=['transliteration', 'translation']).iterrows():\n cln_akk = self.preprocessor.preprocess_batch([str(row['transliteration'])])[0]\n cln_akk = re.sub(r'\\s*\\[Commodities:[^\\]]*\\]', '', cln_akk).strip()\n train_dict[cln_akk] = str(row['translation']).strip()\n\n # Try to load the Michel Old Assyrian Letters Corpus (if attached)\n try:\n michel_path = \"/kaggle/input/datasets/gauravparkhedkar/akkadian/train.csv\"\n michel_df = pd.read_csv(michel_path)\n for _, row in michel_df.dropna(subset=['akkadian', 'english']).iterrows():\n cln_akk = self.preprocessor.preprocess_batch([str(row['akkadian'])])[0]\n cln_akk = re.sub(r'\\s*\\[Commodities:[^\\]]*\\]', '', cln_akk).strip()\n train_dict[cln_akk] = str(row['english']).strip()\n logger.info(\"✅ Michel Corpus Injected!\")\n except Exception:\n logger.info(\"⚠️ Michel corpus not found. Using train.csv only.\")\n\n train_keys = list(train_dict.keys())\n logger.info(f\"Loaded {len(train_keys)} official rows for Legal Override.\")\n # ---------------------------------------------------\n\n logger.info(\"Phase 3/3 — Pool merge + MBR selection + RAG\")\n results = []\n override_count = 0\n\n for idx, sid in enumerate(tqdm(sample_ids, desc=\"MBR & RAG\")):\n raw_text = str(test_df.iloc[idx]['transliteration'])\n cln_input = self.preprocessor.preprocess_batch([raw_text])[0]\n cln_input_no_hints = re.sub(r'\\s*\\[Commodities:[^\\]]*\\]', '', cln_input).strip()\n\n # 1. AGGRESSIVE LEGAL RAG OVERRIDE (Threshold lowered to 88.0)\n # We use token_sort_ratio to ignore word ordering differences\n match = process.extractOne(cln_input_no_hints, train_keys, scorer=fuzz.token_sort_ratio)\n \n # Length safety check: Ensure the matched strings are roughly the same length\n length_ratio = len(cln_input_no_hints) / max(1, len(match[0])) if match else 0\n \n if match and match[1] >= 88.0 and 0.8 <= length_ratio <= 1.2:\n # Hijack the neural network entirely!\n chosen = train_dict[match[0]]\n override_count += 1\n else:\n # 2. Fallback to the 2-Model Ensemble 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 logger.info(f\"🎯 Aggressive RAG Overrides Triggered: {override_count} / {len(sample_ids)}\")\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)\n\n# 10. Environment summary helper\n\ndef 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()\n\n# 11. Run inference and export submission\n\nif __name__ == \"__main__\":\n cfg = EnsembleMBRConfig()\n logger = setup_logging(cfg.output_dir)\n\n print_env(cfg)\n\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 samples: {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 → {out_path} ({len(results_df)} rows)\")\n results_df.head()\n 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 }\n with open(Path(cfg.output_dir) / \"ensemble_mbr_config.json\", \"w\") as f:\n json.dump(cfg_snap, f, indent=2)\n\n print(\"Submission file:\", out_path)\n print(\"Config saved to:\", Path(cfg.output_dir) / \"ensemble_mbr_config.json\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}