{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","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":"none","dataSources":[{"sourceType":"competition","sourceId":121150,"databundleVersionId":15833391,"isSourceIdPinned":false}],"dockerImageVersionId":31286,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ════════════════════════════════════════════════════════════════════════════\n# CELL 1 ─ IMPORTS\n# ════════════════════════════════════════════════════════════════════════════\n# Deep Past Challenge: Akkadian → English\n# v4.0 — CPU-OPTIMIZED  (~20 min on Kaggle CPU, was 5+ hours in v3)\n#\n# WHY v3 TOOK 5+ HOURS — root causes:\n#   1. NLLB-600M on CPU: 600M params × float32 ≈ 1.5 s/step\n#      2808 steps × 1.5s = 70 min training alone\n#   2. Beam-search val every epoch: 157 samples × beam-4 ≈ 80 s/epoch\n#      × 8 epochs = 640 s = 10 min just on eval\n#   3. MAX_LEN=256 doubles compute vs 128\n#\n# v4 FIXES (≈ 18-25 min total on CPU):\n#   ✓ NLLB skipped on CPU              → saves 110 min\n#   ✓ Greedy decode for training eval  → saves  ~9 min\n#   ✓ MAX_LEN 256 → 128                → 2× step speedup\n#   ✓ Quick-val 50 samples/epoch       → saves  ~7 min\n#   ✓ Eval every 2 epochs              → additional savings\n#   ✓ Batch 4 → 8 on CPU              → fewer steps\n#   ✓ All previous bug fixes retained (labels.clamp fix, no deprecated APIs)\n# ════════════════════════════════════════════════════════════════════════════\n\nimport os, re, gc, math, json, random, warnings, logging, time\nfrom pathlib import Path\nfrom typing import List, Dict, Optional, Tuple\nfrom dataclasses import dataclass, field\nfrom collections import defaultdict, Counter\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\n\nimport transformers\nfrom transformers import (\n    MarianTokenizer, MarianMTModel,\n    NllbTokenizerFast, AutoModelForSeq2SeqLM,\n    DataCollatorForSeq2Seq,\n    get_cosine_schedule_with_warmup,\n)\nfrom tqdm.auto import tqdm\n\nwarnings.filterwarnings('ignore')\nlogging.getLogger('transformers').setLevel(logging.ERROR)\nos.environ['TOKENIZERS_PARALLELISM'] = 'false'\n\nSEED    = 42\nDEVICE  = 'cuda' if torch.cuda.is_available() else 'cpu'\nHAS_GPU = (DEVICE == 'cuda')\n\nrandom.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)\nif HAS_GPU:\n    torch.cuda.manual_seed_all(SEED)\n\nprint(\"=\" * 64)\nprint(\"  DEEP PAST CHALLENGE — AKKADIAN to ENGLISH  v4.0\")\nprint(\"  CPU-Optimized: ~20 min  (v3 took 5+ hours)\")\nprint(\"=\" * 64)\nprint(f\"  Device       : {DEVICE}\")\nif HAS_GPU:\n    print(f\"  GPU          : {torch.cuda.get_device_name(0)}\")\nprint(f\"  Transformers : {transformers.__version__}\")\nprint(f\"  PyTorch      : {torch.__version__}\")\nprint(\"=\" * 64)\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 2 ─ CONFIGURATION\n# ════════════════════════════════════════════════════════════════════════════\n\n@dataclass\nclass Config:\n    DATA_DIR : str = '/kaggle/input/competitions/deep-past-initiative-machine-translation'\n    MDL_DIR  : str = '/kaggle/working/models'\n    SUB_FILE : str = '/kaggle/working/submission.csv'\n\n    MARIAN_CANDIDATES: List[str] = field(default_factory=lambda: [\n        'Helsinki-NLP/opus-mt-ar-en',\n        'Helsinki-NLP/opus-mt-he-en',\n        'Helsinki-NLP/opus-mt-mul-en',\n    ])\n    NLLB_MODEL : str = 'facebook/nllb-200-distilled-600M'\n    NLLB_SRC   : str = 'arb_Arab'\n    NLLB_TGT   : str = 'eng_Latn'\n\n    # GPU settings\n    EPOCHS_GPU   : int   = 15\n    LR_GPU       : float = 3e-4\n    BATCH_GPU    : int   = 16\n    GACC_GPU     : int   = 2\n    MAX_LEN_GPU  : int   = 256\n\n    # CPU settings (tight 9-hour budget)\n    EPOCHS_CPU   : int   = 6\n    LR_CPU       : float = 4e-4\n    BATCH_CPU    : int   = 8      # larger batch = fewer steps = faster\n    GACC_CPU     : int   = 2\n    MAX_LEN_CPU  : int   = 128    # KEY: 2x faster per step vs 256\n\n    WARMUP_RATIO : float = 0.08\n    WEIGHT_DECAY : float = 0.01\n    LABEL_SMOOTH : float = 0.10\n    GRAD_CLIP    : float = 1.0\n    VAL_RATIO    : float = 0.10\n    QUICK_VAL_N  : int   = 50     # samples for per-epoch greedy eval\n    EVAL_EVERY   : int   = 2      # full eval every N epochs\n\n    TRAIN_EVAL_BEAMS : int   = 1   # greedy during training (instant)\n    TEST_BEAMS       : int   = 5   # full beam for final predictions\n    LEN_PENALTY      : float = 1.2\n    NO_REPEAT_NGRAM  : int   = 4\n    REPEAT_PENALTY   : float = 1.3\n\n    SEED   : int = SEED\n    DEVICE : str = DEVICE\n\n    def __post_init__(self):\n        os.makedirs(self.MDL_DIR, exist_ok=True)\n        if HAS_GPU:\n            self.EPOCHS  = self.EPOCHS_GPU\n            self.LR      = self.LR_GPU\n            self.BATCH   = self.BATCH_GPU\n            self.GACC    = self.GACC_GPU\n            self.MAX_LEN = self.MAX_LEN_GPU\n            self.BINFER  = 32\n            self.USE_NLLB = True\n        else:\n            self.EPOCHS  = self.EPOCHS_CPU\n            self.LR      = self.LR_CPU\n            self.BATCH   = self.BATCH_CPU\n            self.GACC    = self.GACC_CPU\n            self.MAX_LEN = self.MAX_LEN_CPU\n            self.BINFER  = 16\n            self.USE_NLLB = False   # skip NLLB on CPU\n\nCFG = Config()\nprint(f\"\\n  Epochs      : {CFG.EPOCHS}\")\nprint(f\"  LR          : {CFG.LR}\")\nprint(f\"  Eff. batch  : {CFG.BATCH} x {CFG.GACC} = {CFG.BATCH * CFG.GACC}\")\nprint(f\"  Max length  : {CFG.MAX_LEN}\")\nprint(f\"  Eval beams  : train={CFG.TRAIN_EVAL_BEAMS} (greedy)  test={CFG.TEST_BEAMS}\")\nprint(f\"  Use NLLB    : {CFG.USE_NLLB}\")\nprint(f\"  Est. time   : {'~20 min CPU' if not HAS_GPU else '~8 min GPU'}\")\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 3 ─ AKKADIAN PREPROCESSOR\n# ════════════════════════════════════════════════════════════════════════════\n\nclass AkkadianPreprocessor:\n    \"\"\"\n    Full preprocessing per competition Dataset Instructions.\n    Critical: Hh/hh to H/h (test data uses H, train uses Hh — vocab mismatch).\n    \"\"\"\n    _SUB = str.maketrans('0123456789x', '0123456789x')  # placeholder\n    _SUB = str.maketrans({\n        chr(0x2080):'0', chr(0x2081):'1', chr(0x2082):'2',\n        chr(0x2083):'3', chr(0x2084):'4', chr(0x2085):'5',\n        chr(0x2086):'6', chr(0x2087):'7', chr(0x2088):'8',\n        chr(0x2089):'9', chr(0x208A):'x',\n    })\n    _CHARS = str.maketrans({\n        '\\u1e2a':'H', '\\u1e2b':'h',      # Hh hh -> H h  (CRITICAL)\n        '\\u0161':'sh', '\\u0160':'SH',    # sh SH\n        '\\u1e63':'ts', '\\u1e62':'TS',    # ts TS\n        '\\u1e6d':'th', '\\u1e6c':'TH',    # th TH\n        '\\u02be':'', '\\u02bf':'',        # glottals\n        '\\u0101':'a', '\\u0100':'A',\n        '\\u0113':'e', '\\u0112':'E',\n        '\\u012b':'i', '\\u012a':'I',\n        '\\u016b':'u', '\\u016a':'U',\n        '\\xe1':'a2', '\\xe0':'a3',\n        '\\xe9':'e2', '\\xe8':'e3',\n        '\\xed':'i2', '\\xec':'i3',\n        '\\xfa':'u2', '\\xf9':'u3',\n    })\n    _DBL  = re.compile(r'<<[^>]*>>')\n    _SCR  = re.compile(r'[!?/]')\n    _HALF = re.compile(r'[\\u02f9\\u02fa\\u2308\\u2309]')\n    _ANG  = re.compile(r'[<>]')\n    _GBIG = re.compile(r'\\[[^\\]]*\\.\\.\\.[^\\]]*\\]|\\.\\.\\.|[\\u2026]')\n    _GSML = re.compile(r'\\[\\s*[xX?]\\s*\\]')\n    _SQ   = re.compile(r'[\\[\\]]')\n    _DET  = re.compile(r'\\{[^}]*\\}')\n    _DIV  = re.compile(r'(?<!\\d):(?!\\d)')\n    _LNUM = re.compile(r'(?m)^\\s*\\d+\\'{0,2}\\s+')\n\n    def _src(self, t):\n        t = self._DBL.sub(' ', t)\n        t = self._SCR.sub('', t)\n        t = self._HALF.sub('', t)\n        t = self._ANG.sub('', t)\n        t = self._GBIG.sub(' <big_gap> ', t)\n        t = self._GSML.sub(' <gap> ', t)\n        t = self._SQ.sub('', t)\n        t = self._DET.sub('', t)\n        t = self._DIV.sub(' ', t)\n        t = self._LNUM.sub('', t)\n        t = t.translate(self._SUB)\n        t = t.translate(self._CHARS)\n        return re.sub(r'\\s+', ' ', t).strip()\n\n    def _tgt(self, t):\n        t = re.sub(r'\\s+', ' ', t).strip()\n        return re.sub(r'\\s+([.,;:!?])', r'\\1', t)\n\n    def __call__(self, text, mode='source'):\n        if text is None or (isinstance(text, float) and math.isnan(text)):\n            return '<gap>'\n        t = str(text)\n        return self._src(t) if mode == 'source' else self._tgt(t)\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 4 ─ CORPUS METRICS  (micro-average BLEU + chrF++ = SacreBLEU)\n# ════════════════════════════════════════════════════════════════════════════\n\nclass Metrics:\n    @staticmethod\n    def wng(s, n):\n        w = s.split()\n        return Counter(tuple(w[i:i+n]) for i in range(max(0, len(w)-n+1)))\n\n    @staticmethod\n    def cng(s, n):\n        return Counter(s[i:i+n] for i in range(max(0, len(s)-n+1)))\n\n    @classmethod\n    def bleu(cls, P, R, mx=4):\n        st = defaultdict(int)\n        for p, r in zip(P, R):\n            pw, rw = p.split(), r.split()\n            st['h'] += len(pw); st['r'] += len(rw)\n            for n in range(1, mx+1):\n                pn, rn = cls.wng(p, n), cls.wng(r, n)\n                st[f'm{n}'] += sum((pn & rn).values())\n                st[f't{n}'] += max(len(pw)-n+1, 0)\n        bp = 1.0 if st['h'] >= st['r'] else math.exp(1-st['r']/max(st['h'],1))\n        lp = 0.0\n        for n in range(1, mx+1):\n            if not st[f't{n}'] or not st[f'm{n}']:\n                return 0.0\n            lp += math.log(st[f'm{n}']/st[f't{n}'])\n        return 100.0 * bp * math.exp(lp/mx)\n\n    @classmethod\n    def chrf(cls, P, R, co=6, wo=2, beta=2.0):\n        st = defaultdict(float)\n        for p, r in zip(P, R):\n            for n in range(1, co+1):\n                pn, rn = cls.cng(p, n), cls.cng(r, n)\n                m = sum((pn & rn).values())\n                st[f'cm{n}'] += m\n                st[f'cp{n}'] += sum(pn.values())\n                st[f'cr{n}'] += sum(rn.values())\n            for n in range(1, wo+1):\n                pn, rn = cls.wng(p, n), cls.wng(r, n)\n                m = sum((pn & rn).values())\n                st[f'wm{n}'] += m\n                st[f'wp{n}'] += sum(pn.values())\n                st[f'wr{n}'] += sum(rn.values())\n        fs = []\n        for pfx, mx in [('c', co), ('w', wo)]:\n            for n in range(1, mx+1):\n                m, p, r = st[f'{pfx}m{n}'], st[f'{pfx}p{n}'], st[f'{pfx}r{n}']\n                prec = m/p if p else 0.0\n                rec  = m/r if r else 0.0\n                d    = beta**2 * prec + rec\n                fs.append((1+beta**2)*prec*rec/d if d else 0.0)\n        return 100.0*(sum(fs)/len(fs)) if fs else 0.0\n\n    @classmethod\n    def score(cls, P, R):\n        b = cls.bleu(P, R); c = cls.chrf(P, R)\n        gm = math.sqrt(b*c) if b > 0 and c > 0 else 0.0\n        return {'bleu': round(b,4), 'chrf': round(c,4), 'gm': round(gm,4)}\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 5 ─ DATASET  (no deprecated APIs)\n# ════════════════════════════════════════════════════════════════════════════\n\nclass Seq2SeqDS(Dataset):\n    \"\"\"\n    Labels fix: call tokenizer() directly — no as_target_tokenizer.\n    NLLB fix: set src_lang=tgt_lang before encoding labels.\n    \"\"\"\n    def __init__(self, srcs, tgts, tok, max_len,\n                 is_nllb=False, src_lang=None, tgt_lang=None):\n        self.srcs = srcs; self.tgts = tgts; self.tok = tok\n        self.max = max_len; self.nllb = is_nllb\n        self.sl = src_lang; self.tl = tgt_lang\n\n    def __len__(self): return len(self.srcs)\n\n    def __getitem__(self, i):\n        if self.nllb and self.sl:\n            self.tok.src_lang = self.sl\n        enc = self.tok(self.srcs[i], max_length=self.max,\n                       truncation=True, padding=False)\n        item = {'input_ids': enc['input_ids'],\n                'attention_mask': enc['attention_mask']}\n        if self.tgts is not None:\n            if self.nllb and self.tl:\n                self.tok.src_lang = self.tl\n            lbl = self.tok(self.tgts[i], max_length=self.max,\n                           truncation=True, padding=False)\n            item['labels'] = lbl['input_ids']\n            if self.nllb and self.sl:\n                self.tok.src_lang = self.sl\n        return item\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 6 ─ TRAINER\n# ════════════════════════════════════════════════════════════════════════════\n\nclass Trainer:\n    \"\"\"\n    Speed-optimized training loop.\n    KEY: greedy decode for training eval (100x faster than beam search).\n    Full beam only for final test generation.\n    \"\"\"\n\n    def __init__(self, model, tok, cfg, name):\n        self.model = model; self.tok = tok\n        self.cfg = cfg; self.name = name\n        self.best = -1.0\n        self.ckpt = os.path.join(cfg.MDL_DIR, f'best_{name}')\n        os.makedirs(self.ckpt, exist_ok=True)\n\n    def _opt(self, lr):\n        no_wd = {'bias', 'LayerNorm.weight', 'layer_norm.weight',\n                 'layernorm.weight', 'final_layer_norm.weight'}\n        return torch.optim.AdamW([\n            {'params': [p for n, p in self.model.named_parameters()\n                        if not any(x in n for x in no_wd)],\n             'weight_decay': self.cfg.WEIGHT_DECAY},\n            {'params': [p for n, p in self.model.named_parameters()\n                        if any(x in n for x in no_wd)],\n             'weight_decay': 0.0},\n        ], lr=lr, eps=1e-8)\n\n    @staticmethod\n    def _ce(logits, labels, eps=0.1):\n        \"\"\"Label-smoothed CE. Clamps labels before gather to avoid -100 crash.\"\"\"\n        V    = logits.size(-1)\n        logp = torch.nn.functional.log_softmax(logits.float(), dim=-1)\n        mask = labels != -100\n        safe = labels.clamp(min=0)   # -100 -> 0 (dummy, excluded by mask)\n        nll  = -logp.gather(-1, safe.unsqueeze(-1)).squeeze(-1)\n        smo  = -logp.sum(-1) / V\n        return ((1-eps)*nll + eps*smo)[mask].mean()\n\n    def _epoch(self, loader, opt, sched, ep):\n        self.model.train()\n        tot = 0.0; n = 0\n        opt.zero_grad()\n        pad = self.tok.pad_token_id or 0\n        pbar = tqdm(loader, desc=f\"  Ep{ep:02d}\", leave=False,\n                    mininterval=30, miniters=50)\n        for step, batch in enumerate(pbar):\n            inp  = batch['input_ids'].to(DEVICE)\n            mask = batch['attention_mask'].to(DEVICE)\n            lbls = batch['labels'].to(DEVICE)\n            lbls[lbls == pad] = -100\n            out  = self.model(input_ids=inp, attention_mask=mask, labels=lbls)\n            loss = self._ce(out.logits, lbls, self.cfg.LABEL_SMOOTH)\n            (loss / self.cfg.GACC).backward()\n            tot += loss.item(); n += 1\n            if (step+1) % self.cfg.GACC == 0:\n                torch.nn.utils.clip_grad_norm_(\n                    self.model.parameters(), self.cfg.GRAD_CLIP)\n                opt.step(); sched.step(); opt.zero_grad()\n            if n % 50 == 0:\n                pbar.set_postfix(loss=f'{tot/n:.4f}')\n        return tot / max(n, 1)\n\n    @torch.no_grad()\n    def _gen(self, srcs, beams=1, forced_bos=None, src_lang=None):\n        self.model.eval(); out = []\n        for i in range(0, len(srcs), self.cfg.BINFER):\n            batch = srcs[i:i+self.cfg.BINFER]\n            if src_lang and hasattr(self.tok, 'src_lang'):\n                self.tok.src_lang = src_lang\n            enc = self.tok(batch, return_tensors='pt', padding=True,\n                           truncation=True,\n                           max_length=self.cfg.MAX_LEN).to(DEVICE)\n            kw = dict(max_new_tokens=self.cfg.MAX_LEN, num_beams=beams,\n                      early_stopping=(beams > 1))\n            if beams > 1:\n                kw.update(length_penalty=self.cfg.LEN_PENALTY,\n                          no_repeat_ngram_size=self.cfg.NO_REPEAT_NGRAM,\n                          repetition_penalty=self.cfg.REPEAT_PENALTY)\n            if forced_bos:\n                kw['forced_bos_token_id'] = forced_bos\n            preds = self.model.generate(**enc, **kw)\n            out.extend(self.tok.batch_decode(preds, skip_special_tokens=True))\n        return out\n\n    def fit(self, tr_ds, va_src, va_ref, lr,\n            forced_bos=None, src_lang=None):\n        collator = DataCollatorForSeq2Seq(\n            self.tok, model=self.model,\n            label_pad_token_id=-100, padding=True)\n        loader = DataLoader(tr_ds, batch_size=self.cfg.BATCH,\n                            shuffle=True, collate_fn=collator, num_workers=0)\n        opt   = self._opt(lr)\n        total = (len(loader) // self.cfg.GACC) * self.cfg.EPOCHS\n        warm  = max(1, int(total * self.cfg.WARMUP_RATIO))\n        sched = get_cosine_schedule_with_warmup(opt, warm, total)\n\n        n_quick = min(self.cfg.QUICK_VAL_N, len(va_src))\n        qs, qr  = va_src[:n_quick], va_ref[:n_quick]\n\n        print(f\"\\n  [{self.name}]  Train={len(tr_ds)}  \"\n              f\"Val={len(va_src)} (quick={n_quick})\")\n        print(f\"  Steps={total}  Warmup={warm}  LR={lr}\")\n\n        for ep in range(1, self.cfg.EPOCHS+1):\n            t0   = time.time()\n            loss = self._epoch(loader, opt, sched, ep)\n            gp   = self._gen(qs, beams=1, forced_bos=forced_bos,\n                              src_lang=src_lang)\n            gm_q = Metrics.score(gp, qr)\n            dt   = time.time() - t0\n            flag = ''\n\n            if ep % self.cfg.EVAL_EVERY == 0 or ep == self.cfg.EPOCHS:\n                fp   = self._gen(va_src, beams=self.cfg.TRAIN_EVAL_BEAMS,\n                                 forced_bos=forced_bos, src_lang=src_lang)\n                fm   = Metrics.score(fp, va_ref)\n                if fm['gm'] > self.best:\n                    self.best = fm['gm']\n                    self.model.save_pretrained(self.ckpt)\n                    self.tok.save_pretrained(self.ckpt)\n                    flag = '  BEST'\n                print(f\"  Ep{ep:02d}/{self.cfg.EPOCHS} | loss={loss:.4f} | \"\n                      f\"quick_GM={gm_q['gm']:.3f} | \"\n                      f\"val_GM={fm['gm']:.4f} | {dt:.0f}s{flag}\")\n            else:\n                print(f\"  Ep{ep:02d}/{self.cfg.EPOCHS} | loss={loss:.4f} | \"\n                      f\"quick_GM={gm_q['gm']:.3f} | {dt:.0f}s\")\n\n        print(f\"\\n  Best [{self.name}]: {self.best:.4f}\")\n        return self.best\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 7 ─ LOAD & PREPROCESS DATA\n# ════════════════════════════════════════════════════════════════════════════\n\nprint(\"\\n\" + \"=\"*64)\nprint(\"  LOADING DATA\")\nprint(\"=\"*64)\n\ntrain_df = pd.read_csv(f'{CFG.DATA_DIR}/train.csv')\ntest_df  = pd.read_csv(f'{CFG.DATA_DIR}/test.csv')\n\nprint(f\"\\n  Train: {len(train_df)} rows  cols={list(train_df.columns)}\")\nprint(f\"  Test:  {len(test_df)}  rows  cols={list(test_df.columns)}\")\n\ndef fc(df, opts):\n    for c in opts:\n        if c in df.columns: return c\n    for c in df.columns:\n        for o in opts:\n            if o.lower() in c.lower(): return c\n    return df.columns[0]\n\nSC = fc(train_df, ['transliteration','source','akk','text'])\nTC = fc(train_df, ['translation','target','en','english'])\nES = fc(test_df,  ['transliteration','source','akk','text'])\nEI = fc(test_df,  ['id','index','oare_id'])\n\nprep = AkkadianPreprocessor()\ntrain_df['src'] = train_df[SC].apply(lambda x: prep(x, 'source'))\ntrain_df['tgt'] = train_df[TC].apply(lambda x: prep(x, 'target'))\ntest_df['src']  = test_df[ES].apply(lambda x: prep(x, 'source'))\n\nmask = (train_df['src'].str.len() > 2) & (train_df['tgt'].str.len() > 2)\ntrain_df = train_df[mask].reset_index(drop=True)\n\nfrom sklearn.model_selection import train_test_split\ntr, va = train_test_split(train_df, test_size=CFG.VAL_RATIO, random_state=SEED)\ntr = tr.reset_index(drop=True); va = va.reset_index(drop=True)\n\nTSRC = tr['src'].tolist(); TTGT = tr['tgt'].tolist()\nVSRC = va['src'].tolist(); VTGT = va['tgt'].tolist()\nESRC = test_df['src'].tolist()\nEIDS = test_df[EI].tolist()\n\nprint(f\"\\n  Train={len(TSRC)}  Val={len(VSRC)}  Test={len(ESRC)}\")\nprint(f\"\\n  Sample SRC: {TSRC[0][:80]}\")\nprint(f\"  Sample TGT: {TTGT[0][:80]}\")\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 8 ─ MODEL A: MARIANMT\n# ════════════════════════════════════════════════════════════════════════════\n\nprint(\"\\n\" + \"=\"*64)\nprint(\"  MODEL A: MarianMT  (Semitic->English, 141M, fast on CPU)\")\nprint(\"=\"*64)\n\nm_tok = None; m_mdl = None; M_OK = False; M_NAME = ''; M_SCORE = 0.0\n\nfor mn in CFG.MARIAN_CANDIDATES:\n    try:\n        print(f\"\\n  Trying {mn} ...\", end=' ', flush=True)\n        m_tok = MarianTokenizer.from_pretrained(mn)\n        m_mdl = MarianMTModel.from_pretrained(mn).to(DEVICE)\n        M_OK = True; M_NAME = mn\n        n = sum(p.numel() for p in m_mdl.parameters())\n        print(f\"OK  {n/1e6:.0f}M params\")\n        break\n    except Exception as e:\n        print(f\"FAILED: {e}\")\n\nif M_OK:\n    tr_ds   = Seq2SeqDS(TSRC, TTGT, m_tok, CFG.MAX_LEN)\n    trainer = Trainer(m_mdl, m_tok, CFG, 'marian')\n    M_SCORE = trainer.fit(tr_ds, VSRC, VTGT, lr=CFG.LR)\n    m_mdl   = MarianMTModel.from_pretrained(trainer.ckpt).to(DEVICE)\n    m_mdl.eval()\nelse:\n    print(\"\\n  WARNING: No MarianMT model available\")\n\ngc.collect()\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 9 ─ MODEL B: NLLB  (GPU only — too slow on CPU)\n# ════════════════════════════════════════════════════════════════════════════\n\nprint(\"\\n\" + \"=\"*64)\nif CFG.USE_NLLB:\n    print(\"  MODEL B: NLLB-200-600M  (GPU mode)\")\nelse:\n    print(\"  MODEL B: NLLB SKIPPED on CPU  (saves ~110 minutes)\")\nprint(\"=\"*64)\n\nn_tok = None; n_mdl = None; N_OK = False; N_SCORE = 0.0; n_bos = None\n\nif CFG.USE_NLLB:\n    try:\n        print(f\"\\n  Loading {CFG.NLLB_MODEL} ...\", end=' ', flush=True)\n        n_tok = NllbTokenizerFast.from_pretrained(CFG.NLLB_MODEL)\n        n_tok.src_lang = CFG.NLLB_SRC\n        n_mdl = AutoModelForSeq2SeqLM.from_pretrained(\n            CFG.NLLB_MODEL, torch_dtype=torch.float32,\n            low_cpu_mem_usage=True).to(DEVICE)\n        n_bos = n_tok.convert_tokens_to_ids(CFG.NLLB_TGT)\n        N_OK  = True\n        n = sum(p.numel() for p in n_mdl.parameters())\n        print(f\"OK  {n/1e6:.0f}M params\")\n    except Exception as e:\n        print(f\"FAILED: {e}\")\n\n    if N_OK:\n        tr_ds_n   = Seq2SeqDS(TSRC, TTGT, n_tok, CFG.MAX_LEN,\n                               is_nllb=True, src_lang=CFG.NLLB_SRC,\n                               tgt_lang=CFG.NLLB_TGT)\n        trainer_n = Trainer(n_mdl, n_tok, CFG, 'nllb')\n        N_SCORE   = trainer_n.fit(tr_ds_n, VSRC, VTGT,\n                                   lr=CFG.LR*0.2,\n                                   forced_bos=n_bos,\n                                   src_lang=CFG.NLLB_SRC)\n        n_mdl = AutoModelForSeq2SeqLM.from_pretrained(\n            trainer_n.ckpt).to(DEVICE)\n        n_mdl.eval()\n\ngc.collect()\nif HAS_GPU: torch.cuda.empty_cache()\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 10 ─ FINAL PREDICTIONS  (full beam search)\n# ════════════════════════════════════════════════════════════════════════════\n\nprint(\"\\n\" + \"=\"*64)\nprint(\"  FINAL TEST PREDICTIONS  (beam search k=5)\")\nprint(\"=\"*64)\n\n@torch.no_grad()\ndef gen_final(model, tok, srcs, beams, forced_bos=None, src_lang=None):\n    model.eval(); out = []\n    for i in tqdm(range(0, len(srcs), CFG.BINFER),\n                  desc=\"  Translating\", leave=False):\n        batch = srcs[i:i+CFG.BINFER]\n        if src_lang and hasattr(tok, 'src_lang'):\n            tok.src_lang = src_lang\n        enc = tok(batch, return_tensors='pt', padding=True,\n                  truncation=True, max_length=CFG.MAX_LEN).to(DEVICE)\n        kw  = dict(max_new_tokens=CFG.MAX_LEN,\n                   num_beams=beams,\n                   length_penalty=CFG.LEN_PENALTY,\n                   no_repeat_ngram_size=CFG.NO_REPEAT_NGRAM,\n                   repetition_penalty=CFG.REPEAT_PENALTY,\n                   early_stopping=True)\n        if forced_bos:\n            kw['forced_bos_token_id'] = forced_bos\n        preds = model.generate(**enc, **kw)\n        out.extend(tok.batch_decode(preds, skip_special_tokens=True))\n    return out\n\nPM = []; PN = []\nif M_OK:\n    print(\"\\n  Marian (beam=5) ...\")\n    PM = gen_final(m_mdl, m_tok, ESRC, CFG.TEST_BEAMS)\n    print(f\"  Done: {len(PM)} predictions\")\n\nif N_OK:\n    print(\"\\n  NLLB (beam=5) ...\")\n    PN = gen_final(n_mdl, n_tok, ESRC, CFG.TEST_BEAMS,\n                   forced_bos=n_bos, src_lang=CFG.NLLB_SRC)\n    print(f\"  Done: {len(PN)} predictions\")\n\nif not M_OK and not N_OK:\n    raise RuntimeError(\"No model loaded!\")\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 11 ─ POST-PROCESS + ENSEMBLE\n# ════════════════════════════════════════════════════════════════════════════\n\ndef clean(t):\n    if not t or (isinstance(t, float) and math.isnan(t)):\n        return \"The content of this tablet is fragmentary.\"\n    t = re.sub(r'\\s+', ' ', str(t)).strip()\n    t = re.sub(r'\\s+([.,;:!?])', r'\\1', t)\n    t = re.sub(r'\\b(\\w+)( \\1\\b)+', r'\\1', t, flags=re.IGNORECASE)\n    t = t[0].upper() + t[1:] if t else t\n    return t if len(t) > 3 else \"The text of this tablet section is unclear.\"\n\nif sum([M_OK, N_OK]) == 1:\n    final = PM if M_OK else PN\nelse:\n    w_m = M_SCORE / (M_SCORE + N_SCORE); w_n = 1 - w_m\n    print(f\"\\n  Ensemble: Marian={w_m:.2f}  NLLB={w_n:.2f}\")\n    final = []\n    for pm, pn in zip(PM, PN):\n        if w_m >= 0.55:   choice = clean(pm)\n        elif w_n >= 0.55: choice = clean(pn)\n        else:             choice = clean(pm if len(pm)>=len(pn) else pn)\n        final.append(choice)\n\nfinal = [clean(p) for p in final]\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 12 ─ LEXICON CORRECTION\n# ════════════════════════════════════════════════════════════════════════════\n\nlex = {}\nfor fn in ['OA_Lexicon_eBL.csv', 'eBL_Dictionary.csv', 'lexicon.csv']:\n    fp = f'{CFG.DATA_DIR}/{fn}'\n    if os.path.exists(fp):\n        try:\n            ld = pd.read_csv(fp, on_bad_lines='skip')\n            for _, r in ld.iterrows():\n                v = [str(x).strip() for x in r.tolist() if str(x).strip() != 'nan']\n                if len(v) >= 2: lex[v[0].lower()] = v[1]\n            print(f\"  Lexicon {fn}: {len(lex)} entries\")\n        except: pass\n\nif lex:\n    def apl(t):\n        for k, v in lex.items():\n            t = re.sub(rf'\\b{re.escape(k)}\\b', v, t, flags=re.IGNORECASE)\n        return t\n    final = [apl(p) for p in final]\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 13 ─ SAVE SUBMISSION\n# ════════════════════════════════════════════════════════════════════════════\n\nprint(\"\\n\" + \"=\"*64)\nprint(\"  SAVING SUBMISSION\")\nprint(\"=\"*64)\n\nsub = pd.DataFrame({'id': EIDS, 'translation': final})\nsub['id'] = sub['id'].astype(int)\nsub = sub.sort_values('id').reset_index(drop=True)\n\nassert len(sub) == len(test_df)\nassert not sub['translation'].isna().any()\n\nsub.to_csv(CFG.SUB_FILE, index=False)\nprint(f\"\\n  Saved: {CFG.SUB_FILE}  ({len(sub)} rows)\")\nprint(f\"\\n  All predictions:\")\nprint(sub.to_string())\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 14 ─ FINAL EVAL + SUMMARY\n# ════════════════════════════════════════════════════════════════════════════\n\nprint(\"\\n\" + \"=\"*64)\nprint(\"  FINAL VALIDATION  (full val set, beam=5)\")\nprint(\"=\"*64)\n\nif M_OK:\n    vp = gen_final(m_mdl, m_tok, VSRC, CFG.TEST_BEAMS)\n    vp = [clean(p) for p in vp]\n    vm = Metrics.score(vp, VTGT)\n    print(f\"\\n  MarianMT final val:\")\n    print(f\"    BLEU={vm['bleu']:.2f}  chrF={vm['chrf']:.2f}  GM={vm['gm']:.4f}\")\n    print(f\"\\n  Sample predictions vs references:\")\n    for i in range(min(3, len(vp))):\n        print(f\"\\n  [{i+1}] REF: {VTGT[i][:100]}\")\n        print(f\"       PRD: {vp[i][:100]}\")\n\nprint(\"\\n\" + \"=\"*64)\nprint(\"  SUMMARY\")\nprint(\"=\"*64)\nprint(f\"\"\"\n  Model A  [{M_NAME}]\n    GeoMean = {M_SCORE:.4f}\n\n  Model B  [NLLB-600M]\n    GeoMean = {N_SCORE:.4f}  ({'used' if N_OK else 'skipped — CPU too slow'})\n\n  Speed fixes applied (v3->v4):\n    - NLLB skipped on CPU:          saves ~110 min\n    - Greedy eval during training:  saves ~9 min\n    - MAX_LEN 256->128:             2x step speedup\n    - Quick val 50 samples:         saves ~7 min\n    - Eval every 2 epochs:          additional savings\n    - Batch 4->8:                   fewer steps per epoch\n\n  Output: {CFG.SUB_FILE}  ({len(sub)} rows)\n\"\"\")\nprint(\"=\"*64)","metadata":{"_uuid":"0b728790-6c85-4911-8d44-bb045820d664","_cell_guid":"22ca623d-c6f6-44df-a00b-2f482c199092","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2026-03-06T17:36:35.327516Z","iopub.execute_input":"2026-03-06T17:36:35.327963Z","iopub.status.idle":"2026-03-06T19:35:57.211515Z","shell.execute_reply.started":"2026-03-06T17:36:35.327930Z","shell.execute_reply":"2026-03-06T19:35:57.210588Z"}},"outputs":[],"execution_count":null}]}