{"cells": [{"cell_type": "markdown", "metadata": {}, "source": "# Deep Past Challenge - mT5-base\n\nLarger model for better translation quality"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "import os\nos.environ['PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'] = 'python'\n\nimport re, gc, random\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom torch.utils.data import Dataset\nfrom transformers import AutoTokenizer, AutoModelForSeq2SeqLM, Seq2SeqTrainingArguments, Seq2SeqTrainer, DataCollatorForSeq2Seq\nprint(f\"CUDA: {torch.cuda.is_available()}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "SEED = 42\nMODEL_NAME = \"/kaggle/input/mt5-akkadian-model\"\nMAX_LEN = 128\nBATCH_SIZE = 2\nGRAD_ACCUM = 8\nEPOCHS = 10\nOUTPUT_DIR = \"/kaggle/working/model\"\nSAVE_PATH = f\"{OUTPUT_DIR}/final\"\nDATA_DIR = \"/kaggle/input/deep-past-initiative-machine-translation\"\nrandom.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "train_df = pd.read_csv(f\"{DATA_DIR}/train.csv\").dropna(subset=[\"transliteration\", \"translation\"])\ntest_df = pd.read_csv(f\"{DATA_DIR}/test.csv\")\ndef clean(t):\n    if pd.isna(t): return \"\"\n    return re.sub(r\"\\s+\", \" \", str(t)).strip()\ntrain_df[\"src\"] = train_df[\"transliteration\"].apply(clean)\ntrain_df[\"tgt\"] = train_df[\"translation\"].apply(clean)\ntrain_df = train_df[train_df[\"src\"].str.len() > 10]\ntrain_df = train_df.sample(frac=1, random_state=SEED).reset_index(drop=True)\nval_df = train_df.iloc[:150]\ntrain_df = train_df.iloc[150:]\nprint(f\"Train: {len(train_df)}, Val: {len(val_df)}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "class DS(Dataset):\n    def __init__(self, df, tok):\n        self.s = [f\"translate Akkadian to English: {x}\" for x in df[\"src\"].tolist()]\n        self.t = df[\"tgt\"].tolist()\n        self.tok = tok\n    def __len__(self): return len(self.s)\n    def __getitem__(self, i):\n        a = self.tok(self.s[i], max_length=MAX_LEN, truncation=True)\n        b = self.tok(text_target=self.t[i], max_length=MAX_LEN, truncation=True)\n        return {\"input_ids\": a[\"input_ids\"], \"attention_mask\": a[\"attention_mask\"], \"labels\": b[\"input_ids\"]}"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\nmodel = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)\nmodel.gradient_checkpointing_enable()\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel.to(device)\nprint(f\"Model on {device}\")\nprint(f\"Parameters: {sum(p.numel() for p in model.parameters()):,}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "train_ds = DS(train_df, tokenizer)\nval_ds = DS(val_df, tokenizer)\ncollator = DataCollatorForSeq2Seq(tokenizer, model=model, padding=True)\nargs = Seq2SeqTrainingArguments(\n    output_dir=OUTPUT_DIR, eval_strategy=\"epoch\", save_strategy=\"no\",\n    learning_rate=5e-5,  # Slightly lower for larger model\n    per_device_train_batch_size=BATCH_SIZE,\n    per_device_eval_batch_size=BATCH_SIZE,\n    gradient_accumulation_steps=GRAD_ACCUM,\n    num_train_epochs=EPOCHS,\n    warmup_ratio=0.1,\n    fp16=False,\n    logging_steps=50,\n    report_to=\"none\",\n    gradient_checkpointing=True,\n)\ntrainer = Seq2SeqTrainer(model=model, args=args, train_dataset=train_ds, eval_dataset=val_ds, processing_class=tokenizer, data_collator=collator)\nprint(\"Training...\")\ntrainer.train()\nprint(\"Done!\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "os.makedirs(SAVE_PATH, exist_ok=True)\nmodel.save_pretrained(SAVE_PATH)\ntokenizer.save_pretrained(SAVE_PATH)\nprint(f\"Saved to {SAVE_PATH}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "del model, trainer\ngc.collect()\ntorch.cuda.empty_cache()\nmodel = AutoModelForSeq2SeqLM.from_pretrained(SAVE_PATH)\nmodel.to(device)\nmodel.eval()\nprint(\"Model reloaded\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def translate(text):\n    src = f\"translate Akkadian to English: {clean(text)}\"\n    inputs = tokenizer(src, return_tensors=\"pt\", max_length=MAX_LEN, truncation=True)\n    inputs = {k: v.to(device) for k, v in inputs.items()}\n    with torch.no_grad():\n        out = model.generate(\n            input_ids=inputs[\"input_ids\"],\n            attention_mask=inputs[\"attention_mask\"],\n            max_length=256,\n            num_beams=4,\n            no_repeat_ngram_size=3,\n            early_stopping=True,\n        )\n    r = tokenizer.decode(out[0], skip_special_tokens=True)\n    r = clean(r)\n    if not r: return \"No translation.\"\n    if not r.endswith(\".\"): r += \".\"\n    return r\n\nprint(\"Test:\")\nprint(translate(train_df.iloc[0][\"src\"])[:200])"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "print(\"Generating submission...\")\npreds = []\nfor _, r in test_df.iterrows():\n    p = translate(r[\"transliteration\"])\n    preds.append(p)\n    print(f\"ID {r['id']}: {p[:60]}...\")\nsub = pd.DataFrame({\"id\": test_df[\"id\"], \"translation\": preds})\nsub.to_csv(\"submission.csv\", index=False)\nprint(\"\\nSubmission saved!\")\nsub"}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python", "version": "3.10.0"}}, "nbformat": 4, "nbformat_minor": 4}