{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Finetune Qwen3-VL series on sampled BigEarthNet S2 RGB Dataset\n",
        "\n",
        "This file contains the code to finetune Qwen3-VL series on sampled BigEarthNet S2 RGB Dataset. The dataset has been pointed to the local path where the notebook is located."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "scrolled": true
      },
      "outputs": [],
      "source": [
        "import os, re\n",
        "\n",
        "# Core installs. In Colab we follow Unsloth's pattern for vision finetuning.\n",
        "if \"COLAB_\" not in \"\".join(os.environ.keys()):\n",
        "    !pip install -U -q unsloth\n",
        "    !pip install -U -q datasets==4.3.0 'huggingface-hub>=0.34.0,<1.0' hf_transfer\n",
        "    !pip install -U -q transformers==4.57.0 trl==0.22.2\n",
        "    !pip install -U -q scikit-learn pandas matplotlib tqdm pillow\n",
        "else:\n",
        "    import torch; v = re.match(r\"[0-9\\.]{3,}\", str(torch.__version__)).group(0)\n",
        "    # xformers = \"xformers==\" + (\"0.0.32.post2\" if v == \"2.8.0\" else \"0.0.29.post3\")\n",
        "    !pip install --no-deps bitsandbytes accelerate xformers==0.0.33.post1 peft trl triton cut_cross_entropy unsloth_zoo\n",
        "    !pip install -U sentencepiece protobuf datasets==4.3.0 'huggingface-hub>=0.34.0,<1.0' hf_transfer\n",
        "    !pip install --no-deps unsloth\n",
        "    !pip install -U transformers==4.57.0\n",
        "    !pip install -U scikit-learn pandas matplotlib tqdm pillow\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "collapsed": false,
        "scrolled": true
      },
      "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"
          ]
        }
      ],
      "source": [
        "import os, re, io, json, math, random, datetime, itertools, warnings, copy\n",
        "from dataclasses import dataclass\n",
        "from typing import Dict, List, Any, Optional\n",
        "\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from tqdm.auto import tqdm\n",
        "from PIL import Image\n",
        "\n",
        "import torch\n",
        "from datasets import load_dataset\n",
        "from sklearn.metrics import f1_score, precision_score, recall_score\n",
        "\n",
        "from unsloth import FastVisionModel\n",
        "from trl import SFTTrainer, SFTConfig\n",
        "from unsloth.trainer import UnslothVisionDataCollator"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "scrolled": true
      },
      "outputs": [],
      "source": [
        "# ==== Config ====\n",
        "SEED = 3407\n",
        "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "\n",
        "# Models\n",
        "BASE_MODEL_ID = \"unsloth/Qwen3-VL-4B-Instruct\"   # primary study target\n",
        "# BASE_MODEL_ID = \"unsloth/Qwen3-VL-8B-Instruct\" # 2B, 32B...\n",
        "\n",
        "# LoRA\n",
        "LORA_R = 16\n",
        "LORA_ALPHA = 16\n",
        "LORA_DROPOUT = 0.0\n",
        "BIAS = \"none\"\n",
        "\n",
        "# Data\n",
        "TRAIN_REPO = \"jfang/BigEarthS2-RGB-10k-stratified-balanced\"\n",
        "TEST_REPO  = \"jfang/BigEarthS2-RGB-test-1k-balanced\"\n",
        "\n",
        "VAL_SPLIT = 0.1\n",
        "EVAL_LIMIT = int(os.environ.get(\"EVAL_LIMIT\", \"1000\"))\n",
        "\n",
        "# Training\n",
        "MAX_STEPS = int(os.environ.get(\"TRAIN_STEPS\", \"1000\"))\n",
        "PER_DEVICE_TRAIN_BSZ = 2\n",
        "GRAD_ACC_STEPS = 4\n",
        "LR = 2e-4\n",
        "WARMUP_STEPS = 50\n",
        "WEIGHT_DECAY = 0.001\n",
        "LOG_STEPS = 10\n",
        "\n",
        "\n",
        "RESULTS_DIR = \"./results_ft\"\n",
        "os.makedirs(RESULTS_DIR, exist_ok=True)\n",
        "\n",
        "def seed_everything(seed=SEED):\n",
        "    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)\n",
        "    if torch.cuda.is_available():\n",
        "        torch.cuda.manual_seed_all(seed)\n",
        "\n",
        "seed_everything()\n",
        "print(\"Device:\", DEVICE)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "collapsed": false,
        "scrolled": true
      },
      "outputs": [],
      "source": [
        "# ==== Prompt & class list (same as benchmark) ====\n",
        "PROMPT = (\n",
        "    \"Instructions: Answer the question asked after the given image. Output format: Output the\\n\"\n",
        "    \"option numbers corresponding to the correct answer in the format \\\"(X)\\\" where X is the\\n\"\n",
        "    \"correct digit choice (among 1 to 19). In case of multiple correct answers, output the\\n\"\n",
        "    \"answer choices in the same format separated by commas. For example, if the correct\\n\"\n",
        "    \"answers are choices 1 and 3, output \\\"(1),(3)\\\". Only output the option numbers in the\\n\"\n",
        "    \"above format and DO NOT OUTPUT the full answer text or any other extra words.\\n\"\n",
        "    \"Question: To which of the following class does the given image belong to?\\n\"\n",
        "    \"Select all that apply. Possible answer choices:\\n\"\n",
        "    \"(1)Agro-forestry areas\\n\"\n",
        "    \"(2)Arable land\\n\"\n",
        "    \"(3)Beaches, dunes, sands\\n\"\n",
        "    \"(4)Broad-leaved forest\\n\"\n",
        "    \"(5)Coastal wetlands\\n\"\n",
        "    \"(6)Complex cultivation patterns\\n\"\n",
        "    \"(7)Coniferous forest\\n\"\n",
        "    \"(8)Industrial or commercial units\\n\"\n",
        "    \"(9)Inland waters\\n\"\n",
        "    \"(10)Inland wetlands\\n\"\n",
        "    \"(11)Land principally occupied by agriculture, with significant areas of natural vegetation\\n\"\n",
        "    \"(12)Marine waters\\n\"\n",
        "    \"(13)Mixed forest\\n\"\n",
        "    \"(14)Moors, heathland and sclerophyllous vegetation\\n\"\n",
        "    \"(15)Natural grassland and sparsely vegetated areas\\n\"\n",
        "    \"(16)Pastures\\n\"\n",
        "    \"(17)Permanent crops\\n\"\n",
        "    \"(18)Transitional woodland, shrub\\n\"\n",
        "    \"(19)Urban fabric\\n\"\n",
        "    \"Answer:\"\n",
        ")\n",
        "\n",
        "CLASS_NAMES = [\n",
        "    \"Agro-forestry areas\",\n",
        "    \"Arable land\",\n",
        "    \"Beaches, dunes, sands\",\n",
        "    \"Broad-leaved forest\",\n",
        "    \"Coastal wetlands\",\n",
        "    \"Complex cultivation patterns\",\n",
        "    \"Coniferous forest\",\n",
        "    \"Industrial or commercial units\",\n",
        "    \"Inland waters\",\n",
        "    \"Inland wetlands\",\n",
        "    \"Land principally occupied by agriculture, with significant areas of natural vegetation\",\n",
        "    \"Marine waters\",\n",
        "    \"Mixed forest\",\n",
        "    \"Moors, heathland and sclerophyllous vegetation\",\n",
        "    \"Natural grassland and sparsely vegetated areas\",\n",
        "    \"Pastures\",\n",
        "    \"Permanent crops\",\n",
        "    \"Transitional woodland, shrub\",\n",
        "    \"Urban fabric\",\n",
        "]\n",
        "NUM_CLASSES = len(CLASS_NAMES)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "collapsed": false,
        "scrolled": true
      },
      "outputs": [],
      "source": [
        "# ==== Parsing & metrics ====\n",
        "DIGIT_RE = re.compile(r\"\\((\\d{1,2})\\)\")\n",
        "\n",
        "def parse_prediction_to_indices(text: str, num_classes: int = NUM_CLASSES):\n",
        "    if not isinstance(text, str):\n",
        "        return []\n",
        "    found = DIGIT_RE.findall(text)\n",
        "    if not found:\n",
        "        found = re.findall(r\"\\b(\\d{1,2})\\b\", text)\n",
        "    indices, seen = [], set()\n",
        "    for tok in found:\n",
        "        try:\n",
        "            k = int(tok)\n",
        "        except Exception:\n",
        "            continue\n",
        "        if 1 <= k <= num_classes and k not in seen:\n",
        "            indices.append(k)\n",
        "            seen.add(k)\n",
        "    return indices\n",
        "\n",
        "def indices_to_onehot(indices, num_classes: int = NUM_CLASSES):\n",
        "    y = np.zeros(num_classes, dtype=np.int32)\n",
        "    for k in indices:\n",
        "        if 1 <= k <= num_classes:\n",
        "            y[k-1] = 1\n",
        "    return y\n",
        "\n",
        "def subset_accuracy(y_true: np.ndarray, y_pred: np.ndarray) -> float:\n",
        "    return float(np.mean(np.all(y_true == y_pred, axis=1)))\n",
        "\n",
        "def example_based_f1(y_true: np.ndarray, y_pred: np.ndarray) -> float:\n",
        "    f1s = []\n",
        "    for t, p in zip(y_true, y_pred):\n",
        "        tp = int((t & p).sum())\n",
        "        t_sum = int(t.sum())\n",
        "        p_sum = int(p.sum())\n",
        "        if (t_sum + p_sum) == 0:\n",
        "            f1s.append(1.0)\n",
        "        elif tp == 0:\n",
        "            f1s.append(0.0)\n",
        "        else:\n",
        "            f1s.append(2 * tp / (t_sum + p_sum))\n",
        "    return float(np.mean(f1s))\n",
        "\n",
        "def compute_metrics(y_true: np.ndarray, y_pred: np.ndarray, class_names=CLASS_NAMES):\n",
        "    metrics = {\n",
        "        \"micro/precision\": precision_score(y_true, y_pred, average=\"micro\", zero_division=0),\n",
        "        \"micro/recall\":    recall_score(y_true, y_pred, average=\"micro\", zero_division=0),\n",
        "        \"micro/f1\":        f1_score(y_true, y_pred, average=\"micro\", zero_division=0),\n",
        "        \"macro/precision\": precision_score(y_true, y_pred, average=\"macro\", zero_division=0),\n",
        "        \"macro/recall\":    recall_score(y_true, y_pred, average=\"macro\", zero_division=0),\n",
        "        \"macro/f1\":        f1_score(y_true, y_pred, average=\"macro\", zero_division=0),\n",
        "        \"subset_accuracy\": subset_accuracy(y_true, y_pred),\n",
        "        \"example_f1\":      example_based_f1(y_true, y_pred),\n",
        "    }\n",
        "    per_class = []\n",
        "    for i, name in enumerate(class_names):\n",
        "        y_true_i = y_true[:, i]\n",
        "        y_pred_i = y_pred[:, i]\n",
        "        prec = precision_score(y_true_i, y_pred_i, zero_division=0)\n",
        "        rec  = recall_score(y_true_i, y_pred_i, zero_division=0)\n",
        "        f1   = f1_score(y_true_i, y_pred_i, zero_division=0)\n",
        "        support = int(y_true_i.sum())\n",
        "        per_class.append({\"class_id\": i+1, \"class_name\": name, \"precision\": prec, \"recall\": rec, \"f1\": f1, \"support\": support})\n",
        "    return metrics, pd.DataFrame(per_class)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "scrolled": true
      },
      "outputs": [],
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "# Data\n",
        "TRAIN_REPO = \"jfang/BigEarthS2-RGB-10k-stratified-balanced\"\n",
        "TEST_REPO  = \"jfang/BigEarthS2-RGB-test-1k-balanced\"\n",
        "# ==== Load datasets ====\n",
        "print(\"Loading datasets...\")\n",
        "ds_train = load_dataset(TRAIN_REPO, split=\"train\")\n",
        "ds_test  = load_dataset(TEST_REPO,  split=\"train\")  # repo uses 'train' split for the test set\n",
        "\n",
        "# create val split from ds_train\n",
        "splits = ds_train.train_test_split(test_size=VAL_SPLIT, seed=SEED)\n",
        "train_split, val_split = splits[\"train\"], splits[\"test\"]\n",
        "print(\"Train:\", len(train_split), \"Val:\", len(val_split), \"Test:\", len(ds_test))\n",
        "\n",
        "# Peek\n",
        "print(\"Sample fields:\", list(ds_test[0].keys()))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "collapsed": false,
        "scrolled": true
      },
      "outputs": [],
      "source": [
        "# ==== Inference helpers ====\n",
        "def build_messages_for_inference(prompt_text: str):\n",
        "    # Keep consistent ordering: image first, then text\n",
        "    return [\n",
        "        {\"role\": \"user\", \"content\": [\n",
        "            {\"type\": \"image\"},\n",
        "            {\"type\": \"text\", \"text\": prompt_text},\n",
        "        ]}\n",
        "    ]\n",
        "\n",
        "@torch.inference_mode()\n",
        "def vlm_generate(model, tokenizer, pil_img, prompt_text: str, max_new_tokens: int = 24, do_sample:bool=False):\n",
        "    messages = build_messages_for_inference(prompt_text)\n",
        "    input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True)\n",
        "    inputs = tokenizer(\n",
        "        pil_img,\n",
        "        input_text,\n",
        "        add_special_tokens=False,\n",
        "        return_tensors=\"pt\",\n",
        "    ).to(DEVICE)\n",
        "\n",
        "    out = model.generate(\n",
        "        **inputs,\n",
        "        max_new_tokens=max_new_tokens,\n",
        "        use_cache=True,\n",
        "        do_sample=do_sample,\n",
        "    )\n",
        "    input_len = inputs[\"input_ids\"].shape[1]\n",
        "    gen_tokens = out[0, input_len:]\n",
        "    text = tokenizer.decode(gen_tokens, skip_special_tokens=True).strip()\n",
        "    return text\n",
        "\n",
        "def evaluate_model_on_dataset(model, tokenizer, ds, eval_limit=None, tag=\"run\"):\n",
        "    rows = []\n",
        "    y_true_list, y_pred_list = [], []\n",
        "    n = len(ds) if eval_limit is None else min(eval_limit, len(ds))\n",
        "\n",
        "    pbar = tqdm(total=n, desc=f\"Evaluating [{tag}]\", leave=False)\n",
        "    for i in range(n):\n",
        "        ex = ds[i]\n",
        "        img = ex[\"image\"]\n",
        "        if not isinstance(img, Image.Image):\n",
        "            img = Image.fromarray(img)\n",
        "        gt = np.array(ex[\"labels\"], dtype=np.int32)\n",
        "        pred_text = vlm_generate(model, tokenizer, img, PROMPT)\n",
        "        pred_indices = parse_prediction_to_indices(pred_text, num_classes=NUM_CLASSES)\n",
        "        pred_vec = indices_to_onehot(pred_indices, num_classes=NUM_CLASSES)\n",
        "\n",
        "        rows.append({\n",
        "            \"index\": i,\n",
        "            \"prediction_raw\": pred_text,\n",
        "            \"prediction_indices\": pred_indices,\n",
        "            \"gt_indices\": [k+1 for k, v in enumerate(gt) if v == 1],\n",
        "        })\n",
        "        y_true_list.append(gt)\n",
        "        y_pred_list.append(pred_vec)\n",
        "        pbar.update(1)\n",
        "    pbar.close()\n",
        "\n",
        "    y_true = np.stack(y_true_list, axis=0) if n > 0 else np.zeros((0, NUM_CLASSES), dtype=np.int32)\n",
        "    y_pred = np.stack(y_pred_list, axis=0) if n > 0 else np.zeros((0, NUM_CLASSES), dtype=np.int32)\n",
        "\n",
        "    metrics, per_class_df = compute_metrics(y_true, y_pred, class_names=CLASS_NAMES)\n",
        "    preds_df = pd.DataFrame(rows).sort_values(\"index\").reset_index(drop=True)\n",
        "    return metrics, per_class_df, preds_df, y_true, y_pred"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "collapsed": false,
        "scrolled": true
      },
      "outputs": [],
      "source": [
        "# -------- Fast batched VLM generation + evaluation --------\n",
        "import math\n",
        "import torch\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from tqdm.auto import tqdm\n",
        "\n",
        "# Small perf toggles (safe on Ampere+)\n",
        "torch.backends.cuda.matmul.allow_tf32 = True\n",
        "try:\n",
        "    torch.set_float32_matmul_precision(\"high\")\n",
        "except Exception:\n",
        "    pass\n",
        "\n",
        "def _prepare_batched_inputs(tokenizer, images, texts, device, add_generation_prompt=True):\n",
        "    \"\"\"\n",
        "    Tokenize a batch of (image, text) pairs with padding suitable for generation.\n",
        "    Works with Unsloth's FastVisionModel tokenizer for Qwen-VL.\n",
        "    \"\"\"\n",
        "    # Decoder-only models usually prefer left padding for generation.\n",
        "    pad_side_prev = getattr(tokenizer, \"padding_side\", \"right\")\n",
        "    tokenizer.padding_side = \"left\"\n",
        "    if getattr(tokenizer, \"pad_token_id\", None) is None and getattr(tokenizer, \"eos_token_id\", None) is not None:\n",
        "        tokenizer.pad_token_id = tokenizer.eos_token_id  # safe fallback\n",
        "\n",
        "    # Apply chat template per-example (text is identical here; still explicit)\n",
        "    templated = [tokenizer.apply_chat_template(build_messages_for_inference(t), add_generation_prompt=add_generation_prompt)\n",
        "                 for t in texts]\n",
        "\n",
        "    # Batch tokenize images + texts together\n",
        "    batch_inputs = tokenizer(\n",
        "        images,                       # list of PIL images\n",
        "        templated,                    # list of templated strings\n",
        "        add_special_tokens=False,\n",
        "        padding=True,                 # pad to max in batch\n",
        "        return_tensors=\"pt\",\n",
        "    )\n",
        "\n",
        "    # Restore original padding side\n",
        "    tokenizer.padding_side = pad_side_prev\n",
        "\n",
        "    # Move to device\n",
        "    batch_inputs = {k: v.to(device) if hasattr(v, \"to\") else v for k, v in batch_inputs.items()}\n",
        "    return batch_inputs\n",
        "\n",
        "@torch.inference_mode()\n",
        "def batched_generate_texts(model, tokenizer, images, prompt_texts, max_new_tokens=24, do_sample=False):\n",
        "    \"\"\"\n",
        "    Returns a list of decoded strings, one per (image, text) pair in the batch.\n",
        "    Decodes only new tokens beyond each example's (unpadded) input length.\n",
        "    \"\"\"\n",
        "    model.eval()\n",
        "    inputs = _prepare_batched_inputs(tokenizer, images, prompt_texts, next(model.parameters()).device, add_generation_prompt=True)\n",
        "\n",
        "    # Track true input lengths (exclude pads)\n",
        "    pad_id = tokenizer.pad_token_id if getattr(tokenizer, \"pad_token_id\", None) is not None else 0\n",
        "    in_lens = (inputs[\"input_ids\"] != pad_id).sum(dim=1)  # [B]\n",
        "\n",
        "    # Generate in batch\n",
        "    out = model.generate(\n",
        "        **inputs,\n",
        "        max_new_tokens=max_new_tokens,\n",
        "        do_sample=do_sample,\n",
        "        use_cache=True,\n",
        "    )  # [B, T_out]\n",
        "\n",
        "    # Slice new tokens per example\n",
        "    texts = []\n",
        "    for i in range(out.size(0)):\n",
        "        gen_tokens = out[i, in_lens[i]:]\n",
        "        txt = tokenizer.decode(gen_tokens, skip_special_tokens=True).strip()\n",
        "        texts.append(txt)\n",
        "    return texts\n",
        "\n",
        "def evaluate_model_on_dataset_batched(model, tokenizer, ds, eval_limit=None, tag=\"run\", batch_size=8, max_new_tokens=24):\n",
        "    \"\"\"\n",
        "    Drop-in replacement for your per-item evaluator:\n",
        "    - Same return signature: (metrics, per_class_df, preds_df, y_true, y_pred)\n",
        "    - Same parsing/metrics as your current code\n",
        "    \"\"\"\n",
        "    n = len(ds) if eval_limit is None else min(eval_limit, len(ds))\n",
        "    rows = []\n",
        "    y_true_list, y_pred_list = [], []\n",
        "\n",
        "    pbar = tqdm(total=n, desc=f\"Evaluating[{tag}]-batched(B={batch_size})\", leave=False)\n",
        "    # Constant text for every sample\n",
        "    batch_texts = None  # lazily build once\n",
        "\n",
        "    for start in range(0, n, batch_size):\n",
        "        end = min(n, start + batch_size)\n",
        "        batch = [ds[i] for i in range(start, end)]\n",
        "        imgs = []\n",
        "        gts  = []\n",
        "\n",
        "        for ex in batch:\n",
        "            img = ex[\"image\"]\n",
        "            if not isinstance(img, Image.Image):\n",
        "                img = Image.fromarray(img)\n",
        "            imgs.append(img)\n",
        "            gts.append(np.array(ex[\"labels\"], dtype=np.int32))\n",
        "\n",
        "        if batch_texts is None or len(batch_texts) != len(imgs):\n",
        "            batch_texts = [PROMPT] * len(imgs)\n",
        "\n",
        "        # Generate batched\n",
        "        pred_texts = batched_generate_texts(model, tokenizer, imgs, batch_texts, max_new_tokens=max_new_tokens, do_sample=False)\n",
        "\n",
        "        # Parse + collect\n",
        "        for local_i, (ex, pred_text, gt_vec) in enumerate(zip(batch, pred_texts, gts)):\n",
        "            pred_idx = parse_prediction_to_indices(pred_text, num_classes=len(CLASS_NAMES))\n",
        "            pred_vec = indices_to_onehot(pred_idx, num_classes=len(CLASS_NAMES))\n",
        "\n",
        "            rows.append({\n",
        "                \"index\": start + local_i,\n",
        "                \"prediction_raw\": pred_text,\n",
        "                \"prediction_indices\": pred_idx,\n",
        "                \"gt_indices\": [k+1 for k, v in enumerate(gt_vec) if v == 1],\n",
        "            })\n",
        "            y_true_list.append(gt_vec)\n",
        "            y_pred_list.append(pred_vec)\n",
        "\n",
        "        pbar.update(end - start)\n",
        "\n",
        "    pbar.close()\n",
        "    y_true = np.stack(y_true_list, axis=0) if n > 0 else np.zeros((0, len(CLASS_NAMES)), dtype=np.int32)\n",
        "    y_pred = np.stack(y_pred_list, axis=0) if n > 0 else np.zeros((0, len(CLASS_NAMES)), dtype=np.int32)\n",
        "\n",
        "    metrics, per_class_df = compute_metrics(y_true, y_pred, class_names=CLASS_NAMES)\n",
        "    preds_df = pd.DataFrame(rows).sort_values(\"index\").reset_index(drop=True)\n",
        "    return metrics, per_class_df, preds_df, y_true, y_pred\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "collapsed": false,
        "scrolled": true
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "==((====))==  Unsloth 2025.11.4: Fast Qwen3_Vl patching. Transformers: 4.57.0.\n",
            "   \\\\   /|    NVIDIA A100-SXM4-40GB. Num GPUs = 1. Max memory: 39.494 GB. Platform: Linux.\n",
            "O^O/ \\_/ \\    Torch: 2.9.0+cu128. CUDA: 8.0. CUDA Toolkit: 12.8. Triton: 3.5.0\n",
            "\\        /    Bfloat16 = TRUE. FA [Xformers = 0.0.33.post1. FA2 = False]\n",
            " \"-____-\"     Free license: http://github.com/unslothai/unsloth\n",
            "Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "6809cd8f517f42b4b1c170b537e9e750",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "model.safetensors:   0%|          | 0.00/4.30G [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "33e8cc1e8dd54c4893011b2ed2d19887",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "generation_config.json:   0%|          | 0.00/213 [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "81ccbf387fcf4f42b8852ca257bd32cc",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "preprocessor_config.json:   0%|          | 0.00/782 [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "b46f44f071f24623abe58701dec2ae06",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "tokenizer_config.json: 0.00B [00:00, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "b5f70a90e4574e6ba655e819268c41c3",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "vocab.json: 0.00B [00:00, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "8f550c61037f41a8956aeaa1d1f6a4d3",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "merges.txt: 0.00B [00:00, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "f9cf0fffff38481a9c7b80c44c995814",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "tokenizer.json:   0%|          | 0.00/11.4M [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "d9190792946c4655ae763d359c98342a",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "added_tokens.json:   0%|          | 0.00/707 [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "5049ac547e23477eb6e2f71b9915e755",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "special_tokens_map.json:   0%|          | 0.00/614 [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "9b4868bd028c4366af3317851e02a902",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "chat_template.jinja: 0.00B [00:00, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "b7f9cd39fb2748b7bbe7301ee1d29bf9",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "video_preprocessor_config.json:   0%|          | 0.00/817 [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Running baseline evaluation...\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "a23173253b5e49c3870e492a62cf7ee9",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Evaluating[baseline]-batched(B=64):   0%|          | 0/1000 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{\n",
            "  \"micro/precision\": 0.45239968528717545,\n",
            "  \"micro/recall\": 0.3120759837177748,\n",
            "  \"micro/f1\": 0.369359242010599,\n",
            "  \"macro/precision\": 0.3975275770555858,\n",
            "  \"macro/recall\": 0.25290110114923575,\n",
            "  \"macro/f1\": 0.2570505316125184,\n",
            "  \"subset_accuracy\": 0.016,\n",
            "  \"example_f1\": 0.34380180930180926\n",
            "}\n"
          ]
        }
      ],
      "source": [
        "# ==== Load base model & run baseline ====\n",
        "base_model, base_tokenizer = FastVisionModel.from_pretrained(\n",
        "    BASE_MODEL_ID,\n",
        "    load_in_4bit=True,\n",
        "    use_gradient_checkpointing=\"unsloth\",\n",
        ")\n",
        "FastVisionModel.for_inference(base_model)\n",
        "\n",
        "print(\"Running baseline evaluation...\")\n",
        "base_metrics, base_perclass, base_preds, base_y_true, base_y_pred = evaluate_model_on_dataset_batched(\n",
        "    base_model, base_tokenizer, ds_test, eval_limit=EVAL_LIMIT, tag=\"baseline\", batch_size=64\n",
        ")\n",
        "print(json.dumps(base_metrics, indent=2))\n",
        "\n",
        "# Save baseline artifacts\n",
        "ts = datetime.datetime.now(datetime.UTC).strftime(\"%Y%m%d-%H%M%S\")\n",
        "base_tag = f\"baseline_{ts}\"\n",
        "base_preds.to_csv(os.path.join(RESULTS_DIR, f\"preds_{base_tag}.csv\"), index=False)\n",
        "base_perclass.to_csv(os.path.join(RESULTS_DIR, f\"perclass_{base_tag}.csv\"), index=False)\n",
        "with open(os.path.join(RESULTS_DIR, f\"metrics_{base_tag}.json\"), \"w\") as f:\n",
        "    json.dump({\"model\": BASE_MODEL_ID, \"dataset\": TEST_REPO, \"num_samples\": len(base_preds), \"metrics\": base_metrics}, f, indent=2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "collapsed": false,
        "scrolled": true
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "==((====))==  Unsloth 2025.11.4: Fast Qwen3_Vl patching. Transformers: 4.57.0.\n",
            "   \\\\   /|    NVIDIA A100-SXM4-40GB. Num GPUs = 1. Max memory: 39.494 GB. Platform: Linux.\n",
            "O^O/ \\_/ \\    Torch: 2.9.0+cu128. CUDA: 8.0. CUDA Toolkit: 12.8. Triton: 3.5.0\n",
            "\\        /    Bfloat16 = TRUE. FA [Xformers = 0.0.33.post1. FA2 = False]\n",
            " \"-____-\"     Free license: http://github.com/unslothai/unsloth\n",
            "Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!\n",
            "Unsloth: Model does not have a default image size - using 512\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "[accelerate.utils.other|WARNING]Detected kernel version 4.4.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.\n",
            "The model is already on multiple devices. Skipping the move to device specified in `args`.\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Starting finetuning...\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "==((====))==  Unsloth - 2x faster free finetuning | Num GPUs used = 1\n",
            "   \\\\   /|    Num examples = 9,000 | Num Epochs = 1 | Total steps = 500\n",
            "O^O/ \\_/ \\    Batch size per device = 16 | Gradient accumulation steps = 1\n",
            "\\        /    Data Parallel GPUs = 1 | Total batch size (16 x 1 x 1) = 16\n",
            " \"-____-\"     Trainable parameters = 314,572,800 of 4,752,388,608 (6.62% trained)\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Unsloth: Will smartly offload gradients to save VRAM!\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "\n",
              "    <div>\n",
              "      \n",
              "      <progress value='500' max='500' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
              "      [500/500 14:51, Epoch 0/1]\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>10</td>\n",
              "      <td>2.357000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>20</td>\n",
              "      <td>0.440900</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>30</td>\n",
              "      <td>0.029100</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>40</td>\n",
              "      <td>0.022200</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>50</td>\n",
              "      <td>0.021700</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>60</td>\n",
              "      <td>0.020500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>70</td>\n",
              "      <td>0.020500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>80</td>\n",
              "      <td>0.019000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>90</td>\n",
              "      <td>0.019500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>100</td>\n",
              "      <td>0.018800</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>110</td>\n",
              "      <td>0.019200</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>120</td>\n",
              "      <td>0.018200</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>130</td>\n",
              "      <td>0.018600</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>140</td>\n",
              "      <td>0.018900</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>150</td>\n",
              "      <td>0.018300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>160</td>\n",
              "      <td>0.017800</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>170</td>\n",
              "      <td>0.017100</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>180</td>\n",
              "      <td>0.018300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>190</td>\n",
              "      <td>0.018500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>200</td>\n",
              "      <td>0.017800</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>210</td>\n",
              "      <td>0.018500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>220</td>\n",
              "      <td>0.017000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>230</td>\n",
              "      <td>0.018600</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>240</td>\n",
              "      <td>0.017500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>250</td>\n",
              "      <td>0.017700</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>260</td>\n",
              "      <td>0.017300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>270</td>\n",
              "      <td>0.017200</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>280</td>\n",
              "      <td>0.017700</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>290</td>\n",
              "      <td>0.017200</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>300</td>\n",
              "      <td>0.017500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>310</td>\n",
              "      <td>0.016300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>320</td>\n",
              "      <td>0.017300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>330</td>\n",
              "      <td>0.016800</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>340</td>\n",
              "      <td>0.015500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>350</td>\n",
              "      <td>0.016400</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>360</td>\n",
              "      <td>0.017300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>370</td>\n",
              "      <td>0.016000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>380</td>\n",
              "      <td>0.017500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>390</td>\n",
              "      <td>0.016500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>400</td>\n",
              "      <td>0.017300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>410</td>\n",
              "      <td>0.014800</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>420</td>\n",
              "      <td>0.016500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>430</td>\n",
              "      <td>0.016500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>440</td>\n",
              "      <td>0.016300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>450</td>\n",
              "      <td>0.015700</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>460</td>\n",
              "      <td>0.017300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>470</td>\n",
              "      <td>0.016500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>480</td>\n",
              "      <td>0.016700</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>490</td>\n",
              "      <td>0.015300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>500</td>\n",
              "      <td>0.016500</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table><p>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Saved LoRA adapters to: lora_qwen3vl4b_bigearth_r32\n"
          ]
        }
      ],
      "source": [
        "# LoRA\n",
        "LORA_R = 128\n",
        "LORA_ALPHA = 128\n",
        "LORA_DROPOUT = 0.0\n",
        "BIAS = \"none\"\n",
        "\n",
        "MAX_STEPS = 500\n",
        "PER_DEVICE_TRAIN_BSZ = 16\n",
        "GRAD_ACC_STEPS = 1\n",
        "LR = 2e-4\n",
        "WARMUP_STEPS = 50\n",
        "WEIGHT_DECAY = 0.001\n",
        "\n",
        "BASE_MODEL_ID = \"unsloth/Qwen3-VL-4B-Instruct\"   \n",
        "\n",
        "# ==== Finetune with LoRA ====\n",
        "ft_model, ft_tokenizer = FastVisionModel.from_pretrained(\n",
        "    BASE_MODEL_ID,\n",
        "    load_in_4bit=True,\n",
        "    use_gradient_checkpointing=\"unsloth\",\n",
        ")\n",
        "ft_model = FastVisionModel.get_peft_model(\n",
        "    ft_model,\n",
        "    finetune_vision_layers     = True,\n",
        "    finetune_language_layers   = True,\n",
        "    finetune_attention_modules = True,\n",
        "    finetune_mlp_modules       = True,\n",
        "    r = LORA_R,\n",
        "    lora_alpha = LORA_ALPHA,\n",
        "    lora_dropout = LORA_DROPOUT,\n",
        "    bias = \"none\",\n",
        "    random_state = SEED,\n",
        "    use_rslora = False,\n",
        "    loftq_config = None,  \n",
        ")\n",
        "\n",
        "FastVisionModel.for_training(ft_model)\n",
        "\n",
        "def convert_to_conversation(sample):\n",
        "    # Keep order consistent with inference: image first, then text\n",
        "    answer = \",\".join(f\"({i+1})\" for i, v in enumerate(sample[\"labels\"]) if int(v) == 1)\n",
        "    return {\n",
        "        \"messages\": [\n",
        "            {\"role\": \"user\", \"content\": [\n",
        "                {\"type\": \"image\", \"image\": sample[\"image\"]},\n",
        "                {\"type\": \"text\",  \"text\": PROMPT},\n",
        "            ]},\n",
        "            {\"role\": \"assistant\", \"content\": [{\"type\": \"text\", \"text\": answer}]},\n",
        "        ]\n",
        "    }\n",
        "\n",
        "train_conversations = [convert_to_conversation(sample) for sample in train_split]\n",
        "val_conversations   = [convert_to_conversation(sample) for sample in val_split]\n",
        "\n",
        "trainer = SFTTrainer(\n",
        "    model = ft_model,\n",
        "    tokenizer = ft_tokenizer,\n",
        "    data_collator = UnslothVisionDataCollator(ft_model, ft_tokenizer),\n",
        "    train_dataset = train_conversations,\n",
        "    eval_dataset  = val_conversations[:1024],\n",
        "    args = SFTConfig(\n",
        "        per_device_train_batch_size = PER_DEVICE_TRAIN_BSZ,\n",
        "        gradient_accumulation_steps = GRAD_ACC_STEPS,\n",
        "        warmup_steps = WARMUP_STEPS,\n",
        "        max_steps = MAX_STEPS,\n",
        "        learning_rate = LR,\n",
        "        logging_steps = LOG_STEPS,\n",
        "        optim = \"adamw_8bit\",\n",
        "        weight_decay = WEIGHT_DECAY,\n",
        "        lr_scheduler_type = \"linear\",\n",
        "        seed = SEED,\n",
        "        output_dir = \"outputs_4b_r32\",\n",
        "        report_to = \"none\",\n",
        "        remove_unused_columns = False,\n",
        "        dataset_text_field = \"\",\n",
        "        dataset_kwargs = {\"skip_prepare_dataset\": True},\n",
        "        max_length = 2048,\n",
        "    ),\n",
        ")\n",
        "\n",
        "print(\"Starting finetuning...\")\n",
        "_ = trainer.train()\n",
        "FastVisionModel.for_inference(ft_model)\n",
        "\n",
        "# Save LoRA\n",
        "LORA_DIR = \"lora_qwen3vl4b_bigearth_r32\"\n",
        "os.makedirs(LORA_DIR, exist_ok=True)\n",
        "ft_model.save_pretrained(LORA_DIR)\n",
        "ft_tokenizer.save_pretrained(LORA_DIR)\n",
        "print(\"Saved LoRA adapters to:\", LORA_DIR)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "collapsed": false,
        "scrolled": true
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "f4b60f1ddc7b46dc8cffde74cf333fd0",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Evaluating[finetuned]-batched(B=64):   0%|          | 0/1000 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{\n",
            "  \"micro/precision\": 0.5801567991349014,\n",
            "  \"micro/recall\": 0.5823609226594301,\n",
            "  \"micro/f1\": 0.5812567713976164,\n",
            "  \"macro/precision\": 0.5500226521366196,\n",
            "  \"macro/recall\": 0.5157787307502865,\n",
            "  \"macro/f1\": 0.5079070785761576,\n",
            "  \"subset_accuracy\": 0.11,\n",
            "  \"example_f1\": 0.5962173992673993\n",
            "}\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>run</th>\n",
              "      <th>micro/precision</th>\n",
              "      <th>micro/recall</th>\n",
              "      <th>micro/f1</th>\n",
              "      <th>macro/precision</th>\n",
              "      <th>macro/recall</th>\n",
              "      <th>macro/f1</th>\n",
              "      <th>subset_accuracy</th>\n",
              "      <th>example_f1</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>baseline</td>\n",
              "      <td>0.452400</td>\n",
              "      <td>0.312076</td>\n",
              "      <td>0.369359</td>\n",
              "      <td>0.397528</td>\n",
              "      <td>0.252901</td>\n",
              "      <td>0.257051</td>\n",
              "      <td>0.016</td>\n",
              "      <td>0.343802</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>finetuned</td>\n",
              "      <td>0.580157</td>\n",
              "      <td>0.582361</td>\n",
              "      <td>0.581257</td>\n",
              "      <td>0.550023</td>\n",
              "      <td>0.515779</td>\n",
              "      <td>0.507907</td>\n",
              "      <td>0.110</td>\n",
              "      <td>0.596217</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "         run  micro/precision  micro/recall  micro/f1  macro/precision  \\\n",
              "0   baseline         0.452400      0.312076  0.369359         0.397528   \n",
              "1  finetuned         0.580157      0.582361  0.581257         0.550023   \n",
              "\n",
              "   macro/recall  macro/f1  subset_accuracy  example_f1  \n",
              "0      0.252901  0.257051            0.016    0.343802  \n",
              "1      0.515779  0.507907            0.110    0.596217  "
            ]
          },
          "execution_count": 10,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "ft_metrics, ft_perclass, ft_preds, ft_y_true, ft_y_pred = evaluate_model_on_dataset_batched(\n",
        "    ft_model, ft_tokenizer, ds_test, eval_limit=EVAL_LIMIT, tag=\"finetuned\", batch_size=64\n",
        ")\n",
        "print(json.dumps(ft_metrics, indent=2))\n",
        "\n",
        "ts = datetime.datetime.now(datetime.UTC).strftime(\"%Y%m%d-%H%M%S\")\n",
        "ft_tag = f\"finetuned_{ts}\"\n",
        "ft_preds.to_csv(os.path.join(RESULTS_DIR, f\"preds_{ft_tag}.csv\"), index=False)\n",
        "ft_perclass.to_csv(os.path.join(RESULTS_DIR, f\"perclass_{ft_tag}.csv\"), index=False)\n",
        "with open(os.path.join(RESULTS_DIR, f\"metrics_{ft_tag}.json\"), \"w\") as f:\n",
        "    json.dump({\"model\": BASE_MODEL_ID + \" + LoRA r=16\", \"dataset\": TEST_REPO, \"num_samples\": len(ft_preds), \"metrics\": ft_metrics}, f, indent=2)\n",
        "\n",
        "summary = pd.DataFrame([{\"run\": \"baseline\", **base_metrics}, {\"run\": \"finetuned\", **ft_metrics}])\n",
        "summary.to_csv(os.path.join(RESULTS_DIR, \"summary_baseline_vs_finetuned.csv\"), index=False)\n",
        "summary"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {
        "collapsed": false,
        "scrolled": true
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Showing 12 improved examples\n",
            "Per-class counts (by dominant GT label):\n",
            "   1 - Agro-forestry areas: 1\n",
            "  12 - Marine waters: 1\n",
            "   9 - Inland waters: 1\n",
            "   7 - Coniferous forest: 1\n",
            "  16 - Pastures: 1\n",
            "  15 - Natural grassland and sparsely vegetated areas: 1\n",
            "  10 - Inland wetlands: 1\n",
            "  17 - Permanent crops: 1\n",
            "  13 - Mixed forest: 1\n",
            "   5 - Coastal wetlands: 1\n",
            "  14 - Moors, heathland and sclerophyllous vegetation: 1\n",
            "  19 - Urban fabric: 1\n"
          ]
        },
        {
          "data": {
            "image/png": "[encoded data, 1778000 characters]",
            "text/plain": [
              "<Figure size 1500x2080 with 12 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# ==== Qualitative examples: diverse, biggest gain from fine-tuning ====\n",
        "import math\n",
        "import matplotlib.pyplot as plt\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from collections import Counter\n",
        "import textwrap\n",
        " \n",
        "def multilabel_f1(gt_indices, pred_indices):\n",
        "    \"\"\"Simple multi-label F1 on sets of 1-based class indices.\"\"\"\n",
        "    gt = set(gt_indices or [])\n",
        "    pred = set(pred_indices or [])\n",
        "\n",
        "    if not gt and not pred:\n",
        "        return 1.0\n",
        "    if not gt or not pred:\n",
        "        return 0.0\n",
        "\n",
        "    tp = len(gt & pred)\n",
        "    if tp == 0:\n",
        "        return 0.0\n",
        "\n",
        "    prec = tp / len(pred)\n",
        "    rec  = tp / len(gt)\n",
        "    if prec + rec == 0:\n",
        "        return 0.0\n",
        "    return 2 * prec * rec / (prec + rec)\n",
        "\n",
        "# Merge baseline + fine-tuned predictions by sample index\n",
        "merged = (\n",
        "    base_preds[[\"index\", \"gt_indices\", \"prediction_indices\", \"prediction_raw\"]]\n",
        "    .rename(columns={\n",
        "        \"prediction_indices\": \"base_pred_indices\",\n",
        "        \"prediction_raw\": \"base_pred_raw\",\n",
        "    })\n",
        "    .merge(\n",
        "        ft_preds[[\"index\", \"prediction_indices\", \"prediction_raw\"]]\n",
        "        .rename(columns={\n",
        "            \"prediction_indices\": \"ft_pred_indices\",\n",
        "            \"prediction_raw\": \"ft_pred_raw\",\n",
        "        }),\n",
        "        on=\"index\",\n",
        "        validate=\"one_to_one\",\n",
        "    )\n",
        ")\n",
        "\n",
        "# Per-example F1 and gain\n",
        "merged[\"f1_base\"] = merged.apply(\n",
        "    lambda r: multilabel_f1(r[\"gt_indices\"], r[\"base_pred_indices\"]),\n",
        "    axis=1,\n",
        ")\n",
        "merged[\"f1_ft\"] = merged.apply(\n",
        "    lambda r: multilabel_f1(r[\"gt_indices\"], r[\"ft_pred_indices\"]),\n",
        "    axis=1,\n",
        ")\n",
        "merged[\"gain\"] = merged[\"f1_ft\"] - merged[\"f1_base\"]\n",
        "\n",
        "# Only keep examples where FT improved\n",
        "improved = merged[merged[\"gain\"] > 0].sort_values(\"gain\", ascending=False)\n",
        "\n",
        "# --- Enforce class diversity -------------------------------------------\n",
        "\n",
        "# Estimate per-class support from GT to know which classes are rare\n",
        "support_counter = Counter()\n",
        "for indices in merged[\"gt_indices\"]:\n",
        "    for c in indices or []:\n",
        "        support_counter[c] += 1\n",
        "\n",
        "NUM_SHOW = 12          # total examples to display\n",
        "MAX_PER_CLASS = 1      # at most this many examples \"assigned\" to one class\n",
        "\n",
        "selected_rows = []\n",
        "per_class_counts = Counter()\n",
        "\n",
        "for _, row in improved.iterrows():\n",
        "    if len(selected_rows) >= NUM_SHOW:\n",
        "        break\n",
        "\n",
        "    gt_indices = row[\"gt_indices\"] or []\n",
        "    if not gt_indices:\n",
        "        continue\n",
        "\n",
        "    # Assign this sample to the rarest GT class it contains\n",
        "    dominant_cls = min(gt_indices, key=lambda c: support_counter.get(c, 10**9))\n",
        "\n",
        "    if per_class_counts[dominant_cls] >= MAX_PER_CLASS:\n",
        "        continue\n",
        "\n",
        "    per_class_counts[dominant_cls] += 1\n",
        "    selected_rows.append(row)\n",
        "\n",
        "top_examples = pd.DataFrame(selected_rows)\n",
        "\n",
        "print(f\"Showing {len(top_examples)} improved examples\")\n",
        "print(\"Per-class counts (by dominant GT label):\")\n",
        "for cls_id, cnt in per_class_counts.items():\n",
        "    print(f\"  {cls_id:2d} - {CLASS_NAMES[cls_id-1]}: {cnt}\")\n",
        "\n",
        "# --- Plot grid with non-overlapping captions ---------------------------\n",
        "\n",
        "n = len(top_examples)\n",
        "n_cols = 3\n",
        "n_rows = math.ceil(n / n_cols)\n",
        "\n",
        "fig, axes = plt.subplots(n_rows, n_cols, figsize=(5 * n_cols, 5.2 * n_rows))\n",
        "\n",
        "if n_rows == 1 and n_cols == 1:\n",
        "    axes = np.array([[axes]])\n",
        "elif n_rows == 1:\n",
        "    axes = np.array([axes])\n",
        "elif n_cols == 1:\n",
        "    axes = axes[:, None]\n",
        "\n",
        "# Turn off empty axes\n",
        "for ax in axes.flatten()[n:]:\n",
        "    ax.axis(\"off\")\n",
        "\n",
        "def idx_to_names(indices):\n",
        "    names = []\n",
        "    for i in indices or []:\n",
        "        if 1 <= i <= len(CLASS_NAMES):\n",
        "            names.append(CLASS_NAMES[i - 1])\n",
        "        else:\n",
        "            names.append(f\"idx_{i}\")\n",
        "    return names\n",
        "\n",
        "def make_caption(gt_names, base_names, ft_names, f1_base, f1_ft, width=45):\n",
        "    lines = [\n",
        "        \"GT:  \"   + (\", \".join(gt_names)   if gt_names   else \"∅\"),\n",
        "        \"Base:  \" + (\", \".join(base_names) if base_names else \"∅\") + f\" (F1={f1_base:.2f})\",\n",
        "        \"FT:  \"   + (\", \".join(ft_names)   if ft_names   else \"∅\") + f\" (F1={f1_ft:.2f})\",\n",
        "    ]\n",
        "    wrapped = [textwrap.fill(line, width) for line in lines]\n",
        "    return \"\\n\".join(wrapped)\n",
        "\n",
        "for ax, (_, row) in zip(axes.flatten(), top_examples.iterrows()):\n",
        "    idx        = row[\"index\"]\n",
        "    gt_indices = row[\"gt_indices\"]\n",
        "    base_idx   = row[\"base_pred_indices\"]\n",
        "    ft_idx     = row[\"ft_pred_indices\"]\n",
        "\n",
        "    ex = ds_test[idx]\n",
        "    img = ex[\"image\"]\n",
        "    if not isinstance(img, Image.Image):\n",
        "        img = Image.fromarray(img)\n",
        "\n",
        "    gt_names   = idx_to_names(gt_indices)\n",
        "    base_names = idx_to_names(base_idx)\n",
        "    ft_names   = idx_to_names(ft_idx)\n",
        "\n",
        "    ax.imshow(img)\n",
        "    ax.axis(\"off\")\n",
        "\n",
        "    # Short title on top\n",
        "    ax.set_title(f\"Sample {idx} | ΔF1={row['gain']:.2f}\",\n",
        "                 fontsize=12, pad=2)\n",
        "\n",
        "    # Wrapped caption below the image (inside the axes)\n",
        "    caption = make_caption(gt_names, base_names, ft_names,\n",
        "                           row[\"f1_base\"], row[\"f1_ft\"])\n",
        "    ax.text(0.5, -0.04, caption,\n",
        "            transform=ax.transAxes,\n",
        "            ha=\"center\", va=\"top\",\n",
        "            fontsize=10)\n",
        "\n",
        "# Extra space between rows so captions don't collide\n",
        "plt.subplots_adjust(hspace=0.8, wspace=0.1, bottom=0.05, top=0.75)\n",
        "plt.show()\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python",
      "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"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
