{"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":"gpu","dataSources":[{"sourceId":121150,"databundleVersionId":15061024,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":14217144,"sourceType":"datasetVersion","datasetId":9068929},{"sourceId":14237822,"sourceType":"datasetVersion","datasetId":9083643},{"sourceId":14339497,"sourceType":"datasetVersion","datasetId":9155323},{"sourceId":14374989,"sourceType":"datasetVersion","datasetId":9147887},{"sourceId":14376272,"sourceType":"datasetVersion","datasetId":9181082},{"sourceId":14410665,"sourceType":"datasetVersion","datasetId":9203761},{"sourceId":14599343,"sourceType":"datasetVersion","datasetId":9325369},{"sourceId":292636505,"sourceType":"kernelVersion"},{"sourceId":720901,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":548346,"modelId":480034},{"sourceId":720912,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":548355,"modelId":480034}],"dockerImageVersionId":31260,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"<div class=\"alert alert-block alert-info\">\n<h3>Experiment: Polishing with TranslateGemma</h3>\n\nThis notebook is based on my 34.4 LB submission. I have added TranslateGemma as a second stage to **polish** the outputs.\n\n<b>👇 Try it yourself:</b>\nSince TranslateGemma is optimised for language mapping, I treat this as an **\"English (Broken) $\\to$ English (Fluent)\"** translation task. Feel free to fork this and tweak the prompt logic!\n\n</div>","metadata":{}},{"cell_type":"code","source":"!pip install --no-index --find-links=/kaggle/input/bitsandbytes-wheels/offline_packages bitsandbytes accelerate","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-24T05:09:17.424714Z","iopub.execute_input":"2026-01-24T05:09:17.424958Z","iopub.status.idle":"2026-01-24T05:09:24.195970Z","shell.execute_reply.started":"2026-01-24T05:09:17.424935Z","shell.execute_reply":"2026-01-24T05:09:24.195269Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import re\nimport gc\nimport os\nimport pandas as pd\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nfrom transformers import (\n    AutoTokenizer, \n    AutoModelForSeq2SeqLM, \n    AutoModelForImageTextToText, \n    AutoProcessor, \n    BitsAndBytesConfig\n)\n\nCFG = {\n    \"data_path\": \"/kaggle/input/deep-past-initiative-machine-translation/test.csv\",\n    \"onomasticon_path\": \"/kaggle/input/dpc-onomasticon/onomasticon.csv\", \n    \"models\": [\n        \"/kaggle/input/byt5-base-big-data2\",\n        \"/kaggle/input/byt5-akkadian-model\",\n        \"/kaggle/input/train-gap-all-2/byt5-base-akkadian_gap_setence2\"\n    ],\n    \"weights\": [0.995, 0.98, 0.395],\n    \"device\": torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\"),\n    \"max_len\": 512,\n    \"batch_size\": 8,\n    \"gen_params\": {\n        \"num_beams\": 8,\n        \"max_new_tokens\": 512,\n        \"length_penalty\": 1.10,\n        \"early_stopping\": True\n    },\n    \"translategemma_path\": \"/kaggle/input/translategemma/transformers/translategemma-4b-it/1\" \n}\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-24T05:09:24.197768Z","iopub.execute_input":"2026-01-24T05:09:24.198029Z","iopub.status.idle":"2026-01-24T05:09:50.835151Z","shell.execute_reply.started":"2026-01-24T05:09:24.198000Z","shell.execute_reply":"2026-01-24T05:09:50.834505Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## The foundation","metadata":{}},{"cell_type":"code","source":"def load_name_lookup():\n    lookup = {}\n    if not os.path.exists(CFG['onomasticon_path']):\n        if os.path.exists(\"onomasticon.csv\"):\n            CFG['onomasticon_path'] = \"onomasticon.csv\"\n        else:\n            print(\"Constraints will be disabled.\")\n            return {}\n            \n    try:\n        df = pd.read_csv(CFG['onomasticon_path'])\n        print(f\"📖 Loading Onomasticon from {CFG['onomasticon_path']}...\")\n        for _, row in df.iterrows():\n            if pd.isna(row['Name']) or pd.isna(row['Spellings_semicolon_separated']):\n                continue\n            \n            target = row['Name'].strip()\n            variants = [v.strip() for v in row['Spellings_semicolon_separated'].split(';') if v.strip()]\n            \n            for v in variants:\n                lookup[v] = target\n                \n        print(f\"Loaded {len(lookup)} spelling variants.\")\n        return lookup\n    except Exception as e:\n        print(f\" {e}\")\n        return {}\n\ndef clean_input(text):\n    if pd.isna(text): return \"\"\n    t = str(text)\n    t = re.sub(r'\\(d\\)', '{d}', t)\n    t = re.sub(r'(\\.{3,}|…+|……)', '<big_gap>', t)\n    t = re.sub(r'(xx+|\\s+x\\s+)', '<gap>', t)\n    return t\n\ndef finalize_translation(text, input_text):\n    if not isinstance(text, str) or not text.strip(): return \"hnr\"\n    t = text.replace('ḫ', 'h').replace('Ḫ', 'H')\n    t = t.translate(str.maketrans(\"₀₁₂₃₄₅₆₇₈₉\", \"0123456789\"))\n    t = re.sub(r'(\\[x\\]|\\(x\\)|\\bx\\b)', '<gap>', t, flags=re.I)\n    t = re.sub(r'(\\.{3,}|…|\\[\\.+\\])', '<big_gap>', t)\n    t = re.sub(r'<gap>\\s*<gap>', ' <big_gap> ', t)\n    t = re.sub(r'<big_gap>\\s*<big_gap>', ' <big_gap> ', t)\n    t = t.replace('<gap>', '___GAP___').replace('<big_gap>', '___BIG___')\n    t = t.replace('ד', ' ')          \n    t = t.replace('()', '')          \n    t = re.sub(r'<(.*?)>', r'\\1', t) \n    t = re.sub(r'\\[(.*?)\\]', r'\\1', t) \n    t = re.sub(r'\\((fem|plur|pl|sing|singular|plural|\\?|!)\\.?\\s*\\w*\\)', '', t, flags=re.I)\n    t = t.replace('___GAP___', ' <gap> ').replace('___BIG___', ' <big_gap> ')\n    if \"<big_gap>\" not in input_text: t = t.replace(\"<big_gap>\", \" \")\n    if \"<gap>\" not in input_text: t = t.replace(\"<gap>\", \" \")\n    if t.count('<big_gap>') == 1: t = t.replace('<big_gap>', '')\n    if t.count('<gap>') == 1: t = t.replace('<gap>', '')\n    t = t.translate(str.maketrans('', '', '!?()\"—–<>⌈⌋⌊[]+ʾ/;'))\n    \n    frac_map = {\n        r'\\.5\\b': ' ½', r'\\.25\\b': ' ¼', r'\\.75\\b': ' ¾',\n        r'\\.33+\\d*\\b': ' ⅓', r'\\.66+\\d*\\b': ' ⅔'\n    }\n    for pat, rep in frac_map.items():\n        t = re.sub(r'(\\d+)' + pat, r'\\1' + rep, t)\n        t = re.sub(r'\\b0' + pat, rep.strip(), t)\n\n    t = re.sub(r'\\b(\\w+)(?:\\s+\\1\\b){3,}', r'\\1', t)\n    for n in range(4, 1, -1):\n        pat = r'\\b((?:\\w+\\s+){' + str(n-1) + r'}\\w+)(?:\\s+\\1\\b){3,}'\n        t = re.sub(pat, r'\\1', t)\n\n    return re.sub(r'\\s+', ' ', t).strip().strip('-')\n\ndef load_soup_model():\n    total_score = sum(CFG['weights'])\n    W = [w / total_score for w in CFG['weights']]\n    base = AutoModelForSeq2SeqLM.from_pretrained(CFG['models'][1])\n    sd = base.state_dict()\n    sd1 = AutoModelForSeq2SeqLM.from_pretrained(CFG['models'][0]).state_dict()\n    sd3 = AutoModelForSeq2SeqLM.from_pretrained(CFG['models'][2]).state_dict()\n    for k in sd:\n        val = W[1] * sd[k]\n        norm = W[1]\n        if k in sd1: val += W[0] * sd1[k]; norm += W[0]\n        if k in sd3: val += W[2] * sd3[k]; norm += W[2]\n        sd[k] = val / norm\n    base.load_state_dict(sd)\n    return base.to(CFG['device']).eval().float()\n\nclass AkkadDataset(Dataset):\n    def __init__(self, df):\n        self.ids = df['id'].tolist()\n        self.cleaned_inputs = df['transliteration'].tolist()\n        self.prompts = [\"translate Akkadian to English: \" + str(t) for t in self.cleaned_inputs]\n    def __len__(self): return len(self.ids)\n    def __getitem__(self, i): return self.ids[i], self.prompts[i], self.cleaned_inputs[i]\n\ndef run_byt5_inference():\n    df = pd.read_csv(CFG['data_path'])\n    df['transliteration'] = df['transliteration'].apply(clean_input)\n    \n    model = load_soup_model()\n    tokenizer = AutoTokenizer.from_pretrained(CFG['models'][1])\n    \n    def collate_fn(batch):\n        ids, prompts, clean_ins = zip(*batch)\n        enc = tokenizer(list(prompts), max_length=CFG['max_len'], padding=True, truncation=True, return_tensors=\"pt\")\n        return ids, enc, clean_ins\n\n    loader = DataLoader(AkkadDataset(df), batch_size=CFG['batch_size'], shuffle=False, num_workers=2, collate_fn=collate_fn)\n    \n    results = []\n    with torch.inference_mode():\n        for i, (ids, inputs, clean_ins) in enumerate(loader):\n            outputs = model.generate(input_ids=inputs.input_ids.to(CFG['device']), attention_mask=inputs.attention_mask.to(CFG['device']), **CFG['gen_params'])\n            decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)\n            cleaned_batch = [finalize_translation(pred, src) for pred, src in zip(decoded, clean_ins)]\n            results.extend(zip(ids, cleaned_batch))\n            if (i+1)%50==0: print(f\"   ByT5 batch {i+1}\")\n            \n    sub = pd.DataFrame(results, columns=['id', 'translation'])\n    sub.to_csv(\"hnr_draft.csv\", index=False)\n    del model, tokenizer, loader, inputs, outputs\n    gc.collect()\n    torch.cuda.empty_cache()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-24T05:09:50.836123Z","iopub.execute_input":"2026-01-24T05:09:50.836762Z","iopub.status.idle":"2026-01-24T05:09:50.856336Z","shell.execute_reply.started":"2026-01-24T05:09:50.836725Z","shell.execute_reply":"2026-01-24T05:09:50.855579Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## TranslateGemma Polishing","metadata":{}},{"cell_type":"code","source":"def run_gemma_polishing():\n    try: \n        draft_df = pd.read_csv(\"hnr_draft.csv\")\n        test_df = pd.read_csv(CFG['data_path'])\n        work_df = pd.merge(draft_df, test_df[['id', 'transliteration']], on='id', how='left')\n    except Exception as e:\n        print(f\"Error loading files: {e}\")\n        return\n\n    name_lookup = load_name_lookup()\n\n    bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type=\"nf4\", bnb_4bit_compute_dtype=torch.float16, bnb_4bit_use_double_quant=True)\n\n    print(f\"   -> Loading {CFG['translategemma_path']}...\")\n    processor = AutoProcessor.from_pretrained(CFG['translategemma_path'])\n    model = AutoModelForImageTextToText.from_pretrained(CFG['translategemma_path'], quantization_config=bnb_config, device_map=\"auto\")\n\n    def get_ground_truth_names(source_text):\n        if not name_lookup: return []\n        clean_src = clean_input(source_text)\n        \n        found = []\n        tokens = clean_src.split()\n        for t in tokens:\n            t_clean = t.strip('.,;')\n            if t_clean in name_lookup:\n                found.append(name_lookup[t_clean])\n        \n        return list(set(found))\n\n    def get_numbers_from_draft(text):\n        num_pattern = r'\\b\\d+(?:[\\s\\.]\\d+|[½⅓⅔¼¾])?|\\b[½⅓⅔¼¾]\\b'\n        return list(set(re.findall(num_pattern, text)))\n\n    def make_constrained_prompt(draft, names, nums):\n        name_str = \", \".join(names) if names else \"None\"\n        num_str = \", \".join(nums) if nums else \"None\"\n        \n        system = (\n            \"You are an expert Assyriologist. Refine the Draft Translation.\\n\\n\"\n            \"CONTEXT & CONSTRAINTS:\\n\"\n            f\"- Mandatory Proper Names (Found in Source): [{name_str}] \\n\"\n            f\"- Detected Numbers: [{num_str}] \\n\"\n            \"INSTRUCTIONS:\\n\"\n            \"1. Fix grammar to fluent, archaic business English.\\n\"\n            \"2. YOU MUST INCLUDE all 'Mandatory Proper Names' exactly (Fixes hallucinations).\\n\"\n            \"3. Format decimals as fractions (e.g. 2.5 -> 2 ½).\\n\"\n            \"4. Do not anglicize names (Keep 'Aššur', not 'Assur').\"\n        )\n        return f\"Draft: {draft}\\n{system}\\nCorrected Translation:\"\n\n    polished_results = []\n    print(f\"   -> Processing {len(work_df)} rows...\")\n\n    for idx, row in work_df.iterrows():\n        draft = row['translation']\n        source = row['transliteration']\n        \n        if len(str(draft)) < 3 or draft == \"broken text\":\n            polished_results.append(draft)\n            continue\n\n        true_names = get_ground_truth_names(source)\n        draft_nums = get_numbers_from_draft(draft)\n        \n        prompt = make_constrained_prompt(draft, true_names, draft_nums)\n\n        messages = [{\"role\": \"user\", \"content\": [{\"type\": \"text\", \"source_lang_code\": \"en\", \"target_lang_code\": \"en\", \"text\": prompt}]}]\n\n        try:\n            inputs = processor.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors=\"pt\").to(model.device)\n\n            with torch.inference_mode():\n                generation = model.generate(**inputs, do_sample=False, max_new_tokens=256)\n\n            input_len = len(inputs['input_ids'][0])\n            decoded = processor.decode(generation[0][input_len:], skip_special_tokens=True).strip()\n            polished_results.append(decoded if decoded else draft)\n        except:\n            polished_results.append(draft)\n\n        if (idx+1) % 50 == 0: print(f\"      Polished {idx+1}/{len(work_df)}\")\n\n    final_df = pd.DataFrame({'id': work_df['id'], 'translation': polished_results})\n    final_df.to_csv(\"submission.csv\", index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-24T05:09:50.857456Z","iopub.execute_input":"2026-01-24T05:09:50.857730Z","iopub.status.idle":"2026-01-24T05:09:50.876786Z","shell.execute_reply.started":"2026-01-24T05:09:50.857707Z","shell.execute_reply":"2026-01-24T05:09:50.876272Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if __name__ == \"__main__\":\n    run_byt5_inference()\n    if os.path.exists(CFG['translategemma_path']):\n        run_gemma_polishing()\n    else:\n        print(\"Submitting drafts.\")\n        os.rename(\"hnr_draft.csv\", \"submission.csv\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-24T05:09:50.878353Z","iopub.execute_input":"2026-01-24T05:09:50.878558Z","iopub.status.idle":"2026-01-24T05:13:26.698848Z","shell.execute_reply.started":"2026-01-24T05:09:50.878537Z","shell.execute_reply":"2026-01-24T05:13:26.698175Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Compare the results","metadata":{}},{"cell_type":"code","source":"h = pd.read_csv(\"/kaggle/input/deep-past-initiative-machine-translation/test.csv\")\nn = pd.read_csv(\"/kaggle/working/hnr_draft.csv\")\nr = pd.read_csv(\"/kaggle/working/submission.csv\")\n\npreview = pd.DataFrame({\n    \"id\": h[\"id\"],\n    \"original_transliteration\": h[\"transliteration\"],\n    \"raw_translation\": n[\"translation\"],\n    \"predicted_translation\": r[\"translation\"],\n})\n\nfor _, r in preview.head(10).iterrows():\n    print(\"=\"*90)\n    print(\"id:\", r[\"id\"])\n    print(\"ORIGINAL:\", r[\"original_transliteration\"])\n    print(\"RAW :\", r[\"raw_translation\"])\n    print(\"GEMMA :\", r[\"predicted_translation\"])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-24T05:13:26.699890Z","iopub.execute_input":"2026-01-24T05:13:26.700173Z","iopub.status.idle":"2026-01-24T05:13:26.714058Z","shell.execute_reply.started":"2026-01-24T05:13:26.700135Z","shell.execute_reply":"2026-01-24T05:13:26.713508Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}