{"cells": [{"cell_type": "markdown", "metadata": {}, "source": "# Deep Past - NLLB-200 v2\n\nNLLB-200 + Top Baseline Techniques:\n- Monotone alignment for sentence-level reconstruction\n- Long text chunking\n- Optimized hyperparameters"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "import os, re, math, random, gc\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom torch.utils.data import Dataset\nfrom transformers import (\n    AutoTokenizer,\n    AutoModelForSeq2SeqLM,\n    DataCollatorForSeq2Seq,\n    Seq2SeqTrainingArguments,\n    Seq2SeqTrainer,\n)\n\nSEED = 42\nrandom.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)\n\nMODEL_NAME = \"/kaggle/input/nllb-akkadian-model\"\nDATA_DIR = \"/kaggle/input/deep-past-initiative-machine-translation\"\nOUTPUT_DIR = \"/kaggle/working/model\"\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"Device:\", DEVICE)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Load data\ntrain_df = pd.read_csv(f\"{DATA_DIR}/train.csv\")\ntest_df = pd.read_csv(f\"{DATA_DIR}/test.csv\")\nprint(\"train:\", train_df.shape, \"test:\", test_df.shape)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Text processing functions from top baseline\ndef clean_ws(x: str) -> str:\n    if pd.isna(x): return \"\"\n    return re.sub(r\"\\s+\", \" \", str(x)).strip()\n\ndef split_src_transliteration(doc: str):\n    \"\"\"Split transliteration doc into 'lines' as proxy for sentences.\"\"\"\n    doc = \"\" if pd.isna(doc) else str(doc)\n    parts = re.split(r\"\\n+\", doc)\n    parts = [clean_ws(p) for p in parts if clean_ws(p)]\n    if len(parts) <= 1:\n        parts = re.split(r\"\\s{2,}|;\\s+\", clean_ws(doc))\n        parts = [clean_ws(p) for p in parts if clean_ws(p)]\n    return parts\n\ndef split_tgt_english(doc: str):\n    \"\"\"Split English translation into sentences.\"\"\"\n    doc = clean_ws(doc)\n    if not doc:\n        return []\n    sents = re.split(r\"(?<=[.!?])\\s+\", doc)\n    sents = [clean_ws(s) for s in sents if clean_ws(s)]\n    return sents\n\ndef tok_len(s: str) -> int:\n    return max(1, len(s.split()))\n\ndef align_monotone(src_units, tgt_units, max_src_merge=2, max_tgt_merge=2):\n    \"\"\"\n    Monotone greedy alignment by matching length proportions.\n    Produces pairs (src_text, tgt_text), allowing merges 1-2 units.\n    \"\"\"\n    i = j = 0\n    pairs = []\n    while i < len(src_units) and j < len(tgt_units):\n        best = None\n        for si in range(1, max_src_merge+1):\n            for tj in range(1, max_tgt_merge+1):\n                if i+si <= len(src_units) and j+tj <= len(tgt_units):\n                    src_chunk = \" \".join(src_units[i:i+si])\n                    tgt_chunk = \" \".join(tgt_units[j:j+tj])\n                    ls, lt = tok_len(src_chunk), tok_len(tgt_chunk)\n                    score = abs(math.log(ls/lt))\n                    if (best is None) or (score < best[0]):\n                        best = (score, si, tj, src_chunk, tgt_chunk)\n        if best is None:\n            break\n        _, si, tj, s_chunk, t_chunk = best\n        pairs.append((s_chunk, t_chunk))\n        i += si\n        j += tj\n\n    if pairs:\n        if i < len(src_units):\n            pairs[-1] = (pairs[-1][0] + \" \" + \" \".join(src_units[i:]), pairs[-1][1])\n        if j < len(tgt_units):\n            pairs[-1] = (pairs[-1][0], pairs[-1][1] + \" \" + \" \".join(tgt_units[j:]))\n\n    return pairs"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Build sentence-level pairs\npairs = []\nfor _, row in train_df.iterrows():\n    src_doc = row[\"transliteration\"]\n    tgt_doc = row[\"translation\"]\n\n    src_lines = split_src_transliteration(src_doc)\n    tgt_sents = split_tgt_english(tgt_doc)\n\n    if len(src_lines) == 0 or len(tgt_sents) == 0:\n        continue\n\n    doc_pairs = align_monotone(src_lines, tgt_sents, max_src_merge=2, max_tgt_merge=2)\n    pairs.extend(doc_pairs)\n\npairs_df = pd.DataFrame(pairs, columns=[\"src\", \"tgt\"])\npairs_df = pairs_df[(pairs_df[\"src\"].str.len() > 0) & (pairs_df[\"tgt\"].str.len() > 0)].drop_duplicates()\nprint(f\"Sentence-level pairs: {len(pairs_df)}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Load NLLB model\ntokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, src_lang=\"arb_Arab\", tgt_lang=\"eng_Latn\")\nmodel = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)\n\nENG_TOKEN_ID = tokenizer.convert_tokens_to_ids(\"eng_Latn\")\nprint(f\"English token ID: {ENG_TOKEN_ID}\")\n\nmodel.to(DEVICE)\nprint(f\"Model on {DEVICE}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "MAX_SRC = 256\nMAX_TGT = 128\n\nclass MTDataset(Dataset):\n    def __init__(self, df):\n        self.src = df[\"src\"].tolist()\n        self.tgt = df[\"tgt\"].tolist()\n\n    def __len__(self):\n        return len(self.src)\n\n    def __getitem__(self, idx):\n        x = tokenizer(\n            self.src[idx],\n            max_length=MAX_SRC,\n            truncation=True,\n        )\n        y = tokenizer(\n            text_target=self.tgt[idx],\n            max_length=MAX_TGT,\n            truncation=True,\n        )\n        x[\"labels\"] = y[\"input_ids\"]\n        return x"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Train/val split\nval_frac = 0.08\npairs_df = pairs_df.sample(frac=1.0, random_state=SEED).reset_index(drop=True)\nn_val = max(1, int(len(pairs_df) * val_frac))\nval_df = pairs_df.iloc[:n_val].reset_index(drop=True)\ntrn_df = pairs_df.iloc[n_val:].reset_index(drop=True)\n\nprint(f\"Train: {len(trn_df)}, Val: {len(val_df)}\")\n\ntrain_ds = MTDataset(trn_df)\nval_ds = MTDataset(val_df)\ndata_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Training with optimized hyperparameters\nargs = Seq2SeqTrainingArguments(\n    output_dir=OUTPUT_DIR,\n    eval_strategy=\"epoch\",\n    save_strategy=\"no\",  # Don't save checkpoints to avoid disk space issues\n    logging_steps=50,\n    learning_rate=3e-4,  # Higher LR like top baseline\n    per_device_train_batch_size=8,\n    per_device_eval_batch_size=8,\n    gradient_accumulation_steps=2,\n    num_train_epochs=20,\n    warmup_ratio=0.1,\n    fp16=True,  # Enable fp16\n    report_to=\"none\",\n    remove_unused_columns=False,\n)\n\ntrainer = Seq2SeqTrainer(\n    model=model,\n    args=args,\n    train_dataset=train_ds,\n    eval_dataset=val_ds,\n    data_collator=data_collator,\n)\n\nprint(\"Training...\")\ntrainer.train()\nprint(\"Done!\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Skip saving to disk to avoid space issues\n# Model is already in memory for inference\nprint(\"Training complete, skipping disk save\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Inference functions\ndel trainer\ngc.collect()\ntorch.cuda.empty_cache()\nmodel.eval()\n\ndef make_one_sentence(en: str) -> str:\n    en = clean_ws(en)\n    if not en:\n        return \"...\"\n    parts = re.split(r\"(?<=[.!?])\\s+\", en)\n    parts = [p.strip() for p in parts if p.strip()]\n    parts = [re.sub(r\"[.!?]+\\s*$\", \"\", p).strip() for p in parts]\n    parts = [p for p in parts if p]\n    out = \", \".join(parts).strip()\n    if not out.endswith(\".\"):\n        out += \".\"\n    return out\n\ndef generate_translation(src_text: str):\n    enc = tokenizer(\n        clean_ws(src_text),\n        return_tensors=\"pt\",\n        max_length=MAX_SRC,\n        truncation=True,\n    ).to(DEVICE)\n\n    with torch.no_grad():\n        out_ids = model.generate(\n            **enc,\n            forced_bos_token_id=ENG_TOKEN_ID,\n            num_beams=4,\n            max_new_tokens=MAX_TGT,\n            length_penalty=1.0,\n            early_stopping=True,\n        )\n    pred = tokenizer.decode(out_ids[0], skip_special_tokens=True)\n    return make_one_sentence(pred)\n\ndef long_translation(src_text: str):\n    \"\"\"Handle long texts by chunking.\"\"\"\n    src_text = clean_ws(src_text)\n    if len(src_text) <= 256:\n        return generate_translation(src_text)\n    \n    chunks = []\n    len_num = len(src_text) // 256\n    for i in range(1, len_num + 1):\n        chunks.append(src_text[(i-1)*256:i*256])\n    chunks.append(src_text[len_num*256:])\n    chunks = [c for c in chunks if c.strip()]\n    \n    preds = [generate_translation(c) for c in chunks]\n    out = \" \".join(preds)\n    if not out.endswith(\".\"):\n        out += \".\"\n    return out\n\nprint(\"Ready for inference\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Generate submission\nprint(\"Generating submission...\")\npreds = []\nfor idx, row in test_df.iterrows():\n    p = long_translation(row[\"transliteration\"])\n    preds.append(p)\n    print(f\"ID {row['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}