{"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,"sourceType":"competition"},{"sourceId":14374989,"sourceType":"datasetVersion","datasetId":9147887},{"sourceId":14376272,"sourceType":"datasetVersion","datasetId":9181082},{"sourceId":14398637,"sourceType":"datasetVersion","datasetId":9195725},{"sourceId":14410665,"sourceType":"datasetVersion","datasetId":9203761},{"sourceId":282751,"sourceType":"modelInstanceVersion","modelInstanceId":239470,"modelId":222398}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Complete, fast, Kaggle-safe inference script (no DataLoader multiprocessing hang,\n# fp16 on GPU, in-place model blending, pinned-memory transfers)\n\nimport os\nos.environ[\"TRANSFORMERS_NO_TF\"] = \"1\"\nos.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n\nimport re\nimport gc\nimport pandas as pd\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nfrom transformers import AutoTokenizer, AutoModelForSeq2SeqLM\nfrom tqdm.auto import tqdm\n\nCFG = {\n    \"data_path\": \"/kaggle/input/deep-past-initiative-machine-translation/test.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.99, 0.98, 0.39],\n    \"device\": torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\"),\n    \"max_len\": 496,\n    \"batch_size\": 8,\n    \"gen_params\": {\n        \"num_beams\": 8,\n        \"max_new_tokens\": 496,\n        \"length_penalty\": 1.2,\n        \"early_stopping\": True\n    }\n}\n\ndef clean_input(text):\n    if pd.isna(text):\n        return \"\"\n    t = str(text)\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 hnr(text):\n    if not isinstance(text, str) or not text.strip():\n        return \"hnr\"\n\n    t = text.replace('ḫ', 'h').replace('Ḫ', 'H')\n    sub_map = str.maketrans(\"₀₁₂₃₄₅₆₇₈₉\", \"0123456789\")\n    t = t.translate(sub_map)\n\n    t = re.sub(r'(\\[x\\]|\\(x\\)|\\bx\\b)', '<gap>', t, flags=re.I)\n    t = re.sub(r'(\\.{3,}|…|\\[\\.+\\])', '<big_gap>', t)\n\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\n    t = re.sub(r'\\((fem|plur|pl|sing|singular|plural|\\?|!)\\.?\\s*\\w*\\)', '', t, flags=re.I)\n\n    t = t.replace('<gap>', '\\x00GAP\\x00').replace('<big_gap>', '\\x00BIG\\x00')\n\n    bad_chars = '!?()\"—–<>⌈⌋⌊[]+ʾ/;'\n    t = t.translate(str.maketrans('', '', bad_chars))\n\n    t = t.replace('\\x00GAP\\x00', ' <gap> ').replace('\\x00BIG\\x00', ' <big_gap> ')\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)+', 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)+'\n        t = re.sub(pat, r'\\1', t)\n\n    return re.sub(r'\\s+', ' ', t).strip().strip('-')\n\n@torch.no_grad()\ndef load_blended_model_fp16(model_paths, weights, device):\n    wsum = float(sum(weights))\n    W = [float(w) / wsum for w in weights]\n\n    # Load on CPU in fp16 to reduce RAM/time; move final model to GPU at end.\n    base = AutoModelForSeq2SeqLM.from_pretrained(\n        model_paths[1],\n        torch_dtype=torch.float16,\n        low_cpu_mem_usage=True\n    )\n    final_sd = base.state_dict()  # CPU tensors (fp16)\n\n    sd1 = AutoModelForSeq2SeqLM.from_pretrained(\n        model_paths[0],\n        torch_dtype=torch.float16,\n        low_cpu_mem_usage=True\n    ).state_dict()\n\n    sd3 = AutoModelForSeq2SeqLM.from_pretrained(\n        model_paths[2],\n        torch_dtype=torch.float16,\n        low_cpu_mem_usage=True\n    ).state_dict()\n\n    # In-place blending for speed/memory\n    for k in tqdm(final_sd.keys(), desc=\"Blending models\", leave=False):\n        norm = W[1]\n        final_sd[k].mul_(W[1])\n\n        t1 = sd1.get(k, None)\n        if t1 is not None:\n            final_sd[k].add_(t1, alpha=W[0])\n            norm += W[0]\n\n        t3 = sd3.get(k, None)\n        if t3 is not None:\n            final_sd[k].add_(t3, alpha=W[2])\n            norm += W[2]\n\n        final_sd[k].div_(norm)\n\n    base.load_state_dict(final_sd, strict=True)\n\n    del sd1, sd3, final_sd\n    gc.collect()\n\n    base.eval()\n    base.to(device)\n    return base\n\nclass AkkadDataset(Dataset):\n    def __init__(self, df):\n        self.ids = df[\"id\"].tolist()\n        self.texts = [\"translate Akkadian to English: \" + str(t) for t in df[\"transliteration\"].tolist()]\n\n    def __len__(self):\n        return len(self.ids)\n\n    def __getitem__(self, i):\n        return self.ids[i], self.texts[i]\n\ndef make_collate(tokenizer, max_len):\n    def collate_fn(batch):\n        ids = [x[0] for x in batch]\n        texts = [x[1] for x in batch]\n        tok = tokenizer(\n            texts,\n            max_length=max_len,\n            padding=True,\n            truncation=True,\n            return_tensors=\"pt\"\n        )\n        return ids, tok\n    return collate_fn\n\ndef main():\n    df = pd.read_csv(CFG[\"data_path\"])\n    df[\"transliteration\"] = df[\"transliteration\"].apply(clean_input)\n\n    tokenizer = AutoTokenizer.from_pretrained(CFG[\"models\"][1])\n\n    model = load_blended_model_fp16(\n        model_paths=CFG[\"models\"],\n        weights=CFG[\"weights\"],\n        device=CFG[\"device\"]\n    )\n\n    loader = DataLoader(\n        AkkadDataset(df),\n        batch_size=CFG[\"batch_size\"],\n        shuffle=False,\n        num_workers=0,          # safest (prevents hangs). Generation dominates runtime anyway.\n        pin_memory=True,\n        collate_fn=make_collate(tokenizer, CFG[\"max_len\"])\n    )\n\n    results = []\n    with torch.inference_mode():\n        for ids, inputs in tqdm(loader, desc=\"Generating\", leave=True):\n            input_ids = inputs[\"input_ids\"].to(CFG[\"device\"], non_blocking=True)\n            attention_mask = inputs[\"attention_mask\"].to(CFG[\"device\"], non_blocking=True)\n\n            outputs = model.generate(\n                input_ids=input_ids,\n                attention_mask=attention_mask,\n                **CFG[\"gen_params\"]\n            )\n\n            decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)\n            cleaned = [hnr(x) for x in decoded]\n            results.extend(zip(ids, cleaned))\n\n    sub = pd.DataFrame(results, columns=[\"id\", \"translation\"])\n    sub.to_csv(\"submission.csv\", index=False)\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T10:11:08.519223Z","iopub.execute_input":"2026-01-21T10:11:08.519909Z","iopub.status.idle":"2026-01-21T10:11:19.563627Z","shell.execute_reply.started":"2026-01-21T10:11:08.519870Z","shell.execute_reply":"2026-01-21T10:11:19.561776Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}