{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"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"},"accelerate_config":{"num_processes":2}},"nbformat_minor":5,"nbformat":4,"cells":[{"id":"932ffe2c-915d-4254-b4dd-1d32bfeb87db","cell_type":"markdown","source":"# Java Code Optimization — CodeT5-small Fine-tuning\n**Kaggle T4×2 · Dual-GPU DataParallel · ~6 K train / 680 val pairs**\n\nPipeline:\n1. Load dataset from `dataset_train.jsonl` / `dataset_val.jsonl`\n2. Tokenize with `Salesforce/codet5-small` tokenizer  \n3. Fine-tune with `Seq2SeqTrainer` + both GPUs via `DataParallel`\n4. Evaluate with BLEU + CodeBLEU proxies  \n5. Push to HuggingFace Hub (optional) and save artefacts","metadata":{}},{"id":"10c9c8fb-0f60-4876-9110-7df1fc24b0aa","cell_type":"code","source":"import subprocess, sys\n\npkgs = [\n    \"transformers==4.41.2\",\n    \"datasets==2.20.0\",\n    \"evaluate==0.4.2\",\n    \"sacrebleu==2.4.3\",\n    \"accelerate==0.33.0\",\n    \"peft==0.11.1\",\n    \"sentencepiece==0.2.0\",\n    \"rouge_score==0.1.2\",\n    \"tokenizers==0.19.1\",\n]\n\nsubprocess.run(\n    [sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"--force-reinstall\"] + pkgs\n)\n\nprint(\"✅ Install done. RESTART KERNEL NOW.\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"id":"e8ea6ed7-c9ef-4a0f-8cbb-457c1b599466","cell_type":"code","source":"import os\n\nos.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"   # ✅ FORCE SINGLE GPU\nos.environ[\"CUDA_LAUNCH_BLOCKING\"] = \"1\"   # debug + stability","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-03T12:59:28.032986Z","iopub.execute_input":"2026-04-03T12:59:28.033907Z","iopub.status.idle":"2026-04-03T12:59:28.038196Z","shell.execute_reply.started":"2026-04-03T12:59:28.033868Z","shell.execute_reply":"2026-04-03T12:59:28.037332Z"}},"outputs":[],"execution_count":14},{"id":"5712ecaa-b2af-4446-a009-e210d4897a12","cell_type":"code","source":"import os, json, random, time\nfrom pathlib import Path\nfrom dataclasses import dataclass\n\nimport numpy as np\nimport torch\n\nfrom transformers import (\n    T5ForConditionalGeneration,\n    RobertaTokenizer,\n    Seq2SeqTrainer,\n    Seq2SeqTrainingArguments,\n    DataCollatorForSeq2Seq,\n)\n\nfrom datasets import Dataset as HFDataset, DatasetDict\n\nprint(\"CUDA:\", torch.cuda.is_available())\nprint(\"GPUs:\", torch.cuda.device_count())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-03T12:59:28.247024Z","iopub.execute_input":"2026-04-03T12:59:28.247805Z","iopub.status.idle":"2026-04-03T12:59:28.253552Z","shell.execute_reply.started":"2026-04-03T12:59:28.247769Z","shell.execute_reply":"2026-04-03T12:59:28.252760Z"}},"outputs":[{"name":"stdout","text":"CUDA: True\nGPUs: 1\n","output_type":"stream"}],"execution_count":15},{"id":"ad8ca487-c537-49ba-a1ff-0b124837ced8","cell_type":"code","source":"@dataclass\nclass CFG:\n    model_name: str = \"Salesforce/codet5-small\"\n\n    train_file: str = \"/kaggle/input/datasets/suhaskoheda/java-optimisation/dataset_train.jsonl\"\n    val_file: str = \"/kaggle/input/datasets/suhaskoheda/java-optimisation/dataset_val.jsonl\"\n\n    max_source_length: int = 128\n    max_target_length: int = 128\n\n    num_train_epochs: int = 5  # 🔥 fast\n\n    per_device_train_batch_size: int = 16\n    per_device_eval_batch_size: int = 32\n\n    learning_rate: float = 5e-4\n    weight_decay: float = 0.01\n    warmup_ratio: float = 0.05\n\n    output_dir: str = \"/kaggle/working/codet5-fast\"\n\ncfg = CFG()\nPath(cfg.output_dir).mkdir(parents=True, exist_ok=True)\n\nprint(\"✅ Config ready\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-03T12:59:28.411634Z","iopub.execute_input":"2026-04-03T12:59:28.412174Z","iopub.status.idle":"2026-04-03T12:59:28.419593Z","shell.execute_reply.started":"2026-04-03T12:59:28.412144Z","shell.execute_reply":"2026-04-03T12:59:28.418710Z"}},"outputs":[{"name":"stdout","text":"✅ Config ready\n","output_type":"stream"}],"execution_count":16},{"id":"2a3aae51-9c02-4b13-b4f7-7e95e3c5084e","cell_type":"code","source":"def load_jsonl(path):\n    data = []\n    with open(path, encoding=\"utf-8\") as f:\n        for line in f:\n            if line.strip():\n                data.append(json.loads(line))\n    return data\n\ntrain_raw = load_jsonl(cfg.train_file)\nval_raw = load_jsonl(cfg.val_file)\n\nprint(\"Train:\", len(train_raw))\nprint(\"Val:\", len(val_raw))\n\nprint(\"\\nSample:\")\nprint(train_raw[0][\"input\"][:200])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-03T12:59:28.560595Z","iopub.execute_input":"2026-04-03T12:59:28.561419Z","iopub.status.idle":"2026-04-03T12:59:28.617144Z","shell.execute_reply.started":"2026-04-03T12:59:28.561339Z","shell.execute_reply":"2026-04-03T12:59:28.616172Z"}},"outputs":[{"name":"stdout","text":"Train: 6115\nVal: 680\n\nSample:\nString alertLevel;\nswitch (sensor.getValue()) {\n    case 0:\n    case 1:\n        alertLevel = \"GREEN\";\n        break;\n    case 2:\n        alertLevel = \"YELLOW\";\n        break;\n    case 3:\n        alert\n","output_type":"stream"}],"execution_count":17},{"id":"4ead70f3-795f-42e9-863d-ac14024611d7","cell_type":"code","source":"from transformers import AutoTokenizer\n\ntokenizer = AutoTokenizer.from_pretrained(\n    cfg.model_name,\n    use_fast=False\n)\n\nTASK_PREFIX = \"Optimize Java: \"\n\ndef tokenize_batch(examples):\n    inputs = [TASK_PREFIX + x for x in examples[\"input\"]]\n    targets = examples[\"output\"]\n\n    model_inputs = tokenizer(\n        inputs,\n        max_length=cfg.max_source_length,\n        truncation=True,\n        padding=\"max_length\",\n    )\n\n    labels = tokenizer(\n        targets,\n        max_length=cfg.max_target_length,\n        truncation=True,\n        padding=\"max_length\",\n    )\n\n    # 🔥 CRITICAL FIX\n    cleaned_labels = []\n    for label in labels[\"input_ids\"]:\n        cleaned = []\n        for token in label:\n            if token == tokenizer.pad_token_id:\n                cleaned.append(-100)\n            else:\n                cleaned.append(int(token))   # 🔥 force python int\n        cleaned_labels.append(cleaned)\n\n    model_inputs[\"labels\"] = cleaned_labels\n\n    return model_inputs\n    \nprint(\"✅ Tokenizer ready\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"id":"94508e52-a595-452b-89f2-cf85003d2109","cell_type":"code","source":"def to_hf(rows):\n    return HFDataset.from_dict({\n        \"input\": [r[\"input\"] for r in rows],\n        \"output\": [r[\"output\"] for r in rows],\n    })\n\ndataset = DatasetDict({\n    \"train\": to_hf(train_raw),\n    \"val\": to_hf(val_raw),\n})\n\ntokenized_ds = dataset.map(\n    tokenize_batch,\n    batched=True,\n    remove_columns=[\"input\", \"output\"],\n)\n\ntokenized_ds.set_format(\"torch\")\ntokenized_ds = tokenized_ds.with_format(\"torch\")\nprint(\"✅ Dataset ready\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-03T13:02:56.439164Z","iopub.execute_input":"2026-04-03T13:02:56.440280Z","iopub.status.idle":"2026-04-03T13:03:01.939341Z","shell.execute_reply.started":"2026-04-03T13:02:56.440196Z","shell.execute_reply":"2026-04-03T13:03:01.938429Z"}},"outputs":[{"output_type":"display_data","data":{"text/plain":"Map:   0%|          | 0/6115 [00:00<?, ? examples/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"d5a572e0b3ff4df6a689234b7c46a72d"}},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"Map:   0%|          | 0/680 [00:00<?, ? examples/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"f74452d8df16462ba281f71eb59989bd"}},"metadata":{}},{"name":"stdout","text":"✅ Dataset ready\n","output_type":"stream"}],"execution_count":25},{"id":"508d26e1-5754-41a5-88ae-1b70cb9fc14a","cell_type":"code","source":"model = T5ForConditionalGeneration.from_pretrained(cfg.model_name)\n\nprint(\"Params:\", sum(p.numel() for p in model.parameters()) / 1e6, \"M\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-03T13:03:01.940689Z","iopub.execute_input":"2026-04-03T13:03:01.941138Z","iopub.status.idle":"2026-04-03T13:03:02.573573Z","shell.execute_reply.started":"2026-04-03T13:03:01.941109Z","shell.execute_reply":"2026-04-03T13:03:02.572813Z"}},"outputs":[{"name":"stdout","text":"Params: 60.492288 M\n","output_type":"stream"}],"execution_count":26},{"id":"3a5dfcc8-9f91-414a-932f-62fb9d1f3998","cell_type":"code","source":"from transformers import DataCollatorForSeq2Seq\n\nclass FastCollator(DataCollatorForSeq2Seq):\n    def __call__(self, features):\n        for f in features:\n            if \"labels\" in f:\n                f[\"labels\"] = list(f[\"labels\"])  # 🔥 ensure list\n        return super().__call__(features)\n\ndata_collator = FastCollator(\n    tokenizer=tokenizer,\n    model=model,\n    label_pad_token_id=-100,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-03T13:03:02.574574Z","iopub.execute_input":"2026-04-03T13:03:02.575012Z","iopub.status.idle":"2026-04-03T13:03:02.580743Z","shell.execute_reply.started":"2026-04-03T13:03:02.574963Z","shell.execute_reply":"2026-04-03T13:03:02.580006Z"}},"outputs":[],"execution_count":27},{"id":"ad9b3c55-9cb8-4b67-98a4-a85724e3b6b0","cell_type":"code","source":"training_args = Seq2SeqTrainingArguments(\n    output_dir=cfg.output_dir,\n\n    num_train_epochs=cfg.num_train_epochs,\n\n    per_device_train_batch_size=cfg.per_device_train_batch_size,\n    per_device_eval_batch_size=cfg.per_device_eval_batch_size,\n\n    learning_rate=cfg.learning_rate,\n    weight_decay=cfg.weight_decay,\n    warmup_ratio=cfg.warmup_ratio,\n\n    # 🔥 CRITICAL FIXES\n    fp16=False,                     # ❌ disable fp16\n    optim=\"adamw_torch\",           # ❌ no fused optimizer\n\n    predict_with_generate=False,\n\n    logging_steps=200,\n    report_to=\"none\",\n\n    save_strategy=\"epoch\",\n\n    dataloader_num_workers=2,\n    dataloader_pin_memory=True,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-03T13:10:17.704999Z","iopub.execute_input":"2026-04-03T13:10:17.705790Z","iopub.status.idle":"2026-04-03T13:10:17.747022Z","shell.execute_reply.started":"2026-04-03T13:10:17.705752Z","shell.execute_reply":"2026-04-03T13:10:17.746340Z"}},"outputs":[],"execution_count":31},{"id":"c429c4c1-9626-4889-a1f4-124115c31bdc","cell_type":"code","source":"from transformers import TrainerCallback\nfrom tqdm.auto import tqdm\n\nclass ProgressBarCallback(TrainerCallback):\n    def __init__(self):\n        self.pbar = None\n\n    def on_train_begin(self, args, state, control, **kwargs):\n        self.pbar = tqdm(\n            total=state.max_steps,\n            desc=\"🚀 Training\",\n            unit=\"step\"\n        )\n\n    def on_step_end(self, args, state, control, **kwargs):\n        if self.pbar:\n            self.pbar.update(1)\n\n    def on_train_end(self, args, state, control, **kwargs):\n        if self.pbar:\n            self.pbar.close()\n\ntrainer = Seq2SeqTrainer(\n    model=model,\n    args=training_args,\n    train_dataset=tokenized_ds[\"train\"],\n    eval_dataset=tokenized_ds[\"val\"],\n    tokenizer=tokenizer,\n    data_collator=data_collator,\n    callbacks=[ProgressBarCallback()],  # 🔥 ADD THIS\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-03T13:10:17.884885Z","iopub.execute_input":"2026-04-03T13:10:17.885120Z","iopub.status.idle":"2026-04-03T13:10:17.902139Z","shell.execute_reply.started":"2026-04-03T13:10:17.885097Z","shell.execute_reply":"2026-04-03T13:10:17.901497Z"}},"outputs":[],"execution_count":32},{"id":"97318834-7e71-4716-bdb9-3e00aa27570c","cell_type":"code","source":"\nprint(\"🚀 Training...\")\n\nt0 = time.time()\n\ntrainer.train()\n\nprint(f\"✅ Done in {(time.time()-t0)/60:.1f} min\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-03T13:10:19.398505Z","iopub.execute_input":"2026-04-03T13:10:19.399280Z","iopub.status.idle":"2026-04-03T13:22:11.117488Z","shell.execute_reply.started":"2026-04-03T13:10:19.399231Z","shell.execute_reply":"2026-04-03T13:22:11.116791Z"}},"outputs":[{"name":"stdout","text":"🚀 Training...\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"🚀 Training:   0%|          | 0/1915 [00:00<?, ?step/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"f0fef2302c9149f39803f045ffb08556"}},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"<IPython.core.display.HTML object>","text/html":"\n    <div>\n      \n      <progress value='1915' max='1915' style='width:300px; height:20px; vertical-align: middle;'></progress>\n      [1915/1915 11:50, Epoch 5/5]\n    </div>\n    <table border=\"1\" class=\"dataframe\">\n  <thead>\n <tr style=\"text-align: left;\">\n      <th>Step</th>\n      <th>Training Loss</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <td>200</td>\n      <td>0.145300</td>\n    </tr>\n    <tr>\n      <td>400</td>\n      <td>0.149900</td>\n    </tr>\n    <tr>\n      <td>600</td>\n      <td>0.112700</td>\n    </tr>\n    <tr>\n      <td>800</td>\n      <td>0.092400</td>\n    </tr>\n    <tr>\n      <td>1000</td>\n      <td>0.074900</td>\n    </tr>\n    <tr>\n      <td>1200</td>\n      <td>0.073200</td>\n    </tr>\n    <tr>\n      <td>1400</td>\n      <td>0.070000</td>\n    </tr>\n    <tr>\n      <td>1600</td>\n      <td>0.061400</td>\n    </tr>\n    <tr>\n      <td>1800</td>\n      <td>0.044600</td>\n    </tr>\n  </tbody>\n</table><p>"},"metadata":{}},{"name":"stdout","text":"✅ Done in 11.9 min\n","output_type":"stream"}],"execution_count":33},{"id":"77c6ffc2-efae-437f-9f3d-fb9db0a07c76","cell_type":"code","source":"trainer.save_model(cfg.output_dir)\ntokenizer.save_pretrained(cfg.output_dir)\n\nprint(\"✅ Model saved\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-03T13:22:11.118944Z","iopub.execute_input":"2026-04-03T13:22:11.119273Z","iopub.status.idle":"2026-04-03T13:22:11.697695Z","shell.execute_reply.started":"2026-04-03T13:22:11.119244Z","shell.execute_reply":"2026-04-03T13:22:11.696982Z"}},"outputs":[{"name":"stdout","text":"✅ Model saved\n","output_type":"stream"}],"execution_count":34},{"id":"6537f70f-3c8d-444d-9dd8-f563abc2af0c","cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nmodel = model.to(device)\nmodel.eval()\n\ndef optimize_java(code):\n    inputs = tokenizer(\n        TASK_PREFIX + code,\n        return_tensors=\"pt\",\n        truncation=True,\n        max_length=cfg.max_source_length,\n    ).to(device)\n\n    with torch.no_grad():\n        outputs = model.generate(\n            **inputs,\n            max_new_tokens=128,\n        )\n\n    return tokenizer.decode(outputs[0], skip_special_tokens=True)\n\n\n# TEST\nprint(optimize_java(\"\"\"\nString alertLevel;\nswitch (sensor.getValue()) {\n    case 0:\n    case 1:\n        alertLevel = \"GREEN\";\n        break;\n    case 2:\n        alertLevel = \"YELLOW\";\n        break;\n    case 3:\n        alertLevel = \"RED\";\n        break;\n    default:\n        alertLevel = \"CRITICAL\";\n}\n\"\"\"))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-03T14:05:35.081973Z","iopub.execute_input":"2026-04-03T14:05:35.082240Z","iopub.status.idle":"2026-04-03T14:05:35.749507Z","shell.execute_reply.started":"2026-04-03T14:05:35.082214Z","shell.execute_reply":"2026-04-03T14:05:35.748745Z"}},"outputs":[{"name":"stdout","text":"String alertLevel = switch (sensor.getValue()) {\n    case 0, 1 -> \"1\";\n    case 2 -> \"2\";\n    case 3 -> \"3\";\n    default -> \"4\";\n};\n","output_type":"stream"}],"execution_count":37},{"id":"cb063d9a-5425-40cd-89de-c23edd1a60f3","cell_type":"code","source":"print(optimize_java(\"\"\"\nList<Payment> payments = new ArrayList<>(pendingPayments);\nIterator<Payment> iterator = payments.iterator();\nwhile (iterator.hasNext()) {\n    Payment p = iterator.next();\n    if (p.isExpired()) {\n        iterator.remove();\n    }\n}\"\"\"))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-03T14:06:17.153683Z","iopub.execute_input":"2026-04-03T14:06:17.154494Z","iopub.status.idle":"2026-04-03T14:06:17.546252Z","shell.execute_reply.started":"2026-04-03T14:06:17.154454Z","shell.execute_reply":"2026-04-03T14:06:17.545289Z"}},"outputs":[{"name":"stdout","text":"List<Payment> payments = new ArrayList<>(pendingPayments);\npayments.removeIf(Payment::isExpired);\n","output_type":"stream"}],"execution_count":38},{"id":"ae99af62-1c8a-4adc-bdef-7c74c2da7696","cell_type":"code","source":"print(optimize_java(\"\"\"\nString result = \"\";\nfor (int i = 0; i < 100; i++) {\n    result += \"Value: \" + i + \"\\n\";\n}\"\"\"))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-03T14:08:05.318144Z","iopub.execute_input":"2026-04-03T14:08:05.318944Z","iopub.status.idle":"2026-04-03T14:08:06.065047Z","shell.execute_reply.started":"2026-04-03T14:08:05.318907Z","shell.execute_reply":"2026-04-03T14:08:06.064234Z"}},"outputs":[{"name":"stdout","text":"StringBuilder sb = new StringBuilder(100 * 8);\nfor (int i = 0; i < 100; i++) {\n    sb.append(\"Value: \").append(i).append(\"\n\");\n}\nString result = sb.toString();\n","output_type":"stream"}],"execution_count":39}]}