{"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":14237822,"sourceType":"datasetVersion","datasetId":9083643},{"sourceId":4251,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":3045,"modelId":575}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"V14: Update with 2 stage predicition","metadata":{}},{"cell_type":"code","source":"import os\nimport gc\nimport pandas as pd\nimport torch\nfrom torch.utils.data import DataLoader, Dataset\nfrom transformers import AutoTokenizer, AutoModelForSeq2SeqLM\nfrom tqdm.auto import tqdm\nimport re\n\nos.environ[\"TRANSFORMERS_OFFLINE\"] = \"1\"\nos.environ[\"HF_DATASETS_OFFLINE\"] = \"1\"\n\nCOMP_DIR = \"/kaggle/input/deep-past-initiative-machine-translation\"\nTEST_DATA_PATH = f\"{COMP_DIR}/test.csv\"\nSAMPLE_SUB_PATH = f\"{COMP_DIR}/sample_submission.csv\"\n\n# Model 1: Byt5 Akkadian Translator\nTRANSLATOR_MODEL_DIR = \"/kaggle/input/byt5-base-20/byt5-base-akkadian\"\n\n# Model 2: The LLM \nLLM_MODEL_DIR = \"/kaggle/input/flan-t5/pytorch/base/4\" \n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nTRANSLATOR_PREFIX = \"translate Akkadian to English: \"\n\nBATCH_SIZE = 32\nMAX_SRC = 512 \n\n# Load Data\ntest_df = pd.read_csv(TEST_DATA_PATH)\nsub_df  = pd.read_csv(SAMPLE_SUB_PATH)\nOUT_ID  = sub_df.columns[0]\nOUT_COL = sub_df.columns[1]\nids = test_df[OUT_ID].tolist()\n\nclass InferenceDataset(Dataset):\n    def __init__(self, texts, tokenizer, max_len=512):\n        self.texts = texts\n        self.tok = tokenizer\n        self.max_len = max_len\n        \n    def __len__(self): return len(self.texts)\n    \n    def __getitem__(self, i):\n        enc = self.tok(\n            self.texts[i],\n            max_length=self.max_len,\n            padding=\"max_length\",\n            truncation=True,\n            return_tensors=\"pt\"\n        )\n        return {\n            \"input_ids\": enc[\"input_ids\"].squeeze(0),\n            \"attention_mask\": enc[\"attention_mask\"].squeeze(0),\n        }\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-26T09:47:54.362048Z","iopub.execute_input":"2025-12-26T09:47:54.362335Z","iopub.status.idle":"2025-12-26T09:48:06.455585Z","shell.execute_reply.started":"2025-12-26T09:47:54.362302Z","shell.execute_reply":"2025-12-26T09:48:06.454948Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# STAGE 1: RAW TRANSLATION","metadata":{}},{"cell_type":"code","source":"tokenizer = AutoTokenizer.from_pretrained(TRANSLATOR_MODEL_DIR, local_files_only=True)\nmodel = AutoModelForSeq2SeqLM.from_pretrained(TRANSLATOR_MODEL_DIR, local_files_only=True).to(DEVICE)\nmodel.eval()\n\n# Prepare Inputs\ninput_texts = [TRANSLATOR_PREFIX + str(t) for t in test_df[\"transliteration\"].tolist()]\nds = InferenceDataset(input_texts, tokenizer, max_len=MAX_SRC)\ndl = DataLoader(ds, batch_size=BATCH_SIZE, shuffle=False)\n\ngen_kwargs = dict(\n    num_beams=4, length_penalty=1.0, max_new_tokens=256, early_stopping=True\n)\n\nraw_preds = []\nwith torch.no_grad():\n    for batch in tqdm(dl, desc=\"Translating\"):\n        out = model.generate(\n            input_ids=batch[\"input_ids\"].to(DEVICE),\n            attention_mask=batch[\"attention_mask\"].to(DEVICE),\n            **gen_kwargs\n        )\n        dec = tokenizer.batch_decode(out, skip_special_tokens=True)\n        raw_preds.extend(dec)\n\n\ndel model, tokenizer, ds, dl\ngc.collect()\ntorch.cuda.empty_cache()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-26T09:48:06.456953Z","iopub.execute_input":"2025-12-26T09:48:06.457185Z","iopub.status.idle":"2025-12-26T09:48:43.705730Z","shell.execute_reply.started":"2025-12-26T09:48:06.457162Z","shell.execute_reply":"2025-12-26T09:48:43.704882Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# STAGE 2: POLISHING with Flan-T5","metadata":{}},{"cell_type":"code","source":"if not os.path.exists(LLM_MODEL_DIR):\n    print(f\"WARNING: LLM Path {LLM_MODEL_DIR} not found. Using raw predictions.\")\n    final_preds = raw_preds\nelse:\n    tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_DIR, local_files_only=True)\n    model = AutoModelForSeq2SeqLM.from_pretrained(LLM_MODEL_DIR, local_files_only=True).to(DEVICE)\n    model.eval()\n\n    llm_prompts = [\n        f\"Fix grammar and rewrite as a single coherent English sentence: {text}\" \n        for text in raw_preds\n    ]\n    \n    ds = InferenceDataset(llm_prompts, tokenizer, max_len=512)\n    dl = DataLoader(ds, batch_size=BATCH_SIZE, shuffle=False)\n\n    llm_kwargs = dict(\n        num_beams=4, length_penalty=1.0, max_new_tokens=256, early_stopping=True\n    )\n\n    polished_preds = []\n    with torch.no_grad():\n        for batch in tqdm(dl, desc=\"Polishing\"):\n            out = model.generate(\n                input_ids=batch[\"input_ids\"].to(DEVICE),\n                attention_mask=batch[\"attention_mask\"].to(DEVICE),\n                **llm_kwargs\n            )\n            dec = tokenizer.batch_decode(out, skip_special_tokens=True)\n            polished_preds.extend(dec)\n            \n    final_preds = polished_preds","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-26T09:50:18.487700Z","iopub.execute_input":"2025-12-26T09:50:18.488283Z","iopub.status.idle":"2025-12-26T09:50:20.843537Z","shell.execute_reply.started":"2025-12-26T09:50:18.488253Z","shell.execute_reply":"2025-12-26T09:50:20.842741Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# FINAL POST-PROCESSING","metadata":{}},{"cell_type":"code","source":"def final_cleanup(s: str) -> str:\n    s = re.sub(r\"\\s+\", \" \", (s or \"\")).strip()\n    if not s: return \"broken text\"\n    \n    if len(s) > 280:\n        cut = max(s.rfind(\".\", 0, 280), s.rfind(\"!\", 0, 280), s.rfind(\"?\", 0, 280))\n        if cut > 60:\n            s = s[:cut+1].strip()\n        else:\n            s = s[:280].rstrip() + \".\"\n            \n    if s and s[-1] not in \".!?\":\n        s += \".\"\n    return s\n\nclean_preds = [final_cleanup(p) for p in final_preds]\n\nsubmission = pd.DataFrame({OUT_ID: ids, OUT_COL: clean_preds})\nsubmission.to_csv(\"/kaggle/working/submission.csv\", index=False)\nsubmission.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-26T09:50:23.650940Z","iopub.execute_input":"2025-12-26T09:50:23.651470Z","iopub.status.idle":"2025-12-26T09:50:23.664154Z","shell.execute_reply.started":"2025-12-26T09:50:23.651436Z","shell.execute_reply":"2025-12-26T09:50:23.663532Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"preview = pd.DataFrame({\n    \"id\": test_df[\"id\"],\n    \"original_transliteration\": test_df[\"transliteration\"],\n    \"predicted_translation\": clean_preds\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(\"PREDICT :\", r[\"predicted_translation\"])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-26T09:50:42.491261Z","iopub.execute_input":"2025-12-26T09:50:42.491968Z","iopub.status.idle":"2025-12-26T09:50:42.497861Z","shell.execute_reply.started":"2025-12-26T09:50:42.491936Z","shell.execute_reply":"2025-12-26T09:50:42.497115Z"}},"outputs":[],"execution_count":null}]}