{"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":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":121150,"databundleVersionId":15833391,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":798528,"databundleVersionId":16228118,"modelInstanceId":608180,"modelId":620081,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":796966,"databundleVersionId":16211400,"modelInstanceId":607085,"modelId":619063,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":798622,"databundleVersionId":16228896,"modelInstanceId":605859,"modelId":617861,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":798590,"databundleVersionId":16228606,"modelInstanceId":607175,"modelId":619139,"isSourceIdPinned":false}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Script","metadata":{}},{"cell_type":"code","source":"%%writefile normalize_utils.py\n\nimport re\n\n\ndef normalize_half_brackets(text: str) -> str:\n    \"\"\"Remove half-bracket editorial marks.\"\"\"\n    for ch in \"˹˺⸢⸣⌈⌉⌊⌋\":\n        text = text.replace(ch, \"\")\n    return textcccccbngjchbtgdrfkjfrrblrrttglibcrcrunftndie\n    \n\n\ndef normalize_heth(text: str) -> str:\n    \"\"\"ḫ → h, Ḫ → H.\"\"\"\n    return text.translate(str.maketrans(\"ḫḪ\", \"hH\"))\n\n\ndef normalize_fractions(text: str) -> str:\n    \"\"\"Decimal floats, slash fractions, and superscript/subscript fractions → Unicode.\"\"\"\n    # Superscript/subscript fractions first (before subscript digits get converted)\n    for src, dst in [\n        (\"¹/₂\", \"½\"),\n        (\"¹/₃\", \"⅓\"),\n        (\"²/₃\", \"⅔\"),\n        (\"¹/₄\", \"¼\"),\n        (\"³/₄\", \"¾\"),\n        (\"¹/₆\", \"⅙\"),\n        (\"⁵/₆\", \"⅚\"),\n        (\"⁵/₈\", \"⅝\"),\n    ]:\n        text = text.replace(src, dst)\n\n    # Decimal fractions (longest prefix first to avoid partial matches)\n    for prefix, uni in [\n        (\"0.8333\", \"⅚\"),\n        (\"0.1666\", \"⅙\"),\n        (\"0.6666\", \"⅔\"),\n        (\"0.3333\", \"⅓\"),\n        (\"0.75\", \"¾\"),\n        (\"0.625\", \"⅝\"),\n        (\"0.25\", \"¼\"),\n        (\"0.5\", \"½\"),\n    ]:\n        esc = re.escape(prefix)\n        text = re.sub(rf\"\\b{esc}\\d*\\b\", uni, text)\n        esc_no_zero = re.escape(prefix[1:])\n        text = re.sub(rf\"(\\d+){esc_no_zero}\\d*\\b\", rf\"\\1 {uni}\", text)\n\n    # Slash fractions: 1/3, 1 / 3\n    _frac = {\n        (1, 2): \"½\",\n        (1, 3): \"⅓\",\n        (2, 3): \"⅔\",\n        (1, 4): \"¼\",\n        (3, 4): \"¾\",\n        (1, 6): \"⅙\",\n        (5, 6): \"⅚\",\n        (5, 8): \"⅝\",\n    }\n\n    def _repl(m):\n        return _frac.get((int(m.group(1)), int(m.group(2))), m.group(0))\n\n    text = re.sub(r\"(\\d+)\\s*/\\s*(\\d+)\", _repl, text)\n\n    # Spacing: 1½ → 1 ½\n    text = re.sub(r\"(\\d)([½⅓⅔¼¾⅙⅚⅝])\", r\"\\1 \\2\", text)\n    return text\n\n\ndef normalize_subscripts(text: str) -> str:\n    \"\"\"Vowel subscripts → diacritics, then remaining subscripts → digits.\"\"\"\n    for src, dst in [\n        (\"a₂\", \"á\"),\n        (\"e₂\", \"é\"),\n        (\"i₂\", \"í\"),\n        (\"u₂\", \"ú\"),\n        (\"A₂\", \"Á\"),\n        (\"E₂\", \"É\"),\n        (\"I₂\", \"Í\"),\n        (\"U₂\", \"Ú\"),\n    ]:\n        text = text.replace(src, dst)\n    for src, dst in [\n        (\"a₃\", \"à\"),\n        (\"e₃\", \"è\"),\n        (\"i₃\", \"ì\"),\n        (\"u₃\", \"ù\"),\n        (\"A₃\", \"À\"),\n        (\"E₃\", \"È\"),\n        (\"I₃\", \"Ì\"),\n        (\"U₃\", \"Ù\"),\n    ]:\n        text = text.replace(src, dst)\n    text = text.translate(str.maketrans(\"₀₁₂₃₄₅₆₇₈₉\", \"0123456789\"))\n    return text\n\n\ndef normalize_macrons_tl(text: str) -> str:\n    \"\"\"Remove macron vowels from transliteration (not in test TL charset).\"\"\"\n    return text.translate(str.maketrans(\"āēīūĀĒĪŪ\", \"aeiuAEIU\"))\n\n\ndef normalize_brackets_tl(text: str) -> str:\n    \"\"\"Remove editorial brackets from transliteration, preserving content.\"\"\"\n    # <<text>> erasure removal (content removed entirely)\n    text = re.sub(r\"<<[^>]*>>\", \"\", text)\n    text = re.sub(r\"«[^»]*»\", \"\", text)\n    text = re.sub(r\"≪[^≫]*≫\", \"\", text)\n    # [text] square brackets — remove brackets, keep content\n    text = re.sub(r\"\\[([^\\]]*)\\]\", r\"\\1\", text)\n    # Unmatched brackets\n    text = text.replace(\"[\", \"\").replace(\"]\", \"\")\n\n    # <text> angle brackets — keep content, preserve <gap>/<big_gap>\n    def _strip_angle(m):\n        content = m.group(1)\n        if content in (\"gap\", \"big_gap\"):\n            return m.group(0)\n        return content\n\n    text = re.sub(r\"<([^>]+)>\", _strip_angle, text)\n    return text\n\n\ndef normalize_determinatives(text: str) -> str:\n    \"\"\"Normalize determinative format: (d)→{d}, handle TÚG.\"\"\"\n    # Superscript determinatives\n    for src, dst in [(\"ᵈ\", \"{d}\"), (\"ᵐ\", \"{m}\"), (\"ᶠ\", \"{f}\"), (\"ᵏⁱ\", \"{ki}\")]:\n        text = text.replace(src, dst)\n    # Caret notation → curly\n    for src, dst in [(\"^d\", \"{d}\"), (\"^ki\", \"{ki}\"), (\"^m\", \"{m}\"), (\"^f\", \"{f}\")]:\n        text = text.replace(src, dst)\n    # Parenthetical → curly for core determinatives\n    for src, dst in [(\"(d)\", \"{d}\"), (\"(ki)\", \"{ki}\"), (\"(m)\", \"{m}\"), (\"(f)\", \"{f}\")]:\n        text = text.replace(src, dst)\n    # Variant curly forms → standard\n    for src, dst in [(\"{D}\", \"{d}\"), (\"{DINGIR}\", \"{d}\"), (\"{KI}\", \"{ki}\"), (\"{M}\", \"{m}\"), (\"{1}\", \"{m}\"), (\"{MÍ}\", \"{f}\"), (\"{MI}\", \"{f}\"), (\"{MUNUS}\", \"{f}\")]:\n        text = text.replace(src, dst)\n    # TÚG: host says no brackets in test data\n    for src in [\"{TÚG}\", \"{túg}\", \"{tug}\", \"(TÚG)\", \"(túg)\", \"(tug)\"]:\n        text = text.replace(src, \"TÚG\")\n    return text\n\n\ndef normalize_silver(text: str) -> str:\n    \"\"\"KÙ.B → KÙ.BABBAR, KB → KÙ.BABBAR.\"\"\"\n    text = re.sub(r\"KÙ\\.B\\.\", \"KÙ.BABBAR\", text)\n    text = re.sub(r\"KÙ\\.B\\b\", \"KÙ.BABBAR\", text)\n    text = re.sub(r\"\\bKB\\b\", \"KÙ.BABBAR\", text)\n    return text\n\n\ndef normalize_editorial_marks(text: str) -> str:\n    \"\"\"Strip editorial marks (!, ?, *) from transliteration.\"\"\"\n    for ch in \"!?*\":\n        text = text.replace(ch, \"\")\n    return text\n\n\ndef normalize_dividers(text: str) -> str:\n    \"\"\"Remove colon/slash word dividers and section marks.\"\"\"\n    text = text.replace(\" : \", \" \")\n    text = re.sub(r\"\\s+:$\", \"\", text)\n    text = re.sub(r\"^:\\s+\", \"\", text)\n    text = text.replace(\"§\", \"\")\n    text = text.replace(\"=\", \"-\")\n    text = text.replace(\"--\", \"-\")\n    text = re.sub(r\"-/\\s*\", \"-\", text)\n    text = re.sub(r\"\\s+/\\s+\", \" \", text)\n    return text\n\n\ndef normalize_turkish_chars(text: str) -> str:\n    \"\"\"Map Turkish chars not in test charset to Akkadological equivalents.\"\"\"\n    text = text.replace(\"Ş\", \"Š\")\n    text = text.replace(\"ș\", \"ṣ\")\n    return text\n\n\ndef normalize_dashes(text: str) -> str:\n    \"\"\"Normalize dash variants to hyphen.\"\"\"\n    for ch in \"\\u2014\\u2013\\u2012\":  # em-dash, en-dash, figure dash\n        text = text.replace(ch, \"-\")\n    return text\n\n\ndef normalize_glottal_stops(text: str) -> str:\n    \"\"\"Strip glottal stop and apostrophe variants from transliteration.\"\"\"\n    for ch in \"ˀʾˈʼ᾿ˁˊˋˇꜥꜣʿˮʹꞌ`´\":\n        text = text.replace(ch, \"\")\n    return text\n\n\ndef normalize_gaps(text: str) -> str:\n    \"\"\"Convert damage markers to <gap> and merge adjacent gaps.\"\"\"\n    text = text.replace(\"…\", \"<gap>\")\n    text = re.sub(r\"\\.{3,}\", \"<gap>\", text)\n    text = re.sub(r\"\\[x(\\s+x)*\\]\", \"<gap>\", text)\n    text = re.sub(r\"\\((?:large\\s+)?break\\)\", \"<gap>\", text)\n    text = re.sub(r\"\\(\\d+\\s+broken\\s+lines?\\)\", \"<gap>\", text)\n    text = re.sub(r\"(\\.\\s){2,}\\.?\", \"<gap>\", text)\n    text = re.sub(r\"\\b(?:ca\\.\\s*)?\\d+\\s+lines?\\s+(?:broken|destroyed|lost|missing)(?:\\s+away)?\\b\", \"<gap>\", text, flags=re.IGNORECASE)\n    text = re.sub(r\"\\b(?:one|two|three|four|five|six|seven|eight|nine|ten)\\s+lines?\\s+(?:broken|destroyed|lost|missing)(?:\\s+away)?\\b\", \"<gap>\", text, flags=re.IGNORECASE)\n    text = re.sub(r\"\\bbreak\\s+of\\s+(?:\\w+\\s+)?lines?\\b\", \"<gap>\", text, flags=re.IGNORECASE)\n    text = re.sub(\n        r\"(?<![a-zA-ZšṣṭáàéèíìúùÁÀÉÈÍÌÚÙ<-])x(?![a-zA-ZšṣṭáàéèíìúùÁÀÉÈÍÌÚÙ>-])\",\n        \"<gap>\",\n        text,\n    )\n    text = text.replace(\"<big_gap>\", \"<gap>\")\n    text = text.replace(\"(<gap>)\", \"<gap>\")\n    # Merge adjacent\n    prev = None\n    while prev != text:\n        prev = text\n        text = re.sub(r\"<gap>\\s+<gap>\", \"<gap>\", text)\n        text = text.replace(\"<gap>-<gap>\", \"<gap>\")\n    return text\n\n\ndef normalize_whitespace(text: str) -> str:\n    \"\"\"Collapse multiple spaces, strip.\"\"\"\n    return re.sub(r\"  +\", \" \", text).strip()\n\n\n# ── Pipeline ─────────────────────────────────────────────────────────\n\n\ndef normalize_transliteration(text: str) -> str:\n    \"\"\"Normalize transliteration for pre-training.\"\"\"\n    try:\n        text = normalize_half_brackets(text)\n        text = normalize_brackets_tl(text)\n        text = normalize_determinatives(text)\n        text = normalize_silver(text)\n        text = normalize_editorial_marks(text)\n        text = normalize_dividers(text)\n        text = normalize_turkish_chars(text)\n        text = normalize_dashes(text)\n        text = normalize_heth(text)\n        text = normalize_fractions(text)\n        text = normalize_subscripts(text)\n        text = normalize_macrons_tl(text)\n        text = normalize_glottal_stops(text)\n        text = normalize_gaps(text)\n        text = normalize_whitespace(text)\n        return text\n    except Exception as e:\n        print(e)\n        return text\n\n\n# ── Translation-specific normalizers ─────────────────────────────────\n\n\ndef normalize_linebreaks_tr(text: str) -> str:\n    \"\"\"Rejoin OCR line-break hyphens in translation: mer- chant → merchant.\"\"\"\n    return re.sub(r\"(?<=\\s)([a-z]+)- ([a-z]+)\", r\"\\1\\2\", text)\n\n\ndef normalize_brackets_tr(text: str) -> str:\n    \"\"\"Remove brackets from translation, keep content.\"\"\"\n    # <<text>> erasure removal\n    text = re.sub(r\"<<[^>]*>>\", \"\", text)\n\n    # <text> angle brackets — keep content, preserve <gap>/<big_gap>\n    def _strip_angle(m):\n        content = m.group(1)\n        if content in (\"gap\", \"big_gap\"):\n            return m.group(0)\n        return content\n\n    text = re.sub(r\"<([^>]+)>\", _strip_angle, text)\n    # [text] square brackets — keep content\n    text = re.sub(r\"\\[([^\\]]*)\\]\", r\"\\1\", text)\n    text = text.replace(\"[\", \"\").replace(\"]\", \"\")\n    return text\n\n\ndef normalize_quotes(text: str) -> str:\n    \"\"\"Curly/smart quotes → straight quotes.\"\"\"\n    for src, dst in [(\"\\u201c\", '\"'), (\"\\u201d\", '\"'), (\"\\u201e\", '\"'), (\"\\u2018\", \"'\"), (\"\\u2019\", \"'\"), (\"\\u201a\", \"'\")]:\n        text = text.replace(src, dst)\n    return text\n\n\ndef normalize_macrons_tr(text: str) -> str:\n    \"\"\"Capital macrons → plain. Keep lowercase ā ē ī ū (in test charset).\"\"\"\n    return text.translate(str.maketrans(\"ĀĒĪŪ\", \"AEIU\"))\n\n\ndef normalize_circumflex_tr(text: str) -> str:\n    \"\"\"Circumflex → macron for û ê î. Keep â (in test charset).\"\"\"\n    for src, dst in [(\"û\", \"ū\"), (\"ê\", \"ē\"), (\"î\", \"ī\"), (\"Û\", \"Ū\"), (\"Ê\", \"Ē\"), (\"Î\", \"Ī\")]:\n        text = text.replace(src, dst)\n    return text\n\n\ndef normalize_glottal_tr(text: str) -> str:\n    \"\"\"Glottal stop variants → straight apostrophe in translation.\"\"\"\n    for ch in \"ˀʾˈʼ᾿ˁꜥꜣʿˮʹꞌ\":\n        text = text.replace(ch, \"'\")\n    return text\n\n\ndef normalize_grammar(text: str) -> str:\n    \"\"\"Remove grammatical annotations from translation.\"\"\"\n    for pattern in [\n        r\"\\(fem\\.\\s*plur\\.\\)\",\n        r\"\\(fem\\.\\s*pl\\.\\)\",\n        r\"\\(fem\\.\\s*sing\\.\\)\",\n        r\"\\(fem\\.\\)\",\n        r\"\\(plur\\.\\)\",\n        r\"\\(sing\\.\\)\",\n        r\"\\(masc\\.\\)\",\n        r\"\\(pl\\.\\)\",\n        r\"\\(m\\.\\)\",\n        r\"\\(f\\.\\)\",\n        r\"\\(\\?\\)\",\n        r\"\\bfem\\.\\s\",\n        r\"\\bsing\\.\\s\",\n        r\"\\bpl\\.\\s\",\n        r\"\\bplural\\b\",\n    ]:\n        text = re.sub(pattern, \"\", text)\n    return text\n\n\ndef normalize_punctuation(text: str) -> str:\n    \"\"\"Remove space before punctuation.\"\"\"\n    return re.sub(r\"\\s+([.,;:?!])\", r\"\\1\", text)\n\n\ndef normalize_double_dots(text: str) -> str:\n    \"\"\"Remove stray double dots (not part of longer sequences).\"\"\"\n    return re.sub(r\"(?<!\\.)\\.\\.(?!\\.)\", \"\", text)\n\n\ndef normalize_hyphen_terms(text: str) -> str:\n    \"\"\"Replace abbreviated hyphenated terms per host.\"\"\"\n    text = text.replace(\" -gold\", \" pašallum gold\")\n    text = text.replace(\" -tax\", \" šadduātum tax\")\n    text = text.replace(\" -textiles\", \" kutānum textiles\")\n    return text\n\n\ndef normalize_pn(text: str) -> str:\n    \"\"\"Literal PN token → <gap>.\"\"\"\n    return re.sub(r\"\\bPN\\b\", \"<gap>\", text)\n\n\ndef normalize_months(text: str) -> str:\n    \"\"\"Month Roman numeral → integer.\"\"\"\n    _roman = {\"XII\": 12, \"XI\": 11, \"X\": 10, \"IX\": 9, \"VIII\": 8, \"VII\": 7, \"VI\": 6, \"V\": 5, \"IV\": 4, \"III\": 3, \"II\": 2, \"I\": 1}\n\n    def _repl(m):\n        prefix = m.group(1)\n        roman = m.group(2)\n        return f\"{prefix} {_roman[roman]}\" if roman in _roman else m.group(0)\n\n    return re.sub(r\"\\b(month|Month)\\s+(XII|XI|X|IX|VIII|VII|VI|V|IV|III|II|I)\\b\", _repl, text)\n\n\n# ── Translation pipeline ─────────────────────────────────────────────\n\n\ndef normalize_translation(text: str) -> str:\n    \"\"\"Normalize translation for pre-training.\"\"\"\n    try:\n        text = normalize_half_brackets(text)\n        text = normalize_linebreaks_tr(text)\n        text = normalize_brackets_tr(text)\n        text = normalize_heth(text)\n        text = normalize_quotes(text)\n        text = normalize_macrons_tr(text)\n        text = normalize_circumflex_tr(text)\n        text = normalize_glottal_tr(text)\n        text = normalize_grammar(text)\n        text = normalize_hyphen_terms(text)\n        text = normalize_pn(text)\n        text = normalize_double_dots(text)\n        text = normalize_months(text)\n        text = normalize_fractions(text)\n        text = normalize_subscripts(text)\n        text = normalize_gaps(text)\n        text = text.replace(\"--\", \"-\")\n        text = normalize_punctuation(text)\n        text = normalize_whitespace(text)\n        return text\n    \n    except Exception as e:\n        print(e)\n        return text\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T19:44:05.352998Z","iopub.execute_input":"2026-03-22T19:44:05.353796Z","iopub.status.idle":"2026-03-22T19:44:05.364204Z","shell.execute_reply.started":"2026-03-22T19:44:05.353766Z","shell.execute_reply":"2026-03-22T19:44:05.363422Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile run_inference.py\n\nimport argparse\nimport math\nfrom dataclasses import dataclass\nfrom pathlib import Path\n\nimport pandas as pd\nimport sacrebleu\nimport torch\nfrom accelerate import Accelerator\nfrom omegaconf import OmegaConf\nfrom torch.utils.data import DataLoader, Dataset\nfrom tqdm.auto import tqdm\nfrom transformers import AutoTokenizer, DataCollatorWithPadding, GenerationConfig, T5ForConditionalGeneration\n\nimport sys\nsys.path.insert(0, \"./\")\nfrom normalize_utils import normalize_translation, normalize_transliteration\n\n# ── MBR Decoding ──────────────────────────────────────────────────────\n\n\ndef comp_metric_utility(hypothesis, reference):\n    \"\"\"Sentence-level sqrt(BLEU * chrF++) — same as competition metric.\"\"\"\n    bleu = sacrebleu.sentence_bleu(hypothesis, [reference]).score\n    chrf = sacrebleu.sentence_chrf(hypothesis, [reference], word_order=2).score\n    return math.sqrt(max(bleu, 0.0) * max(chrf, 0.0))\n\n\ndef mbr_select(candidates, utility_fn=None):\n    \"\"\"Return index of candidate with highest average pairwise utility.\"\"\"\n    if utility_fn is None:\n        utility_fn = comp_metric_utility\n    n = len(candidates)\n    if n == 1:\n        return 0\n    scores = []\n    for i in range(n):\n        total = sum(utility_fn(candidates[i], candidates[j]) for j in range(n) if j != i)\n        scores.append(total / (n - 1))\n    return max(range(n), key=lambda i: scores[i])\n\n\n@torch.no_grad()\ndef generate_mbr_predictions(model, tokenizer, eval_dataloader, accelerator, generation_config, mbr_n=16):\n    model.eval()\n    accelerator.print(f\"MBR decoding: N={mbr_n} candidates per input\")\n    accelerator.wait_for_everyone()\n\n    unwrapped_model = accelerator.unwrap_model(model)\n    pad_id = tokenizer.pad_token_id or 0\n    all_predictions = []\n\n    progress_bar = tqdm(eval_dataloader, desc=\"MBR Decoding\", disable=not accelerator.is_local_main_process)\n\n    for batch in progress_bar:\n        local_batch_size = batch[\"input_ids\"].shape[0]\n\n        generated_ids = unwrapped_model.generate(\n            input_ids=batch[\"input_ids\"],\n            attention_mask=batch[\"attention_mask\"],\n            generation_config=generation_config,\n        )\n\n        all_candidates = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)\n\n        best_token_ids = []\n        for i in range(local_batch_size):\n            group = all_candidates[i * mbr_n : (i + 1) * mbr_n]\n            best_idx = mbr_select(group)\n            best_token_ids.append(generated_ids[i * mbr_n + best_idx])\n\n        selected = torch.stack(best_token_ids, dim=0)\n        selected = accelerator.pad_across_processes(selected, dim=1, pad_index=pad_id)\n        selected = accelerator.gather_for_metrics(selected)\n\n        predictions = tokenizer.batch_decode(selected, skip_special_tokens=True)\n        all_predictions.extend(predictions)\n\n    accelerator.wait_for_everyone()\n    return all_predictions\n\n\n# ── Dataset & Collator ───────────────────────────────────────────────\n\n\ndef get_tokenizer(cfg):\n    tokenizer = AutoTokenizer.from_pretrained(cfg.model.backbone_path)\n    return tokenizer\n\n\nclass DPCDataset(Dataset):\n    \"\"\"\n    Dataset class for Deep Past Challenge - Translate Akkadian to English\n    \"\"\"\n\n    def __init__(self, cfg, df, id_col=\"oare_id\"):\n        self.cfg = cfg\n        self.df = df\n        self.id_col = id_col\n\n        self.tokenizer = get_tokenizer(cfg)\n\n    def __len__(self):\n        return len(self.df)\n\n    def tokenize(self, texts):\n        return self.tokenizer(texts, padding=False, truncation=True, max_length=self.cfg.model.max_length, return_length=True, add_special_tokens=True)\n\n    def pre_process(self, row):\n        ak_en_template = \"Translate Akkadian to English: {source}\"\n\n        return [ak_en_template.format(source=row[\"transliteration\"])], [None]\n\n    def __getitem__(self, idx):\n        data = self.df.iloc[idx].to_dict()\n        this_id = data[self.id_col]\n        sources, targets = self.pre_process(data)\n        tx_src = self.tokenize(sources)\n\n        return dict(id=this_id, input_ids=tx_src[\"input_ids\"], attention_mask=tx_src[\"attention_mask\"])\n\n\n@dataclass\nclass DPCCollator(DataCollatorWithPadding):\n    \"\"\"\n    Data collector for DPC dataset\n    \"\"\"\n\n    tokenizer = None\n    padding = True\n    max_length = None\n    pad_to_multiple_of = None\n\n    def __call__(self, features):\n        encoder_features = {\"input_ids\": [], \"attention_mask\": []}\n        labels = []\n\n        for feature in features:\n            encoder_features[\"input_ids\"].extend(feature[\"input_ids\"])\n            encoder_features[\"attention_mask\"].extend(feature[\"attention_mask\"])\n\n            if \"labels\" in feature:\n                labels.extend(feature[\"labels\"])\n\n        if len(labels) > 0:\n            assert len(labels) == len(encoder_features[\"input_ids\"])\n\n        batch = self.tokenizer.pad(\n            encoder_features,\n            padding=self.padding,\n            max_length=self.max_length,\n            pad_to_multiple_of=self.pad_to_multiple_of,\n            return_tensors=\"pt\",\n        )\n\n        # labels\n        if len(labels) > 0:\n            labels_batch = self.tokenizer.pad({\"input_ids\": labels}, padding=self.padding, max_length=self.max_length, pad_to_multiple_of=self.pad_to_multiple_of, return_tensors=None)\n            labels = torch.tensor(labels_batch[\"input_ids\"], dtype=torch.int64)\n            # Mask padding with -100 (ignored in loss)\n            labels[labels == self.tokenizer.pad_token_id] = -100\n            batch[\"labels\"] = labels\n        return batch\n\n\ndef show_batch(batch, tokenizer, print_fn=print, **kwargs):\n    bs = batch[\"input_ids\"].size(0)\n    print_fn(f\"batch size: {bs}\")\n\n    print_fn(f\"shape of input_ids: {batch['input_ids'].shape}\")\n    print_fn(f\"shape of attention_mask: {batch['attention_mask'].shape}\")\n    if \"labels\" in batch.keys():\n        print_fn(f\"shape of labels: {batch['labels'].shape}\")\n\n    print_fn(\"\\n\\n\")\n    max_n = kwargs.get(\"n\", 5)\n\n    for idx in range(bs):\n        print_fn(f\"\\n\\n===== Example {idx + 1} =====\\n\\n\")\n        print_fn(f\"## Input:\\n\\n{tokenizer.decode(batch['input_ids'][idx], skip_special_tokens=False)}\\n\\n\")\n\n        if \"labels\" in batch.keys():\n            labels = batch[\"labels\"][idx]\n            labels[labels == -100] = tokenizer.pad_token_id\n            print_fn(f\"## Output:\\n\\n{tokenizer.decode(labels, skip_special_tokens=False)}\")\n\n        if idx >= max_n:\n            break\n\n\n# ── Generation ────────────────────────────────────────────────────────\n\n\n@torch.no_grad()\ndef generate_predictions(model, tokenizer, eval_dataloader, accelerator, generation_config):\n    model.eval()\n    accelerator.print(\"Generating predictions...\")\n    accelerator.wait_for_everyone()\n\n    unwrapped_model = accelerator.unwrap_model(model)\n\n    all_predictions = []\n\n    progress_bar = tqdm(eval_dataloader, desc=\"Evaluating\", disable=not accelerator.is_local_main_process)\n\n    for batch in progress_bar:\n        # Generate predictions\n        generated_ids = unwrapped_model.generate(input_ids=batch[\"input_ids\"], attention_mask=batch[\"attention_mask\"], generation_config=generation_config)\n\n        # Pad and gather across processes\n        generated_ids = accelerator.pad_across_processes(generated_ids, dim=1, pad_index=tokenizer.pad_token_id or 0)\n        generated_ids = accelerator.gather_for_metrics(generated_ids)\n\n        # Decode on main process\n        predictions = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)\n        all_predictions.extend(predictions)\n\n    accelerator.wait_for_everyone()\n\n    return all_predictions\n\n\ndef get_generation_config(cfg, tokenizer):\n    return GenerationConfig(\n        max_new_tokens=cfg.generation.max_new_tokens,\n        do_sample=cfg.generation.do_sample,\n        top_k=cfg.generation.top_k,\n        top_p=cfg.generation.top_p,\n        temperature=cfg.generation.temperature,\n        use_cache=True,\n        pad_token_id=tokenizer.pad_token_id,\n        eos_token_id=tokenizer.eos_token_id,\n    )\n\n\n# ── Main ──────────────────────────────────────────────────────────────\n\n\ndef main(cfg, save_path):\n    accelerator = Accelerator()\n\n    test_df = pd.read_csv(\"/kaggle/input/deep-past-initiative-machine-translation/test.csv\")\n    test_df[\"transliteration\"] = test_df[\"transliteration\"].apply(normalize_transliteration)\n\n    infer_ds = DPCDataset(cfg, test_df, id_col=\"id\")\n    tokenizer = infer_ds.tokenizer\n    data_collator = DPCCollator(tokenizer=tokenizer, pad_to_multiple_of=16)\n\n    infer_dl = DataLoader(\n        infer_ds,\n        batch_size=cfg.predict_params.per_device_eval_batch_size,\n        shuffle=False,\n        collate_fn=data_collator,\n    )\n\n    accelerator.print(\"showing batch...\")\n    for idx, b in enumerate(infer_dl):\n        show_batch(b, tokenizer, task=\"infer\", print_fn=accelerator.print, n=2)\n        break\n\n    model = T5ForConditionalGeneration.from_pretrained(cfg.model.backbone_path)\n    model, infer_dl = accelerator.prepare(model, infer_dl)\n\n    use_mbr = cfg.generation.get(\"mbr\", False)\n    mbr_n = cfg.generation.get(\"mbr_n\", 16)\n\n    if use_mbr:\n        mbr_num_beams = cfg.generation.get(\"mbr_num_beams\", 0)\n        if mbr_num_beams > 0:\n            generation_config = GenerationConfig(\n                max_new_tokens=cfg.generation.max_new_tokens,\n                use_cache=True,\n                do_sample=False,\n                num_beams=mbr_num_beams,\n                num_return_sequences=mbr_n,\n                pad_token_id=tokenizer.pad_token_id,\n                eos_token_id=tokenizer.eos_token_id,\n            )\n        else:\n            generation_config = GenerationConfig(\n                max_new_tokens=cfg.generation.max_new_tokens,\n                use_cache=True,\n                do_sample=True,\n                temperature=cfg.generation.get(\"mbr_temperature\", 0.5),\n                top_p=cfg.generation.get(\"mbr_top_p\", 0.9),\n                top_k=cfg.generation.get(\"mbr_top_k\", 0),\n                num_return_sequences=mbr_n,\n                pad_token_id=tokenizer.pad_token_id,\n                eos_token_id=tokenizer.eos_token_id,\n            )\n        predictions = generate_mbr_predictions(model=model, tokenizer=tokenizer, eval_dataloader=infer_dl, accelerator=accelerator, generation_config=generation_config, mbr_n=mbr_n)\n    else:\n        generation_config = get_generation_config(cfg, tokenizer)\n        predictions = generate_predictions(model=model, tokenizer=tokenizer, eval_dataloader=infer_dl, accelerator=accelerator, generation_config=generation_config)\n\n    predictions_df = pd.DataFrame({\"id\": test_df[\"id\"], \"translation\": predictions, \"transliteration\": test_df[\"transliteration\"]})\n    predictions_df[\"translation\"] = predictions_df[\"translation\"].apply(normalize_translation)\n    \n    predictions_df.to_parquet(save_path)\n\n\nif __name__ == \"__main__\":\n    ap = argparse.ArgumentParser()\n    ap.add_argument(\"--config_path\", type=str, required=True)\n    ap.add_argument(\"--save_path\", type=str, required=True)\n    args = ap.parse_args()\n\n    cfg = OmegaConf.load(args.config_path)\n\n    Path(args.save_path).parent.mkdir(parents=True, exist_ok=True)\n\n    # run\n    main(cfg, args.save_path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T19:44:16.961006Z","iopub.execute_input":"2026-03-22T19:44:16.961845Z","iopub.status.idle":"2026-03-22T19:44:16.971194Z","shell.execute_reply.started":"2026-03-22T19:44:16.961816Z","shell.execute_reply":"2026-03-22T19:44:16.970507Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile run_inference_xl.py\n\nimport argparse\nimport math\nfrom dataclasses import dataclass\nfrom pathlib import Path\n\nimport pandas as pd\nimport sacrebleu\nimport torch\nfrom omegaconf import OmegaConf\nfrom torch.utils.data import DataLoader, Dataset\nfrom tqdm.auto import tqdm\nfrom transformers import AutoTokenizer, DataCollatorWithPadding, GenerationConfig, T5ForConditionalGeneration\n\nimport sys\nsys.path.insert(0, \"./\")\nfrom normalize_utils import normalize_translation, normalize_transliteration\n\n# ── MBR Decoding ──────────────────────────────────────────────────────\n\n\ndef comp_metric_utility(hypothesis, reference):\n    \"\"\"Sentence-level sqrt(BLEU * chrF++) — same as competition metric.\"\"\"\n    bleu = sacrebleu.sentence_bleu(hypothesis, [reference]).score\n    chrf = sacrebleu.sentence_chrf(hypothesis, [reference], word_order=2).score\n    return math.sqrt(max(bleu, 0.0) * max(chrf, 0.0))\n\n\ndef mbr_select(candidates, utility_fn=None):\n    \"\"\"Return index of candidate with highest average pairwise utility.\"\"\"\n    if utility_fn is None:\n        utility_fn = comp_metric_utility\n    n = len(candidates)\n    if n == 1:\n        return 0\n    scores = []\n    for i in range(n):\n        total = sum(utility_fn(candidates[i], candidates[j]) for j in range(n) if j != i)\n        scores.append(total / (n - 1))\n    return max(range(n), key=lambda i: scores[i])\n\n\n@torch.no_grad()\ndef generate_mbr_predictions(model, tokenizer, eval_dataloader, generation_config, mbr_n=16):\n    model.eval()\n    print(f\"MBR decoding: N={mbr_n} candidates per input\")\n\n    all_predictions = []\n    progress_bar = tqdm(eval_dataloader, desc=\"MBR Decoding\")\n\n    for batch in progress_bar:\n        input_ids = batch[\"input_ids\"].to(model.device)\n        attention_mask = batch[\"attention_mask\"].to(model.device)\n        local_batch_size = input_ids.shape[0]\n\n        generated_ids = model.generate(\n            input_ids=input_ids,\n            attention_mask=attention_mask,\n            generation_config=generation_config,\n        )\n\n        all_candidates = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)\n\n        for i in range(local_batch_size):\n            group = all_candidates[i * mbr_n : (i + 1) * mbr_n]\n            best_idx = mbr_select(group)\n            all_predictions.append(group[best_idx])\n\n    return all_predictions\n\n\n# ── Dataset & Collator ───────────────────────────────────────────────\n\n\ndef get_tokenizer(cfg):\n    tokenizer = AutoTokenizer.from_pretrained(cfg.model.backbone_path)\n    return tokenizer\n\n\nclass DPCDataset(Dataset):\n    \"\"\"\n    Dataset class for Deep Past Challenge - Translate Akkadian to English\n    \"\"\"\n\n    def __init__(self, cfg, df, id_col=\"oare_id\"):\n        self.cfg = cfg\n        self.df = df\n        self.id_col = id_col\n\n        self.tokenizer = get_tokenizer(cfg)\n\n    def __len__(self):\n        return len(self.df)\n\n    def tokenize(self, texts):\n        return self.tokenizer(texts, padding=False, truncation=True, max_length=self.cfg.model.max_length, return_length=True, add_special_tokens=True)\n\n    def pre_process(self, row):\n        ak_en_template = \"Translate Akkadian to English: {source}\"\n\n        return [ak_en_template.format(source=row[\"transliteration\"])], [None]\n\n    def __getitem__(self, idx):\n        data = self.df.iloc[idx].to_dict()\n        this_id = data[self.id_col]\n        sources, targets = self.pre_process(data)\n        tx_src = self.tokenize(sources)\n\n        return dict(id=this_id, input_ids=tx_src[\"input_ids\"], attention_mask=tx_src[\"attention_mask\"])\n\n\n@dataclass\nclass DPCCollator(DataCollatorWithPadding):\n    \"\"\"\n    Data collector for DPC dataset\n    \"\"\"\n\n    tokenizer = None\n    padding = True\n    max_length = None\n    pad_to_multiple_of = None\n\n    def __call__(self, features):\n        encoder_features = {\"input_ids\": [], \"attention_mask\": []}\n        labels = []\n\n        for feature in features:\n            encoder_features[\"input_ids\"].extend(feature[\"input_ids\"])\n            encoder_features[\"attention_mask\"].extend(feature[\"attention_mask\"])\n\n            if \"labels\" in feature:\n                labels.extend(feature[\"labels\"])\n\n        if len(labels) > 0:\n            assert len(labels) == len(encoder_features[\"input_ids\"])\n\n        batch = self.tokenizer.pad(\n            encoder_features,\n            padding=self.padding,\n            max_length=self.max_length,\n            pad_to_multiple_of=self.pad_to_multiple_of,\n            return_tensors=\"pt\",\n        )\n\n        # labels\n        if len(labels) > 0:\n            labels_batch = self.tokenizer.pad({\"input_ids\": labels}, padding=self.padding, max_length=self.max_length, pad_to_multiple_of=self.pad_to_multiple_of, return_tensors=None)\n            labels = torch.tensor(labels_batch[\"input_ids\"], dtype=torch.int64)\n            # Mask padding with -100 (ignored in loss)\n            labels[labels == self.tokenizer.pad_token_id] = -100\n            batch[\"labels\"] = labels\n        return batch\n\n\ndef show_batch(batch, tokenizer, print_fn=print, **kwargs):\n    bs = batch[\"input_ids\"].size(0)\n    print_fn(f\"batch size: {bs}\")\n\n    print_fn(f\"shape of input_ids: {batch['input_ids'].shape}\")\n    print_fn(f\"shape of attention_mask: {batch['attention_mask'].shape}\")\n    if \"labels\" in batch.keys():\n        print_fn(f\"shape of labels: {batch['labels'].shape}\")\n\n    print_fn(\"\\n\\n\")\n    max_n = kwargs.get(\"n\", 5)\n\n    for idx in range(bs):\n        print_fn(f\"\\n\\n===== Example {idx + 1} =====\\n\\n\")\n        print_fn(f\"## Input:\\n\\n{tokenizer.decode(batch['input_ids'][idx], skip_special_tokens=False)}\\n\\n\")\n\n        if \"labels\" in batch.keys():\n            labels = batch[\"labels\"][idx]\n            labels[labels == -100] = tokenizer.pad_token_id\n            print_fn(f\"## Output:\\n\\n{tokenizer.decode(labels, skip_special_tokens=False)}\")\n\n        if idx >= max_n:\n            break\n\n\n# ── Generation ────────────────────────────────────────────────────────\n\n\n@torch.no_grad()\ndef generate_predictions(model, tokenizer, eval_dataloader, generation_config):\n    model.eval()\n    print(\"Generating predictions...\")\n\n    all_predictions = []\n    progress_bar = tqdm(eval_dataloader, desc=\"Evaluating\")\n\n    for batch in progress_bar:\n        input_ids = batch[\"input_ids\"].to(model.device)\n        attention_mask = batch[\"attention_mask\"].to(model.device)\n\n        generated_ids = model.generate(\n            input_ids=input_ids,\n            attention_mask=attention_mask,\n            generation_config=generation_config,\n        )\n\n        predictions = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)\n        all_predictions.extend(predictions)\n\n    return all_predictions\n\n\ndef get_generation_config(cfg, tokenizer):\n    return GenerationConfig(\n        max_new_tokens=cfg.generation.max_new_tokens,\n        do_sample=cfg.generation.do_sample,\n        top_k=cfg.generation.top_k,\n        top_p=cfg.generation.top_p,\n        temperature=cfg.generation.temperature,\n        use_cache=True,\n        pad_token_id=tokenizer.pad_token_id,\n        eos_token_id=tokenizer.eos_token_id,\n    )\n\n\n# ── Main ──────────────────────────────────────────────────────────────\n\n\ndef main(cfg, save_path):\n    test_df = pd.read_csv(\"/kaggle/input/deep-past-initiative-machine-translation/test.csv\")\n    # test_df = pd.read_csv(\"/kaggle/input/deep-past-initiative-machine-translation/train.csv\").head(128).rename(columns={'oare_id': 'id'})\n    # test_df['transliteration'] = test_df['transliteration'].apply(lambda x: x[:80])\n    test_df[\"transliteration\"] = test_df[\"transliteration\"].apply(normalize_transliteration)\n\n    infer_ds = DPCDataset(cfg, test_df, id_col=\"id\")\n    tokenizer = infer_ds.tokenizer\n    data_collator = DPCCollator(tokenizer=tokenizer, pad_to_multiple_of=16)\n\n    infer_dl = DataLoader(\n        infer_ds,\n        batch_size=cfg.predict_params.per_device_eval_batch_size,\n        shuffle=False,\n        collate_fn=data_collator,\n    )\n\n    print(\"showing batch...\")\n    for idx, b in enumerate(infer_dl):\n        show_batch(b, tokenizer, n=2)\n        break\n\n    model = T5ForConditionalGeneration.from_pretrained(\n        cfg.model.backbone_path,\n        dtype=torch.float16,\n        device_map=\"auto\",\n    )\n\n    use_mbr = cfg.generation.get(\"mbr\", False)\n    mbr_n = cfg.generation.get(\"mbr_n\", 16)\n\n    if use_mbr:\n        mbr_num_beams = cfg.generation.get(\"mbr_num_beams\", 0)\n        if mbr_num_beams > 0:\n            generation_config = GenerationConfig(\n                max_new_tokens=cfg.generation.max_new_tokens,\n                use_cache=True,\n                do_sample=False,\n                num_beams=mbr_num_beams,\n                num_return_sequences=mbr_n,\n                pad_token_id=tokenizer.pad_token_id,\n                eos_token_id=tokenizer.eos_token_id,\n            )\n        else:\n            generation_config = GenerationConfig(\n                max_new_tokens=cfg.generation.max_new_tokens,\n                use_cache=True,\n                do_sample=True,\n                temperature=cfg.generation.get(\"mbr_temperature\", 0.5),\n                top_p=cfg.generation.get(\"mbr_top_p\", 0.9),\n                top_k=cfg.generation.get(\"mbr_top_k\", 0),\n                num_return_sequences=mbr_n,\n                pad_token_id=tokenizer.pad_token_id,\n                eos_token_id=tokenizer.eos_token_id,\n            )\n        predictions = generate_mbr_predictions(model=model, tokenizer=tokenizer, eval_dataloader=infer_dl, generation_config=generation_config, mbr_n=mbr_n)\n    else:\n        generation_config = get_generation_config(cfg, tokenizer)\n        predictions = generate_predictions(model=model, tokenizer=tokenizer, eval_dataloader=infer_dl, generation_config=generation_config)\n\n    predictions_df = pd.DataFrame({\"id\": test_df[\"id\"], \"translation\": predictions})\n    predictions_df[\"translation\"] = predictions_df[\"translation\"].apply(normalize_translation)\n\n    predictions_df.to_parquet(save_path)\n\n\nif __name__ == \"__main__\":\n    ap = argparse.ArgumentParser()\n    ap.add_argument(\"--config_path\", type=str, required=True)\n    ap.add_argument(\"--save_path\", type=str, required=True)\n    args = ap.parse_args()\n\n    cfg = OmegaConf.load(args.config_path)\n\n    Path(args.save_path).parent.mkdir(parents=True, exist_ok=True)\n\n    # run\n    main(cfg, args.save_path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T20:26:07.158736Z","iopub.execute_input":"2026-03-22T20:26:07.159352Z","iopub.status.idle":"2026-03-22T20:26:07.168579Z","shell.execute_reply.started":"2026-03-22T20:26:07.159324Z","shell.execute_reply":"2026-03-22T20:26:07.167827Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Config","metadata":{}},{"cell_type":"code","source":"%%writefile conf_m1.yaml\n\nmodel:\n    backbone_path: \"/kaggle/input/models/conjuring92/dpc_byt5_b2_ns/transformers/default/1\"\n    max_length: 1024\n    num_proc: 2\n\npredict_params:\n    per_device_eval_batch_size: 1\n\ngeneration:\n    max_new_tokens: 512\n    do_sample: true\n    top_k: 50\n    top_p: 1.0\n    temperature: 0.0\n    mbr: true\n    mbr_n: 1\n    mbr_num_beams: 8","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T19:46:14.294403Z","iopub.execute_input":"2026-03-22T19:46:14.294910Z","iopub.status.idle":"2026-03-22T19:46:14.299931Z","shell.execute_reply.started":"2026-03-22T19:46:14.294882Z","shell.execute_reply":"2026-03-22T19:46:14.299150Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile conf_m2.yaml\n\nmodel:\n    backbone_path: \"/kaggle/input/models/conjuring92/dpc_a05_14k_soup/transformers/default/1\"\n    max_length: 1024\n    num_proc: 2\n\npredict_params:\n    per_device_eval_batch_size: 1\n\ngeneration:\n    max_new_tokens: 512\n    do_sample: true\n    top_k: 50\n    top_p: 1.0\n    temperature: 0.0\n    mbr: true\n    mbr_n: 1\n    mbr_num_beams: 8","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T19:46:14.671381Z","iopub.execute_input":"2026-03-22T19:46:14.672070Z","iopub.status.idle":"2026-03-22T19:46:14.676690Z","shell.execute_reply.started":"2026-03-22T19:46:14.672040Z","shell.execute_reply":"2026-03-22T19:46:14.675936Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile conf_m3.yaml\n\nmodel:\n    backbone_path: \"/kaggle/input/models/conjuring92/dpc_byt5_xl/transformers/default/4\"\n    max_length: 1024\n    num_proc: 2\n\npredict_params:\n    per_device_eval_batch_size: 1\n\ngeneration:\n    max_new_tokens: 512\n    do_sample: true\n    top_k: 50\n    top_p: 1.0\n    temperature: 0.0\n    mbr: true\n    mbr_n: 1\n    mbr_num_beams: 8","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T20:28:48.476394Z","iopub.execute_input":"2026-03-22T20:28:48.477272Z","iopub.status.idle":"2026-03-22T20:28:48.482239Z","shell.execute_reply.started":"2026-03-22T20:28:48.477242Z","shell.execute_reply":"2026-03-22T20:28:48.481535Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Infer","metadata":{}},{"cell_type":"code","source":"%%time\n!accelerate launch --multi_gpu --num_processes=2 run_inference.py --config_path \"./conf_m1.yaml\" --save_path \"./model_outputs/pred_m1.parquet\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T19:46:17.637999Z","iopub.execute_input":"2026-03-22T19:46:17.638579Z","iopub.status.idle":"2026-03-22T19:49:25.447839Z","shell.execute_reply.started":"2026-03-22T19:46:17.638551Z","shell.execute_reply":"2026-03-22T19:49:25.446847Z"},"scrolled":true,"_kg_hide-input":false,"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# %%time\n# !accelerate launch --multi_gpu --num_processes=2 run_inference.py --config_path \"./conf_m2.yaml\" --save_path \"./model_outputs/pred_m2.parquet\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T20:31:28.126921Z","iopub.execute_input":"2026-03-22T20:31:28.127213Z","iopub.status.idle":"2026-03-22T20:31:28.131327Z","shell.execute_reply.started":"2026-03-22T20:31:28.127184Z","shell.execute_reply":"2026-03-22T20:31:28.130436Z"},"scrolled":true,"_kg_hide-input":false,"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%time\n!python run_inference_xl.py --config_path \"./conf_m3.yaml\" --save_path \"./model_outputs/pred_m3.parquet\"","metadata":{"trusted":true,"scrolled":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2026-03-22T20:29:41.561569Z","iopub.execute_input":"2026-03-22T20:29:41.562192Z","iopub.status.idle":"2026-03-22T20:31:28.124703Z","shell.execute_reply.started":"2026-03-22T20:29:41.562155Z","shell.execute_reply":"2026-03-22T20:31:28.123509Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# RM","metadata":{}},{"cell_type":"code","source":"%%writefile run_rm.py\n\nimport argparse\nimport os\nfrom pathlib import Path\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport vllm\nfrom omegaconf import OmegaConf\n\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\nos.environ[\"VLLM_WORKER_MULTIPROC_METHOD\"] = \"spawn\"\nos.environ[\"VLLM_ATTENTION_BACKEND\"] = \"TRITON_ATTN\"\n\n# Token IDs for Qwen3-8B (from training)\nTOK_A = 32\nTOK_B = 33\nTOK_EQUAL = 96024\n\nSYS_PROMPT = \"\"\"You are an expert judge comparing two Akkadian-to-English translations of the same source text. The transliteration is the ground truth — a syllable-by-syllable representation of an Old Assyrian cuneiform tablet (circa 1950-1750 BCE).\n\nYour task: determine which translation is more faithful to the transliteration, or whether they are effectively equal.\n\nKey evaluation criteria (in order of importance):\n1. SEMANTIC FIDELITY: Who did what to whom, under what conditions, with what consequences. Reversing a debt relationship or dropping a negation is catastrophic.\n2. NAME ACCURACY: Personal names must be derivable from the syllables in the transliteration. Minor spelling variants (macrons, gemination) are acceptable.\n3. NUMBERS AND QUANTITIES: Weights, counts, and amounts must be preserved exactly. 1 mina = 60 shekels, 1 talent = 60 minas.\n4. COMPLETENESS: A translation covering all content is better than one that stops midway. But fabrication is worse than omission.\n5. ALIGNMENT: Every element in the transliteration should have a corresponding element in the translation.\n\nPick A if Translation A is more faithful. Pick B if Translation B is more faithful. Pick EQUAL if both are substantively equivalent.\"\"\"\n\nUSER_TEMPLATE = \"\"\"Transliteration: {transliteration}\n\nTranslation A: {translation_a}\n\nTranslation B: {translation_b}\n\nWhich translation is more faithful to the transliteration? Respond with only: A, B, or EQUAL.\"\"\"\n\n\ndef get_probs(logprobs):\n    \"\"\"Extract softmax probabilities for A/B/EQUAL from vLLM logprobs.\"\"\"\n    tok_map = {TOK_A: 0, TOK_B: 1, TOK_EQUAL: 2}\n    logits = np.full(3, -20.0)\n\n    for token_id, logprob_obj in logprobs.items():\n        if token_id in tok_map:\n            logits[tok_map[token_id]] = logprob_obj.logprob\n\n    exp_logits = np.exp(logits - np.max(logits))\n    return exp_logits / np.sum(exp_logits)\n\n\ndef build_prompts(df, tokenizer):\n    prompts = []\n    for _, row in df.iterrows():\n        user_message = USER_TEMPLATE.format(\n            transliteration=row[\"transliteration\"],\n            translation_a=row[\"translation_a\"],\n            translation_b=row[\"translation_b\"],\n        )\n        conversation = [\n            {\"role\": \"system\", \"content\": SYS_PROMPT},\n            {\"role\": \"user\", \"content\": user_message},\n        ]\n        text = tokenizer.apply_chat_template(\n            conversation,\n            tokenize=False,\n            add_generation_prompt=True,\n            enable_thinking=False,\n        )\n        prompts.append(text)\n    return prompts\n\n\ndef main():\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--config_path\", type=str, required=True)\n    args = parser.parse_args()\n\n    cfg = OmegaConf.load(args.config_path)\n\n    # Load predictions\n    preds_a = pd.read_parquet(cfg.preds_a_path)\n    preds_b = pd.read_parquet(cfg.preds_b_path)\n    print(f\"Preds A: {len(preds_a):,} | Preds B: {len(preds_b):,}\")\n\n    id_col = cfg.get(\"id_col\", \"id\")\n    df = preds_a[[id_col, \"transliteration\", \"translation\"]].merge(\n        preds_b[[id_col, \"translation\"]],\n        on=id_col,\n        suffixes=(\"_a\", \"_b\"),\n    )\n    print(f\"Merged: {len(df):,} rows\")\n\n    # Load model\n    print(f\"Loading model: {cfg.model_path}\")\n    llm = vllm.LLM(\n        model=cfg.model_path,\n        max_model_len=cfg.get(\"max_model_len\", 2048),\n        max_num_seqs=cfg.get(\"max_num_seqs\", 32),\n        dtype=cfg.get(\"dtype\", \"half\"),\n        cpu_offload_gb=cfg.get(\"cpu_offload_gb\", 0),\n        swap_space=cfg.get(\"swap_space\", 0),\n        trust_remote_code=True,\n        tensor_parallel_size=torch.cuda.device_count(),\n        enforce_eager=True,\n    )\n\n    tokenizer = llm.get_tokenizer()\n    prompts = build_prompts(df, tokenizer)\n    print(f\"Built {len(prompts):,} prompts\")\n\n    # Generate with logprobs\n    outputs = llm.generate(\n        prompts,\n        vllm.SamplingParams(temperature=0.0, max_tokens=1, logprobs=20),\n        use_tqdm=True,\n    )\n\n    # Extract probabilities\n    probs_list = []\n    for output in outputs:\n        lps = output.outputs[0].logprobs[0]\n        probs_list.append(get_probs(lps))\n\n    probs_arr = np.array(probs_list)\n    df[\"prob_A\"] = probs_arr[:, 0]\n    df[\"prob_B\"] = probs_arr[:, 1]\n    df[\"prob_EQUAL\"] = probs_arr[:, 2]\n\n    # Pick: switch to B only if prob_B >= threshold, otherwise keep A\n    b_threshold = cfg.b_threshold\n    df[\"pick\"] = np.where(df[\"prob_B\"] >= b_threshold, \"B\", \"A\")\n    df[\"translation\"] = np.where(df[\"pick\"] == \"B\", df[\"translation_b\"], df[\"translation_a\"])\n    print(f\"\\nb_threshold={b_threshold}\")\n\n    # Stats\n    n = len(df)\n    for label in [\"A\", \"B\"]:\n        cnt = (df[\"pick\"] == label).sum()\n        print(f\"  {label}: {cnt:,} ({cnt / n * 100:.1f}%)\")\n\n    print(f\"\\n{df[['pick', 'prob_A', 'prob_B', 'prob_EQUAL']].head(20).to_string()}\")\n\n    for col in [\"prob_A\", \"prob_B\", \"prob_EQUAL\"]:\n        v = df[col]\n        print(f\"  {col}: mean={v.mean():.3f}  std={v.std():.3f}  min={v.min():.3f}  max={v.max():.3f}\")\n\n    # Save\n    output_path = Path(cfg.output_path)\n    output_path.parent.mkdir(parents=True, exist_ok=True)\n    df.to_parquet(output_path, index=False)\n    print(f\"Saved {len(df):,} rows -> {output_path}\")\n\n    # Submission\n    sub_df = df[[id_col, \"translation\"]].copy()\n    sub_path = output_path.parent / \"submission.csv\"\n    sub_df.to_csv(sub_path, index=False)\n    print(f\"Submission -> {sub_path}\")\n\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T20:31:44.002469Z","iopub.execute_input":"2026-03-22T20:31:44.003204Z","iopub.status.idle":"2026-03-22T20:31:44.011519Z","shell.execute_reply.started":"2026-03-22T20:31:44.003177Z","shell.execute_reply":"2026-03-22T20:31:44.010594Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile conf_rm.yaml\n\nmodel_path: /kaggle/input/models/conjuring92/dpc_qwen8b_rm/transformers/default/2\npreds_a_path: ./model_outputs/pred_m1.parquet\npreds_b_path: ./model_outputs/pred_m3.parquet\noutput_path: ./model_outputs/reward_picks.parquet\nid_col: id\n\nmax_model_len: 2048\nmax_num_seqs: 32\ncpu_offload_gb: 0\nswap_space: 0\ndtype: half\n\nb_threshold: 0.65","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T20:31:53.077586Z","iopub.execute_input":"2026-03-22T20:31:53.078350Z","iopub.status.idle":"2026-03-22T20:31:53.084372Z","shell.execute_reply.started":"2026-03-22T20:31:53.078318Z","shell.execute_reply":"2026-03-22T20:31:53.083183Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%time\n!python run_rm.py --config_path conf_rm.yaml","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T20:31:54.479692Z","iopub.execute_input":"2026-03-22T20:31:54.480339Z"},"scrolled":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Submission","metadata":{}},{"cell_type":"code","source":"import pandas as pd\n\ntest_df = pd.read_csv(\"/kaggle/input/deep-past-initiative-machine-translation/test.csv\")\npred_df = pd.read_parquet(\"./model_outputs/reward_picks.parquet\")\n\n# df_m1 = pd.read_parquet(\"./model_outputs/pred_m1.parquet\")\n# df_m2 = pd.read_parquet(\"./model_outputs/pred_m2.parquet\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T20:04:18.736147Z","iopub.execute_input":"2026-03-22T20:04:18.736749Z","iopub.status.idle":"2026-03-22T20:04:18.746221Z","shell.execute_reply.started":"2026-03-22T20:04:18.736718Z","shell.execute_reply":"2026-03-22T20:04:18.745525Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission = pd.DataFrame({\"id\": test_df[\"id\"].values.tolist(), \"translation\": pred_df['translation'].values.tolist()})\nsubmission[\"translation\"] = submission[\"translation\"].apply(lambda x: x if len(x) > 0 else \"broken text\")\nsubmission[\"translation\"] = submission[\"translation\"].apply(lambda x: x.replace(\"<unknown>\", \"<gap>\"))\n\nsubmission.to_csv(\"submission.csv\", index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T20:04:20.228155Z","iopub.execute_input":"2026-03-22T20:04:20.228821Z","iopub.status.idle":"2026-03-22T20:04:20.237617Z","shell.execute_reply.started":"2026-03-22T20:04:20.228792Z","shell.execute_reply":"2026-03-22T20:04:20.236883Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pd.options.display.max_colwidth = None\nsubmission.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T20:04:20.563137Z","iopub.execute_input":"2026-03-22T20:04:20.563767Z","iopub.status.idle":"2026-03-22T20:04:20.580957Z","shell.execute_reply.started":"2026-03-22T20:04:20.563743Z","shell.execute_reply":"2026-03-22T20:04:20.580387Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# df_m1","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T20:06:04.276056Z","iopub.execute_input":"2026-03-22T20:06:04.276557Z","iopub.status.idle":"2026-03-22T20:06:04.280398Z","shell.execute_reply.started":"2026-03-22T20:06:04.276527Z","shell.execute_reply":"2026-03-22T20:06:04.279538Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# df_m2","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T20:06:07.125861Z","iopub.execute_input":"2026-03-22T20:06:07.126518Z","iopub.status.idle":"2026-03-22T20:06:07.129891Z","shell.execute_reply.started":"2026-03-22T20:06:07.126487Z","shell.execute_reply":"2026-03-22T20:06:07.129093Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}