{"cells":[{"cell_type":"markdown","metadata":{},"source":"# 🏺 Improved Akkadian Translation - 3-Model Ensemble with Geometric MBR\n\n**Key improvements over LB 35.9 baseline:**\n1. **Three-model ensemble** instead of two\n2. **More candidate diversity** - 6 beam + 3 sampling per model\n3. **Geometric MBR** - Uses geometric mean of BLEU and chrF++ (competition metric)\n4. **Enhanced preprocessing** with Old Assyrian commodity hints\n5. **Improved postprocessing** with better fraction/commodity normalization\n\n## Models Used:\n- **Model A**: `byt5-akkadian-optimized-34x` (high quality base)\n- **Model B**: `byt5-akkadian-mbr-v2` (MBR-tuned)\n- **Model C**: mattiaangeli's MBR model (diverse training)\n\n### Credits\n- @mattiaangeli and @assiaben for pre-trained models\n- @giovannyrodrguez for the 35.9 baseline architecture"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"import os\nimport gc\nimport re\nimport json\nimport math\nimport random\nimport logging\nimport warnings\nfrom pathlib import Path\nfrom dataclasses import dataclass\nfrom contextlib import nullcontext\nfrom typing import List, Tuple, Dict, Optional\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\")"},{"cell_type":"markdown","metadata":{},"source":"## Configuration"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"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()\n\n@dataclass\nclass ImprovedEnsembleConfig:\n test_data_path: str = \"/kaggle/input/deep-past-initiative-machine-translation/test.csv\"\n output_dir: str = \"/kaggle/working/\"\n \n # Three models for improved diversity\n # Model A: assiaben/final-byt5 - high quality base model\n model_a_path: str = \"/kaggle/input/final-byt5/byt5-akkadian-optimized-34x\"\n # Model B: llkh0a/byt5-akkadian-model - diverse training\n model_b_path: str = \"/kaggle/input/byt5-akkadian-model\"\n # Model C: mattiaangeli MBR model - MBR-optimized\n model_c_path: str = \"/kaggle/input/models/mattiaangeli/byt5-akkadian-mbr/pytorch/default/6\"\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 # Increased candidate generation for better MBR\n num_beam_cands: int = 6\n num_beams: int = 10\n length_penalty: float = 1.25\n early_stopping: bool = True\n repetition_penalty: float = 1.15\n \n # More sampling candidates\n num_sample_cands: int = 3\n mbr_top_p: float = 0.90\n mbr_temperature: float = 0.70\n mbr_pool_cap: int = 48 # Increased pool cap for 3 models\n \n use_mixed_precision: bool = True\n use_better_transformer: bool = True\n use_bucket_batching: bool = True\n use_adaptive_beams: bool = True\n checkpoint_freq: int = 200\n \n # Use geometric mean for MBR (matching competition metric)\n use_geometric_mbr: 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 \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 assert self.num_beams >= self.num_beam_cands"},{"cell_type":"markdown","metadata":{},"source":"## Logging"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"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 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) / \"improved_ensemble.log\"),\n ],\n )\n return logging.getLogger(\"improved_ensemble\")"},{"cell_type":"markdown","metadata":{},"source":"## Enhanced Preprocessing with OA Hints"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"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# Old Assyrian commodity hints (improves translation of business terms)\nOA_HINTS = {\n \"KÙ.BABBAR\": \"silver\",\n \"KU.BABBAR\": \"silver\",\n \"AN.NA\": \"tin\",\n \"MA.NA\": \"mina\",\n \"GÍN\": \"shekel\",\n \"GIN\": \"shekel\",\n \"TÚG\": \"textile\",\n \"TUG\": \"textile\",\n \"GADA\": \"linen\",\n \"ŠE\": \"barley\",\n \"SE\": \"barley\",\n \"URUDU\": \"copper\",\n \"ANŠE\": \"donkey\",\n \"É\": \"house\",\n \"DUMU\": \"son\",\n \"DUMU.MUNUS\": \"daughter\",\n}\n\n_ALLOWED_FRACS = [\n (1/6, \"0.16666\"), (1/4, \"0.25\"), (1/3, \"0.33333\"),\n (1/2, \"0.5\"), (2/3, \"0.66666\"), (3/4, \"0.75\"), (5/6, \"0.83333\"),\n]\n_FRAC_TOL = 2e-3\n_FLOAT_RE = re.compile(r\"(? str:\n ip = int(math.floor(x + 1e-12))\n frac = x - ip\n best = min(_ALLOWED_FRACS, key=lambda t: abs(frac - t[0]))\n 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_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\n_DET_UPPER_RE = re.compile(r\"\\(([\" + _UNICODE_UPPER + r\"0-9]{1,6})\\)\")\n_DET_LOWER_RE = re.compile(r\"\\(([\" + _UNICODE_LOWER + r\"]{1,4})\\)\")\n\n_PN_RE = re.compile(r\"\\bPN\\b\")\n_KUBABBAR_RE = re.compile(r\"KÙ\\.B\\.\")\n\n_EXACT_FRAC_RE = re.compile(r\"0\\.8333|0\\.6666|0\\.3333|0\\.1666|0\\.625|0\\.75|0\\.25|0\\.5\")\n_EXACT_FRAC_MAP = {\n \"0.8333\": \"⅚\", \"0.6666\": \"⅔\", \"0.3333\": \"⅓\", \"0.1666\": \"⅙\",\n \"0.625\": \"⅝\", \"0.75\": \"¾\", \"0.25\": \"¼\", \"0.5\": \"½\",\n}\n\ndef _frac_repl(m: re.Match) -> str:\n return _EXACT_FRAC_MAP[m.group(0)]\n\n_WS_RE = re.compile(r\"\\s+\")\n\nclass EnhancedPreprocessor:\n \"\"\"Enhanced preprocessor with OA commodity hints.\"\"\"\n \n def _extract_hints(self, text: str) -> str:\n \"\"\"Extract Old Assyrian commodity hints from text.\"\"\"\n found = []\n upper_text = text.upper()\n for logogram, meaning in OA_HINTS.items():\n if logogram.upper() in upper_text:\n found.append(meaning)\n if found:\n # Remove duplicates and limit\n unique = list(dict.fromkeys(found))[:4]\n return f\" [Commodities: {', '.join(unique)}]\"\n return \"\"\n \n def preprocess_batch(self, texts: List[str]) -> List[str]:\n ser = pd.Series(texts).fillna(\"\").astype(str)\n ser = ser.apply(_ascii_to_diacritics)\n \n # Uppercase determinatives are unwrapped, lowercase ones converted to braces\n ser = ser.str.replace(_DET_UPPER_RE, r\"\\1\", regex=True)\n ser = ser.str.replace(_DET_LOWER_RE, r\"{\\1}\", regex=True)\n \n ser = _normalize_gaps_vec(ser)\n ser = ser.str.translate(_CHAR_TRANS)\n ser = ser.str.replace(_SUB_X, \"\", regex=False)\n ser = ser.str.replace(_KUBABBAR_RE, \"KÙ.BABBAR\", regex=True)\n ser = ser.str.replace(_EXACT_FRAC_RE, _frac_repl, regex=True)\n ser = ser.str.replace(_FLOAT_RE, lambda m: _canon_decimal(float(m.group(1))), regex=True)\n ser = ser.str.replace(_WS_RE, \" \", regex=True).str.strip()\n \n # Add commodity hints\n results = []\n for text in ser.tolist():\n hints = self._extract_hints(text)\n results.append(text + hints)\n \n return results"},{"cell_type":"markdown","metadata":{},"source":"## Enhanced Postprocessing"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"_SOFT_GRAM_RE = re.compile(\n r\"\\(\\s*(?:fem|plur|pl|sing|singular|plural|\\?|\\!)\"\n r\"(?:\\.\\s*(?:plur|plural|sing|singular))?\"\n r\"\\.?\\s*[^)]*\\)\", re.I\n)\n_BARE_GRAM_RE = re.compile(r\"(?⌈⌋⌊[]+ʾ;')\n\n# Enhanced commodity replacements\n_COMMODITY_RE = re.compile(r'-(gold|tax|textiles)\\b')\n_COMMODITY_REPL = {\n \"gold\": \"pašallum gold\",\n \"tax\": \"šadduātum tax\",\n \"textiles\": \"kutānum textiles\",\n}\n\ndef _commodity_repl(m: re.Match) -> str:\n return _COMMODITY_REPL[m.group(1)]\n\n# Enhanced shekel replacements\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 (re.compile(r'1\\s*/\\s*3\\s+shekels?', re.I), '⅓ shekel'),\n (re.compile(r'2\\s*/\\s*3\\s+shekels?', re.I), '⅔ shekel'),\n (re.compile(r'1\\s*/\\s*2\\s+shekels?', re.I), '½ shekel'),\n (re.compile(r'1\\s*/\\s*4\\s+shekels?', re.I), '¼ shekel'),\n (re.compile(r'3\\s*/\\s*4\\s+shekels?', re.I), '¾ shekel'),\n]\n\n# Additional numeric normalizations\n_MINA_REPLS = [\n (re.compile(r'1\\s*/\\s*3\\s+minas?', re.I), '⅓ mina'),\n (re.compile(r'2\\s*/\\s*3\\s+minas?', re.I), '⅔ mina'),\n (re.compile(r'1\\s*/\\s*2\\s+minas?', re.I), '½ mina'),\n]\n\n_SLASH_ALT_RE = re.compile(r'(?]*>>|<(?!gap\\b)[^>]*>')\n_MULTI_GAP_RE = re.compile(r'(?:\\s*){2,}')\n\n# Remove hint brackets from output\n_HINT_BRACKET_RE = re.compile(r'\\s*\\[Commodities:[^\\]]*\\]', re.I)\n\ndef _month_repl(m: re.Match) -> str:\n return f\"Month {_ROMAN2INT.get(m.group(1).upper(), m.group(1))}\"\n\nclass EnhancedPostprocessor:\n \"\"\"Enhanced postprocessor with better normalization.\"\"\"\n \n def postprocess_batch(self, translations: List[str]) -> List[str]:\n s = pd.Series(translations).fillna(\"\").astype(str)\n \n # Remove any leaked hints\n s = s.str.replace(_HINT_BRACKET_RE, \"\", regex=True)\n \n s = _normalize_gaps_vec(s)\n s = s.str.replace(_PN_RE, \"\", regex=True)\n s = s.str.replace(_COMMODITY_RE, _commodity_repl, regex=True)\n \n for pat, repl in _SHEKEL_REPLS:\n s = s.str.replace(pat, repl, regex=True)\n \n for pat, repl in _MINA_REPLS:\n s = s.str.replace(pat, repl, regex=True)\n \n s = s.str.replace(_EXACT_FRAC_RE, _frac_repl, regex=True)\n s = s.str.replace(_FLOAT_RE, lambda m: _canon_decimal(float(m.group(1))), regex=True)\n \n s = s.str.replace(_SOFT_GRAM_RE, \" \", regex=True)\n s = s.str.replace(_BARE_GRAM_RE, \" \", regex=True)\n s = s.str.replace(_UNCERTAIN_RE, \"\", regex=True)\n \n s = s.str.replace(_STRAY_MARKS_RE, \"\", regex=True)\n s = s.str.replace(_SLASH_ALT_RE, \"\", regex=True)\n \n s = s.str.replace(_CURLY_QUOTES_RE, \"\", regex=True)\n \n s = s.str.replace(_MONTH_RE, _month_repl, regex=True)\n s = s.str.replace(_MULTI_GAP_RE, \"\", regex=True)\n \n s = s.str.replace(\"\", \"\\x00GAP\\x00\", regex=False)\n s = s.str.translate(_FORBIDDEN_TRANS)\n s = s.str.replace(\"\\x00GAP\\x00\", \" \", regex=False)\n \n s = s.str.replace(_REPEAT_WORD_RE, r\"\\1\", regex=True)\n for n in range(4, 1, -1):\n pat = r\"\\b((?:\\w+\\s+){\" + str(n-1) + r\"}\\w+)(?:\\s+\\1\\b)+\"\n s = s.str.replace(pat, r\"\\1\", regex=True)\n \n s = s.str.replace(_PUNCT_SPACE_RE, r\"\\1\", regex=True)\n s = s.str.replace(_REPEAT_PUNCT_RE, r\"\\1\", regex=True)\n s = s.str.replace(_WS_RE, \" \", regex=True).str.strip()\n \n return s.tolist()"},{"cell_type":"markdown","metadata":{},"source":"## Dataset and Bucketed Batching"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"class AkkadianDataset(Dataset):\n def __init__(self, df: pd.DataFrame, preprocessor: EnhancedPreprocessor, logger: logging.Logger):\n self.sample_ids = df[\"id\"].tolist()\n proc = preprocessor.preprocess_batch(df[\"transliteration\"].tolist())\n self.input_texts = [\"translate Akkadian to English: \" + t for t in proc]\n logger.info(f\"Dataset: {len(self.sample_ids)} samples\")\n \n def __len__(self):\n return len(self.sample_ids)\n \n def __getitem__(self, idx):\n return self.sample_ids[idx], self.input_texts[idx]\n\n\nclass BucketBatchSampler(Sampler):\n def __init__(self, dataset, batch_size, num_buckets, logger, shuffle=False):\n self.batch_size = batch_size\n self.shuffle = shuffle\n \n lengths = [len(t.split()) for _, t in dataset]\n sorted_idx = sorted(range(len(lengths)), key=lambda i: lengths[i])\n \n bsize = max(1, len(sorted_idx) // max(1, num_buckets))\n self.buckets = [\n sorted_idx[i * bsize : None if i == num_buckets - 1 else (i + 1) * bsize]\n for i in range(num_buckets)\n ]\n \n for i, b in enumerate(self.buckets):\n if b:\n bl = [lengths[x] for x in b]\n logger.info(f\" Bucket {i}: {len(b)} samples, len [{min(bl)}, {max(bl)}]\")\n \n def __iter__(self):\n for bucket in self.buckets:\n b = list(bucket)\n if self.shuffle:\n random.shuffle(b)\n for i in range(0, len(b), self.batch_size):\n yield b[i:i+self.batch_size]\n \n def __len__(self):\n return sum(math.ceil(len(b) / self.batch_size) for b in self.buckets)"},{"cell_type":"markdown","metadata":{},"source":"## Model Wrapper"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"class ModelWrapper:\n def __init__(self, model_path: str, cfg: ImprovedEnsembleConfig, logger: logging.Logger, label: str):\n self.cfg = cfg\n self.logger = logger\n self.label = label\n self.tokenizer = None\n self.model = None\n \n logger.info(f\"[{label}] Loading from {model_path}\")\n self.tokenizer = AutoTokenizer.from_pretrained(model_path)\n self.model = AutoModelForSeq2SeqLM.from_pretrained(model_path).to(cfg.device).eval()\n \n if cfg.device.type == \"cuda\":\n try:\n torch.set_float32_matmul_precision(\"high\")\n except Exception:\n pass\n \n n = sum(p.numel() for p in self.model.parameters())\n logger.info(f\"[{label}] {n:,} parameters\")\n \n if cfg.device.type == \"cuda\":\n used = torch.cuda.memory_allocated() / 1e9\n logger.info(f\"[{label}] GPU mem used: {used:.2f} GB\")\n \n if cfg.use_better_transformer and cfg.device.type == \"cuda\":\n try:\n from optimum.bettertransformer import BetterTransformer\n self.model = BetterTransformer.transform(self.model)\n logger.info(f\"[{label}] BetterTransformer applied\")\n except Exception as e:\n logger.warning(f\"[{label}] BetterTransformer skipped: {e}\")\n \n def collate(self, batch_samples):\n ids = [s[0] for s in batch_samples]\n texts = [s[1] for s in batch_samples]\n \n enc = self.tokenizer(\n texts,\n max_length=self.cfg.max_input_length,\n padding=True,\n truncation=True,\n return_tensors=\"pt\",\n )\n return ids, enc\n \n def generate_candidates(self, input_ids, attention_mask, beam_size: int) -> List[List[str]]:\n cfg = self.cfg\n B = input_ids.shape[0]\n ctx = _bf16_ctx(cfg.device, cfg.use_bf16_amp)\n \n with ctx:\n nb = max(beam_size, cfg.num_beam_cands)\n beam_out = self.model.generate(\n input_ids=input_ids,\n attention_mask=attention_mask,\n do_sample=False,\n num_beams=nb,\n num_return_sequences=cfg.num_beam_cands,\n max_new_tokens=cfg.max_new_tokens,\n length_penalty=cfg.length_penalty,\n early_stopping=cfg.early_stopping,\n repetition_penalty=cfg.repetition_penalty,\n use_cache=True,\n )\n beam_texts = self.tokenizer.batch_decode(beam_out, skip_special_tokens=True)\n \n samp_texts = []\n if cfg.num_sample_cands > 0:\n samp_out = self.model.generate(\n input_ids=input_ids,\n attention_mask=attention_mask,\n do_sample=True,\n num_beams=1,\n top_p=cfg.mbr_top_p,\n temperature=cfg.mbr_temperature,\n num_return_sequences=cfg.num_sample_cands,\n max_new_tokens=cfg.max_new_tokens,\n repetition_penalty=cfg.repetition_penalty,\n use_cache=True,\n )\n samp_texts = self.tokenizer.batch_decode(samp_out, skip_special_tokens=True)\n \n Rb, Rs = cfg.num_beam_cands, cfg.num_sample_cands\n pools = []\n for i in range(B):\n p = list(beam_texts[i*Rb:(i+1)*Rb])\n if Rs > 0:\n p += list(samp_texts[i*Rs:(i+1)*Rs])\n pools.append(p)\n \n return pools\n \n def unload(self):\n label = self.label\n \n try:\n from optimum.bettertransformer import BetterTransformer\n self.model = BetterTransformer.reverse(self.model)\n except Exception:\n pass\n \n del self.model\n del self.tokenizer\n self.model = None\n self.tokenizer = None\n \n gc.collect()\n if torch.cuda.is_available():\n torch.cuda.empty_cache()\n torch.cuda.synchronize()\n free = (torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_allocated()) / 1e9\n self.logger.info(f\"[{label}] Unloaded. GPU free: {free:.2f} GB\")"},{"cell_type":"markdown","metadata":{},"source":"## Geometric MBR Selector\n\nUses the geometric mean of BLEU and chrF++ to match competition metric."},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"class GeometricMBRSelector:\n \"\"\"MBR selector using geometric mean of BLEU and chrF++ (competition metric).\"\"\"\n \n def __init__(self, pool_cap: int = 48, use_geometric: bool = True):\n self._chrf = sacrebleu.metrics.CHRF(word_order=2)\n self._bleu = sacrebleu.metrics.BLEU(smooth_method='floor', smooth_value=0.1)\n self.pool_cap = pool_cap\n self.use_geometric = use_geometric\n \n def _chrfpp(self, a: str, b: str) -> float:\n a, b = (a or \"\").strip(), (b or \"\").strip()\n if not a or not b:\n return 0.0\n return float(self._chrf.sentence_score(a, [b]).score)\n \n def _bleu_score(self, a: str, b: str) -> float:\n a, b = (a or \"\").strip(), (b or \"\").strip()\n if not a or not b:\n return 0.0\n try:\n return float(self._bleu.sentence_score(a, [b]).score)\n except:\n return 0.0\n \n def _geometric_score(self, a: str, b: str) -> float:\n \"\"\"Compute geometric mean of BLEU and chrF++.\"\"\"\n chrf = self._chrfpp(a, b)\n bleu = self._bleu_score(a, b)\n # Add small epsilon to avoid zero\n return math.sqrt((chrf + 1e-6) * (bleu + 1e-6))\n \n @staticmethod\n def _dedup(xs: List[str]) -> List[str]:\n seen, out = set(), []\n for x in xs:\n x = str(x).strip()\n if x and x not in seen:\n out.append(x)\n seen.add(x)\n return out\n \n def pick(self, candidates: List[str]) -> str:\n cands = self._dedup(candidates)\n \n if self.pool_cap:\n cands = cands[:self.pool_cap]\n \n n = len(cands)\n if n == 0:\n return \"\"\n if n == 1:\n return cands[0]\n \n # Choose scoring function\n score_fn = self._geometric_score if self.use_geometric else self._chrfpp\n \n best_i, best_s = 0, -1e9\n for i in range(n):\n s = sum(score_fn(cands[i], cands[j]) for j in range(n) if j != i)\n s /= max(1, n - 1)\n if s > best_s:\n best_s, best_i = s, i\n \n return cands[best_i]"},{"cell_type":"markdown","metadata":{},"source":"## Ensemble Engine (3 Models)"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"class ImprovedEnsembleEngine:\n \"\"\"Three-model ensemble with geometric MBR.\"\"\"\n \n def __init__(self, cfg: ImprovedEnsembleConfig, logger: logging.Logger):\n self.cfg = cfg\n self.logger = logger\n self.preprocessor = EnhancedPreprocessor()\n self.postprocessor = EnhancedPostprocessor()\n self.mbr = GeometricMBRSelector(\n pool_cap=cfg.mbr_pool_cap,\n use_geometric=cfg.use_geometric_mbr\n )\n \n def _adaptive_beams(self, attn: torch.Tensor) -> int:\n if not self.cfg.use_adaptive_beams:\n return self.cfg.num_beams\n med = float(attn.sum(dim=1).float().median().item())\n short = max(self.cfg.num_beam_cands, self.cfg.num_beams // 2)\n return short if med < 100 else self.cfg.num_beams\n \n def _build_dataloader(self, dataset: AkkadianDataset, wrapper: ModelWrapper) -> DataLoader:\n if self.cfg.use_bucket_batching:\n sampler = BucketBatchSampler(\n dataset,\n self.cfg.batch_size,\n self.cfg.num_buckets,\n self.logger,\n )\n return DataLoader(\n dataset,\n batch_sampler=sampler,\n num_workers=self.cfg.num_workers,\n collate_fn=wrapper.collate,\n pin_memory=(self.cfg.device.type == \"cuda\"),\n )\n \n return DataLoader(\n dataset,\n batch_size=self.cfg.batch_size,\n shuffle=False,\n num_workers=self.cfg.num_workers,\n collate_fn=wrapper.collate,\n pin_memory=(self.cfg.device.type == \"cuda\"),\n )\n \n def _run_one_model(self, wrapper: ModelWrapper, dataset: AkkadianDataset) -> dict:\n dl = self._build_dataloader(dataset, wrapper)\n pools_by_id = {}\n \n with torch.inference_mode():\n for batch_ids, enc in tqdm(dl, desc=f\"[{wrapper.label}]\"):\n input_ids = enc.input_ids.to(self.cfg.device, non_blocking=True)\n attn = enc.attention_mask.to(self.cfg.device, non_blocking=True)\n beam_size = self._adaptive_beams(attn)\n \n try:\n batch_pools = wrapper.generate_candidates(input_ids, attn, beam_size)\n for sid, pool in zip(batch_ids, batch_pools):\n pools_by_id[str(sid)] = pool\n \n except RuntimeError as e:\n if \"out of memory\" in str(e).lower():\n self.logger.error(f\"OOM in [{wrapper.label}] — skipping batch\")\n torch.cuda.empty_cache()\n for sid in batch_ids:\n pools_by_id.setdefault(str(sid), [])\n else:\n raise\n \n except Exception as e:\n self.logger.error(f\"[{wrapper.label}] batch error: {e}\")\n for sid in batch_ids:\n pools_by_id.setdefault(str(sid), [])\n \n if torch.cuda.is_available():\n torch.cuda.empty_cache()\n \n return pools_by_id\n \n def run(self, test_df: pd.DataFrame) -> pd.DataFrame:\n cfg, logger = self.cfg, self.logger\n \n logger.info(\"=\" * 60)\n logger.info(\"Improved 3-Model Ensemble × Geometric MBR\")\n logger.info(f\" Model A : {cfg.model_a_path}\")\n logger.info(f\" Model B : {cfg.model_b_path}\")\n logger.info(f\" Model C : {cfg.model_c_path}\")\n logger.info(f\" Pool per model: beam×{cfg.num_beam_cands} + sample×{cfg.num_sample_cands}\")\n logger.info(f\" MBR pool cap : {cfg.mbr_pool_cap}\")\n logger.info(f\" Geometric MBR : {cfg.use_geometric_mbr}\")\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 # Phase 1/3 — Model A\n logger.info(\"Phase 1/3 — 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/3 — Model B\n logger.info(\"Phase 2/3 — 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 # Phase 3/3 — Model C\n logger.info(\"Phase 3/3 — Model C inference\")\n wrapper_c = ModelWrapper(cfg.model_c_path, cfg, logger, \"Model-C\")\n pools_c = self._run_one_model(wrapper_c, dataset)\n wrapper_c.unload()\n del wrapper_c\n \n logger.info(\"Pool merge + Geometric MBR selection\")\n results: List[Tuple[str, str]] = []\n \n for sid in tqdm(sample_ids, desc=\"MBR\"):\n # Combine all three model pools\n combined = pools_a.get(sid, []) + pools_b.get(sid, []) + pools_c.get(sid, [])\n pp = self.postprocessor.postprocess_batch(combined) if combined else []\n chosen = self.mbr.pick(pp)\n \n if not chosen or not chosen.strip():\n chosen = \"The tablet is too damaged to translate.\"\n \n results.append((sid, chosen))\n \n if len(results) % cfg.checkpoint_freq == 0:\n ckpt = Path(cfg.output_dir) / f\"checkpoint_{len(results)}.csv\"\n pd.DataFrame(results, columns=[\"id\", \"translation\"]).to_csv(ckpt, index=False)\n logger.info(f\" Checkpoint: {len(results)} rows → {ckpt}\")\n \n result_df = pd.DataFrame(results, columns=[\"id\", \"translation\"])\n self._validate(result_df)\n return result_df\n \n def _validate(self, df: pd.DataFrame):\n logger = self.logger\n logger.info(\"=\" * 60)\n \n empty = df[\"translation\"].str.strip().eq(\"\").sum()\n lens = df[\"translation\"].str.len()\n \n logger.info(f\"Empty : {empty} ({100 * empty / max(1, len(df)):.2f}%)\")\n logger.info(\n f\"Len mean : {lens.mean():.1f} median: {lens.median():.1f} min: {lens.min()} max: {lens.max()}\"\n )\n \n for idx in [0, len(df)//4, len(df)//2, 3*len(df)//4, len(df)-1]:\n row = df.iloc[idx]\n logger.info(f\" ID {row['id']}: {str(row['translation'])[:80]}\")\n \n logger.info(\"=\" * 60)"},{"cell_type":"markdown","metadata":{},"source":"## Runtime Info"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"def print_env(cfg: ImprovedEnsembleConfig):\n print(f\"PyTorch : {torch.__version__}\")\n print(f\"CUDA : {torch.cuda.is_available()}\")\n \n if torch.cuda.is_available():\n print(f\"GPU : {torch.cuda.get_device_name(0)}\")\n mem = torch.cuda.get_device_properties(0).total_memory / 1e9\n print(f\"GPU Mem : {mem:.2f} GB\")\n print(f\"BF16 : {_cuda_bf16_supported()}\")\n \n print(f\"BF16 AMP : {cfg.use_bf16_amp}\")\n print(f\"Geometric MBR: {cfg.use_geometric_mbr}\")"},{"cell_type":"markdown","metadata":{},"source":"## Run Inference"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"cfg = ImprovedEnsembleConfig()\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 = ImprovedEnsembleEngine(cfg, logger)\nresults_df = engine.run(test_df)\n\nout_path = Path(cfg.output_dir) / \"submission.csv\"\nresults_df.to_csv(out_path, index=False)\nlogger.info(f\"Saved → {out_path} ({len(results_df)} rows)\")\n\ncfg_snap = {\n \"model_a\": cfg.model_a_path,\n \"model_b\": cfg.model_b_path,\n \"model_c\": cfg.model_c_path,\n \"num_beam_cands\": cfg.num_beam_cands,\n \"num_beams\": cfg.num_beams,\n \"num_sample_cands\": cfg.num_sample_cands,\n \"length_penalty\": cfg.length_penalty,\n \"repetition_penalty\": cfg.repetition_penalty,\n \"mbr_top_p\": cfg.mbr_top_p,\n \"mbr_temperature\": cfg.mbr_temperature,\n \"mbr_pool_cap\": cfg.mbr_pool_cap,\n \"max_new_tokens\": cfg.max_new_tokens,\n \"use_bf16_amp\": cfg.use_bf16_amp,\n \"use_geometric_mbr\": cfg.use_geometric_mbr,\n \"batch_size\": cfg.batch_size,\n}\n\nwith open(Path(cfg.output_dir) / \"improved_ensemble_config.json\", \"w\") as f:\n json.dump(cfg_snap, f, indent=2)\n\nprint(\"\\n\" + \"=\" * 60)\nprint(\"Improved 3-Model Ensemble × Geometric MBR complete\")\nprint(f\"Submission : {out_path}\")\nprint(f\"Total rows : {len(results_df)}\")\nprint(\"=\" * 60)"},{"cell_type":"markdown","metadata":{},"source":"## Notes\n\n### Improvements over LB 35.9:\n1. **3 models** instead of 2 for more diversity\n2. **More candidates** - 6 beam + 3 sampling per model (vs 4+2)\n3. **Geometric MBR** - Uses geometric mean of BLEU and chrF++ matching competition metric\n4. **OA commodity hints** - Helps model focus on business terms\n5. **Enhanced postprocessing** - Better fraction and commodity normalization\n\n### Model Paths (adjust for your setup):\n- Model A: `assiaben/final-byt5` → `byt5-akkadian-optimized-34x`\n- Model B: `giovannyrodrguez/nomarl36` → `byt5-akkadian-model_final`\n- Model C: `mattiaangeli/byt5-akkadian-mbr-v2/pytorch/default/1`"}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.12.0"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":121150,"databundleVersionId":15833391},{"sourceType":"datasetVersion","datasetId":9366631},{"sourceType":"datasetVersion","datasetId":9203761},{"sourceType":"modelInstanceVersion","modelId":585575}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat":4,"nbformat_minor":4}