{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":121150,"databundleVersionId":15745837},{"sourceType":"modelInstanceVersion","sourceId":521642,"databundleVersionId":13388065,"modelInstanceId":410134,"modelId":222398}],"dockerImageVersionId":31236,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Note:\n#\n# increase rank -> Better translation?\n#\n# scale up -> 270M already overfits\n#\n# Long inputs: Long inputs should be beneficial to the model\n# However, it might interfere with eos tokens learning, since translations\n# are truncated.","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:11:49.414981Z","iopub.execute_input":"2026-02-20T16:11:49.415798Z","iopub.status.idle":"2026-02-20T16:11:49.419341Z","shell.execute_reply.started":"2026-02-20T16:11:49.415759Z","shell.execute_reply":"2026-02-20T16:11:49.418627Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"debug = False","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:11:49.420458Z","iopub.execute_input":"2026-02-20T16:11:49.420733Z","iopub.status.idle":"2026-02-20T16:11:49.433996Z","shell.execute_reply.started":"2026-02-20T16:11:49.420648Z","shell.execute_reply":"2026-02-20T16:11:49.433261Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install -Uq bitsandbytes\n!pip install -q sacrebleu trl","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:11:49.434898Z","iopub.execute_input":"2026-02-20T16:11:49.435139Z","iopub.status.idle":"2026-02-20T16:12:01.208658Z","shell.execute_reply.started":"2026-02-20T16:11:49.435112Z","shell.execute_reply":"2026-02-20T16:12:01.207747Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import warnings\nif not debug:\n    warnings.filterwarnings(\"ignore\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:12:01.210077Z","iopub.execute_input":"2026-02-20T16:12:01.210319Z","iopub.status.idle":"2026-02-20T16:12:01.215263Z","shell.execute_reply.started":"2026-02-20T16:12:01.210295Z","shell.execute_reply":"2026-02-20T16:12:01.214218Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%capture\nimport pandas as pd\nimport torch\nfrom transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM, BitsAndBytesConfig\nfrom datasets import Dataset\nfrom transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling\nfrom transformers import TrainingArguments\nfrom torch.utils.data import DataLoader\nfrom torch.optim import AdamW\nfrom torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, CosineAnnealingLR\nfrom tqdm.auto import tqdm\nfrom transformers.optimization import Adafactor,AdafactorSchedule\nfrom peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training\nimport math\nimport re\nimport heapq\n\nimport wandb\nfrom kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nwandb_API = user_secrets.get_secret(\"wandb_API\")\n!wandb login $wandb_API\n\nSEED = 123\n\nimport sacrebleu\nimport trl","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:12:01.217507Z","iopub.execute_input":"2026-02-20T16:12:01.217819Z","iopub.status.idle":"2026-02-20T16:12:38.023599Z","shell.execute_reply.started":"2026-02-20T16:12:01.217795Z","shell.execute_reply":"2026-02-20T16:12:38.022452Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df = pd.read_csv(\"/kaggle/input/deep-past-initiative-machine-translation/train.csv\")\ndf = df.drop(columns=[\"oare_id\"])\nif debug:\n    df = df[:100]\ndataset = Dataset.from_pandas(df)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:12:38.024920Z","iopub.execute_input":"2026-02-20T16:12:38.025721Z","iopub.status.idle":"2026-02-20T16:12:38.093228Z","shell.execute_reply.started":"2026-02-20T16:12:38.025652Z","shell.execute_reply":"2026-02-20T16:12:38.092609Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# replace gaps no longer useful\ndef preprocessing(samp):\n    text = samp['transliteration']\n    \n\n    text = re.sub(r'\\.{3,}', '<gap>', text) # ... -> <gap>\n    text = re.sub(r'…', '<gap>', text) # … -> <gap>\n    text = re.sub(\n        r'(?:\\bx\\b\\s*){2,}',\n        '<big_gap> ',\n        text\n    ) # x x x x... -> <big_gap>\n\n    text = re.sub(r'(x{4,})', '<big_gap>', text) # xxxx... -> <big_gap>\n    \n    text = re.sub(r'(x{2,3})', '<big_gap>', text) # xx and xxx -> <big_gap>\n    \n    text = re.sub(r'\\b(x)\\b', '<gap>', text) # x -> <gap>\n    \n    text = re.sub(r'\\s+', ' ', text).strip() # Collapse extra whitespace\n    \n    text = re.sub(\n        r'(?:<gap>\\s*){2,}',\n        '<big_gap> ',\n        text\n    ) # <gap> <gap> -> <big_gap>\n    \n    text = re.sub(\n        r'(<gap>\\s*<big_gap>|<big_gap>\\s*<gap>)',\n        '<big_gap>',\n        text\n    ) # <gap> <big_gap> and <big_gap> <gap> -> <big_gap>\n    \n    text = re.sub(r'\\s+', ' ', text).strip() # Collapse extra whitespace\n    \n    samp['transliteration'] = text\n    return samp","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:12:38.094245Z","iopub.execute_input":"2026-02-20T16:12:38.094571Z","iopub.status.idle":"2026-02-20T16:12:38.100511Z","shell.execute_reply.started":"2026-02-20T16:12:38.094537Z","shell.execute_reply":"2026-02-20T16:12:38.099812Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"prep_dataset = dataset.map(preprocessing)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:12:38.101565Z","iopub.execute_input":"2026-02-20T16:12:38.101926Z","iopub.status.idle":"2026-02-20T16:12:38.170098Z","shell.execute_reply.started":"2026-02-20T16:12:38.101890Z","shell.execute_reply":"2026-02-20T16:12:38.169167Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def formatting_func(example):\n    return {\n        \"prompt\": \"You are an Akkadian expert. Translate the following Akkadian transliteration to English.\\n\\n\"\n        f\"{example['transliteration']}\\n\\n\",\n        \"completion\": f\"{example['translation']}\"\n    }","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:12:38.171345Z","iopub.execute_input":"2026-02-20T16:12:38.171880Z","iopub.status.idle":"2026-02-20T16:12:38.175745Z","shell.execute_reply.started":"2026-02-20T16:12:38.171854Z","shell.execute_reply":"2026-02-20T16:12:38.174947Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"formatted_ds = prep_dataset.map(formatting_func,\n           remove_columns=dataset.column_names)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:12:38.176807Z","iopub.execute_input":"2026-02-20T16:12:38.177111Z","iopub.status.idle":"2026-02-20T16:12:38.216460Z","shell.execute_reply.started":"2026-02-20T16:12:38.177078Z","shell.execute_reply":"2026-02-20T16:12:38.215455Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"q_config = BitsAndBytesConfig(\n    load_in_4bit=True,\n    bnb_4bit_quant_type=\"nf4\",\n    bnb_4bit_use_double_quant=True,\n    bnb_4bit_compute_dtype=torch.bfloat16,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:12:38.217520Z","iopub.execute_input":"2026-02-20T16:12:38.217950Z","iopub.status.idle":"2026-02-20T16:12:38.223694Z","shell.execute_reply.started":"2026-02-20T16:12:38.217910Z","shell.execute_reply":"2026-02-20T16:12:38.222827Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"tokenizer = AutoTokenizer.from_pretrained(\"/kaggle/input/gemma-3/transformers/gemma-3-270m-it/1\")\ntokenizer.save_pretrained('./gemma3_tokenizer')\nmodel = AutoModelForCausalLM.from_pretrained(\n    \"/kaggle/input/gemma-3/transformers/gemma-3-270m-it/1\",\n    dtype=torch.bfloat16,\n    device_map=\"auto\",\n    quantization_config=q_config,\n)\nmodel = prepare_model_for_kbit_training(model)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:12:38.224625Z","iopub.execute_input":"2026-02-20T16:12:38.225000Z","iopub.status.idle":"2026-02-20T16:12:54.417996Z","shell.execute_reply.started":"2026-02-20T16:12:38.224959Z","shell.execute_reply":"2026-02-20T16:12:54.417087Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def filter_fn(samp):\n    prompt = tokenizer(samp['prompt'])\n    completion = tokenizer(samp['completion'])\n    length = len(prompt['input_ids']) + len(completion['input_ids'])\n    diff = abs(len(prompt['input_ids']) - len(completion['input_ids']))\n    if length > 1024:\n        return False\n    if diff > 500:\n        return False\n    return True","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:12:54.419020Z","iopub.execute_input":"2026-02-20T16:12:54.419888Z","iopub.status.idle":"2026-02-20T16:12:54.425192Z","shell.execute_reply.started":"2026-02-20T16:12:54.419862Z","shell.execute_reply":"2026-02-20T16:12:54.424203Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"filtered_ds = formatted_ds.filter(filter_fn)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:12:54.428014Z","iopub.execute_input":"2026-02-20T16:12:54.428329Z","iopub.status.idle":"2026-02-20T16:12:54.975548Z","shell.execute_reply.started":"2026-02-20T16:12:54.428303Z","shell.execute_reply":"2026-02-20T16:12:54.974567Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(len(formatted_ds))\nprint(len(filtered_ds))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:12:54.976581Z","iopub.execute_input":"2026-02-20T16:12:54.976936Z","iopub.status.idle":"2026-02-20T16:12:54.981503Z","shell.execute_reply.started":"2026-02-20T16:12:54.976911Z","shell.execute_reply":"2026-02-20T16:12:54.980799Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# EDA\nlengths,prompts,completions,p_lengths,c_lengths = [],[],[],[],[]\nfor samp in formatted_ds:\n    prompt = tokenizer(samp['prompt'])\n    completion = tokenizer(samp['completion'])\n    lengths.append(len(prompt['input_ids']) + len(completion['input_ids']))\n    prompts.append(prompt)\n    completions.append(completion)\n    p_lengths.append(len(prompt['input_ids']))\n    c_lengths.append(len(completion['input_ids']))\n\nprint(heapq.nlargest(10,lengths))\n# [3852, 3159, 2876, 2750, 2658, 2603, 2553, 2512, 2476, 2325]\nprint(heapq.nlargest(10,p_lengths))\n# [2490, 2036, 1864, 1783, 1762, 1734, 1689, 1625, 1600, 1549]\nprint(heapq.nlargest(10,c_lengths))\n# [1362, 1123, 1094, 1093, 1091, 1022, 1021, 924, 912, 886]\n\ncount = 0\nfor x in lengths:\n    if x > 1024:\n        count += 1\n\nprint(count) # 119\n\ndiffs = []\nfor p,c in zip(prompts,completions):\n    diff = abs(len(p['input_ids'])-len(c['input_ids']))\n    diffs.append(diff)\n\nprint(heapq.nlargest(10,diffs))\n# [1128, 1042, 978, 962, 939, 938, 925, 921, 913, 902]\nprint(heapq.nsmallest(10,c_lengths))\n# [5, 6, 7, 7, 8, 9, 10, 10, 11, 12] Further investigations\n\ncount = 0\nfor x in diffs:\n    if x > 500:\n        count += 1\n\nprint(count) # 44","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:12:54.982558Z","iopub.execute_input":"2026-02-20T16:12:54.983239Z","iopub.status.idle":"2026-02-20T16:12:55.082920Z","shell.execute_reply.started":"2026-02-20T16:12:54.983212Z","shell.execute_reply":"2026-02-20T16:12:55.082001Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"ds = filtered_ds.train_test_split(test_size=0.2, seed=SEED)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:12:55.083797Z","iopub.execute_input":"2026-02-20T16:12:55.084090Z","iopub.status.idle":"2026-02-20T16:12:55.094779Z","shell.execute_reply.started":"2026-02-20T16:12:55.084050Z","shell.execute_reply":"2026-02-20T16:12:55.093867Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"lora_config = LoraConfig(\n    task_type=TaskType.CAUSAL_LM,\n    r=32,\n    lora_alpha=48,\n    lora_dropout=0.05,\n    target_modules=\"all-linear\", # all linear shows modest improvements compared to smaller ones\n    bias=\"none\"\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:12:55.095954Z","iopub.execute_input":"2026-02-20T16:12:55.096262Z","iopub.status.idle":"2026-02-20T16:12:55.102149Z","shell.execute_reply.started":"2026-02-20T16:12:55.096237Z","shell.execute_reply":"2026-02-20T16:12:55.101467Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from trl import SFTConfig\n\nsft_config = SFTConfig(\n    output_dir=\"./gemma_sft\",\n    num_train_epochs=1 if debug else 3, # debug\n    per_device_train_batch_size=2,\n    per_device_eval_batch_size=4,\n    gradient_accumulation_steps=2,\n    eval_strategy=\"epoch\",\n    save_strategy=\"epoch\",\n    logging_strategy=\"steps\",\n    logging_steps=10,\n    learning_rate=3e-4,\n    optim=\"adamw_torch\",\n    lr_scheduler_type=\"cosine\",\n    warmup_ratio=0.05,\n    bf16=True,\n    report_to=None if debug else \"wandb\", # debug\n    load_best_model_at_end=True,\n    max_length=1024, # Longer input translation got truncated\n    completion_only_loss=True,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:12:55.103108Z","iopub.execute_input":"2026-02-20T16:12:55.103360Z","iopub.status.idle":"2026-02-20T16:12:56.545574Z","shell.execute_reply.started":"2026-02-20T16:12:55.103337Z","shell.execute_reply":"2026-02-20T16:12:56.544511Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from trl import SFTTrainer\ntrainer = SFTTrainer(\n    model=model,\n    args=sft_config,\n    train_dataset=ds[\"train\"],\n    eval_dataset=ds[\"test\"],\n    peft_config=lora_config,\n)\n\nif not debug:\n    wandb.init(project=\"MyProject\")\n    wandb.watch(trainer.model, log=\"gradients\", log_freq=10)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:12:56.546340Z","iopub.execute_input":"2026-02-20T16:12:56.546809Z","iopub.status.idle":"2026-02-20T16:13:00.362469Z","shell.execute_reply.started":"2026-02-20T16:12:56.546778Z","shell.execute_reply":"2026-02-20T16:13:00.361740Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"trainer.train()\ntrainer.save_model(\"./gemma_sft_lora\")\ntokenizer.save_pretrained(\"./gemma_sft_lora\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:13:00.363521Z","iopub.execute_input":"2026-02-20T16:13:00.363831Z","iopub.status.idle":"2026-02-20T16:13:52.087262Z","shell.execute_reply.started":"2026-02-20T16:13:00.363806Z","shell.execute_reply":"2026-02-20T16:13:52.086360Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"@torch.inference_mode()\ndef evaluate(\n    model,\n    tokenizer,\n    dataset,\n    batch_size=8,\n    max_new_tokens=1024,\n    num_worst=50,\n    wandb_run=None,\n):\n    model.eval()\n\n    records = []\n\n    for i in tqdm(range(0, len(dataset), batch_size), desc=\"Evaluating\"):\n        batch = dataset[i : i + batch_size]\n\n        prompts = batch['prompt']\n        refs = batch['completion']\n\n        inputs = tokenizer(\n            prompts,\n            return_tensors=\"pt\",\n            padding=True,\n        ).to(model.device)\n\n        outputs = model.generate(\n            **inputs,\n            max_new_tokens=max_new_tokens,\n            do_sample=False,\n            eos_token_id=tokenizer.eos_token_id,\n            pad_token_id=tokenizer.eos_token_id,\n        )\n\n        # Gemma3 tokenizer uses left padding\n        decoded = []\n\n        for i in range(len(prompts)):\n            attn = inputs[\"attention_mask\"][i]\n            prompt_len = attn.sum().item()\n            prompt_start = attn.nonzero(as_tuple=False)[0].item()\n            prompt_end = prompt_start + prompt_len\n        \n            gen_tokens = outputs[i, prompt_end:]\n            text = tokenizer.decode(gen_tokens, skip_special_tokens=True)\n            decoded.append(text)\n\n        \n        for prompt, pred_text, ref in zip(prompts, decoded, refs):\n            pred = pred_text.strip()\n\n            # ---- Sentence-level metrics ----\n            bleu = sacrebleu.sentence_bleu(pred, [ref]).score\n            chrf = sacrebleu.sentence_chrf(pred, [ref]).score\n            \n            geo_mean = math.sqrt(bleu * chrf)\n\n            records.append(\n                {\n                    \"prompt\": prompt,\n                    \"prediction\": pred,\n                    \"reference\": ref,\n                    \"bleu\": bleu,\n                    \"chrf\": chrf,\n                    \"geo_mean\": geo_mean,\n                }\n            )\n\n    # ---- Corpus metrics ----\n    all_preds = [r[\"prediction\"] for r in records]\n    all_refs = [r[\"reference\"] for r in records]\n\n    corpus_bleu = sacrebleu.corpus_bleu(all_preds, [all_refs]).score\n    corpus_chrf = sacrebleu.corpus_chrf(all_preds, [all_refs]).score\n    corpus_geo = math.sqrt(corpus_bleu * corpus_chrf)\n\n    # ---- Worst samples ----\n    records_sorted = sorted(records, key=lambda x: x[\"geo_mean\"])\n    worst = records_sorted[:num_worst]\n\n    # ---- W&B logging ----\n    if wandb_run is not None:\n        table = wandb.Table(\n            columns=[\n                \"prompt\",\n                \"prediction\",\n                \"reference\",\n                \"BLEU\",\n                \"chrF\",\n                \"GeoMean\",\n            ]\n        )\n        for r in worst:\n            table.add_data(\n                r[\"prompt\"],\n                r[\"prediction\"],\n                r[\"reference\"],\n                r[\"bleu\"],\n                r[\"chrf\"],\n                r[\"geo_mean\"],\n            )\n\n        wandb.log(\n            {\n                \"eval/bleu\": corpus_bleu,\n                \"eval/chrf\": corpus_chrf,\n                \"eval/geo_mean\": corpus_geo,\n                \"eval/worst_samples\": table,\n            }\n        )\n\n    return {\n        \"bleu\": corpus_bleu,\n        \"chrf\": corpus_chrf,\n        \"geo_mean\": corpus_geo,\n        \"worst_samples\": worst,\n    }\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:13:52.089345Z","iopub.execute_input":"2026-02-20T16:13:52.089866Z","iopub.status.idle":"2026-02-20T16:13:52.106175Z","shell.execute_reply.started":"2026-02-20T16:13:52.089814Z","shell.execute_reply":"2026-02-20T16:13:52.105498Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"history = evaluate(\n    trainer.model,\n    tokenizer,\n    ds['test'], # test\n    batch_size=16,\n    max_new_tokens=256,\n    num_worst=50,\n    wandb_run=None if debug else True,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:13:52.107170Z","iopub.execute_input":"2026-02-20T16:13:52.107500Z","execution_failed":"2026-02-20T16:16:37.747Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if not debug:\n    wandb.finish()","metadata":{"trusted":true,"execution":{"execution_failed":"2026-02-20T16:16:37.749Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"history","metadata":{"trusted":true,"execution":{"execution_failed":"2026-02-20T16:16:37.750Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}