{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e2381cb4",
   "metadata": {},
   "source": [
    "# Fine-tune DiffusionGemma on PubMedQA with Before and After Evaluation\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "533a0d2a",
   "metadata": {},
   "source": [
    "## 1. Check GPU\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bbeca747",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CUDA available: True\n",
      "GPU: NVIDIA H100 80GB HBM3\n",
      "Free GPU memory: 84.5 GB\n",
      "Total GPU memory: 85.0 GB\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "print(\"CUDA available:\", torch.cuda.is_available())\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    print(\"GPU:\", torch.cuda.get_device_name(0))\n",
    "    free_gb, total_gb = torch.cuda.mem_get_info()\n",
    "    print(f\"Free GPU memory: {free_gb / 1e9:.1f} GB\")\n",
    "    print(f\"Total GPU memory: {total_gb / 1e9:.1f} GB\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49f1a81b",
   "metadata": {},
   "source": [
    "## 2. Install Packages\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6b58f698",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%capture\n",
    "# %pip install --upgrade pip wheel setuptools packaging ninja\n",
    "# %pip install unsloth\n",
    "# %pip install --no-deps --upgrade --force-reinstall git+https://github.com/unslothai/unsloth-zoo.git git+https://github.com/unslothai/unsloth.git\n",
    "# %pip install sentencepiece protobuf \"datasets==4.3.0\" \"huggingface_hub>=0.34.0\" hf_transfer\n",
    "# %pip install --no-deps bitsandbytes accelerate peft trl triton\n",
    "# %pip install --no-deps --upgrade \"torchao>=0.16.0\"\n",
    "# %pip install --no-deps transformers==5.11.0 \"tokenizers>=0.22.0,<=0.23.0\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "871befcc",
   "metadata": {},
   "source": [
    "## 3. Import Libraries\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5f38209d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.\n",
      "🦥 Unsloth Zoo will now patch everything to make training faster!\n",
      "Torch: 2.10.0+cu128\n",
      "CUDA available: True\n",
      "GPU: NVIDIA H100 80GB HBM3\n"
     ]
    }
   ],
   "source": [
    "import copy\n",
    "import os\n",
    "import random\n",
    "import time\n",
    "\n",
    "import torch\n",
    "from datasets import load_dataset\n",
    "from unsloth import FastModel\n",
    "\n",
    "os.environ[\"HF_HUB_ENABLE_HF_TRANSFER\"] = \"1\"\n",
    "torch._dynamo.config.recompile_limit = 64\n",
    "\n",
    "print(\"Torch:\", torch.__version__)\n",
    "print(\"CUDA available:\", torch.cuda.is_available())\n",
    "print(\"GPU:\", torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"None\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "157d7536",
   "metadata": {},
   "source": [
    "## 4. Set Config\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5d20e819",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = \"unsloth/diffusiongemma-26B-A4B-it\"\n",
    "DATASET_NAME = \"qiaojin/PubMedQA\"\n",
    "\n",
    "TRAIN_SUBSET = \"pqa_artificial\"\n",
    "EVAL_SUBSET = \"pqa_labeled\"\n",
    "\n",
    "N_TRAIN = 3000\n",
    "N_EVAL = 200\n",
    "\n",
    "MAX_CONTEXT_CHARS = 2500\n",
    "\n",
    "STEPS = 60\n",
    "GRAD_ACCUM = 4\n",
    "LR = 1e-4\n",
    "T_LO = 0.1\n",
    "\n",
    "EVAL_TOTAL = 50\n",
    "EVAL_DENOISING_STEPS = 16\n",
    "\n",
    "OUTPUT_DIR = \"diffusiongemma_pubmedqa_lora\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ccd2a0b",
   "metadata": {},
   "source": [
    "## 5. Load DiffusionGemma\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "13e33bdf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==((  Unsloth: FastDiffusionModel (slow / transformers-only path)  ))==\n",
      "   Model: unsloth/diffusiongemma-26B-A4B-it  | class: DiffusionGemmaForBlockDiffusion  | model_type: diffusion_gemma\n",
      "   dtype: torch.bfloat16 | 4bit: False | 8bit: False | attn: eager\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5c4a128418044a0dbb45bd077a6ac9d2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading weights:   0%|          | 0/1047 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vocab size: 262144\n",
      "Canvas length: 256\n",
      "Model device: cuda:0\n"
     ]
    }
   ],
   "source": [
    "model, tokenizer = FastModel.from_pretrained(\n",
    "    model_name=MODEL_NAME,\n",
    "    dtype=torch.bfloat16,\n",
    "    load_in_4bit=False,\n",
    ")\n",
    "\n",
    "processor = tokenizer\n",
    "tok = processor.tokenizer if hasattr(processor, \"tokenizer\") else processor\n",
    "\n",
    "vocab = model.config.text_config.vocab_size\n",
    "canvas_len = model.config.canvas_length\n",
    "\n",
    "dev = next(\n",
    "    (p.device for p in model.parameters() if p.device.type != \"meta\"),\n",
    "    torch.device(\"cuda\"),\n",
    ")\n",
    "\n",
    "print(\"Vocab size:\", vocab)\n",
    "print(\"Canvas length:\", canvas_len)\n",
    "print(\"Model device:\", dev)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02869124",
   "metadata": {},
   "source": [
    "## 6. Add LoRA\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "590ec647",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 149,630,976 || all params: 25,973,409,840 || trainable%: 0.5761\n"
     ]
    }
   ],
   "source": [
    "model = FastModel.get_peft_model(\n",
    "    model,\n",
    "    r=64,\n",
    "    lora_alpha=128,\n",
    "    use_gradient_checkpointing=False,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "546f0664",
   "metadata": {},
   "source": [
    "## 7. Load PubMedQA\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e1f4b158",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train size: 211269\n",
      "Eval size: 1000\n",
      "{'pubid': 25429730, 'question': 'Are group 2 innate lymphoid cells ( ILC2s ) increased in chronic rhinosinusitis with nasal polyps or eosinophilia?', 'context': {'contexts': ['Chronic rhinosinusitis (CRS) is a heterogeneous disease with an uncertain pathogenesis. Group 2 innate lymphoid cells (ILC2s) represent a recently discovered cell population which has been implicated in driving Th2 inflammation in CRS; however, their relationship with clinical disease characteristics has yet to be investigated.', 'The aim of this study was to identify ILC2s in sinus mucosa in patients with CRS and controls and compare ILC2s across characteristics of disease.', 'A cross-sectional study of patients with CRS undergoing endoscopic sinus surgery was conducted. Sinus mucosal biopsies were obtained during surgery and control tissue from patients undergoing pituitary tumour resection through transphenoidal approach. ILC2s were identified as CD45(+) Lin(-) CD127(+) CD4(-) CD8(-) CRTH2(CD294)(+) CD161(+) cells in single cell suspensions through flow cytometry. ILC2 frequencies, measured as a percentage of CD45(+) cells, were compared across CRS phenotype, endotype, inflammatory CRS subtype and other disease characteristics including blood eosinophils, serum IgE, asthma status and nasal symptom score.', '35 patients (40% female, age 48 ± 17 years) including 13 with eosinophilic CRS (eCRS), 13 with non-eCRS and 9 controls were recruited. ILC2 frequencies were associated with the presence of nasal polyps (P = 0.002) as well as high tissue eosinophilia (P = 0.004) and eosinophil-dominant CRS (P = 0.001) (Mann-Whitney U). They were also associated with increased blood eosinophilia (P = 0.005). There were no significant associations found between ILC2s and serum total IgE and allergic disease. In the CRS with nasal polyps (CRSwNP) population, ILC2s were increased in patients with co-existing asthma (P = 0.03). ILC2s were also correlated with worsening nasal symptom score in CRS (P = 0.04).'], 'labels': ['BACKGROUND', 'OBJECTIVE', 'METHODS', 'RESULTS'], 'meshes': ['Adult', 'Aged', 'Antigens, Surface', 'Case-Control Studies', 'Chronic Disease', 'Eosinophilia', 'Female', 'Humans', 'Hypersensitivity', 'Immunity, Innate', 'Immunoglobulin E', 'Immunophenotyping', 'Leukocyte Count', 'Lymphocyte Subsets', 'Male', 'Middle Aged', 'Nasal Mucosa', 'Nasal Polyps', 'Neutrophil Infiltration', 'Patient Outcome Assessment', 'Rhinitis', 'Sinusitis', 'Young Adult']}, 'long_answer': 'As ILC2s are elevated in patients with CRSwNP, they may drive nasal polyp formation in CRS. ILC2s are also linked with high tissue and blood eosinophilia and have a potential role in the activation and survival of eosinophils during the Th2 immune response. The association of innate lymphoid cells in CRS provides insights into its pathogenesis.', 'final_decision': 'yes'}\n"
     ]
    }
   ],
   "source": [
    "train_data = load_dataset(DATASET_NAME, TRAIN_SUBSET, split=\"train\")\n",
    "eval_data = load_dataset(DATASET_NAME, EVAL_SUBSET, split=\"train\")\n",
    "\n",
    "print(\"Train size:\", len(train_data))\n",
    "print(\"Eval size:\", len(eval_data))\n",
    "print(train_data[0])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f7ff8e7",
   "metadata": {},
   "source": [
    "## 8. Convert Dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d8c697a7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prepared train examples: 3000\n",
      "Prepared eval examples: 200\n",
      "Answer the biomedical research question using only the context.\n",
      "\n",
      "Context:\n",
      "Chronic rhinosinusitis (CRS) is a heterogeneous disease with an uncertain pathogenesis. Group 2 innate lymphoid cells (ILC2s) represent a recently discovered cell population which has been implicated in driving Th2 inflammation in CRS; however, their relationship with clinical disease characteristics has yet to be investigated. The aim of this study was to identify ILC2s in sinus mucosa in patients with CRS and controls and compare ILC2s across characteristics of disease. A cross-sectional study of patients with CRS undergoing endoscopic sinus surgery was conducted. Sinus mucosal biopsies were obtained during surgery and control tissue from patients undergoing pituitary tumour resection through transphenoidal approach. ILC2s were identified as CD45(+) Lin(-) CD127(+) CD4(-) CD8(-) CRTH2(CD294)(+) CD161(+) cells in single cell suspensions through flow cytometry. ILC2 frequencies, measured as a percentage of CD45(+) cells, were compared across CRS phenotype, endotype, inflammatory CRS subtype and other disease characteristics including blood eosinophils, serum IgE, asthma status and nasal symptom score. 35 patients (40% female, age 48 ± 17 years) including 13 with eosinophilic CRS (eCRS), 13 with non-eCRS and 9 controls were recruited. ILC2 frequencies were associated with the presence of nasal polyps (P = 0.002) as well as high tissue eosinophilia (P = 0.004) and eosinophil-dominant CRS (P = 0.001) (Mann-Whitney U). They were also associated with increased blood eosinophilia (P = 0.005). There were no significant associations found between ILC2s and serum total IgE and allergic disease. In the CRS with nasal polyps (CRSwNP) population, ILC2s were increased in patients with co-existing asthma (P = 0.03). ILC2s were also correlated with worsening nasal symptom score in CRS (P = 0.04).\n",
      "\n",
      "Question:\n",
      "Are group 2 innate lymphoid cells ( ILC2s ) increased in chronic rhinosinusitis with nasal polyps or eosinophilia?\n",
      "\n",
      "Answer with only one word: yes, no, or maybe.\n",
      "Answer: yes\n"
     ]
    }
   ],
   "source": [
    "def make_prompt(row):\n",
    "    context = \" \".join(row[\"context\"][\"contexts\"])\n",
    "    context = context[:MAX_CONTEXT_CHARS]\n",
    "    question = row[\"question\"]\n",
    "\n",
    "    return f\"\"\"Answer the biomedical research question using only the context.\n",
    "\n",
    "Context:\n",
    "{context}\n",
    "\n",
    "Question:\n",
    "{question}\n",
    "\n",
    "Answer with only one word: yes, no, or maybe.\"\"\"\n",
    "\n",
    "\n",
    "def make_answer(row):\n",
    "    return row[\"final_decision\"].strip().lower()\n",
    "\n",
    "\n",
    "def convert_row(row):\n",
    "    answer = make_answer(row)\n",
    "\n",
    "    if answer not in [\"yes\", \"no\", \"maybe\"]:\n",
    "        return None\n",
    "\n",
    "    return {\n",
    "        \"messages\": [\n",
    "            {\"role\": \"user\", \"content\": make_prompt(row)},\n",
    "            {\"role\": \"assistant\", \"content\": answer},\n",
    "        ]\n",
    "    }\n",
    "\n",
    "\n",
    "train_rows = []\n",
    "for row in train_data.select(range(N_TRAIN)):\n",
    "    item = convert_row(row)\n",
    "    if item is not None:\n",
    "        train_rows.append(item)\n",
    "\n",
    "eval_rows = []\n",
    "for row in eval_data.select(range(N_EVAL)):\n",
    "    item = convert_row(row)\n",
    "    if item is not None:\n",
    "        eval_rows.append(item)\n",
    "\n",
    "print(\"Prepared train examples:\", len(train_rows))\n",
    "print(\"Prepared eval examples:\", len(eval_rows))\n",
    "print(train_rows[0][\"messages\"][0][\"content\"])\n",
    "print(\"Answer:\", train_rows[0][\"messages\"][1][\"content\"])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b448b84",
   "metadata": {},
   "source": [
    "## 9. Build Diffusion Training Examples\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "df6f8b40",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Usable training examples: 3000\n"
     ]
    }
   ],
   "source": [
    "eos = model.generation_config.eos_token_id or [1]\n",
    "eos = eos[0] if isinstance(eos, (list, tuple)) else eos\n",
    "\n",
    "pad = tok.pad_token_id if tok.pad_token_id is not None else eos\n",
    "\n",
    "\n",
    "def build_examples(rows):\n",
    "    examples = []\n",
    "\n",
    "    for row in rows:\n",
    "        user_message = row[\"messages\"][0]\n",
    "        assistant_message = row[\"messages\"][1]\n",
    "\n",
    "        prompt_ids = processor.apply_chat_template(\n",
    "            [user_message],\n",
    "            tokenize=True,\n",
    "            add_generation_prompt=True,\n",
    "            return_tensors=\"pt\",\n",
    "        )[0]\n",
    "\n",
    "        answer_ids = tok.encode(\n",
    "            assistant_message[\"content\"],\n",
    "            add_special_tokens=False,\n",
    "        )\n",
    "\n",
    "        content = answer_ids + [eos]\n",
    "        n = len(content)\n",
    "\n",
    "        if n > canvas_len:\n",
    "            continue\n",
    "\n",
    "        x0 = torch.tensor(\n",
    "            content + [pad] * (canvas_len - n),\n",
    "            dtype=torch.long,\n",
    "        )\n",
    "\n",
    "        loss_mask = torch.zeros(canvas_len, dtype=torch.bool)\n",
    "        loss_mask[:n] = True\n",
    "\n",
    "        examples.append((prompt_ids, x0, loss_mask))\n",
    "\n",
    "    return examples\n",
    "\n",
    "\n",
    "examples = build_examples(train_rows)\n",
    "print(\"Usable training examples:\", len(examples))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "077c118d",
   "metadata": {},
   "source": [
    "## 10. Inference and Evaluation Functions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "107d2d21",
   "metadata": {},
   "outputs": [],
   "source": [
    "def answer_question(prompt, steps=64):\n",
    "    input_ids = processor.apply_chat_template(\n",
    "        [{\"role\": \"user\", \"content\": prompt}],\n",
    "        tokenize=True,\n",
    "        add_generation_prompt=True,\n",
    "        return_tensors=\"pt\",\n",
    "    ).to(dev)\n",
    "\n",
    "    gen_config = copy.deepcopy(model.generation_config)\n",
    "    gen_config.max_denoising_steps = steps\n",
    "    gen_config.max_new_tokens = canvas_len\n",
    "\n",
    "    model.eval()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        output = model.generate(\n",
    "            input_ids=input_ids,\n",
    "            generation_config=gen_config,\n",
    "        )\n",
    "\n",
    "    generated = output.sequences[0, input_ids.shape[1]:]\n",
    "    text = tok.decode(generated.tolist(), skip_special_tokens=True)\n",
    "    return text.strip().lower()\n",
    "\n",
    "\n",
    "def clean_prediction(text):\n",
    "    text = text.lower().strip()\n",
    "\n",
    "    if text.startswith(\"yes\"):\n",
    "        return \"yes\"\n",
    "    if text.startswith(\"no\"):\n",
    "        return \"no\"\n",
    "    if text.startswith(\"maybe\"):\n",
    "        return \"maybe\"\n",
    "\n",
    "    words = text.replace(\".\", \" \").replace(\",\", \" \").split()\n",
    "\n",
    "    for word in words:\n",
    "        if word in [\"yes\", \"no\", \"maybe\"]:\n",
    "            return word\n",
    "\n",
    "    return \"unknown\"\n",
    "\n",
    "\n",
    "def evaluate_model(rows, total=50, steps=64, title=\"Evaluation\"):\n",
    "    correct = 0\n",
    "    results = []\n",
    "    total = min(total, len(rows))\n",
    "\n",
    "    print(title)\n",
    "    print(\"-\" * len(title))\n",
    "\n",
    "    for i, row in enumerate(rows[:total], start=1):\n",
    "        prompt = row[\"messages\"][0][\"content\"]\n",
    "        gold = row[\"messages\"][1][\"content\"]\n",
    "\n",
    "        raw_pred = answer_question(prompt, steps=steps)\n",
    "        pred = clean_prediction(raw_pred)\n",
    "\n",
    "        is_correct = pred == gold\n",
    "        correct += int(is_correct)\n",
    "\n",
    "        results.append({\n",
    "            \"index\": i,\n",
    "            \"gold\": gold,\n",
    "            \"prediction\": pred,\n",
    "            \"raw_prediction\": raw_pred,\n",
    "            \"correct\": is_correct,\n",
    "        })\n",
    "\n",
    "        print(f\"{i:02d}. Gold: {gold} | Pred: {pred} | Correct: {is_correct}\")\n",
    "\n",
    "    accuracy = correct / total if total else 0\n",
    "\n",
    "    print()\n",
    "    print(\"Accuracy:\", accuracy)\n",
    "    print()\n",
    "\n",
    "    return {\n",
    "        \"accuracy\": accuracy,\n",
    "        \"correct\": correct,\n",
    "        \"total\": total,\n",
    "        \"results\": results,\n",
    "    }\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e779832f",
   "metadata": {},
   "source": [
    "## 11. Evaluate Before Fine-Tuning\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4ac8ca3d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before Fine-Tuning Evaluation\n",
      "-----------------------------\n",
      "01. Gold: yes | Pred: yes | Correct: True\n",
      "02. Gold: no | Pred: yes | Correct: False\n",
      "03. Gold: yes | Pred: maybe | Correct: False\n",
      "04. Gold: no | Pred: maybe | Correct: False\n",
      "05. Gold: yes | Pred: yes | Correct: True\n",
      "06. Gold: yes | Pred: maybe | Correct: False\n",
      "07. Gold: maybe | Pred: yes | Correct: False\n",
      "08. Gold: no | Pred: yes | Correct: False\n",
      "09. Gold: no | Pred: no | Correct: True\n",
      "10. Gold: yes | Pred: yes | Correct: True\n",
      "11. Gold: yes | Pred: yes | Correct: True\n",
      "12. Gold: no | Pred: maybe | Correct: False\n",
      "13. Gold: yes | Pred: yes | Correct: True\n",
      "14. Gold: no | Pred: no | Correct: True\n",
      "15. Gold: yes | Pred: yes | Correct: True\n",
      "16. Gold: yes | Pred: yes | Correct: True\n",
      "17. Gold: yes | Pred: maybe | Correct: False\n",
      "18. Gold: yes | Pred: yes | Correct: True\n",
      "19. Gold: yes | Pred: yes | Correct: True\n",
      "20. Gold: yes | Pred: yes | Correct: True\n",
      "21. Gold: yes | Pred: maybe | Correct: False\n",
      "22. Gold: yes | Pred: yes | Correct: True\n",
      "23. Gold: yes | Pred: yes | Correct: True\n",
      "24. Gold: yes | Pred: yes | Correct: True\n",
      "25. Gold: yes | Pred: maybe | Correct: False\n",
      "26. Gold: no | Pred: no | Correct: True\n",
      "27. Gold: yes | Pred: no | Correct: False\n",
      "28. Gold: maybe | Pred: yes | Correct: False\n",
      "29. Gold: yes | Pred: yes | Correct: True\n",
      "30. Gold: yes | Pred: yes | Correct: True\n",
      "31. Gold: no | Pred: maybe | Correct: False\n",
      "32. Gold: maybe | Pred: yes | Correct: False\n",
      "33. Gold: no | Pred: maybe | Correct: False\n",
      "34. Gold: yes | Pred: yes | Correct: True\n",
      "35. Gold: yes | Pred: yes | Correct: True\n",
      "36. Gold: no | Pred: no | Correct: True\n",
      "37. Gold: no | Pred: maybe | Correct: False\n",
      "38. Gold: no | Pred: yes | Correct: False\n",
      "39. Gold: yes | Pred: yes | Correct: True\n",
      "40. Gold: no | Pred: no | Correct: True\n",
      "41. Gold: yes | Pred: yes | Correct: True\n",
      "42. Gold: yes | Pred: yes | Correct: True\n",
      "43. Gold: maybe | Pred: maybe | Correct: True\n",
      "44. Gold: no | Pred: maybe | Correct: False\n",
      "45. Gold: yes | Pred: maybe | Correct: False\n",
      "46. Gold: yes | Pred: yes | Correct: True\n",
      "47. Gold: yes | Pred: yes | Correct: True\n",
      "48. Gold: yes | Pred: yes | Correct: True\n",
      "49. Gold: yes | Pred: yes | Correct: True\n",
      "50. Gold: no | Pred: maybe | Correct: False\n",
      "\n",
      "Accuracy: 0.6\n",
      "\n"
     ]
    }
   ],
   "source": [
    "before_eval = evaluate_model(\n",
    "    eval_rows,\n",
    "    total=EVAL_TOTAL,\n",
    "    steps=EVAL_DENOISING_STEPS,\n",
    "    title=\"Before Fine-Tuning Evaluation\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adb680c8",
   "metadata": {},
   "source": [
    "## 12. Set Up Training\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a9ccf9b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.config.use_cache = True\n",
    "model.train()\n",
    "\n",
    "opt = torch.optim.AdamW(\n",
    "    [p for p in model.parameters() if p.requires_grad],\n",
    "    lr=LR,\n",
    "    betas=(0.9, 0.95),\n",
    "    weight_decay=0.0,\n",
    ")\n",
    "\n",
    "sched = torch.optim.lr_scheduler.OneCycleLR(\n",
    "    opt,\n",
    "    max_lr=LR,\n",
    "    total_steps=STEPS,\n",
    "    pct_start=0.03,\n",
    "    anneal_strategy=\"cos\",\n",
    ")\n",
    "\n",
    "\n",
    "def corrupt(x0):\n",
    "    noise_level = random.uniform(T_LO, 1.0)\n",
    "    xt = x0.to(dev).clone()\n",
    "    noise_mask = torch.rand(canvas_len, device=dev) < noise_level\n",
    "    xt[noise_mask] = torch.randint(0, vocab, (canvas_len,), device=dev)[noise_mask]\n",
    "    return xt.unsqueeze(0)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d44af4fa",
   "metadata": {},
   "source": [
    "## 13. Train\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ea5e6bcf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 20/60 | loss 0.0019 | 43s\n",
      "step 40/60 | loss 0.0003 | 85s\n",
      "step 60/60 | loss 0.0001 | 126s\n"
     ]
    }
   ],
   "source": [
    "order = list(range(len(examples)))\n",
    "ptr = 0\n",
    "start_time = time.time()\n",
    "\n",
    "opt.zero_grad(set_to_none=True)\n",
    "\n",
    "for step in range(1, STEPS + 1):\n",
    "    step_loss = 0.0\n",
    "\n",
    "    for _ in range(GRAD_ACCUM):\n",
    "        if ptr >= len(order):\n",
    "            random.shuffle(order)\n",
    "            ptr = 0\n",
    "\n",
    "        prompt_ids, x0, loss_mask = examples[order[ptr]]\n",
    "        ptr += 1\n",
    "\n",
    "        output = model(\n",
    "            input_ids=prompt_ids.unsqueeze(0).to(dev),\n",
    "            canvas_ids=corrupt(x0),\n",
    "            self_conditioning_logits=None,\n",
    "        )\n",
    "\n",
    "        logits = output.logits[0].float()\n",
    "        mask = loss_mask.to(dev)\n",
    "\n",
    "        loss = torch.nn.functional.cross_entropy(\n",
    "            logits[mask],\n",
    "            x0.to(dev)[mask],\n",
    "        )\n",
    "\n",
    "        (loss / GRAD_ACCUM).backward()\n",
    "        step_loss += loss.item() / GRAD_ACCUM\n",
    "\n",
    "    torch.nn.utils.clip_grad_norm_(\n",
    "        [p for p in model.parameters() if p.requires_grad],\n",
    "        1.0,\n",
    "    )\n",
    "\n",
    "    opt.step()\n",
    "    sched.step()\n",
    "    opt.zero_grad(set_to_none=True)\n",
    "\n",
    "    if step % 20 == 0:\n",
    "        elapsed = time.time() - start_time\n",
    "        print(f\"step {step}/{STEPS} | loss {step_loss:.4f} | {elapsed:.0f}s\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce4f69d9",
   "metadata": {},
   "source": [
    "## 14. Evaluate After Fine-Tuning\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2117d937",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "After Fine-Tuning Evaluation\n",
      "----------------------------\n",
      "01. Gold: yes | Pred: yes | Correct: True\n",
      "02. Gold: no | Pred: yes | Correct: False\n",
      "03. Gold: yes | Pred: yes | Correct: True\n",
      "04. Gold: no | Pred: yes | Correct: False\n",
      "05. Gold: yes | Pred: yes | Correct: True\n",
      "06. Gold: yes | Pred: yes | Correct: True\n",
      "07. Gold: maybe | Pred: yes | Correct: False\n",
      "08. Gold: no | Pred: yes | Correct: False\n",
      "09. Gold: no | Pred: no | Correct: True\n",
      "10. Gold: yes | Pred: yes | Correct: True\n",
      "11. Gold: yes | Pred: yes | Correct: True\n",
      "12. Gold: no | Pred: no | Correct: True\n",
      "13. Gold: yes | Pred: yes | Correct: True\n",
      "14. Gold: no | Pred: no | Correct: True\n",
      "15. Gold: yes | Pred: yes | Correct: True\n",
      "16. Gold: yes | Pred: yes | Correct: True\n",
      "17. Gold: yes | Pred: yes | Correct: True\n",
      "18. Gold: yes | Pred: yes | Correct: True\n",
      "19. Gold: yes | Pred: yes | Correct: True\n",
      "20. Gold: yes | Pred: yes | Correct: True\n",
      "21. Gold: yes | Pred: yes | Correct: True\n",
      "22. Gold: yes | Pred: yes | Correct: True\n",
      "23. Gold: yes | Pred: yes | Correct: True\n",
      "24. Gold: yes | Pred: yes | Correct: True\n",
      "25. Gold: yes | Pred: yes | Correct: True\n",
      "26. Gold: no | Pred: no | Correct: True\n",
      "27. Gold: yes | Pred: yes | Correct: True\n",
      "28. Gold: maybe | Pred: yes | Correct: False\n",
      "29. Gold: yes | Pred: yes | Correct: True\n",
      "30. Gold: yes | Pred: yes | Correct: True\n",
      "31. Gold: no | Pred: no | Correct: True\n",
      "32. Gold: maybe | Pred: yes | Correct: False\n",
      "33. Gold: no | Pred: no | Correct: True\n",
      "34. Gold: yes | Pred: yes | Correct: True\n",
      "35. Gold: yes | Pred: yes | Correct: True\n",
      "36. Gold: no | Pred: yes | Correct: False\n",
      "37. Gold: no | Pred: yes | Correct: False\n",
      "38. Gold: no | Pred: yes | Correct: False\n",
      "39. Gold: yes | Pred: yes | Correct: True\n",
      "40. Gold: no | Pred: no | Correct: True\n",
      "41. Gold: yes | Pred: yes | Correct: True\n",
      "42. Gold: yes | Pred: yes | Correct: True\n",
      "43. Gold: maybe | Pred: no | Correct: False\n",
      "44. Gold: no | Pred: no | Correct: True\n",
      "45. Gold: yes | Pred: yes | Correct: True\n",
      "46. Gold: yes | Pred: yes | Correct: True\n",
      "47. Gold: yes | Pred: yes | Correct: True\n",
      "48. Gold: yes | Pred: yes | Correct: True\n",
      "49. Gold: yes | Pred: yes | Correct: True\n",
      "50. Gold: no | Pred: no | Correct: True\n",
      "\n",
      "Accuracy: 0.8\n",
      "\n"
     ]
    }
   ],
   "source": [
    "after_eval = evaluate_model(\n",
    "    eval_rows,\n",
    "    total=EVAL_TOTAL,\n",
    "    steps=EVAL_DENOISING_STEPS,\n",
    "    title=\"After Fine-Tuning Evaluation\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4a8273e",
   "metadata": {},
   "source": [
    "## 15. Compare Before and After\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "aa62e0f3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before fine-tuning accuracy: 0.6\n",
      "After fine-tuning accuracy: 0.8\n",
      "Improvement: 0.20000000000000007\n"
     ]
    }
   ],
   "source": [
    "before_accuracy = before_eval[\"accuracy\"]\n",
    "after_accuracy = after_eval[\"accuracy\"]\n",
    "improvement = after_accuracy - before_accuracy\n",
    "\n",
    "print(\"Before fine-tuning accuracy:\", before_accuracy)\n",
    "print(\"After fine-tuning accuracy:\", after_accuracy)\n",
    "print(\"Improvement:\", improvement)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9800fdcd",
   "metadata": {},
   "source": [
    "## 16. Save Adapter\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9e30c8e7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved LoRA adapter to: diffusiongemma_pubmedqa_lora\n"
     ]
    }
   ],
   "source": [
    "model.save_pretrained(OUTPUT_DIR)\n",
    "processor.save_pretrained(OUTPUT_DIR)\n",
    "\n",
    "print(f\"Saved LoRA adapter to: {OUTPUT_DIR}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c900a6a",
   "metadata": {},
   "source": [
    "## 17. Push to Hugging Face\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "d63ce717",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No files have been modified since last commit. Skipping to prevent empty commit.\n",
      "[huggingface_hub.hf_api|WARNING]No files have been modified since last commit. Skipping to prevent empty commit.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/kingabzpro/diffusiongemma_pubmedqa/commit/080c9609c8ac80355bb27e6ce1e9e6478b297fcb', commit_message='Upload processor', commit_description='', oid='080c9609c8ac80355bb27e6ce1e9e6478b297fcb', pr_url=None, repo_url=RepoUrl('https://huggingface.co/kingabzpro/diffusiongemma_pubmedqa', endpoint='https://huggingface.co', repo_type='model', repo_id='kingabzpro/diffusiongemma_pubmedqa'), pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "model.push_to_hub(\"kingabzpro/diffusiongemma_pubmedqa\")\n",
    "processor.push_to_hub(\"kingabzpro/diffusiongemma_pubmedqa\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}