{"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":14410665,"datasetId":9203761,"databundleVersionId":15226414},{"sourceType":"datasetVersion","sourceId":14376272,"datasetId":9181082,"databundleVersionId":15188111},{"sourceType":"datasetVersion","sourceId":14899351,"datasetId":9366631,"databundleVersionId":15763950},{"sourceType":"datasetVersion","sourceId":14374989,"datasetId":9147887,"databundleVersionId":15186684},{"sourceType":"modelInstanceVersion","sourceId":282751,"databundleVersionId":11375937,"modelInstanceId":239470,"modelId":222398}],"dockerImageVersionId":31260,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## 1️⃣ System-Level Optimization\n\nSet optimal environment variables **before** importing PyTorch.","metadata":{}},{"cell_type":"code","source":"# ============================================================\n# Robust ByT5 Inference + Two-pass Self-Ensemble (Merged)\n# - Safe → Conditional Aggressive postprocess\n# - Two runs (cfg A/B) + fuse (badness-first, quality tie-break)\n# - Bucket batching, adaptive beams, optional BetterTransformer, optional FP16\n# - Checkpointing + basic diagnostics\n# ============================================================\n\nimport os\nimport re\nimport gc\nimport json\nimport math\nimport time\nimport logging\nfrom dataclasses import dataclass, asdict\nfrom typing import List, Tuple, Optional, Dict\n\nimport numpy as np\nimport pandas as pd\n\nimport torch\nfrom torch.utils.data import Dataset, DataLoader, Sampler\nfrom transformers import AutoTokenizer, AutoModelForSeq2SeqLM\n\n# ---------------------------\n# 0) Environment knobs\n# ---------------------------\nos.environ.setdefault(\"OMP_NUM_THREADS\", \"4\")\nos.environ.setdefault(\"MKL_NUM_THREADS\", \"4\")\nos.environ.setdefault(\"CUDA_LAUNCH_BLOCKING\", \"0\")\nos.environ.setdefault(\"TORCH_CUDNN_V8_API_ENABLED\", \"1\")\nos.environ.setdefault(\"TOKENIZERS_PARALLELISM\", \"true\")\n\nlogging.basicConfig(\n    level=logging.INFO,\n    format=\"%(asctime)s | %(levelname)s | %(message)s\",\n)\nlogger = logging.getLogger(\"byt-robust\")\n\ndef print_environment_info():\n    logger.info(f\"torch={torch.__version__}\")\n    logger.info(f\"cuda_available={torch.cuda.is_available()}\")\n    if torch.cuda.is_available():\n        logger.info(f\"gpu={torch.cuda.get_device_name(0)}\")\n        props = torch.cuda.get_device_properties(0)\n        logger.info(f\"gpu_mem_total_gb={props.total_memory/1024**3:.2f}\")\n    # BetterTransformer availability (optional)\n    try:\n        from optimum.bettertransformer import BetterTransformer  # noqa: F401\n        logger.info(\"optimum.bettertransformer is available (will try to use if enabled).\")\n    except Exception:\n        logger.info(\"optimum.bettertransformer not available (will skip).\")\n\nprint_environment_info()\n\n# ---------------------------\n# 1) Config\n# ---------------------------\n@dataclass\nclass UltraConfig:\n    # paths\n    test_data_path: str = \"/kaggle/input/deep-past-initiative-machine-translation/test.csv\"   # <-- adjust to your dataset\n    model_path: str = \"/kaggle/input/final-byt5/byt5-akkadian-optimized-34x\"\n    output_dir: str = \"/kaggle/working\"\n\n    # dataloader/tokenization\n    max_length: int = 512\n    batch_size: int = 8\n    num_workers: int = 2\n    pin_memory: bool = True\n\n    # generation\n    num_beams: int = 8\n    max_new_tokens: int = 512\n    length_penalty: float = 1.30\n    early_stopping: bool = True\n    no_repeat_ngram_size: int = 0\n    repetition_penalty: Optional[float] = None\n\n    # performance\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    use_vectorized_postproc: bool = True\n\n    # robustness/ops\n    checkpoint_freq: int = 2000  # save partial predictions periodically\n    empty_cache_every: int = 10  # batches\n    validate_samples: int = 6\n\n    # postprocess behavior\n    postprocess_mode: str = \"safe_then_conditional_aggressive\"  # or \"safe_only\" / \"aggressive_only\"\n    aggressive_trigger_badness: float = 2.5  # if badness >= threshold -> run aggressive refinement\n    min_words_fallback: int = 3\n\n    # device\n    device: str = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n    def __post_init__(self):\n        if self.device == \"cpu\":\n            self.use_mixed_precision = False\n            self.use_better_transformer = False\n            self.pin_memory = False\n\n# ---------------------------\n# 2) Preprocess (input)\n# ---------------------------\nclass OptimizedPreprocessor:\n    def __init__(self):\n        # big gaps: ... … …… etc\n        self.re_big_gap = re.compile(r\"(\\.\\.\\.+|…+)\")\n        # small gaps: occurrences of 'x' / 'xx' in isolation-ish patterns\n        # This is intentionally conservative; you can expand if needed.\n        self.re_small_gap = re.compile(r\"(?:(?<=\\s)|^)(x{1,3})(?=(\\s|$))\", flags=re.IGNORECASE)\n\n    def preprocess_batch(self, texts: List[str]) -> List[str]:\n        s = pd.Series(texts, dtype=\"object\").fillna(\"\")\n        s = s.str.replace(self.re_big_gap, \" <big_gap> \", regex=True)\n        s = s.str.replace(self.re_small_gap, \" <gap> \", regex=True)\n        s = s.str.replace(r\"\\s+\", \" \", regex=True).str.strip()\n        return s.tolist()\n\n# ---------------------------\n# 3) Postprocess (output) – Safe + Aggressive\n# ---------------------------\nclass SafePostprocessor:\n    \"\"\"\n    Conservative cleanup.\n    - Preserve common punctuation and formatting.\n    - Protect <gap>/<big_gap>.\n    - Normalize a few known characters.\n    \"\"\"\n    SUBSCRIPT_TO_NORMAL = str.maketrans({\n        \"₀\":\"0\",\"₁\":\"1\",\"₂\":\"2\",\"₃\":\"3\",\"₄\":\"4\",\"₅\":\"5\",\"₆\":\"6\",\"₇\":\"7\",\"₈\":\"8\",\"₉\":\"9\"\n    })\n\n    def __init__(self, use_unicode_fractions: bool = False, strip_dash_old: bool = False):\n        self.use_unicode_fractions = use_unicode_fractions\n        self.strip_dash_old = strip_dash_old\n\n        self.forbidden_chars = re.compile(r\"[⌈⌋⌊⌉]\")  # very conservative\n        self.multi_space = re.compile(r\"\\s+\")\n        self.space_before_punct = re.compile(r\"\\s+([,.;:!?])\")\n        self.multi_punct = re.compile(r\"([!?.,])\\1{2,}\")\n        self.dashes = re.compile(r\"[–—]\")\n\n        # Protect tokens\n        self.prot_gap = \"\\uFFF0\"\n        self.prot_big = \"\\uFFF1\"\n\n        # Gaps normalization (mild)\n        self.bracket_x = re.compile(r\"\\[\\s*x\\s*\\]|\\(\\s*x\\s*\\)\", re.IGNORECASE)\n        self.bare_x = re.compile(r\"(?:(?<=\\s)|^)x(?=(\\s|$))\", re.IGNORECASE)\n        self.ellipsis = re.compile(r\"(\\.\\.\\.+|…+)\")\n\n        # optional fractions\n        self.frac_map = {\n            \"0.5\": \"½\", \"0.25\": \"¼\", \"0.75\": \"¾\",\n            \"1/2\": \"½\", \"1/4\": \"¼\", \"3/4\": \"¾\",\n        }\n        self.frac_re = re.compile(r\"\\b(0\\.5|0\\.25|0\\.75|1/2|1/4|3/4)\\b\")\n\n    def postprocess_one(self, text: str) -> str:\n        if text is None:\n            text = \"\"\n        t = str(text)\n\n        # Normalize known chars\n        t = t.replace(\"ḫ\", \"h\").replace(\"Ḫ\", \"H\")\n        t = t.translate(self.SUBSCRIPT_TO_NORMAL)\n        t = self.dashes.sub(\"-\", t)\n\n        # Normalize gaps lightly\n        t = self.bracket_x.sub(\"<gap>\", t)\n        t = self.bare_x.sub(\"<gap>\", t)\n        t = self.ellipsis.sub(\"<big_gap>\", t)\n\n        # Protect tokens before stripping\n        t = t.replace(\"<gap>\", self.prot_gap).replace(\"<big_gap>\", self.prot_big)\n\n        # Remove only very rare/unsafe glyphs\n        t = self.forbidden_chars.sub(\"\", t)\n\n        # Restore protected tokens\n        t = t.replace(self.prot_gap, \"<gap>\").replace(self.prot_big, \"<big_gap>\")\n\n        # Optional: old dash stripping behavior (risky; kept for ablation)\n        if self.strip_dash_old:\n            t = t.strip(\" -\")\n\n        # Optional unicode fractions\n        if self.use_unicode_fractions:\n            t = self.frac_re.sub(lambda m: self.frac_map.get(m.group(1), m.group(1)), t)\n\n        # Whitespace & punctuation tidying (safe)\n        t = self.space_before_punct.sub(r\"\\1\", t)\n        t = self.multi_punct.sub(r\"\\1\", t)\n        t = self.multi_space.sub(\" \", t).strip()\n\n        return t\n\n    def postprocess_batch(self, texts: List[str]) -> List[str]:\n        return [self.postprocess_one(x) for x in texts]\n\n\nclass AggressivePostprocessor:\n    \"\"\"\n    More aggressive cleanup used only when output looks \"bad\".\n    - Remove common bracketed grammatical notes\n    - Dedupe repeated words and repeated ngrams\n    - Stronger gap consolidation\n    \"\"\"\n    def __init__(self, ngram_dedupe_max_n: int = 4):\n        self.ngram_dedupe_max_n = ngram_dedupe_max_n\n\n        self.multi_space = re.compile(r\"\\s+\")\n        self.space_before_punct = re.compile(r\"\\s+([,.;:!?])\")\n        self.multi_punct = re.compile(r\"([!?.,])\\1{2,}\")\n\n        self.remove_notes = re.compile(\n            r\"\\((?:plur\\.?|sing\\.?|fem\\.?|masc\\.?|uncertain|\\?|\\!|damaged|broken)\\)\",\n            flags=re.IGNORECASE\n        )\n        self.remove_weird = re.compile(r\"[<>⌈⌋⌊⌉⌊⌋\\+ʾ/;]\")  # somewhat aggressive (still not too crazy)\n\n        self.gap_runs = re.compile(r\"(?:<gap>\\s*){2,}\")\n        self.biggap_runs = re.compile(r\"(?:<big_gap>\\s*){2,}\")\n\n        self.repeat_word = re.compile(r\"\\b(\\w+)(\\s+\\1){1,}\\b\", flags=re.IGNORECASE)\n\n    def _dedupe_ngrams(self, text: str) -> str:\n        # simple greedy n-gram dedupe for 2..N\n        tokens = text.split()\n        if len(tokens) < 12:\n            return text\n        for n in range(2, self.ngram_dedupe_max_n + 1):\n            i = 0\n            out = []\n            while i < len(tokens):\n                if i + 2*n <= len(tokens) and tokens[i:i+n] == tokens[i+n:i+2*n]:\n                    # skip one repetition\n                    out.extend(tokens[i:i+n])\n                    i += 2*n\n                else:\n                    out.append(tokens[i])\n                    i += 1\n            tokens = out\n        return \" \".join(tokens)\n\n    def postprocess_one(self, text: str) -> str:\n        t = str(text or \"\")\n\n        # remove notes\n        t = self.remove_notes.sub(\"\", t)\n\n        # stronger gap consolidation\n        t = self.gap_runs.sub(\"<big_gap> \", t)   # 2+ gaps -> big gap\n        t = self.biggap_runs.sub(\"<big_gap> \", t)\n\n        # remove weird chars\n        t = self.remove_weird.sub(\" \", t)\n\n        # word repeat dedupe\n        t = self.repeat_word.sub(r\"\\1\", t)\n\n        # ngram dedupe\n        t = self._dedupe_ngrams(t)\n\n        # normalize spaces/punct\n        t = self.space_before_punct.sub(r\"\\1\", t)\n        t = self.multi_punct.sub(r\"\\1\", t)\n        t = self.multi_space.sub(\" \", t).strip()\n\n        return t\n\n    def postprocess_batch(self, texts: List[str]) -> List[str]:\n        return [self.postprocess_one(x) for x in texts]\n\n\n# ---------------------------\n# 4) Scoring (badness + quality)\n# ---------------------------\ndef badness_score(text: str) -> float:\n    \"\"\"\n    Lower is better.\n    Penalize:\n      - very short\n      - very long\n      - excessive gaps\n      - heavy repetition\n    \"\"\"\n    if text is None:\n        text = \"\"\n    t = str(text).strip()\n    if not t:\n        return 10.0\n\n    words = t.split()\n    n = len(words)\n\n    score = 0.0\n    if n < 5:\n        score += 3.0\n    if n < 3:\n        score += 3.0\n    if n > 500:\n        score += 2.0\n    if n > 650:\n        score += 3.0\n\n    gaps = t.count(\"<gap>\") + t.count(\"<big_gap>\")\n    if gaps > 6:\n        score += (gaps - 6) * 0.35\n\n    # 3+ consecutive repeats (same word)\n    rep = 0\n    for i in range(2, n):\n        if words[i].lower() == words[i-1].lower() == words[i-2].lower():\n            rep += 1\n    score += rep * 0.75\n\n    # long repeated bigrams (very rough)\n    if n >= 20:\n        bigrams = list(zip(words, words[1:]))\n        uniq = len(set(bigrams))\n        if uniq > 0:\n            repetitiveness = 1.0 - (uniq / max(1, len(bigrams)))\n            if repetitiveness > 0.35:\n                score += (repetitiveness - 0.35) * 6.0\n\n    return score\n\n\nKEYWORDS = {\n    \"tablet\", \"king\", \"city\", \"god\", \"silver\", \"gold\", \"temple\", \"house\", \"palace\",\n    \"year\", \"son\", \"daughter\", \"brother\", \"mother\", \"father\", \"gave\", \"took\", \"sent\",\n    \"received\", \"grain\", \"sheep\", \"oil\", \"wine\"\n}\n\ndef quality_score(text: str) -> float:\n    \"\"\"\n    Higher is better (simple heuristic).\n    \"\"\"\n    if text is None:\n        text = \"\"\n    t = str(text).strip()\n    if not t:\n        return 0.0\n\n    score = 0.0\n    words = t.split()\n\n    # length preference\n    n = len(words)\n    if 8 <= n <= 120:\n        score += 2.0\n    elif 5 <= n <= 200:\n        score += 1.0\n\n    # formatting\n    if t[0].isupper():\n        score += 0.5\n    if t.endswith((\".\", \"!\", \"?\")):\n        score += 0.5\n\n    # keyword hints\n    lw = {w.strip(\".,;:!?\\\"'()[]\").lower() for w in words}\n    hit = len(lw.intersection(KEYWORDS))\n    score += min(2.0, 0.25 * hit)\n\n    # penalize too many gaps\n    gaps = t.count(\"<gap>\") + t.count(\"<big_gap>\")\n    score -= min(2.0, 0.15 * gaps)\n\n    # penalize heavy repetition by badness\n    score -= 0.5 * max(0.0, badness_score(t) - 1.0)\n\n    return score\n\n\ndef fuse_texts(a: str, b: str, prefer: str = \"a\", tie_badness_thresh: float = 0.5,\n               w_a: float = 0.60, w_b: float = 0.40) -> str:\n    \"\"\"\n    Robust fuse:\n      1) pick lower badness if clearly better\n      2) if close, pick higher weighted quality\n    \"\"\"\n    ba = badness_score(a)\n    bb = badness_score(b)\n\n    if ba + tie_badness_thresh < bb:\n        return a\n    if bb + tie_badness_thresh < ba:\n        return b\n\n    qa = quality_score(a) * w_a\n    qb = quality_score(b) * w_b\n    if qa > qb:\n        return a\n    if qb > qa:\n        return b\n    return a if prefer == \"a\" else b\n\n\n# ---------------------------\n# 5) Dataset + Bucket sampler\n# ---------------------------\nclass AkkadianDataset(Dataset):\n    def __init__(self, df: pd.DataFrame, preprocessor: OptimizedPreprocessor):\n        self.ids = df[\"id\"].astype(str).tolist()\n        raw = df[\"transliteration\"].astype(\"object\").fillna(\"\").tolist()\n        proc = preprocessor.preprocess_batch(raw)\n        self.inputs = [f\"translate Akkadian to English: {x}\" for x in proc]\n\n    def __len__(self):\n        return len(self.ids)\n\n    def __getitem__(self, i):\n        return self.ids[i], self.inputs[i]\n\n\nclass BucketBatchSampler(Sampler[List[int]]):\n    def __init__(self, texts: List[str], batch_size: int, num_buckets: int = 32, shuffle: bool = False):\n        self.batch_size = batch_size\n        self.shuffle = shuffle\n\n        lengths = np.array([max(1, len(t.split())) for t in texts], dtype=np.int32)\n        order = np.argsort(lengths)\n        self.indices = order.tolist()\n\n        # split into buckets\n        self.num_buckets = max(1, num_buckets)\n        self.buckets = np.array_split(self.indices, self.num_buckets)\n\n        self.batches = []\n        rng = np.random.default_rng(12345)\n        for b in self.buckets:\n            b = list(b)\n            if shuffle:\n                rng.shuffle(b)\n            for i in range(0, len(b), batch_size):\n                chunk = b[i:i+batch_size]\n                if len(chunk) > 0:\n                    self.batches.append(chunk)\n\n        if shuffle:\n            rng.shuffle(self.batches)\n\n    def __iter__(self):\n        for batch in self.batches:\n            yield batch\n\n    def __len__(self):\n        return len(self.batches)\n\n\n# ---------------------------\n# 6) Inference Engine\n# ---------------------------\nclass UltraInferenceEngine:\n    def __init__(self, cfg: UltraConfig):\n        self.cfg = cfg\n        self.device = torch.device(cfg.device)\n        self.preprocessor = OptimizedPreprocessor()\n\n        # postprocessors\n        self.safe_pp = SafePostprocessor(\n            use_unicode_fractions=False,\n            strip_dash_old=False,\n        )\n        self.safe_pp_stripdash = SafePostprocessor(\n            use_unicode_fractions=False,\n            strip_dash_old=True,\n        )\n        self.agg_pp = AggressivePostprocessor()\n\n        self.tokenizer = AutoTokenizer.from_pretrained(cfg.model_path)\n        self.model = AutoModelForSeq2SeqLM.from_pretrained(cfg.model_path).to(self.device)\n        self.model.eval()\n\n        # optional BetterTransformer\n        if cfg.use_better_transformer and cfg.device == \"cuda\":\n            try:\n                from optimum.bettertransformer import BetterTransformer\n                self.model = BetterTransformer.transform(self.model)\n                logger.info(\"BetterTransformer enabled.\")\n            except Exception as e:\n                logger.warning(f\"BetterTransformer failed, continue without it: {e}\")\n\n    def _collate_fn(self, batch):\n        ids, texts = zip(*batch)\n        enc = self.tokenizer(\n            list(texts),\n            padding=True,\n            truncation=True,\n            max_length=self.cfg.max_length,\n            return_tensors=\"pt\",\n        )\n        return list(ids), enc\n\n    def _adaptive_beams(self, attention_mask: torch.Tensor) -> int:\n        if not self.cfg.use_adaptive_beams:\n            return self.cfg.num_beams\n        # use median token count in batch for robustness\n        lens = attention_mask.sum(dim=1).detach().cpu().numpy()\n        med = float(np.median(lens)) if len(lens) else 0.0\n        if med < 100:\n            return max(4, self.cfg.num_beams // 2)\n        return self.cfg.num_beams\n\n    def _postprocess(self, texts: List[str]) -> List[str]:\n        mode = self.cfg.postprocess_mode\n\n        if mode == \"safe_only\":\n            out = self.safe_pp.postprocess_batch(texts)\n            return out\n\n        if mode == \"aggressive_only\":\n            out = self.safe_pp.postprocess_batch(texts)\n            out = self.agg_pp.postprocess_batch(out)\n            return out\n\n        # default: safe then conditional aggressive for bad samples\n        safe = self.safe_pp.postprocess_batch(texts)\n        refined = []\n        thr = self.cfg.aggressive_trigger_badness\n        for t in safe:\n            if badness_score(t) >= thr:\n                refined.append(self.agg_pp.postprocess_one(t))\n            else:\n                refined.append(t)\n        return refined\n\n    @staticmethod\n    def _final_fix_one(t: str, min_words: int = 3) -> str:\n        tt = (t or \"\").strip()\n        if not tt:\n            return \"The tablet contains fragmentary text.\"\n        words = tt.split()\n        if len(words) < min_words:\n            return \"The tablet contains an incomplete inscription.\"\n        # capitalize first letter if it starts with alpha\n        if tt and tt[0].isalpha() and tt[0].islower():\n            tt = tt[0].upper() + tt[1:]\n        # add ending punctuation if missing\n        if not tt.endswith((\".\", \"!\", \"?\")):\n            tt = tt + \".\"\n        # clean spaces before punctuation just in case\n        tt = re.sub(r\"\\s+([,.;:!?])\", r\"\\1\", tt)\n        tt = re.sub(r\"\\s+\", \" \", tt).strip()\n        return tt\n\n    def run_inference(self, test_df: pd.DataFrame, run_tag: str = \"run\") -> pd.DataFrame:\n        cfg = self.cfg\n        ds = AkkadianDataset(test_df, self.preprocessor)\n\n        if cfg.use_bucket_batching:\n            sampler = BucketBatchSampler(ds.inputs, batch_size=cfg.batch_size, num_buckets=32, shuffle=False)\n            dl = DataLoader(\n                ds,\n                batch_sampler=sampler,\n                num_workers=cfg.num_workers,\n                pin_memory=cfg.pin_memory,\n                collate_fn=self._collate_fn,\n            )\n        else:\n            dl = DataLoader(\n                ds,\n                batch_size=cfg.batch_size,\n                shuffle=False,\n                num_workers=cfg.num_workers,\n                pin_memory=cfg.pin_memory,\n                collate_fn=self._collate_fn,\n            )\n\n        results: List[Tuple[str, str]] = []\n        t0 = time.time()\n\n        # autocast context\n        if cfg.use_mixed_precision and cfg.device == \"cuda\":\n            autocast_ctx = torch.cuda.amp.autocast\n        else:\n            # no-op context\n            class _NullCtx:\n                def __enter__(self): return None\n                def __exit__(self, exc_type, exc, tb): return False\n            autocast_ctx = lambda: _NullCtx()\n\n        for step, (ids, enc) in enumerate(dl):\n            input_ids = enc[\"input_ids\"].to(self.device, non_blocking=True)\n            attn = enc[\"attention_mask\"].to(self.device, non_blocking=True)\n\n            beams = self._adaptive_beams(attn)\n\n            gen_kwargs = dict(\n                num_beams=beams,\n                max_new_tokens=cfg.max_new_tokens,\n                length_penalty=cfg.length_penalty,\n                early_stopping=cfg.early_stopping,\n                no_repeat_ngram_size=cfg.no_repeat_ngram_size,\n            )\n            if cfg.repetition_penalty is not None:\n                gen_kwargs[\"repetition_penalty\"] = cfg.repetition_penalty\n\n            with torch.inference_mode():\n                with autocast_ctx():\n                    out_ids = self.model.generate(\n                        input_ids=input_ids,\n                        attention_mask=attn,\n                        **gen_kwargs\n                    )\n\n            decoded = self.tokenizer.batch_decode(out_ids, skip_special_tokens=True)\n            processed = self._postprocess(decoded)\n            processed = [self._final_fix_one(x, cfg.min_words_fallback) for x in processed]\n\n            results.extend(list(zip(ids, processed)))\n\n            # periodic checkpoint\n            if cfg.checkpoint_freq and (len(results) % cfg.checkpoint_freq == 0):\n                ck = pd.DataFrame(results, columns=[\"id\", \"translation\"])\n                ck_path = os.path.join(cfg.output_dir, f\"checkpoint_{run_tag}_{len(results)}.csv\")\n                ck.to_csv(ck_path, index=False)\n                logger.info(f\"Saved checkpoint: {ck_path}\")\n\n            if cfg.device == \"cuda\" and cfg.empty_cache_every and (step + 1) % cfg.empty_cache_every == 0:\n                torch.cuda.empty_cache()\n\n            if (step + 1) % 50 == 0:\n                elapsed = time.time() - t0\n                logger.info(f\"[{run_tag}] step={step+1}/{len(dl)} | done={len(results)} | elapsed={elapsed:.1f}s\")\n\n        df_out = pd.DataFrame(results, columns=[\"id\", \"translation\"])\n\n        # quick validation\n        self._validate(df_out, run_tag=run_tag)\n\n        return df_out\n\n    def _validate(self, df: pd.DataFrame, run_tag: str):\n        if df.empty:\n            logger.warning(f\"[{run_tag}] Empty output dataframe.\")\n            return\n        lens = df[\"translation\"].fillna(\"\").map(lambda x: len(str(x).split()))\n        empties = (df[\"translation\"].fillna(\"\").str.strip() == \"\").mean() * 100\n        short = (lens < 5).sum()\n        logger.info(f\"[{run_tag}] rows={len(df)} | empty%={empties:.2f} | len(mean/med/min/max)={lens.mean():.1f}/{lens.median():.1f}/{lens.min()}/{lens.max()} | short(<5)={short}\")\n\n        k = min(self.cfg.validate_samples, len(df))\n        if k > 0:\n            sample = df.sample(k, random_state=123)\n            for _, r in sample.iterrows():\n                logger.info(f\"[{run_tag}] sample id={r['id']} | {str(r['translation'])[:160]}\")\n\n# ---------------------------\n# 7) Helpers to build configs (two-pass)\n# ---------------------------\ndef make_cfg(base: UltraConfig,\n             preset: str,\n             batch_size: Optional[int] = None,\n             num_beams: Optional[int] = None,\n             length_penalty: Optional[float] = None,\n             repetition_penalty: Optional[float] = None,\n             no_repeat_ngram_size: Optional[int] = None,\n             strip_dash_old: Optional[bool] = None) -> UltraConfig:\n    cfg = UltraConfig(**asdict(base))\n\n    # presets\n    if preset == \"baseline\":\n        cfg.length_penalty = 1.30\n        cfg.repetition_penalty = None\n    elif preset == \"len115\":\n        cfg.length_penalty = 1.15\n        cfg.repetition_penalty = None\n    elif preset == \"rep_pen\":\n        cfg.length_penalty = 1.30\n        cfg.repetition_penalty = 1.08\n    elif preset == \"len115_rep\":\n        cfg.length_penalty = 1.15\n        cfg.repetition_penalty = 1.08\n    else:\n        raise ValueError(f\"Unknown preset: {preset}\")\n\n    # overrides\n    if batch_size is not None: cfg.batch_size = int(batch_size)\n    if num_beams is not None: cfg.num_beams = int(num_beams)\n    if length_penalty is not None: cfg.length_penalty = float(length_penalty)\n    if repetition_penalty is not None: cfg.repetition_penalty = float(repetition_penalty)\n    if no_repeat_ngram_size is not None: cfg.no_repeat_ngram_size = int(no_repeat_ngram_size)\n\n    # strip_dash_old is implemented inside SafePostprocessor; we control it via engine if desired\n    # here we store it in cfg as metadata only\n    cfg._strip_dash_old = bool(strip_dash_old) if strip_dash_old is not None else False  # type: ignore\n\n    return cfg\n\ndef fuse_submissions(df_a: pd.DataFrame, df_b: pd.DataFrame,\n                     prefer: str = \"a\", tie_badness_thresh: float = 0.5,\n                     w_a: float = 0.60, w_b: float = 0.40) -> pd.DataFrame:\n    a = df_a.set_index(\"id\")[\"translation\"]\n    b = df_b.set_index(\"id\")[\"translation\"]\n    ids = a.index\n\n    out = []\n    for _id in ids:\n        ta = a.loc[_id]\n        tb = b.loc[_id]\n        fused = fuse_texts(ta, tb, prefer=prefer, tie_badness_thresh=tie_badness_thresh, w_a=w_a, w_b=w_b)\n        out.append((_id, fused))\n    return pd.DataFrame(out, columns=[\"id\", \"translation\"])\n\n\n# ---------------------------\n# 8) Main: load test, run 2 passes, fuse, save\n# ---------------------------\nbase_cfg = UltraConfig()\n\n# Ensure output dir exists\nos.makedirs(base_cfg.output_dir, exist_ok=True)\n\n# Load test\ntest_df = pd.read_csv(base_cfg.test_data_path)\nassert \"id\" in test_df.columns and \"transliteration\" in test_df.columns, \"test.csv must have columns: id, transliteration\"\n\n# Save config snapshot\nwith open(os.path.join(base_cfg.output_dir, \"run_config_base.json\"), \"w\", encoding=\"utf-8\") as f:\n    json.dump(asdict(base_cfg), f, ensure_ascii=False, indent=2)\n\n# Pass A / Pass B presets (recommended defaults)\ncfg_a = make_cfg(base_cfg, preset=\"len115\")\ncfg_b = make_cfg(base_cfg, preset=\"rep_pen\")\n\nlogger.info(f\"PassA: len_pen={cfg_a.length_penalty}, rep_pen={cfg_a.repetition_penalty}, beams={cfg_a.num_beams}, bs={cfg_a.batch_size}\")\nlogger.info(f\"PassB: len_pen={cfg_b.length_penalty}, rep_pen={cfg_b.repetition_penalty}, beams={cfg_b.num_beams}, bs={cfg_b.batch_size}\")\n\n# Run inference A\nengine_a = UltraInferenceEngine(cfg_a)\nsub_a = engine_a.run_inference(test_df, run_tag=\"A_len115\")\npath_a = os.path.join(cfg_a.output_dir, \"submission_A.csv\")\nsub_a.to_csv(path_a, index=False)\nlogger.info(f\"Saved: {path_a}\")\n\n# Free GPU memory a bit before second engine (optional)\ndel engine_a\ngc.collect()\nif torch.cuda.is_available():\n    torch.cuda.empty_cache()\n\n# Run inference B\nengine_b = UltraInferenceEngine(cfg_b)\nsub_b = engine_b.run_inference(test_df, run_tag=\"B_rep_pen\")\npath_b = os.path.join(cfg_b.output_dir, \"submission_B.csv\")\nsub_b.to_csv(path_b, index=False)\nlogger.info(f\"Saved: {path_b}\")\n\n# Fuse\nfused = fuse_submissions(sub_a, sub_b, prefer=\"a\", tie_badness_thresh=0.5, w_a=0.60, w_b=0.40)\n\n# Final sanity (no empties)\nfused[\"translation\"] = fused[\"translation\"].fillna(\"\").map(lambda x: UltraInferenceEngine._final_fix_one(str(x), min_words=base_cfg.min_words_fallback))\n\nout_path = os.path.join(base_cfg.output_dir, \"submission.csv\")\nfused.to_csv(out_path, index=False)\nlogger.info(f\"Saved FINAL submission: {out_path}\")\n\n# Quick summary\nlens = fused[\"translation\"].map(lambda x: len(str(x).split()))\nlogger.info(f\"FINAL | rows={len(fused)} | len(mean/med/min/max)={lens.mean():.1f}/{lens.median():.1f}/{lens.min()}/{lens.max()}\")\nprint(fused.head())\nprint(fused.tail())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-27T00:12:18.816843Z","iopub.execute_input":"2026-02-27T00:12:18.817188Z","iopub.status.idle":"2026-02-27T00:13:23.516473Z","shell.execute_reply.started":"2026-02-27T00:12:18.817159Z","shell.execute_reply":"2026-02-27T00:13:23.515631Z"}},"outputs":[],"execution_count":null}]}