{"cells": [{"cell_type": "markdown", "metadata": {}, "source": "# Deep Past Challenge - ByT5 + Sentence Alignment\n\nByte-level T5 handles noisy Akkadian transliterations at character level"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "import os\nimport gc\nimport re\nimport 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\n\nprint(f\"CUDA: {torch.cuda.is_available()}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "SEED = 42\nMODEL_NAME = \"/kaggle/input/byt5-small\"\nMAX_LEN = 512  # ByT5 needs longer sequences (byte-level)\nBATCH_SIZE = 4\nGRAD_ACCUM = 2\nEPOCHS = 30\nLEARNING_RATE = 2e-4\nOUTPUT_DIR = \"/kaggle/working/model\"\nDATA_DIR = \"/kaggle/input/deep-past-initiative-machine-translation\"\n\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\ntorch.cuda.manual_seed(SEED)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def simple_sentence_aligner(df):\n    \"\"\"\n    Split document-level pairs into sentence-level pairs.\n    Only split when English sentence count matches Akkadian line count.\n    \"\"\"\n    aligned_data = []\n    \n    for idx, row in df.iterrows():\n        src = str(row['transliteration'])\n        tgt = str(row['translation'])\n        \n        # Split English by sentence-ending punctuation\n        tgt_sents = [t.strip() for t in re.split(r'(?<=[.!?])\\s+', tgt) if t.strip()]\n        \n        # Split Akkadian by newlines\n        src_lines = [s.strip() for s in src.split('\\n') if s.strip()]\n        \n        # Only split if counts match (safe alignment)\n        if len(tgt_sents) > 1 and len(tgt_sents) == len(src_lines):\n            for s, t in zip(src_lines, tgt_sents):\n                if len(s) > 3 and len(t) > 3:\n                    aligned_data.append({'transliteration': s, 'translation': t})\n        else:\n            aligned_data.append({'transliteration': src, 'translation': tgt})\n            \n    return pd.DataFrame(aligned_data)\n\ndef clean(t):\n    if pd.isna(t): return \"\"\n    return re.sub(r\"\\s+\", \" \", str(t)).strip()"}, {"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\")\n\nprint(f\"Original: {len(train_df)} docs\")\n\n# Apply sentence alignment\ntrain_df = simple_sentence_aligner(train_df)\nprint(f\"After alignment: {len(train_df)} sentences\")\n\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() > 3]\ntrain_df = train_df.sample(frac=1, random_state=SEED).reset_index(drop=True)\n\nval_size = int(len(train_df) * 0.1)\nval_df = train_df.iloc[:val_size]\ntrain_df = train_df.iloc[val_size:]\nprint(f\"Train: {len(train_df)}, Val: {len(val_df)}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\nmodel = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)\n\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": "PREFIX = \"translate Akkadian to English: \"\n\nclass DS(Dataset):\n    def __init__(self, df, tok):\n        self.s = [PREFIX + 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": "train_ds = DS(train_df, tokenizer)\nval_ds = DS(val_df, tokenizer)\ncollator = DataCollatorForSeq2Seq(tokenizer, model=model, padding=True)\n\nargs = Seq2SeqTrainingArguments(\n    output_dir=OUTPUT_DIR,\n    eval_strategy=\"epoch\",\n    save_strategy=\"epoch\",\n    learning_rate=LEARNING_RATE,\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    weight_decay=0.01,\n    fp16=False,  # FP32 to avoid NaN\n    logging_steps=10,\n    save_total_limit=2,\n    report_to=\"none\",\n)\n\ntrainer = Seq2SeqTrainer(\n    model=model,\n    args=args,\n    train_dataset=train_ds,\n    eval_dataset=val_ds,\n    data_collator=collator,\n    tokenizer=tokenizer,\n)\n\nprint(\"Training...\")\ntrainer.train()\nprint(\"Done!\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "trainer.save_model(OUTPUT_DIR)\ntokenizer.save_pretrained(OUTPUT_DIR)\nprint(f\"Model saved to {OUTPUT_DIR}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "del trainer\ngc.collect()\ntorch.cuda.empty_cache()\nmodel.eval()\nprint(\"Ready for inference\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def translate(text):\n    src = PREFIX + 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]}...\")\n\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}