{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.11.0"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":121150,"databundleVersionId":15833391}],"dockerImageVersionId":31287,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"1sz3cutvzwk","cell_type":"markdown","source":"# DPC Training - ByT5-large (full fine-tune)\n\nByT5-large (1.2B params) on train.csv only. Comparing against ByT5-base (580M) baseline which scored 28.1.","metadata":{}},{"id":"i08x9slgb7n","cell_type":"code","source":"# our code starts here\n!pip install -q sacrebleu wandb","metadata":{},"outputs":[],"execution_count":null},{"id":"kqxnja1c08","cell_type":"code","source":"import os\n\n# Single GPU to prevent DataParallel OOM on T4x2\nos.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n\nimport gc\nimport math\nimport re\nimport numpy as np\nimport pandas as pd\nimport torch\nimport sacrebleu\nfrom datasets import Dataset, disable_progress_bars\nfrom transformers import (\n    AutoTokenizer,\n    AutoModelForSeq2SeqLM,\n    DataCollatorForSeq2Seq,\n    Seq2SeqTrainingArguments,\n    Seq2SeqTrainer,\n)\n\ndisable_progress_bars()\n\n# Load W&B API key from Kaggle Secrets\ntry:\n    from kaggle_secrets import UserSecretsClient\n    os.environ[\"WANDB_API_KEY\"] = UserSecretsClient().get_secret(\"WANDB_API_KEY\")\n    print(\"W&B API key loaded from Kaggle Secrets\")\nexcept Exception:\n    print(\"Not on Kaggle or WANDB_API_KEY secret not set, W&B will use local login\")\n\ndef seed_everything(seed=42):\n    torch.manual_seed(seed)\n    np.random.seed(seed)\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed_all(seed)\n\nseed_everything()","metadata":{},"outputs":[],"execution_count":null},{"id":"e9dtg9odhvg","cell_type":"markdown","source":"## Config","metadata":{}},{"id":"qf5z72r115s","cell_type":"code","source":"MODEL_NAME = \"google/byt5-large\"\nMAX_LENGTH = 512\nEPOCHS = 15\nLR = 1e-4\nBATCH_SIZE = 1\nGRAD_ACCUM = 8  # effective batch = 8\nOUTPUT_DIR = \"./byt5-akkadian-large\"\nPREFIX = \"translate Akkadian to English: \"\nSKIP_OVERFIT = True  # set False to run overfit check\nEVAL_SAMPLES = 30   # subset of val set for competition metric eval","metadata":{},"outputs":[],"execution_count":null},{"id":"u0ptw6802zr","cell_type":"markdown","source":"## Load Data","metadata":{}},{"id":"4rcvz5ys9je","cell_type":"code","source":"INPUT_DIR = \"/kaggle/input/competitions/deep-past-initiative-machine-translation\"\ntrain_df = pd.read_csv(f\"{INPUT_DIR}/train.csv\")\nprint(f\"Raw training data: {len(train_df)} rows\")\ntrain_df.head(3)","metadata":{},"outputs":[],"execution_count":null},{"id":"bogpko3zq3f","cell_type":"markdown","source":"## Preprocessing","metadata":{}},{"id":"yc78ahvtuih","cell_type":"code","source":"def clean_transliteration(text):\n    \"\"\"Clean Akkadian transliteration per competition host guidelines.\"\"\"\n    if pd.isna(text):\n        return \"\"\n    text = str(text)\n\n    # Test data only has H/h, not Ḫ/ḫ\n    text = text.replace(\"Ḫ\", \"H\").replace(\"ḫ\", \"h\")\n\n    # Gap markers (before removing brackets)\n    text = re.sub(r\"\\[x\\]\", \"<gap>\", text)\n    text = re.sub(r\"\\[…\\]\", \"<big_gap>\", text)\n    text = re.sub(r\"\\[\\.\\.\\.\\]\", \"<big_gap>\", text)\n    text = re.sub(r\"…\", \"<big_gap>\", text)\n    text = re.sub(r\"\\.{3,}\", \"<big_gap>\", text)\n\n    # Remove scribal notation: ! ? / : .\n    text = re.sub(r\"[!?/:.]\", \"\", text)\n\n    # Remove bracket/brace characters but keep content inside\n    text = re.sub(r\"[˹˺\\[\\]<>]\", \"\", text)\n\n    # Collapse whitespace\n    text = re.sub(r\"\\s+\", \" \", text).strip()\n    return text\n\n\ndef clean_translation(text):\n    \"\"\"Clean English translation per competition host guidelines.\"\"\"\n    if pd.isna(text):\n        return \"\"\n    text = str(text)\n\n    # Remove parenthetical annotations like (plur.), (fem.), (lit. ...), etc.\n    text = re.sub(r\"\\([^)]*\\)\", \"\", text)\n\n    # Collapse whitespace\n    text = re.sub(r\"\\s+\", \" \", text).strip()\n    return text\n\n\ntrain_df[\"transliteration\"] = train_df[\"transliteration\"].apply(clean_transliteration)\ntrain_df[\"translation\"] = train_df[\"translation\"].apply(clean_translation)\n\n# Drop empty rows\ntrain_df = train_df[\n    (train_df[\"transliteration\"].str.len() > 0) & (train_df[\"translation\"].str.len() > 0)\n].reset_index(drop=True)\n\nprint(f\"After cleaning: {len(train_df)} rows\")\ntrain_df.head(3)","metadata":{},"outputs":[],"execution_count":null},{"id":"ykp7z2apsj","cell_type":"markdown","source":"## Tokenize & Split","metadata":{}},{"id":"xjnt8j6y8g","cell_type":"code","source":"# Add task prefix and split\ntrain_df[\"input_text\"] = PREFIX + train_df[\"transliteration\"]\ntrain_df[\"target_text\"] = train_df[\"translation\"]\n\ndataset = Dataset.from_pandas(train_df[[\"input_text\", \"target_text\"]])\nsplit = dataset.train_test_split(test_size=0.1, seed=42)\n\nprint(f\"Train: {len(split['train'])}, Val: {len(split['test'])}\")\n\n# Tokenize\ntokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n\ndef tokenize(examples):\n    inputs = tokenizer(examples[\"input_text\"], max_length=MAX_LENGTH, truncation=True)\n    labels = tokenizer(examples[\"target_text\"], max_length=MAX_LENGTH, truncation=True)\n    inputs[\"labels\"] = labels[\"input_ids\"]\n    return inputs\n\ntok_train = split[\"train\"].map(tokenize, batched=True, remove_columns=[\"input_text\", \"target_text\"])\ntok_val = split[\"test\"].map(tokenize, batched=True, remove_columns=[\"input_text\", \"target_text\"])","metadata":{},"outputs":[],"execution_count":null},{"id":"8ww0403kug3","cell_type":"code","source":"# Competition metric: sqrt(BLEU * chrF++)\ndef compute_metrics(eval_preds):\n    preds, labels = eval_preds\n    if isinstance(preds, tuple):\n        preds = preds[0]\n    if preds.ndim == 3:\n        preds = np.argmax(preds, axis=-1)\n\n    preds = np.clip(preds.astype(np.int64), 0, tokenizer.vocab_size - 1)\n    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)\n\n    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)\n    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n\n    bleu = sacrebleu.corpus_bleu(decoded_preds, [decoded_labels]).score\n    chrf = sacrebleu.corpus_chrf(decoded_preds, [decoded_labels], word_order=2).score\n    geo = math.sqrt(bleu * chrf) if bleu > 0 and chrf > 0 else 0.0\n    return {\"bleu\": bleu, \"chrf\": chrf, \"geo_mean\": geo}","metadata":{},"outputs":[],"execution_count":null},{"id":"5d5eyxx4ho","cell_type":"markdown","source":"## Overfit a single batch","metadata":{}},{"id":"e3rx48s9fl7","cell_type":"code","source":"if not SKIP_OVERFIT:\n    # Overfit on 8 examples, loss should drop to near 0\n    gc.collect()\n    torch.cuda.empty_cache()\n\n    tiny_train = tok_train.select(range(8))\n\n    overfit_model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)\n    overfit_collator = DataCollatorForSeq2Seq(tokenizer, model=overfit_model)\n\n    overfit_args = Seq2SeqTrainingArguments(\n        output_dir=\"./overfit-test\",\n        num_train_epochs=100,\n        learning_rate=3e-4,\n        per_device_train_batch_size=1,\n        gradient_accumulation_steps=8,\n        gradient_checkpointing=True,\n        fp16=False,\n        optim=\"adafactor\",\n        logging_steps=10,\n        save_strategy=\"no\",\n        report_to=\"none\",\n    )\n\n    overfit_trainer = Seq2SeqTrainer(\n        model=overfit_model,\n        args=overfit_args,\n        train_dataset=tiny_train,\n        data_collator=overfit_collator,\n        processing_class=tokenizer,\n    )\n\n    print(\"Overfitting on 8 examples, loss should approach 0...\")\n    overfit_trainer.train()\nelse:\n    print(\"Skipping overfit check (SKIP_OVERFIT=True)\")","metadata":{},"outputs":[],"execution_count":null},{"id":"i05i13yx8mf","cell_type":"code","source":"if not SKIP_OVERFIT:\n    # Re-run this cell to train more epochs without reloading the model\n    overfit_args.num_train_epochs = 50\n    overfit_trainer.train(resume_from_checkpoint=False)\nelse:\n    print(\"Skipping (SKIP_OVERFIT=True)\")","metadata":{},"outputs":[],"execution_count":null},{"id":"ilhewsq3ync","cell_type":"code","source":"if not SKIP_OVERFIT:\n    # Check predictions on the overfit set\n    overfit_model.eval()\n    all_preds = []\n    all_trues = []\n\n    for i in range(len(tiny_train)):\n        row = split[\"train\"][i]\n        inp = tokenizer(row[\"input_text\"], return_tensors=\"pt\", max_length=MAX_LENGTH, truncation=True).to(overfit_model.device)\n        with torch.no_grad():\n            out = overfit_model.generate(**inp, max_length=MAX_LENGTH, num_beams=1)\n        pred = tokenizer.decode(out[0], skip_special_tokens=True)\n        true = row[\"target_text\"]\n        all_preds.append(pred)\n        all_trues.append(true)\n        print(f\"\\nTRUE: {true[:200]}\")\n        print(f\"PRED: {pred[:200]}\")\n\n    bleu = sacrebleu.corpus_bleu(all_preds, [all_trues])\n    chrf = sacrebleu.corpus_chrf(all_preds, [all_trues], word_order=2)\n    geo = math.sqrt(bleu.score * chrf.score) if bleu.score > 0 and chrf.score > 0 else 0.0\n\n    print(f\"\\n{'='*50}\")\n    print(f\"BLEU:                {bleu.score:.2f}\")\n    print(f\"chrF++:              {chrf.score:.2f}\")\n    print(f\"Competition metric:  {geo:.2f}  (sqrt(BLEU * chrF++))\")\n    print(f\"{'='*50}\")\nelse:\n    print(\"Skipping (SKIP_OVERFIT=True)\")","metadata":{},"outputs":[],"execution_count":null},{"id":"z3xijbbh45e","cell_type":"code","source":"if not SKIP_OVERFIT:\n    del overfit_model, overfit_trainer\n    gc.collect()\n    torch.cuda.empty_cache()\n    print(\"Overfit check complete. Moving on to full training.\")\nelse:\n    print(\"Skipping (SKIP_OVERFIT=True)\")","metadata":{},"outputs":[],"execution_count":null},{"id":"i08npykanqq","cell_type":"markdown","source":"## Full Training with W&B","metadata":{}},{"id":"lw4ch2hr6d","cell_type":"code","source":"import wandb\n\nwandb.init(\n    project=\"dpc-akkadian\",\n    name=f\"byt5large-lr{LR}-ep{EPOCHS}\",\n    config={\n        \"model\": MODEL_NAME,\n        \"epochs\": EPOCHS,\n        \"lr\": LR,\n        \"batch_size\": BATCH_SIZE,\n        \"grad_accum\": GRAD_ACCUM,\n        \"effective_batch\": BATCH_SIZE * GRAD_ACCUM,\n        \"max_length\": MAX_LENGTH,\n        \"train_rows\": len(tok_train),\n        \"val_rows\": len(tok_val),\n        \"method\": \"full_finetune\",\n    },\n)\n\ngc.collect()\ntorch.cuda.empty_cache()\n\nmodel = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)\ncollator = DataCollatorForSeq2Seq(tokenizer, model=model)\n\nargs = Seq2SeqTrainingArguments(\n    output_dir=OUTPUT_DIR,\n    num_train_epochs=EPOCHS,\n    learning_rate=LR,\n    per_device_train_batch_size=BATCH_SIZE,\n    per_device_eval_batch_size=BATCH_SIZE,\n    gradient_accumulation_steps=GRAD_ACCUM,\n    gradient_checkpointing=True,\n    fp16=False,\n    optim=\"adafactor\",\n    warmup_ratio=0.1,\n    eval_strategy=\"epoch\",\n    save_strategy=\"epoch\",\n    save_total_limit=1,\n    logging_steps=10,\n    report_to=\"wandb\",\n    load_best_model_at_end=True,\n    metric_for_best_model=\"eval_loss\",\n)\n\ntrainer = Seq2SeqTrainer(\n    model=model,\n    args=args,\n    train_dataset=tok_train,\n    eval_dataset=tok_val,\n    data_collator=collator,\n    processing_class=tokenizer,\n)\n\ntrainer.train()","metadata":{},"outputs":[],"execution_count":null},{"id":"mashbzhe0zr","cell_type":"markdown","source":"## Sample Predictions","metadata":{}},{"id":"e4plcnu4ydo","cell_type":"code","source":"import random\n\nval_df = split[\"test\"].to_pandas()\n\n# Generate predictions on a subset for competition metrics\neval_indices = list(range(min(EVAL_SAMPLES, len(val_df))))\nmodel.eval()\nall_preds = []\nall_trues = []\n\nfor i in eval_indices:\n    row = val_df.iloc[i]\n    inp = tokenizer(row[\"input_text\"], return_tensors=\"pt\", max_length=MAX_LENGTH, truncation=True).to(model.device)\n    with torch.no_grad():\n        out = model.generate(**inp, max_length=256, num_beams=4, early_stopping=True)\n    pred = tokenizer.decode(out[0], skip_special_tokens=True)\n    all_preds.append(pred)\n    all_trues.append(row[\"target_text\"])\n\nbleu = sacrebleu.corpus_bleu(all_preds, [all_trues])\nchrf = sacrebleu.corpus_chrf(all_preds, [all_trues], word_order=2)\ngeo = math.sqrt(bleu.score * chrf.score) if bleu.score > 0 and chrf.score > 0 else 0.0\n\nprint(f\"{'='*50}\")\nprint(f\"Val subset ({len(eval_indices)} examples)\")\nprint(f\"BLEU:                {bleu.score:.2f}\")\nprint(f\"chrF++:              {chrf.score:.2f}\")\nprint(f\"Competition metric:  {geo:.2f}  (sqrt(BLEU * chrF++))\")\nprint(f\"{'='*50}\")\n\n# Show 10 random samples\nsamples = random.sample(eval_indices, min(10, len(eval_indices)))\nfor i in samples:\n    print(f\"\\n--- {i} ---\")\n    print(f\"AKK:  {val_df.iloc[i]['input_text'][:200]}\")\n    print(f\"PRED: {all_preds[eval_indices.index(i)][:200]}\")\n    print(f\"TRUE: {all_trues[eval_indices.index(i)][:200]}\")","metadata":{},"outputs":[],"execution_count":null},{"id":"gls8fzxev7","cell_type":"markdown","source":"## Save Model","metadata":{}},{"id":"6zyul3512jb","cell_type":"code","source":"# Log final competition metrics and sample translations to W&B\nwandb.log({\"final/bleu\": bleu.score, \"final/chrf\": chrf.score, \"final/geo_mean\": geo})\n\n# Save translation examples as a W&B Table\ntable = wandb.Table(columns=[\"akkadian\", \"prediction\", \"ground_truth\"])\nfor i in range(len(all_preds)):\n    table.add_data(\n        val_df.iloc[eval_indices[i]][\"input_text\"][:500],\n        all_preds[i][:500],\n        all_trues[i][:500],\n    )\nwandb.log({\"translations\": table})\n\ntrainer.save_model(OUTPUT_DIR)\ntokenizer.save_pretrained(OUTPUT_DIR)\nwandb.finish()\nprint(f\"Model saved to {OUTPUT_DIR}\")\n# our code ends here","metadata":{},"outputs":[],"execution_count":null}]}