Upload opengraph_image_gemma4_e4b_finetune_starter.ipynb

opengraph_image_gemma4_e4b_finetune_starter.ipynb

@@ -0,0 +1,539 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "e681502a",
+   "metadata": {},
+   "source": [
+    "# OpenGraph-Image — Fine-Tuning Gemma 4 E4B (QLoRA) — Starter Notebook\n",
+    "\n",
+    "This notebook fine-tunes **Gemma 4 E4B** to turn an image into a valid `graph.json` knowledge graph, using QLoRA — following [Google's official Gemma 4 vision QLoRA guide](https://ai.google.dev/gemma/docs/core/huggingface_vision_finetune_qlora).\n",
+    "\n",
+    "**Before you run anything, confirm the guidebook prerequisites are done:**\n",
+    "\n",
+    "1. You created / joined the `OpenGraphAI` Hugging Face org.\n",
+    "2. You accepted the **Gemma 4 license** on the [E4B model page](https://huggingface.co/google/gemma-4-E4B-it) while logged in.\n",
+    "3. This Colab session has an **L4 GPU** selected: `Runtime → Change runtime type → L4`.\n",
+    "4. Your **write-scope HF token** is stored in Colab Secrets (🔑 icon) as `HF_TOKEN`.\n",
+    "\n",
+    "> **Placeholders:** anything written as `<<LIKE_THIS>>` or `[More Information Needed]` is for the team to fill in. Search the notebook for `<<` to find them all.\n",
+    "\n",
+    "> **Note on hyperparameters:** the values here are Google's published defaults. **Do not change them on your first real run** — get a baseline number first, then tune."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "58dcb06f",
+   "metadata": {},
+   "source": [
+    "## 1. Install dependencies\n",
+    "\n",
+    "Installs the training stack. `transformers>=5.10.1` is required "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2403b734",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install -q -U \"transformers>=5.10.1\" \"trl\" \"peft\" \"bitsandbytes\" \"accelerate\" \"datasets\" \"Pillow\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0d970981",
+   "metadata": {},
+   "source": [
+    "## 2. Authenticate with Hugging Face\n",
+    "\n",
+    "Pulls your `HF_TOKEN` from Colab Secrets and logs in. This is what lets you (a) download the gated Gemma model and (b) push your trained adapter back to the org. If you haven't added the secret yet: click the 🔑 icon in the left sidebar, add `HF_TOKEN`, and toggle \"Notebook access\" on."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ed661ea8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from google.colab import userdata\n",
+    "from huggingface_hub import login\n",
+    "\n",
+    "login(token=userdata.get(\"HF_TOKEN\"))\n",
+    "print(\"Logged in to Hugging Face.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9ebeb277",
+   "metadata": {},
+   "source": [
+    "## 3. Smoke test — is this GPU actually usable?\n",
+    "\n",
+    "Gemma 4 QLoRA training needs a **bfloat16-capable GPU** (NVIDIA Ampere or newer). Compute capability `>= 8` means Ampere+ (L4 = 8.9, A100 = 8.0). A T4 (7.5) will fail this check — reconnect to an L4 runtime if it does."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "00c2d692",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "\n",
+    "assert torch.cuda.is_available(), \"No GPU. Runtime -> Change runtime type -> L4.\"\n",
+    "name = torch.cuda.get_device_name(0)\n",
+    "cc = torch.cuda.get_device_capability()\n",
+    "assert cc[0] >= 8, f\"GPU {name} (cc {cc}) lacks bf16. Reconnect for an L4 or A100.\"\n",
+    "print(f\"GPU OK: {name}  (compute capability {cc[0]}.{cc[1]})\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a46ddabc",
+   "metadata": {},
+   "source": [
+    "## 4. Configuration\n",
+    "\n",
+    "All the knobs in one place. Fill the placeholders, then the rest of the notebook reads from these constants."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e26190c6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# --- Models -----------------------------------------------------------------\n",
+    "BASE_MODEL_ID = \"google/gemma-4-E4B\"       # the model we TRAIN (base, not -it)\n",
+    "PROCESSOR_ID  = \"google/gemma-4-E4B-it\"    # processor gives us the official chat template\n",
+    "\n",
+    "# --- Data -------------------------------------------------------------------\n",
+    "DATASET_ID = \"OpenGraphAI/opengraph-image-gold-v1\"   # <<CONFIRM_OR_CHANGE>> the gold-pair dataset on the Hub\n",
+    "\n",
+    "# --- Output -----------------------------------------------------------------\n",
+    "OUTPUT_REPO = \"OpenGraphAI/opengraph-image-gemma4-e4b-v1\"   # <<CONFIRM>> where the adapter is pushed\n",
+    "\n",
+    "# --- The extraction system prompt (must match production / extract.py) ------\n",
+    "SYSTEM_PROMPT = \"\"\"<<PASTE_THE_OPENGRAPH_EXTRACTION_SYSTEM_PROMPT_HERE>>\"\"\"\n",
+    "\n",
+    "USER_INSTRUCTION = \"Extract the knowledge graph from this image.\"\n",
+    "\n",
+    "# --- Quick-start / debug switch ---------------------------------------------\n",
+    "# Set True for the \"throwaway plumbing run\" (guidebook step 4.1): trains on a\n",
+    "# tiny slice for a few minutes to shake out bugs before the real run.\n",
+    "SMOKE_TEST = True\n",
+    "SMOKE_TEST_N = 100"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "13bf02c8",
+   "metadata": {},
+   "source": [
+    "## 5. Load the dataset\n",
+    "\n",
+    "Loads the gold pairs from the Hub. Each row is expected to be a chat-format example with an image in the user turn and the gold `graph.json` (as a string) in the assistant turn — the structure defined in **guidebook step 2.5**.\n",
+    "\n",
+    "If `SMOKE_TEST` is on, we shrink the training split so a full pass takes minutes, not hours."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4c366b48",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import load_dataset\n",
+    "\n",
+    "dataset = load_dataset(DATASET_ID)\n",
+    "print(dataset)\n",
+    "\n",
+    "train_ds = dataset[\"train\"]\n",
+    "val_ds   = dataset[\"validation\"]\n",
+    "\n",
+    "if SMOKE_TEST:\n",
+    "    train_ds = train_ds.select(range(min(SMOKE_TEST_N, len(train_ds))))\n",
+    "    val_ds   = val_ds.select(range(min(16, len(val_ds))))\n",
+    "    print(f\"SMOKE_TEST on: training on {len(train_ds)} examples, validating on {len(val_ds)}.\")\n",
+    "\n",
+    "# Peek at one example's structure (without dumping a whole base64 image)\n",
+    "example = train_ds[0]\n",
+    "print(\"\\nExample keys:\", list(example.keys()))\n",
+    "print(\"Roles in messages:\", [m[\"role\"] for m in example[\"messages\"]])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2d89d108",
+   "metadata": {},
+   "source": [
+    "## 6. Load the model + processor (4-bit)\n",
+    "\n",
+    "Loads Gemma 4 E4B in **4-bit NF4** quantization — this is the \"Q\" in QLoRA, and it's what makes an ~8B-parameter model fit in an L4's 24 GB. The frozen 4-bit model is what we attach LoRA adapters to in the next step."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "210c7259",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import AutoModelForImageTextToText, AutoProcessor, BitsAndBytesConfig\n",
+    "import torch\n",
+    "\n",
+    "bnb_config = BitsAndBytesConfig(\n",
+    "    load_in_4bit=True,\n",
+    "    bnb_4bit_use_double_quant=True,\n",
+    "    bnb_4bit_quant_type=\"nf4\",\n",
+    "    bnb_4bit_compute_dtype=torch.bfloat16,\n",
+    ")\n",
+    "\n",
+    "model = AutoModelForImageTextToText.from_pretrained(\n",
+    "    BASE_MODEL_ID,\n",
+    "    quantization_config=bnb_config,\n",
+    "    dtype=torch.bfloat16,\n",
+    "    device_map=\"auto\",\n",
+    "    attn_implementation=\"eager\",   # recommended for Gemma-family stability; try \"flash_attention_2\" later for speed\n",
+    ")\n",
+    "\n",
+    "processor = AutoProcessor.from_pretrained(PROCESSOR_ID)\n",
+    "print(\"Model and processor loaded.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "40028d2f",
+   "metadata": {},
+   "source": [
+    "## 7. The vision collate function\n",
+    "\n",
+    "This is the one piece you should **copy verbatim from the [official guide](https://ai.google.dev/gemma/docs/core/huggingface_vision_finetune_qlora)** rather than trust from memory — the exact image-token IDs used for masking are Gemma-version-specific, and getting them wrong silently corrupts training.\n",
+    "\n",
+    "What it does: turns a batch of chat-format examples (text + PIL images) into model tensors, then builds `labels` and **masks out** the pad tokens and image-placeholder tokens so the model is only graded on generating the graph text, never on \"predicting\" image pixels or padding."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d95d1af6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from PIL import Image\n",
+    "\n",
+    "def process_vision_info(messages):\n",
+    "    \"\"\"Pull the PIL image(s) out of a chat-format message list.\"\"\"\n",
+    "    images = []\n",
+    "    for msg in messages:\n",
+    "        content = msg.get(\"content\")\n",
+    "        if isinstance(content, list):\n",
+    "            for part in content:\n",
+    "                if part.get(\"type\") == \"image\":\n",
+    "                    img = part[\"image\"]\n",
+    "                    if not isinstance(img, Image.Image):\n",
+    "                        img = Image.open(img)\n",
+    "                    images.append(img.convert(\"RGB\"))\n",
+    "    return images\n",
+    "\n",
+    "\n",
+    "def collate_fn(examples):\n",
+    "    texts, images = [], []\n",
+    "    for ex in examples:\n",
+    "        texts.append(\n",
+    "            processor.apply_chat_template(\n",
+    "                ex[\"messages\"], add_generation_prompt=False, tokenize=False\n",
+    "            ).strip()\n",
+    "        )\n",
+    "        images.append(process_vision_info(ex[\"messages\"]))\n",
+    "\n",
+    "    batch = processor(text=texts, images=images, return_tensors=\"pt\", padding=True)\n",
+    "\n",
+    "    labels = batch[\"input_ids\"].clone()\n",
+    "    labels[labels == processor.tokenizer.pad_token_id] = -100\n",
+    "\n",
+    "    # --- <<VERIFY AGAINST THE GUIDE>> ---------------------------------------\n",
+    "    # Mask the image placeholder tokens so they don't contribute to the loss.\n",
+    "    # The exact special token differs by Gemma version; resolve it from the\n",
+    "    # processor rather than hardcoding an ID. If training loss looks wrong,\n",
+    "    # this block is the first place to check against the official guide.\n",
+    "    special = getattr(processor.tokenizer, \"special_tokens_map\", {})\n",
+    "    boi = special.get(\"boi_token\")\n",
+    "    if boi is not None:\n",
+    "        image_token_id = processor.tokenizer.convert_tokens_to_ids(boi)\n",
+    "        labels[labels == image_token_id] = -100\n",
+    "    # ------------------------------------------------------------------------\n",
+    "\n",
+    "    batch[\"labels\"] = labels\n",
+    "    return batch\n",
+    "\n",
+    "\n",
+    "# Sanity check on one batch\n",
+    "_batch = collate_fn([train_ds[0]])\n",
+    "print(\"Batch keys:\", list(_batch.keys()))\n",
+    "print(\"input_ids shape:\", _batch[\"input_ids\"].shape)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "806ab79f",
+   "metadata": {},
+   "source": [
+    "## 8. LoRA configuration\n",
+    "\n",
+    "Defines the small trainable adapter. Straight from the guide: rank 16, attached to every linear layer, plus the embedding/head layers (`modules_to_save`) which Gemma needs for its special tokens."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "195178d8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from peft import LoraConfig\n",
+    "\n",
+    "peft_config = LoraConfig(\n",
+    "    r=16,\n",
+    "    lora_alpha=16,\n",
+    "    lora_dropout=0.05,\n",
+    "    bias=\"none\",\n",
+    "    target_modules=\"all-linear\",\n",
+    "    task_type=\"CAUSAL_LM\",\n",
+    "    modules_to_save=[\"lm_head\", \"embed_tokens\"],\n",
+    "    ensure_weight_tying=True,\n",
+    ")\n",
+    "print(\"LoRA config ready.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "65324603",
+   "metadata": {},
+   "source": [
+    "## 9. Training arguments\n",
+    "\n",
+    "The training recipe (`SFTConfig`). Defaults from the QLoRA paper / Gemma guide. `push_to_hub=True` auto-uploads the adapter to `OUTPUT_REPO`. `report_to=\"tensorboard\"` gives you the loss curves for the next step."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6800710e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from trl import SFTConfig\n",
+    "\n",
+    "training_args = SFTConfig(\n",
+    "    output_dir=OUTPUT_REPO.split(\"/\")[-1],\n",
+    "    num_train_epochs=1 if SMOKE_TEST else 3,\n",
+    "    per_device_train_batch_size=1,\n",
+    "    gradient_accumulation_steps=4,\n",
+    "    learning_rate=2e-4,\n",
+    "    max_grad_norm=0.3,\n",
+    "    lr_scheduler_type=\"constant\",\n",
+    "    warmup_ratio=0.03,\n",
+    "    bf16=True,\n",
+    "    logging_steps=5,\n",
+    "    save_strategy=\"epoch\",\n",
+    "    eval_strategy=\"epoch\",\n",
+    "    push_to_hub=not SMOKE_TEST,          # don't pollute the Hub with smoke-test runs\n",
+    "    hub_model_id=OUTPUT_REPO,\n",
+    "    report_to=\"tensorboard\",\n",
+    "    gradient_checkpointing=True,\n",
+    "    gradient_checkpointing_kwargs={\"use_reentrant\": False},\n",
+    "    dataset_kwargs={\"skip_prepare_dataset\": True},\n",
+    "    remove_unused_columns=False,\n",
+    ")\n",
+    "print(\"Training args ready. push_to_hub =\", training_args.push_to_hub)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b22d808f",
+   "metadata": {},
+   "source": [
+    "## 10. Train\n",
+    "\n",
+    "Wires everything into `SFTTrainer` and runs the loop. On a smoke test this is a few minutes; a full run on the real dataset is 1–2 hours on an L4.\n",
+    "\n",
+    "While it runs, launch TensorBoard (next cell) to watch **eval loss** — it should trend down. If it doesn't, the usual culprit is the dataset format (guidebook step 2.5) or the image-token masking in the collate function, **not** the learning rate."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5a42432c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from trl import SFTTrainer\n",
+    "\n",
+    "trainer = SFTTrainer(\n",
+    "    model=model,\n",
+    "    args=training_args,\n",
+    "    train_dataset=train_ds,\n",
+    "    eval_dataset=val_ds,\n",
+    "    peft_config=peft_config,\n",
+    "    processing_class=processor,\n",
+    "    data_collator=collate_fn,\n",
+    ")\n",
+    "\n",
+    "trainer.train()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1ec4c350",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Optional: view loss curves (run in a separate cell during/after training)\n",
+    "%load_ext tensorboard\n",
+    "%tensorboard --logdir {training_args.output_dir}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3dcdaa9f",
+   "metadata": {},
+   "source": [
+    "## 11. Save & push the adapter\n",
+    "\n",
+    "Saves the trained LoRA adapter. For a real (non-smoke) run with `push_to_hub=True`, this also uploads it to `OUTPUT_REPO` on the Hub."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6132caa7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer.save_model()\n",
+    "\n",
+    "if training_args.push_to_hub:\n",
+    "    trainer.push_to_hub()\n",
+    "    print(f\"Adapter pushed to https://huggingface.co/{OUTPUT_REPO}\")\n",
+    "else:\n",
+    "    print(f\"Adapter saved locally to ./{training_args.output_dir} (smoke test — not pushed).\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "308bc33d",
+   "metadata": {},
+   "source": [
+    "## 12. Quick inference test\n",
+    "\n",
+    "Loads the fine-tuned adapter and runs it on one image so you can eyeball the output before committing to a full evaluation. **This is a smell test, not evaluation** — real numbers come from the eval harness (Phase 3), not from staring at one output."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d2eb3a7c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from PIL import Image\n",
+    "\n",
+    "# Reload base + adapter for a clean-ish inference path\n",
+    "model.eval()\n",
+    "\n",
+    "# <<REPLACE>> with a path to a real test image, or pull one from the test split\n",
+    "test_image = Image.open(\"<<PATH_TO_A_TEST_IMAGE.jpg>>\").convert(\"RGB\")\n",
+    "\n",
+    "messages = [\n",
+    "    {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
+    "    {\"role\": \"user\", \"content\": [\n",
+    "        {\"type\": \"image\", \"image\": test_image},\n",
+    "        {\"type\": \"text\", \"text\": USER_INSTRUCTION},\n",
+    "    ]},\n",
+    "]\n",
+    "\n",
+    "inputs = processor.apply_chat_template(\n",
+    "    messages, add_generation_prompt=True, tokenize=True,\n",
+    "    return_dict=True, return_tensors=\"pt\",\n",
+    ").to(model.device)\n",
+    "\n",
+    "import torch\n",
+    "with torch.no_grad():\n",
+    "    out = model.generate(**inputs, max_new_tokens=2048, do_sample=False)\n",
+    "\n",
+    "generated = processor.decode(out[0][inputs[\"input_ids\"].shape[-1]:], skip_special_tokens=True)\n",
+    "print(generated)\n",
+    "\n",
+    "# Validate against the real schema (guidebook: reuse validate_extraction from the repo)\n",
+    "# from opengraph_image.schema import validate_extraction\n",
+    "# errors = validate_extraction(generated)\n",
+    "# print(\"Schema-valid!\" if not errors else f\"Schema errors: {errors}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "afe564a8",
+   "metadata": {},
+   "source": [
+    "## 13. (Optional) Merge adapter into base weights\n",
+    "\n",
+    "Only needed if you want a single standalone model for serving with vLLM (guidebook step 4.4). Needs a lot of CPU RAM; skip it if you're serving the adapter directly. On Colab this may require a high-RAM runtime."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2f38ae59",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# from peft import PeftModel\n",
+    "# from transformers import AutoModelForImageTextToText\n",
+    "#\n",
+    "# base = AutoModelForImageTextToText.from_pretrained(BASE_MODEL_ID, low_cpu_mem_usage=True)\n",
+    "# merged = PeftModel.from_pretrained(base, training_args.output_dir).merge_and_unload()\n",
+    "# merged.save_pretrained(\"merged_model\", safe_serialization=True, max_shard_size=\"2GB\")\n",
+    "# processor.save_pretrained(\"merged_model\")\n",
+    "# print(\"Merged model saved to ./merged_model\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "eec6bae1",
+   "metadata": {},
+   "source": [
+    "## Next steps\n",
+    "\n",
+    "1. **Turn off `SMOKE_TEST`** and run on the full dataset once the plumbing works end to end.\n",
+    "2. **Run the eval harness** (guidebook Phase 3) on the fine-tuned model against the held-out `test` split — that produces the real numbers for the model card, not this one-image smell test.\n",
+    "3. **Fill the model card** (`README.md`) with those measured numbers.\n",
+    "\n",
+    "Reminder from the success-metrics discussion: the honest quality gate is *matching the incumbent pipeline* on the same held-out set, with F1 measured under a **fixed matching function**, node and edge F1 reported **separately**, and every metric **sliced by data source**. Don't publish a headline number off this notebook — publish it off the harness."
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "L4",
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}