{"cells": [{"cell_type": "markdown", "metadata": {}, "source": "# Deep Past Challenge - NLLB-200 + Sentence Alignment\n\nUsing Meta's No Language Left Behind model with sentence-level data augmentation"}, {"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/nllb-akkadian-model\"\nMAX_LEN = 128\nBATCH_SIZE = 4\nGRAD_ACCUM = 4\nEPOCHS = 15\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": "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() > 10]\ntrain_df = train_df.sample(frac=1, random_state=SEED).reset_index(drop=True)\n\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": "# Load tokenizer with language settings\n# Using Arabic as proxy for Akkadian (both Semitic languages)\ntokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, src_lang=\"arb_Arab\", tgt_lang=\"eng_Latn\")\nmodel = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)\n\n# Get English language token ID for generation\nENG_TOKEN_ID = tokenizer.convert_tokens_to_ids(\"eng_Latn\")\nprint(f\"English token ID: {ENG_TOKEN_ID}\")\n\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": "class DS(Dataset):\n    def __init__(self, df, tok):\n        self.s = 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, eval_strategy=\"epoch\", save_strategy=\"no\",\n    learning_rate=5e-5,\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)\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 = 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            forced_bos_token_id=ENG_TOKEN_ID,\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!\")\nprint(sub.to_string())"}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python", "version": "3.10.0"}}, "nbformat": 4, "nbformat_minor": 4}