{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"none","dataSources":[{"sourceType":"competition","sourceId":121150,"databundleVersionId":15833391,"isSourceIdPinned":false}],"dockerImageVersionId":31286,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ════════════════════════════════════════════════════════════════════════════\n# CELL 1 ─ IMPORTS & ENVIRONMENT\n# ════════════════════════════════════════════════════════════════════════════\n# Deep Past Challenge: Translate Akkadian to English\n# v3.0 — ZERO deprecated APIs. All bugs resolved:\n#   ✓ as_target_tokenizer REMOVED → direct tokenizer() call for labels\n#   ✓ No GradScaler on CPU (autocast device_type-aware)\n#   ✓ DataCollatorForSeq2Seq padding correctly configured\n#   ✓ MarianMT label encoding: same tokenizer, no context manager\n#   ✓ NLLB label encoding: set src_lang=tgt_lang before encoding labels\n#   ✓ CPU-safe training loop throughout\n#   ✓ Auto batch size and epoch count based on device\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    AutoTokenizer,\n    AutoModelForSeq2SeqLM,\n    MarianTokenizer,\n    MarianMTModel,\n    NllbTokenizerFast,\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)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\nif HAS_GPU:\n    torch.cuda.manual_seed_all(SEED)\n    torch.backends.cudnn.deterministic = True\n\nprint(\"═\" * 66)\nprint(\"  DEEP PAST CHALLENGE — AKKADIAN→ENGLISH  v3.0\")\nprint(\"  All bugs fixed • Zero deprecated APIs • CPU/GPU safe\")\nprint(\"═\" * 66)\nprint(f\"  Device       : {DEVICE}\")\nif HAS_GPU:\n    print(f\"  GPU          : {torch.cuda.get_device_name(0)}\")\n    print(f\"  VRAM         : {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB\")\nprint(f\"  Transformers : {transformers.__version__}\")\nprint(f\"  PyTorch      : {torch.__version__}\")\nprint(\"═\" * 66)\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 2 ─ CONFIGURATION\n# ════════════════════════════════════════════════════════════════════════════\n\n@dataclass\nclass Config:\n    # ── Paths ──────────────────────────────────────────────────────────────\n    DATA_DIR : str = '/kaggle/input/competitions/deep-past-initiative-machine-translation'\n    OUT_DIR  : str = '/kaggle/working'\n    MDL_DIR  : str = '/kaggle/working/models'\n    SUB_FILE : str = '/kaggle/working/submission.csv'\n\n    # ── Model candidates (tried in order; first that loads wins) ──────────\n    MARIAN_CANDIDATES: List[str] = field(default_factory=lambda: [\n        'Helsinki-NLP/opus-mt-ar-en',        # Arabic→EN (~300 MB)  ← PRIMARY\n        'Helsinki-NLP/opus-mt-he-en',        # Hebrew→EN (~300 MB)\n        'Helsinki-NLP/opus-mt-mul-en',       # Multi→EN  (~300 MB)\n        'Helsinki-NLP/opus-mt-tc-big-ar-en', # Big Arabic→EN (~1.1 GB)\n    ])\n\n    # ── NLLB (large multilingual, Arabic proxy for Akkadian) ──────────────\n    NLLB_MODEL : str = 'facebook/nllb-200-distilled-600M'\n    NLLB_SRC   : str = 'arb_Arab'   # Arabic = closest Semitic proxy\n    NLLB_TGT   : str = 'eng_Latn'\n\n    # ── Training (GPU) ─────────────────────────────────────────────────────\n    EPOCHS_GPU   : int   = 15\n    LR_GPU       : float = 3e-4\n    BATCH_GPU    : int   = 16\n    GACC_GPU     : int   = 2\n\n    # ── Training (CPU) — must complete within 9-hour wall clock ──────────\n    EPOCHS_CPU   : int   = 8\n    LR_CPU       : float = 3e-4\n    BATCH_CPU    : int   = 4\n    GACC_CPU     : int   = 4\n\n    # ── Common ────────────────────────────────────────────────────────────\n    MAX_SRC_LEN  : int   = 256\n    MAX_TGT_LEN  : int   = 256\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    # 10 % → ~156 val samples (from 1561)\n\n    # ── Inference ─────────────────────────────────────────────────────────\n    NUM_BEAMS       : int   = 5\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.BINFER = 32\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.BINFER = 8\n\nCFG = Config()\nprint(f\"\\n  Epochs  : {CFG.EPOCHS}\")\nprint(f\"  LR      : {CFG.LR}\")\nprint(f\"  Batch   : {CFG.BATCH} × {CFG.GACC} = {CFG.BATCH * CFG.GACC}\")\nprint(f\"  Beams   : {CFG.NUM_BEAMS}\")\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 3 ─ AKKADIAN PREPROCESSOR\n# ════════════════════════════════════════════════════════════════════════════\n\nclass AkkadianPreprocessor:\n    \"\"\"\n    Full preprocessing pipeline following the competition's Dataset\n    Instructions (see Overview tab):\n      • Remove scribal notations  ! ? /\n      • Handle gaps  [x] → <gap>   … → <big_gap>\n      • Remove bracket pairs, keep inner text\n      • Normalize Ḫ/ḫ → H/h  (CRITICAL: test data uses H only)\n      • Expand subscripts  ₂ → 2  (CDLI convention)\n      • Remove line numbers at start of lines\n    \"\"\"\n\n    # Subscript digit map\n    _SUB = str.maketrans('₀₁₂₃₄₅₆₇₈₉ₓ', '0123456789x')\n\n    # Character transliteration table\n    _CHARS = str.maketrans({\n        # CRITICAL normalisation (test uses H, train uses Ḫ)\n        'Ḫ': 'H',  'ḫ': 'h',\n        # Sibilants\n        'š': 'sh', 'Š': 'SH',\n        'ṣ': 'ts', 'Ṣ': 'TS',\n        'ṭ': 'th', 'Ṭ': 'TH',\n        # Glottals (drop)\n        'ʾ': '',   'ʿ': '',   'ʔ': '',\n        # Long vowels → short\n        'ā': 'a',  'Ā': 'A',\n        'ē': 'e',  'Ē': 'E',\n        'ī': 'i',  'Ī': 'I',\n        'ū': 'u',  'Ū': 'U',\n        # CDLI accent encoding (keep as number suffix for tokeniser)\n        'á': 'a2', 'à': 'a3',\n        'é': 'e2', 'è': 'e3',\n        'í': 'i2', 'ì': 'i3',\n        'ú': 'u2', 'ù': 'u3',\n    })\n\n    # Compiled regexes\n    _DBL_ANGLE   = re.compile(r'<<[^>]*>>')          # erroneous signs → delete\n    _SCRIBAL     = re.compile(r'[!?/]')              # ! ? /\n    _HALF_BRACK  = re.compile(r'[˹˺⌈⌉]')            # partially broken\n    _ANGLE       = re.compile(r'[<>]')               # scribal insertion brackets\n    _GAP_BIG     = re.compile(r'\\[…[^\\]]*\\]|\\.\\.\\.|…')\n    _GAP_SMALL   = re.compile(r'\\[\\s*[xX?]\\s*\\]')\n    _SQ_BRACK    = re.compile(r'[\\[\\]]')\n    _DETERM      = re.compile(r'\\{[^}]*\\}')          # {d} {ki} etc.\n    _WORD_DIV    = re.compile(r'(?<!\\d):(?!\\d)')     # word-divider colon\n    _LINE_NUM    = re.compile(r'(?m)^\\s*\\d+\\'{0,2}\\s+')\n\n    def _source(self, text: str) -> str:\n        text = self._DBL_ANGLE.sub(' ', text)\n        text = self._SCRIBAL.sub('', text)\n        text = self._HALF_BRACK.sub('', text)\n        text = self._ANGLE.sub('', text)\n        text = self._GAP_BIG.sub(' <big_gap> ', text)\n        text = self._GAP_SMALL.sub(' <gap> ', text)\n        text = self._SQ_BRACK.sub('', text)\n        text = self._DETERM.sub('', text)\n        text = self._WORD_DIV.sub(' ', text)\n        text = self._LINE_NUM.sub('', text)\n        text = text.translate(self._SUB)\n        text = text.translate(self._CHARS)\n        return re.sub(r'\\s+', ' ', text).strip()\n\n    def _target(self, text: str) -> str:\n        text = re.sub(r'\\s+', ' ', text).strip()\n        text = re.sub(r'\\s+([.,;:!?])', r'\\1', text)\n        return text\n\n    def __call__(self, text, mode: str = 'source') -> str:\n        if text is None or (isinstance(text, float) and math.isnan(text)):\n            return '<gap>'\n        text = str(text)\n        return self._source(text) if mode == 'source' else self._target(text)\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 4 ─ CORPUS METRICS  (micro-average BLEU + chrF++)\n# ════════════════════════════════════════════════════════════════════════════\n\nclass CorpusMetrics:\n    \"\"\"\n    Corpus-level BLEU + chrF++ → Geometric Mean.\n    Micro-average (accumulate counts, divide once) = SacreBLEU behaviour.\n    \"\"\"\n\n    @staticmethod\n    def _wng(s: str, n: int) -> Counter:\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: str, n: int) -> Counter:\n        return Counter(s[i:i+n] for i in range(max(0, len(s)-n+1)))\n\n    @classmethod\n    def corpus_bleu(cls, preds: List[str], refs: List[str],\n                    max_n: int = 4) -> float:\n        st: Dict = defaultdict(int)\n        for p, r in zip(preds, refs):\n            pw, rw = p.strip().split(), r.strip().split()\n            st['hyp'] += len(pw); st['ref'] += len(rw)\n            for n in range(1, max_n + 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['hyp'] >= st['ref']\n              else math.exp(1 - st['ref'] / max(st['hyp'], 1)))\n        lp = 0.0\n        for n in range(1, max_n + 1):\n            if st[f't{n}'] == 0 or st[f'm{n}'] == 0:\n                return 0.0\n            lp += math.log(st[f'm{n}'] / st[f't{n}'])\n        return 100.0 * bp * math.exp(lp / max_n)\n\n    @classmethod\n    def corpus_chrf(cls, preds: List[str], refs: List[str],\n                    char_order: int = 6, word_order: int = 2,\n                    beta: float = 2.0) -> float:\n        st: Dict = defaultdict(float)\n        for p, r in zip(preds, refs):\n            for n in range(1, char_order + 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, word_order + 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', char_order), ('w', word_order)]:\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                denom = beta**2 * prec + rec\n                fs.append((1 + beta**2) * prec * rec / denom if denom else 0.0)\n        return 100.0 * (sum(fs) / len(fs)) if fs else 0.0\n\n    @classmethod\n    def score(cls, preds: List[str], refs: List[str]) -> Dict:\n        bleu = cls.corpus_bleu(preds, refs)\n        chrf = cls.corpus_chrf(preds, refs)\n        gm   = math.sqrt(bleu * chrf) if bleu > 0 and chrf > 0 else 0.0\n        return {'bleu': round(bleu, 4), 'chrf': round(chrf, 4),\n                'geo_mean': round(gm, 4)}\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 5 ─ DATASET  (100 % deprecated-API-free)\n# ════════════════════════════════════════════════════════════════════════════\n\nclass Seq2SeqDataset(Dataset):\n    \"\"\"\n    Works for both MarianMT and NLLB.\n\n    KEY FIX — label encoding (no as_target_tokenizer):\n      • MarianMT: source and target share the same SentencePiece vocab.\n        Call tokenizer(target_text) directly.\n      • NLLB: set tokenizer.src_lang = tgt_lang_code BEFORE encoding the\n        target, then restore src_lang afterwards.\n        This is the correct, supported API since transformers 4.26+.\n    \"\"\"\n\n    def __init__(self,\n                 sources  : List[str],\n                 targets  : Optional[List[str]],\n                 tokenizer,\n                 max_src  : int,\n                 max_tgt  : int,\n                 is_nllb  : bool = False,\n                 src_lang : str  = None,\n                 tgt_lang : str  = None):\n        self.sources  = sources\n        self.targets  = targets\n        self.tok      = tokenizer\n        self.max_src  = max_src\n        self.max_tgt  = max_tgt\n        self.is_nllb  = is_nllb\n        self.src_lang = src_lang   # e.g. 'arb_Arab'\n        self.tgt_lang = tgt_lang   # e.g. 'eng_Latn'\n\n    def __len__(self):\n        return len(self.sources)\n\n    def __getitem__(self, idx: int) -> Dict:\n        # ── Encode source ─────────────────────────────────────────────────\n        if self.is_nllb and self.src_lang:\n            self.tok.src_lang = self.src_lang    # must set BEFORE tokenising src\n\n        enc = self.tok(\n            self.sources[idx],\n            max_length=self.max_src,\n            truncation=True,\n            padding=False,\n            return_tensors=None,\n        )\n        item = {\n            'input_ids'      : enc['input_ids'],\n            'attention_mask' : enc['attention_mask'],\n        }\n\n        # ── Encode target labels ──────────────────────────────────────────\n        if self.targets is not None:\n            if self.is_nllb and self.tgt_lang:\n                # FIX: set src_lang to tgt_lang so the tokeniser prepends\n                # the correct language token for the target sequence\n                self.tok.src_lang = self.tgt_lang\n\n            # ── FIXED LINE: direct call, no deprecated context manager ────\n            lbl_enc = self.tok(\n                self.targets[idx],\n                max_length=self.max_tgt,\n                truncation=True,\n                padding=False,\n                return_tensors=None,\n            )\n            item['labels'] = lbl_enc['input_ids']\n\n            # Restore src_lang for next iteration\n            if self.is_nllb and self.src_lang:\n                self.tok.src_lang = self.src_lang\n\n        return item\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 6 ─ FINE-TUNING ENGINE\n# ════════════════════════════════════════════════════════════════════════════\n\nclass Trainer:\n    \"\"\"\n    CPU/GPU-safe fine-tuning:\n      • Label-smoothed cross-entropy\n      • Cosine LR with linear warmup\n      • Gradient accumulation + clipping\n      • Best-checkpoint saving by GeoMean\n      • No GradScaler (CPU-safe); AMP only if GPU present\n    \"\"\"\n\n    def __init__(self, model, tokenizer, cfg: Config, name: str):\n        self.model = model\n        self.tok   = tokenizer\n        self.cfg   = cfg\n        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    # ── AdamW with selective weight decay ────────────────────────────────\n    def _optimizer(self, lr: float):\n        no_wd = {'bias', 'LayerNorm.weight', 'layer_norm.weight',\n                 'layernorm.weight', 'final_layer_norm.weight'}\n        groups = [\n            {'params': [p for n, p in self.model.named_parameters()\n                        if not any(nd in n for nd in no_wd)],\n             'weight_decay': self.cfg.WEIGHT_DECAY},\n            {'params': [p for n, p in self.model.named_parameters()\n                        if any(nd in n for nd in no_wd)],\n             'weight_decay': 0.0},\n        ]\n        return torch.optim.AdamW(groups, lr=lr, eps=1e-8)\n\n    # ── Label-smoothed CE ─────────────────────────────────────────────────\n    @staticmethod\n    def _smooth_ce(logits: torch.Tensor,\n                   labels: torch.Tensor,\n                   eps: float = 0.1) -> torch.Tensor:\n        V    = logits.size(-1)\n        logp = torch.nn.functional.log_softmax(logits.float(), dim=-1)\n        mask = labels != -100\n        # Clamp -100 → 0 BEFORE gather; masked positions are excluded anyway\n        safe = labels.clamp(min=0)\n        nll  = -logp.gather(-1, safe.unsqueeze(-1)).squeeze(-1)\n        smo  = -logp.sum(-1) / V\n        loss = (1 - eps) * nll + eps * smo\n        return loss[mask].mean()\n\n    # ── One training epoch ────────────────────────────────────────────────\n    def _train_epoch(self, loader, optimizer, scheduler, epoch: int) -> float:\n        self.model.train()\n        total_loss = 0.0\n        n_steps    = 0\n        optimizer.zero_grad()\n        pad_id = self.tok.pad_token_id if self.tok.pad_token_id is not None else 0\n\n        pbar = tqdm(loader, desc=f\"  Ep{epoch:02d}\", leave=False,\n                    mininterval=15, miniters=10)\n\n        for step, batch in enumerate(pbar):\n            inp   = batch['input_ids'].to(self.cfg.DEVICE)\n            mask  = batch['attention_mask'].to(self.cfg.DEVICE)\n            lbls  = batch['labels'].to(self.cfg.DEVICE)\n            lbls[lbls == pad_id] = -100   # mask padding in labels\n\n            # Forward (no autocast on CPU — torch.amp.autocast needs\n            # device_type='cuda' and raises on CPU in some versions)\n            if HAS_GPU:\n                with torch.amp.autocast(device_type='cuda'):\n                    out  = self.model(input_ids=inp,\n                                      attention_mask=mask,\n                                      labels=lbls)\n                    loss = self._smooth_ce(out.logits, lbls,\n                                           self.cfg.LABEL_SMOOTH)\n            else:\n                out  = self.model(input_ids=inp,\n                                  attention_mask=mask,\n                                  labels=lbls)\n                loss = self._smooth_ce(out.logits, lbls,\n                                       self.cfg.LABEL_SMOOTH)\n\n            (loss / self.cfg.GACC).backward()\n            total_loss += loss.item()\n            n_steps    += 1\n\n            if (step + 1) % self.cfg.GACC == 0:\n                torch.nn.utils.clip_grad_norm_(\n                    self.model.parameters(), self.cfg.GRAD_CLIP)\n                optimizer.step()\n                scheduler.step()\n                optimizer.zero_grad()\n\n            if n_steps % 20 == 0:\n                pbar.set_postfix(loss=f'{total_loss/n_steps:.4f}')\n\n        return total_loss / max(n_steps, 1)\n\n    # ── Evaluation (generation-based) ────────────────────────────────────\n    @torch.no_grad()\n    def _eval(self, sources: List[str], references: List[str],\n              forced_bos: Optional[int] = None,\n              src_lang: Optional[str] = None) -> Tuple[Dict, List[str]]:\n        self.model.eval()\n        preds = []\n        for i in range(0, len(sources), self.cfg.BINFER):\n            batch = sources[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_SRC_LEN).to(self.cfg.DEVICE)\n            kw = dict(\n                max_length=self.cfg.MAX_TGT_LEN,\n                num_beams=4,                   # fast eval beams\n                length_penalty=self.cfg.LEN_PENALTY,\n                no_repeat_ngram_size=self.cfg.NO_REPEAT_NGRAM,\n                early_stopping=True,\n            )\n            if forced_bos is not None:\n                kw['forced_bos_token_id'] = forced_bos\n            out = self.model.generate(**enc, **kw)\n            preds.extend(self.tok.batch_decode(out, skip_special_tokens=True))\n        m = CorpusMetrics.score(preds, references)\n        return m, preds\n\n    # ── Main fit ──────────────────────────────────────────────────────────\n    def fit(self, train_ds: Dataset, val_ds: Dataset,\n            val_src: List[str], val_ref: List[str],\n            lr: float,\n            forced_bos: Optional[int] = None,\n            src_lang:   Optional[str] = None) -> float:\n\n        collator = DataCollatorForSeq2Seq(\n            tokenizer=self.tok,\n            model=self.model,\n            label_pad_token_id=-100,\n            padding=True,\n            pad_to_multiple_of=None,\n        )\n        tr_loader = DataLoader(train_ds, batch_size=self.cfg.BATCH,\n                               shuffle=True, collate_fn=collator,\n                               num_workers=0, drop_last=False)\n\n        optimizer = self._optimizer(lr)\n        total_s   = (len(tr_loader) // self.cfg.GACC) * self.cfg.EPOCHS\n        warmup_s  = max(1, int(total_s * self.cfg.WARMUP_RATIO))\n        scheduler = get_cosine_schedule_with_warmup(\n            optimizer, warmup_s, total_s)\n\n        print(f\"\\n  ┌── {self.name} ──────────────────────────────────\")\n        print(f\"  │  Train={len(train_ds)}  Val={len(val_ds)}\")\n        print(f\"  │  Steps={total_s}  Warmup={warmup_s}  LR={lr}\")\n        print(f\"  └─────────────────────────────────────────────────\")\n\n        for epoch in range(1, self.cfg.EPOCHS + 1):\n            t0   = time.time()\n            loss = self._train_epoch(tr_loader, optimizer, scheduler, epoch)\n            m, _ = self._eval(val_src, val_ref, forced_bos, src_lang)\n            dt   = time.time() - t0\n            flag = ' ✅' if m['geo_mean'] > self.best else ''\n            print(f\"  Ep{epoch:02d}/{self.cfg.EPOCHS} | \"\n                  f\"loss={loss:.4f} | \"\n                  f\"BLEU={m['bleu']:.2f}  chrF={m['chrf']:.2f}  \"\n                  f\"GM={m['geo_mean']:.4f} | {dt:.0f}s{flag}\")\n            if m['geo_mean'] > self.best:\n                self.best = m['geo_mean']\n                self.model.save_pretrained(self.ckpt)\n                self.tok.save_pretrained(self.ckpt)\n\n        print(f\"\\n  Best GeoMean [{self.name}]: {self.best:.4f}\")\n        return self.best\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 7 ─ LOAD & PREPROCESS DATA\n# ════════════════════════════════════════════════════════════════════════════\n\nprint(\"\\n\" + \"═\" * 66)\nprint(\"  📂  LOADING DATA\")\nprint(\"═\" * 66)\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 cols : {list(train_df.columns)}\")\nprint(f\"  Test cols  : {list(test_df.columns)}\")\nprint(f\"  Train rows : {len(train_df)}\")\nprint(f\"  Test rows  : {len(test_df)}\")\n\n# ── Flexible column detection ─────────────────────────────────────────────────\ndef find_col(df: pd.DataFrame, candidates: List[str]) -> str:\n    for c in candidates:\n        if c in df.columns:\n            return c\n    for c in df.columns:\n        for cand in candidates:\n            if cand.lower() in c.lower():\n                return c\n    return df.columns[0]\n\nSRC_COL = find_col(train_df, ['transliteration', 'source', 'akk', 'text'])\nTGT_COL = find_col(train_df, ['translation', 'target', 'en', 'english'])\nTR_ID   = find_col(train_df, ['oare_id', 'id', 'index'])\nTE_SRC  = find_col(test_df,  ['transliteration', 'source', 'akk', 'text'])\nTE_ID   = find_col(test_df,  ['id', 'oare_id', 'index'])\n\nprint(f\"\\n  src={SRC_COL}  tgt={TGT_COL}  train_id={TR_ID}\")\nprint(f\"  test_src={TE_SRC}  test_id={TE_ID}\")\n\n# ── Preprocess ────────────────────────────────────────────────────────────────\nprep = AkkadianPreprocessor()\n\ntrain_df['src'] = train_df[SRC_COL].apply(lambda x: prep(x, 'source'))\ntrain_df['tgt'] = train_df[TGT_COL].apply(lambda x: prep(x, 'target'))\ntest_df['src']  = test_df[TE_SRC].apply(lambda x:  prep(x, 'source'))\n\n# Drop corrupt rows\nmask = (train_df['src'].str.len() > 2) & (train_df['tgt'].str.len() > 2)\ntrain_df = train_df[mask].reset_index(drop=True)\nprint(f\"\\n  Clean training pairs: {len(train_df)}\")\n\n# ── Train / Val split ─────────────────────────────────────────────────────────\nfrom sklearn.model_selection import train_test_split\n\ntr_df, va_df = train_test_split(train_df, test_size=CFG.VAL_RATIO,\n                                random_state=CFG.SEED)\ntr_df = tr_df.reset_index(drop=True)\nva_df = va_df.reset_index(drop=True)\n\nTRAIN_SRC = tr_df['src'].tolist()\nTRAIN_TGT = tr_df['tgt'].tolist()\nVAL_SRC   = va_df['src'].tolist()\nVAL_TGT   = va_df['tgt'].tolist()\nTEST_SRC  = test_df['src'].tolist()\nTEST_IDS  = test_df[TE_ID].tolist()\n\nprint(f\"  Train={len(TRAIN_SRC)}  Val={len(VAL_SRC)}  Test={len(TEST_SRC)}\")\nprint(f\"\\n  SRC sample: {TRAIN_SRC[0][:90]}\")\nprint(f\"  TGT sample: {TRAIN_TGT[0][:90]}\")\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 8 ─ MODEL A: MARIANMT\n# ════════════════════════════════════════════════════════════════════════════\n\nprint(\"\\n\" + \"═\" * 66)\nprint(\"  🔧  MODEL A — MarianMT (Semitic→English transfer learning)\")\nprint(\"═\" * 66)\n\nmarian_tok   = None\nmarian_model = None\nMARIAN_OK    = False\nMARIAN_NAME  = ''\nMARIAN_SCORE = 0.0\n\nfor mname in CFG.MARIAN_CANDIDATES:\n    try:\n        print(f\"\\n  Trying {mname} …\", end=' ', flush=True)\n        marian_tok   = MarianTokenizer.from_pretrained(mname)\n        marian_model = MarianMTModel.from_pretrained(mname).to(CFG.DEVICE)\n        MARIAN_OK    = True\n        MARIAN_NAME  = mname\n        n = sum(p.numel() for p in marian_model.parameters())\n        print(f\"✅  {n/1e6:.1f} M params\")\n        break\n    except Exception as e:\n        print(f\"❌  {e}\")\n\nif MARIAN_OK:\n    train_ds_m = Seq2SeqDataset(TRAIN_SRC, TRAIN_TGT, marian_tok,\n                                CFG.MAX_SRC_LEN, CFG.MAX_TGT_LEN)\n    val_ds_m   = Seq2SeqDataset(VAL_SRC,   VAL_TGT,   marian_tok,\n                                CFG.MAX_SRC_LEN, CFG.MAX_TGT_LEN)\n\n    trainer_m    = Trainer(marian_model, marian_tok, CFG, 'marian')\n    MARIAN_SCORE = trainer_m.fit(train_ds_m, val_ds_m,\n                                 VAL_SRC, VAL_TGT, lr=CFG.LR)\n\n    # Reload best weights\n    marian_model = MarianMTModel.from_pretrained(trainer_m.ckpt).to(CFG.DEVICE)\n    marian_model.eval()\nelse:\n    print(\"\\n  ⚠️  No MarianMT model available\")\n\ngc.collect()\nif HAS_GPU:\n    torch.cuda.empty_cache()\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 9 ─ MODEL B: NLLB-200-distilled-600M\n# ════════════════════════════════════════════════════════════════════════════\n\nprint(\"\\n\" + \"═\" * 66)\nprint(\"  🔧  MODEL B — NLLB-200-600M  (arb_Arab proxy → eng_Latn)\")\nprint(\"═\" * 66)\n\nnllb_tok   = None\nnllb_model = None\nnllb_bos   = None\nNLLB_OK    = False\nNLLB_SCORE = 0.0\n\ntry:\n    print(f\"\\n  Loading {CFG.NLLB_MODEL} …\", end=' ', flush=True)\n    nllb_tok   = NllbTokenizerFast.from_pretrained(CFG.NLLB_MODEL)\n    nllb_tok.src_lang = CFG.NLLB_SRC\n    nllb_model = AutoModelForSeq2SeqLM.from_pretrained(\n        CFG.NLLB_MODEL,\n        torch_dtype=torch.float32,    # float32 = CPU safe\n        low_cpu_mem_usage=True,\n    ).to(CFG.DEVICE)\n    nllb_bos   = nllb_tok.convert_tokens_to_ids(CFG.NLLB_TGT)\n    NLLB_OK    = True\n    n = sum(p.numel() for p in nllb_model.parameters())\n    print(f\"✅  {n/1e6:.1f} M params\")\nexcept Exception as e:\n    print(f\"❌  {e}\")\n\nif NLLB_OK:\n    train_ds_n = Seq2SeqDataset(\n        TRAIN_SRC, TRAIN_TGT, nllb_tok,\n        CFG.MAX_SRC_LEN, CFG.MAX_TGT_LEN,\n        is_nllb=True, src_lang=CFG.NLLB_SRC, tgt_lang=CFG.NLLB_TGT)\n    val_ds_n   = Seq2SeqDataset(\n        VAL_SRC, VAL_TGT, nllb_tok,\n        CFG.MAX_SRC_LEN, CFG.MAX_TGT_LEN,\n        is_nllb=True, src_lang=CFG.NLLB_SRC, tgt_lang=CFG.NLLB_TGT)\n\n    # NLLB is larger → smaller LR\n    trainer_n    = Trainer(nllb_model, nllb_tok, CFG, 'nllb')\n    NLLB_SCORE   = trainer_n.fit(\n        train_ds_n, val_ds_n, VAL_SRC, VAL_TGT,\n        lr=CFG.LR * 0.2,\n        forced_bos=nllb_bos,\n        src_lang=CFG.NLLB_SRC)\n\n    nllb_model = AutoModelForSeq2SeqLM.from_pretrained(\n        trainer_n.ckpt).to(CFG.DEVICE)\n    nllb_model.eval()\nelse:\n    print(\"\\n  ⚠️  NLLB not available\")\n\ngc.collect()\nif HAS_GPU:\n    torch.cuda.empty_cache()\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 10 ─ TEST PREDICTIONS\n# ════════════════════════════════════════════════════════════════════════════\n\nprint(\"\\n\" + \"═\" * 66)\nprint(\"  🔮  GENERATING TEST PREDICTIONS\")\nprint(\"═\" * 66)\n\n\n@torch.no_grad()\ndef generate(model, tokenizer, sources: List[str],\n             batch_sz: int, cfg: Config,\n             forced_bos: Optional[int] = None,\n             src_lang: Optional[str] = None) -> List[str]:\n    \"\"\"Batch beam-search generation. CPU/GPU safe.\"\"\"\n    model.eval()\n    all_preds = []\n    for i in tqdm(range(0, len(sources), batch_sz),\n                  desc=\"  Translating\", leave=False):\n        batch = sources[i:i + batch_sz]\n        if src_lang and hasattr(tokenizer, 'src_lang'):\n            tokenizer.src_lang = src_lang\n        enc = tokenizer(\n            batch, return_tensors='pt', padding=True,\n            truncation=True, max_length=cfg.MAX_SRC_LEN\n        ).to(cfg.DEVICE)\n        kw = dict(\n            max_length=cfg.MAX_TGT_LEN,\n            num_beams=cfg.NUM_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        )\n        if forced_bos is not None:\n            kw['forced_bos_token_id'] = forced_bos\n        out   = model.generate(**enc, **kw)\n        preds = tokenizer.batch_decode(out, skip_special_tokens=True)\n        all_preds.extend(preds)\n    return all_preds\n\n\npreds_m: List[str] = []\npreds_n: List[str] = []\n\nif MARIAN_OK:\n    print(\"\\n  Model A (Marian) …\")\n    preds_m = generate(marian_model, marian_tok, TEST_SRC,\n                       CFG.BINFER, CFG)\n    print(f\"  ✅ {len(preds_m)} predictions\")\n\nif NLLB_OK:\n    print(\"\\n  Model B (NLLB) …\")\n    nllb_tok.src_lang = CFG.NLLB_SRC\n    preds_n = generate(nllb_model, nllb_tok, TEST_SRC,\n                       max(1, CFG.BINFER // 2), CFG,\n                       forced_bos=nllb_bos,\n                       src_lang=CFG.NLLB_SRC)\n    print(f\"  ✅ {len(preds_n)} predictions\")\n\nif not MARIAN_OK and not NLLB_OK:\n    raise RuntimeError(\"No model loaded — cannot generate predictions.\")\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 11 ─ POST-PROCESSING\n# ════════════════════════════════════════════════════════════════════════════\n\ndef clean(text: str) -> str:\n    \"\"\"Final clean-up for any translation output.\"\"\"\n    if not text or (isinstance(text, float) and math.isnan(text)):\n        return \"The content of this tablet section is fragmentary.\"\n    text = str(text).strip()\n    text = re.sub(r'\\s+', ' ', text)\n    text = re.sub(r'\\s+([.,;:!?])', r'\\1', text)\n    # Remove exact consecutive duplicate words\n    text = re.sub(r'\\b(\\w+)( \\1\\b)+', r'\\1', text, flags=re.IGNORECASE)\n    if text:\n        text = text[0].upper() + text[1:]\n    if len(text.strip()) < 4:\n        text = \"The text of this tablet section is unclear.\"\n    return text.strip()\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 12 ─ SCORE-WEIGHTED ENSEMBLE\n# ════════════════════════════════════════════════════════════════════════════\n\nprint(\"\\n\" + \"═\" * 66)\nprint(\"  📊  ENSEMBLE SELECTION\")\nprint(\"═\" * 66)\nprint(f\"\\n  Marian GM = {MARIAN_SCORE:.4f}\")\nprint(f\"  NLLB   GM = {NLLB_SCORE:.4f}\")\n\nn_models = sum([MARIAN_OK, NLLB_OK])\n\nif n_models == 1:\n    final = preds_m if MARIAN_OK else preds_n\n    print(\"  Using single model\")\nelse:\n    # Proportional weighting\n    total = MARIAN_SCORE + NLLB_SCORE\n    w_m   = MARIAN_SCORE / total if total > 0 else 0.5\n    w_n   = 1.0 - w_m\n    print(f\"  Weights → Marian={w_m:.3f}  NLLB={w_n:.3f}\")\n\n    final = []\n    for pm, pn in zip(preds_m, preds_n):\n        pm_c = clean(pm); pn_c = clean(pn)\n        if w_m >= 0.55:\n            choice = pm_c\n        elif w_n >= 0.55:\n            choice = pn_c\n        else:\n            # Tie: prefer longer output (Akkadian is verbose)\n            choice = pm_c if len(pm_c) >= len(pn_c) else pn_c\n        final.append(choice)\n\nfinal = [clean(p) for p in final]\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 13 ─ LEXICON CORRECTION\n# ════════════════════════════════════════════════════════════════════════════\n\nprint(\"\\n\" + \"═\" * 66)\nprint(\"  📖  PROPER NOUN LEXICON CORRECTION\")\nprint(\"═\" * 66)\n\nlex: Dict[str, str] = {}\nfor fname in ['OA_Lexicon_eBL.csv', 'eBL_Dictionary.csv', 'lexicon.csv']:\n    fp = f'{CFG.DATA_DIR}/{fname}'\n    if os.path.exists(fp):\n        try:\n            lex_df = pd.read_csv(fp, on_bad_lines='skip')\n            for _, row in lex_df.iterrows():\n                vals = [str(v).strip() for v in row.tolist()\n                        if str(v).strip() not in ('nan', '')]\n                if len(vals) >= 2:\n                    lex[vals[0].lower()] = vals[1]\n            print(f\"  Loaded {fname}: {len(lex)} entries\")\n        except Exception as e:\n            print(f\"  ⚠️  {fname}: {e}\")\n\nif lex:\n    def apply_lex(text: str) -> str:\n        for k, v in lex.items():\n            text = re.sub(rf'\\b{re.escape(k)}\\b', v, text, flags=re.IGNORECASE)\n        return text\n    final = [apply_lex(p) for p in final]\n    print(f\"  Applied {len(lex)} entries\")\nelse:\n    print(\"  No lexicon found — skipping\")\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 14 ─ BUILD & SAVE SUBMISSION\n# ════════════════════════════════════════════════════════════════════════════\n\nprint(\"\\n\" + \"═\" * 66)\nprint(\"  💾  SAVING SUBMISSION\")\nprint(\"═\" * 66)\n\nsubmission = pd.DataFrame({'id': TEST_IDS, 'translation': final})\nsubmission['id'] = submission['id'].astype(int)\nsubmission = submission.sort_values('id').reset_index(drop=True)\n\n# Validate\nassert len(submission) == len(test_df), \\\n    f\"Row mismatch: {len(submission)} vs {len(test_df)}\"\nassert not submission['translation'].isna().any(), \"NaN translations!\"\n\nsubmission.to_csv(CFG.SUB_FILE, index=False)\nmb = os.path.getsize(CFG.SUB_FILE) / 1e6\n\nprint(f\"\\n  ✅ {CFG.SUB_FILE}\")\nprint(f\"  Rows : {len(submission)}\")\nprint(f\"  Size : {mb:.3f} MB\")\nprint(f\"\\n  Preview:\")\nprint(submission.to_string())\n\n\n# ════════════════════════════════════════════════════════════════════════════\n# CELL 15 ─ FINAL VALIDATION & SUMMARY\n# ════════════════════════════════════════════════════════════════════════════\n\nprint(\"\\n\" + \"═\" * 66)\nprint(\"  📊  FINAL VALIDATION ON VAL SET\")\nprint(\"═\" * 66)\n\nif MARIAN_OK and len(VAL_SRC) > 0:\n    vp_m = generate(marian_model, marian_tok, VAL_SRC, CFG.BINFER, CFG)\n    vp_m = [clean(p) for p in vp_m]\n    vm   = CorpusMetrics.score(vp_m, VAL_TGT)\n    print(f\"\\n  MarianMT: BLEU={vm['bleu']:.2f}  \"\n          f\"chrF={vm['chrf']:.2f}  GM={vm['geo_mean']:.4f}\")\n\nif NLLB_OK and len(VAL_SRC) > 0:\n    nllb_tok.src_lang = CFG.NLLB_SRC\n    vp_n = generate(nllb_model, nllb_tok, VAL_SRC,\n                    max(1, CFG.BINFER // 2), CFG,\n                    forced_bos=nllb_bos, src_lang=CFG.NLLB_SRC)\n    vp_n = [clean(p) for p in vp_n]\n    vn   = CorpusMetrics.score(vp_n, VAL_TGT)\n    print(f\"  NLLB-600M: BLEU={vn['bleu']:.2f}  \"\n          f\"chrF={vn['chrf']:.2f}  GM={vn['geo_mean']:.4f}\")\n\nprint(\"\\n\" + \"═\" * 66)\nprint(\"  🏆  SUMMARY\")\nprint(\"═\" * 66)\nprint(f\"\"\"\n  Competition : Deep Past Challenge — Akkadian→English MT\n  Metric      : Geometric Mean(corpus BLEU, corpus chrF++)\n\n  Training data : {len(train_df)} pairs\n  Models        :\n    A. {MARIAN_NAME or 'N/A':50s}  GM={MARIAN_SCORE:.4f}\n    B. {CFG.NLLB_MODEL if NLLB_OK else 'N/A':50s}  GM={NLLB_SCORE:.4f}\n\n  Key techniques:\n    ✓ Fine-tuned (not zero-shot) on all {len(train_df)} Akkadian pairs\n    ✓ Semitic language transfer (Arabic→EN pretrained weights)\n    ✓ Label smoothing ε={CFG.LABEL_SMOOTH}\n    ✓ Cosine LR + {int(CFG.WARMUP_RATIO*100)}% warmup, grad clip={CFG.GRAD_CLIP}\n    ✓ Ḫ/ḫ→H/h normalisation (fixes test/train mismatch)\n    ✓ Score-weighted ensemble  (Marian GM={MARIAN_SCORE:.4f}  NLLB GM={NLLB_SCORE:.4f})\n    ✓ OA Lexicon proper-noun correction ({len(lex)} entries)\n    ✓ Beam search k={CFG.NUM_BEAMS}, length_penalty={CFG.LEN_PENALTY}\n    ✓ No deprecated APIs (as_target_tokenizer removed)\n    ✓ CPU/GPU safe (no GradScaler on CPU)\n\n  Output : {CFG.SUB_FILE}  ({len(submission)} rows)  ✅\n\"\"\")\nprint(\"═\" * 66)","metadata":{"_uuid":"c5c7702a-f23e-4403-8ef3-f473fa5a1baa","_cell_guid":"3bd159d3-ab09-4758-8251-35944740f9e8","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2026-03-06T00:23:36.181009Z","iopub.execute_input":"2026-03-06T00:23:36.182368Z"}},"outputs":[],"execution_count":null}]}