{"cells": [{"cell_type": "markdown", "metadata": {}, "source": "# Deep Past - ByT5 v1\n\nByT5-base for Akkadian to English translation"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "source": "# Fix protobuf first\nimport subprocess, sys\nsubprocess.check_call([sys.executable, '-m', 'pip', 'install', 'protobuf==3.20.3', '-q'])\nprint('protobuf fixed')\n\nimport os, re, 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)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\n\n# Use byt5-base (larger than small)\nMODEL_NAME = \"google/byt5-base\"\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(f\"Device: {DEVICE}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "source": "train_df = pd.read_csv(f\"{DATA_DIR}/train.csv\")\ntest_df = pd.read_csv(f\"{DATA_DIR}/test.csv\")\nprint(f\"Train: {train_df.shape}, Test: {test_df.shape}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "source": "def clean_text(text):\n    if pd.isna(text):\n        return \"\"\n    return re.sub(r\"\\s+\", \" \", str(text)).strip()\n\ntrain_df[\"src\"] = train_df[\"transliteration\"].apply(clean_text)\ntrain_df[\"tgt\"] = train_df[\"translation\"].apply(clean_text)\ntrain_df = train_df[train_df[\"src\"].str.len() > 0]\ntrain_df = train_df[train_df[\"tgt\"].str.len() > 0]\nprint(f\"Clean train: {len(train_df)}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "source": "print(f\"Loading {MODEL_NAME}...\")\ntokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\nmodel = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)\nmodel.to(DEVICE)\nprint(f\"Model loaded on {DEVICE}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "source": "PREFIX = \"translate Akkadian to English: \"\nMAX_SRC = 512\nMAX_TGT = 512\n\nclass TranslationDataset(Dataset):\n    def __init__(self, df):\n        self.src = [PREFIX + s for s in 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        src_enc = tokenizer(self.src[idx], max_length=MAX_SRC, truncation=True, padding=\"max_length\")\n        tgt_enc = tokenizer(text_target=self.tgt[idx], max_length=MAX_TGT, truncation=True, padding=\"max_length\")\n        src_enc[\"labels\"] = tgt_enc[\"input_ids\"]\n        return src_enc"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "source": "val_frac = 0.08\ntrain_df = train_df.sample(frac=1.0, random_state=SEED).reset_index(drop=True)\nn_val = max(1, int(len(train_df) * val_frac))\nval_df = train_df.iloc[:n_val].reset_index(drop=True)\ntrn_df = train_df.iloc[n_val:].reset_index(drop=True)\nprint(f\"Train: {len(trn_df)}, Val: {len(val_df)}\")\n\ntrain_ds = TranslationDataset(trn_df)\nval_ds = TranslationDataset(val_df)\ndata_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model, padding=True)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "source": "args = Seq2SeqTrainingArguments(\n    output_dir=OUTPUT_DIR,\n    eval_strategy=\"epoch\",\n    save_strategy=\"no\",\n    logging_steps=50,\n    learning_rate=1e-4,\n    per_device_train_batch_size=2,\n    per_device_eval_batch_size=2,\n    gradient_accumulation_steps=8,\n    num_train_epochs=50,\n    warmup_ratio=0.1,\n    weight_decay=0.01,\n    fp16=True,\n    report_to=\"none\",\n    remove_unused_columns=False,\n    predict_with_generate=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 ByT5-base for 50 epochs...\")\ntrainer.train()\nprint(\"Done!\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "source": "del trainer\ngc.collect()\ntorch.cuda.empty_cache()\nmodel.eval()\n\ndef translate(text):\n    text = clean_text(text)\n    if not text:\n        return \"...\"\n    inputs = tokenizer(PREFIX + text, return_tensors=\"pt\", max_length=MAX_SRC, truncation=True).to(DEVICE)\n    with torch.no_grad():\n        outputs = model.generate(\n            **inputs,\n            max_length=MAX_TGT,\n            num_beams=4,\n            early_stopping=True,\n            no_repeat_ngram_size=3,\n        )\n    return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()\n\nprint(\"Generating translations...\")\npredictions = [translate(row[\"transliteration\"]) for _, row in test_df.iterrows()]\nfor i, p in enumerate(predictions):\n    print(f\"ID {i}: {p[:100]}...\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "source": "submission = pd.DataFrame({\"id\": test_df[\"id\"], \"translation\": predictions})\nsubmission[\"translation\"] = submission[\"translation\"].apply(lambda x: x if len(x) > 0 else \"...\")\nsubmission.to_csv(\"submission.csv\", index=False)\nprint(\"Submission saved!\")\nprint(submission)"}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python", "version": "3.10.0"}}, "nbformat": 4, "nbformat_minor": 4}