{"cells": [{"cell_type": "markdown", "metadata": {}, "source": "# \ud83c\udffa Simplified Akkadian Translation - 3-Model Ensemble\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. **Simplified MBR** - Uses character F-score (chrF-like) without sacrebleu dependency\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-model` (diverse training) \n- **Model C**: mattiaangeli's MBR model (MBR-tuned)\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\nfrom collections import Counter\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\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 SimplifiedEnsembleConfig:\n    test_data_path: str = \"/kaggle/input/competitions/deep-past-initiative-machine-translation/test.csv\"\n    output_dir: str = \"/kaggle/working/\"\n    \n    # Three models for improved diversity\n    model_a_path: str = \"/kaggle/input/datasets/assiaben/final-byt5/byt5-akkadian-optimized-34x\"\n    model_b_path: str = \"/kaggle/input/datasets/llkh0a/byt5-akkadian-model\"\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    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) / \"simplified_ensemble.log\"),\n        ],\n    )\n    return logging.getLogger(\"simplified_ensemble\")"}, {"cell_type": "markdown", "metadata": {}, "source": "## Simple CharF-like Metric (No sacrebleu dependency)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def char_f_score(pred: str, ref: str, n: int = 6, beta: float = 2.0) -> float:\n    \"\"\"Simple character-based F-score similar to chrF++.\"\"\"\n    pred = pred.strip().lower()\n    ref = ref.strip().lower()\n    \n    if not pred or not ref:\n        return 0.0\n    \n    # Extract character n-grams\n    def get_ngrams(text: str, n: int) -> Counter:\n        if len(text) < n:\n            return Counter([text])\n        return Counter([text[i:i+n] for i in range(len(text) - n + 1)])\n    \n    pred_ngrams = get_ngrams(pred, n)\n    ref_ngrams = get_ngrams(ref, n)\n    \n    # Calculate precision and recall\n    common = sum((pred_ngrams & ref_ngrams).values())\n    pred_total = sum(pred_ngrams.values())\n    ref_total = sum(ref_ngrams.values())\n    \n    if pred_total == 0:\n        precision = 0.0\n    else:\n        precision = common / pred_total\n    \n    if ref_total == 0:\n        recall = 0.0\n    else:\n        recall = common / ref_total\n    \n    # F-beta score\n    if precision + recall == 0:\n        return 0.0\n    \n    f_score = (1 + beta**2) * precision * recall / (beta**2 * precision + recall)\n    return f_score\n\ndef bleu_score_simple(pred: str, ref: str, n: int = 4) -> float:\n    \"\"\"Simple BLEU-like score using word n-grams.\"\"\"\n    pred_words = pred.strip().lower().split()\n    ref_words = ref.strip().lower().split()\n    \n    if not pred_words or not ref_words:\n        return 0.0\n    \n    # Calculate n-gram precisions\n    precisions = []\n    for i in range(1, n + 1):\n        pred_ngrams = Counter([tuple(pred_words[j:j+i]) for j in range(len(pred_words) - i + 1)])\n        ref_ngrams = Counter([tuple(ref_words[j:j+i]) for j in range(len(ref_words) - i + 1)])\n        \n        common = sum((pred_ngrams & ref_ngrams).values())\n        total = sum(pred_ngrams.values())\n        \n        if total == 0:\n            precisions.append(0.0)\n        else:\n            precisions.append(common / total)\n    \n    # Geometric mean of precisions\n    if all(p > 0 for p in precisions):\n        bleu = (np.prod(precisions) ** (1.0 / len(precisions)))\n    else:\n        bleu = 0.0\n    \n    # Length penalty\n    bp = min(1.0, math.exp(1 - len(ref_words) / max(1, len(pred_words))))\n    \n    return bp * bleu\n\ndef geometric_score(pred: str, ref: str) -> float:\n    \"\"\"Compute geometric mean of BLEU and chrF-like scores.\"\"\"\n    chrf = char_f_score(pred, ref)\n    bleu = bleu_score_simple(pred, ref)\n    # Add small epsilon to avoid zero\n    return math.sqrt((chrf + 1e-6) * (bleu + 1e-6))"}, {"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|\u2082)\")\n_V3 = re.compile(r\"([aAeEiIuU])(?:3|\u2083)\")\n_ACUTE = str.maketrans({\"a\":\"\u00e1\",\"e\":\"\u00e9\",\"i\":\"\u00ed\",\"u\":\"\u00fa\",\"A\":\"\u00c1\",\"E\":\"\u00c9\",\"I\":\"\u00cd\",\"U\":\"\u00da\"})\n_GRAVE = str.maketrans({\"a\":\"\u00e0\",\"e\":\"\u00e8\",\"i\":\"\u00ec\",\"u\":\"\u00f9\",\"A\":\"\u00c0\",\"E\":\"\u00c8\",\"I\":\"\u00cc\",\"U\":\"\u00d9\"})\n\ndef _ascii_to_diacritics(s: str) -> str:\n    s = s.replace(\"sz\", \"\u0161\").replace(\"SZ\", \"\u0160\")\n    s = s.replace(\"s,\", \"\u1e63\").replace(\"S,\", \"\u1e62\")\n    s = s.replace(\"t,\", \"\u1e6d\").replace(\"T,\", \"\u1e6c\")\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\nOA_HINTS = {\n    \"K\u00d9.BABBAR\": \"silver\", \"KU.BABBAR\": \"silver\", \"AN.NA\": \"tin\", \"MA.NA\": \"mina\",\n    \"G\u00cdN\": \"shekel\", \"GIN\": \"shekel\", \"T\u00daG\": \"textile\", \"TUG\": \"textile\",\n    \"GADA\": \"linen\", \"\u0160E\": \"barley\", \"SE\": \"barley\", \"URUDU\": \"copper\",\n    \"AN\u0160E\": \"donkey\", \"\u00c9\": \"house\", \"DUMU\": \"son\", \"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\"(?<![\\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_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,}|\u2026+|\\[\\.+\\]\"\n    r\"|\\[\\s*x\\s*\\]|\\(\\s*x\\s*\\)\"\n    r\"|(?<!\\w)x{2,}(?!\\w)\"\n    r\"|(?<!\\w)x(?!\\w)\"\n    r\"|\\(\\s*large\\s+break\\s*\\)\"\n    r\"|\\(\\s*break\\s*\\)\"\n    r\"|\\(\\s*\\d+\\s+broken\\s+lines?\\s*\\)\",\n    re.I\n)\n\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    \"\u1e2b\":\"h\",\"\u1e2a\":\"H\",\"\u02be\":\"\",\n    \"\u2080\":\"0\",\"\u2081\":\"1\",\"\u2082\":\"2\",\"\u2083\":\"3\",\"\u2084\":\"4\",\n    \"\u2085\":\"5\",\"\u2086\":\"6\",\"\u2087\":\"7\",\"\u2088\":\"8\",\"\u2089\":\"9\",\n    \"\u2014\":\"-\",\"\u2013\":\"-\",\n})\n_SUB_X = \"\u2093\"\n_UNICODE_UPPER = r\"A-Z\u0160\u1e6c\u1e62\u1e2a\\u00C0-\\u00D6\\u00D8-\\u00DE\\u0160\\u1E00-\\u1EFF\"\n_UNICODE_LOWER = r\"a-z\u0161\u1e6d\u1e63\u1e2b\\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_PN_RE = re.compile(r\"\\bPN\\b\")\n_KUBABBAR_RE = re.compile(r\"K\u00d9\\.B\\.\")\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\": \"\u215a\", \"0.6666\": \"\u2154\", \"0.3333\": \"\u2153\", \"0.1666\": \"\u2159\",\n    \"0.625\": \"\u215d\", \"0.75\": \"\u00be\", \"0.25\": \"\u00bc\", \"0.5\": \"\u00bd\",\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    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            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        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\u00d9.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\"(?<!\\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(\"\", \"\", '()\u2014\u2014<>\u2308\u230b\u230a[]+\u02be;')\n_COMMODITY_RE = re.compile(r'-(gold|tax|textiles)\\b')\n_COMMODITY_REPL = {\n    \"gold\": \"pa\u0161allum gold\", \"tax\": \"\u0161addu\u0101tum tax\", \"textiles\": \"kut\u0101num textiles\",\n}\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), '\u2154 shekel 15 grains'),\n    (re.compile(r'7\\s*/\\s*12\\s+shekels?', re.I), '\u00bd 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), '\u2153 shekel'),\n    (re.compile(r'2\\s*/\\s*3\\s+shekels?', re.I), '\u2154 shekel'),\n    (re.compile(r'1\\s*/\\s*2\\s+shekels?', re.I), '\u00bd shekel'),\n    (re.compile(r'1\\s*/\\s*4\\s+shekels?', re.I), '\u00bc shekel'),\n    (re.compile(r'3\\s*/\\s*4\\s+shekels?', re.I), '\u00be shekel'),\n]\n_MINA_REPLS = [\n    (re.compile(r'1\\s*/\\s*3\\s+minas?', re.I), '\u2153 mina'),\n    (re.compile(r'2\\s*/\\s*3\\s+minas?', re.I), '\u2154 mina'),\n    (re.compile(r'1\\s*/\\s*2\\s+minas?', re.I), '\u00bd mina'),\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_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    def postprocess_batch(self, translations: List[str]) -> List[str]:\n        s = pd.Series(translations).fillna(\"\").astype(str)\n        \n        s = s.str.replace(_HINT_BRACKET_RE, \"\", regex=True)  # Remove leaked hints\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        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        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        \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\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: SimplifiedEnsembleConfig, 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": "## Simplified MBR Selector (No sacrebleu dependency)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "class SimplifiedMBRSelector:\n    \"\"\"MBR selector using simple chrF-like and BLEU-like metrics.\"\"\"\n    \n    def __init__(self, pool_cap: int = 48):\n        self.pool_cap = pool_cap\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        best_i, best_s = 0, -1e9\n        for i in range(n):\n            s = sum(geometric_score(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 SimplifiedEnsembleEngine:\n    def __init__(self, cfg: SimplifiedEnsembleConfig, logger: logging.Logger):\n        self.cfg = cfg\n        self.logger = logger\n        self.preprocessor = EnhancedPreprocessor()\n        self.postprocessor = EnhancedPostprocessor()\n        self.mbr = SimplifiedMBRSelector(pool_cap=cfg.mbr_pool_cap)\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, self.cfg.batch_size, self.cfg.num_buckets, self.logger,\n            )\n            return DataLoader(\n                dataset, batch_sampler=sampler, num_workers=self.cfg.num_workers,\n                collate_fn=wrapper.collate, pin_memory=(self.cfg.device.type == \"cuda\"),\n            )\n        \n        return DataLoader(\n            dataset, batch_size=self.cfg.batch_size, shuffle=False,\n            num_workers=self.cfg.num_workers, 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}] \u2014 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(\"Simplified 3-Model Ensemble \u00d7 Simplified 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\u00d7{cfg.num_beam_cands} + sample\u00d7{cfg.num_sample_cands}\")\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        # Phase 1/3 \u2014 Model A\n        logger.info(\"Phase 1/3 \u2014 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 \u2014 Model B\n        logger.info(\"Phase 2/3 \u2014 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 \u2014 Model C\n        logger.info(\"Phase 3/3 \u2014 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 + Simplified 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 \u2192 {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            if idx < len(df):\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: SimplifiedEnsembleConfig):\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\"Simplified MBR: Character F-score + BLEU-like geometric mean\")"}, {"cell_type": "markdown", "metadata": {}, "source": "## Run Inference"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "cfg = SimplifiedEnsembleConfig()\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 = SimplifiedEnsembleEngine(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 \u2192 {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    \"batch_size\": cfg.batch_size,\n}\n\nwith open(Path(cfg.output_dir) / \"simplified_ensemble_config.json\", \"w\") as f:\n    json.dump(cfg_snap, f, indent=2)\n\nprint(\"\\n\" + \"=\" * 60)\nprint(\"Simplified 3-Model Ensemble \u00d7 Simplified 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. **No sacrebleu dependency** - implements chrF-like and BLEU-like metrics directly\n4. **OA commodity hints** - helps model focus on business terms\n5. **Enhanced postprocessing** - better fraction and commodity normalization\n6. **Geometric mean selection** - combines character and word-level metrics\n\n### Model Paths:\n- Model A: `assiaben/final-byt5` \u2192 `byt5-akkadian-optimized-34x`\n- Model B: `llkh0a/byt5-akkadian-model`\n- Model C: `mattiaangeli/byt5-akkadian-mbr/pytorch/default/6`\n\n### Simplified MBR:\n- Uses character-level F-score (similar to chrF++)\n- Uses word-level BLEU-like score\n- Combines them with geometric mean for final selection\n- No external dependencies required"}], "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": 9630558}, {"sourceType": "datasetVersion", "datasetId": 9203761}, {"sourceType": "modelInstanceVersion", "modelId": 585575}], "isInternetEnabled": false, "language": "python", "sourceType": "notebook", "isGpuEnabled": true}}, "nbformat": 4, "nbformat_minor": 4}