{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GzMT0d7XRdQ3"
      },
      "source": [
        "# Supervised Fine-tuning Gemini 2.5 Flash for Visual Defect Detection\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/tuning/sft_gemini_visual_defect_detection.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Ftuning%2Fsft_gemini_visual_defect_detection.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/tuning/sft_gemini_visual_defect_detection.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/tuning/sft_gemini_visual_defect_detection.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_uco5wDNcIRq"
      },
      "source": [
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/tuning/sft_gemini_visual_defect_detection.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/tuning/sft_gemini_visual_defect_detection.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/tuning/sft_gemini_visual_defect_detection.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/tuning/sft_gemini_visual_defect_detection.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/tuning/sft_gemini_visual_defect_detection.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M04y-KnqcSCq"
      },
      "source": [
        "| Author |\n",
        "| --- |\n",
        "| [Aniket Agrawal](https://github.com/aniketagrawal2012) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WURYK3ZRRdQ5"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how to perform **supervised fine-tuning** on a Gemini model for a **visual defect detection** task within a manufacturing context. We will use the `google-genai` SDK integrated with Vertex AI to train the model to classify product images and identify flaws.\n",
        "\n",
        "### Use Case: Classifying Product Quality from Images\n",
        "\n",
        "We'll fine-tune Gemini to analyze an image of a product from a manufacturing line and classify its quality (e.g., \"Pass\", \"Defect\") and provide a short description of the issue if one is found. This is a multimodal task combining image analysis (vision) with text generation (classification and description).\n",
        "\n",
        "**Workflow:**\n",
        "1.  **Generate Data**: Create simulated product images (Pass/Defect), upload them to GCS, and create a manifest DataFrame.\n",
        "2.  **Prepare Tuning Data (JSONL)**: Convert image GCS URIs and corresponding labels (e.g., \"Status: Defect - Scratch detected\") into the JSON Lines format required for Gemini supervised tuning.\n",
        "3.  **Upload to GCS**: Store the formatted tuning data (JSONL files) in a Google Cloud Storage bucket.\n",
        "4.  **Launch Fine-tuning Job**: Use the `google-genai` SDK client (configured for Vertex AI) to start the supervised tuning job.\n",
        "5.  **Monitor Job**: Track the progress of the fine-tuning job.\n",
        "6.  **Evaluate Tuned Model**: Make predictions on new product images using the fine-tuned model endpoint and compare qualitatively.\n",
        "7.  **Integrate Gemini for Reporting**: Use a base Gemini model to summarize the tuning job results."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NTmKZlmIRdQ5"
      },
      "source": [
        "## Setup\n",
        "\n",
        "### Install required packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PdrVBfYMRdQ5"
      },
      "outputs": [],
      "source": [
        "import sys  # noqa: F401\n",
        "\n",
        "# Install necessary libraries\n",
        "# gcsfs is added to allow pandas to write directly to GCS\n",
        "# Pillow (PIL) is needed for image generation\n",
        "!{sys.executable} -m pip install --upgrade --user --quiet pandas numpy google-cloud-aiplatform google-genai google-cloud-storage gcsfs Pillow"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yqy4-3cBRdQ6"
      },
      "source": [
        "**⚠️ Important:** Restart the kernel after installation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KS9rUufMRdQ6"
      },
      "source": [
        "### Authenticate and Initialize Vertex AI\n",
        "\n",
        "Set your project, region, and GCS bucket information. We configure the notebook for Vertex AI fine-tuning and reporting."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PyQnTYGlRdQ6"
      },
      "outputs": [],
      "source": [
        "import subprocess\n",
        "\n",
        "import vertexai\n",
        "from google.genai import (\n",
        "    Client as VertexClient,  # This is for Vertex AI tuning/models client\n",
        ")\n",
        "\n",
        "# --- Vertex AI Configuration (Required for Fine-tuning Job) ---\n",
        "PROJECT_ID = \"\"  # @param {type: \"string\", placeholder: \"your-gcp-project-id\"}\n",
        "REGION = \"\"  # @param {type:\"string\"}\n",
        "BUCKET_NAME = \"\"  # @param {type:\"string\", placeholder: \"your-gcs-bucket-name\"}\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}\"\n",
        "\n",
        "\n",
        "def ensure_gcs_bucket_exists(project_id: str, region: str, bucket_uri: str) -> None:\n",
        "    \"\"\"Ensures the specified GCS bucket exists, creating it if necessary.\"\"\"\n",
        "    print(f\"Checking/Creating bucket: {bucket_uri}\")\n",
        "\n",
        "    # 1. Check if bucket exists. If yes, return immediately (Guard Clause).\n",
        "    try:\n",
        "        subprocess.run(\n",
        "            [\"gsutil\", \"ls\", bucket_uri],\n",
        "            check=True,\n",
        "            stdout=subprocess.PIPE,\n",
        "            stderr=subprocess.PIPE,\n",
        "        )\n",
        "        print(f\"Bucket {bucket_uri} already exists.\")\n",
        "        return\n",
        "    except subprocess.CalledProcessError:\n",
        "        # Fall through only if check failed\n",
        "        print(f\"Bucket {bucket_uri} not found. Attempting to create it.\")\n",
        "\n",
        "    # 2. Create bucket (No longer indented inside an 'else' or 'except' block)\n",
        "    try:\n",
        "        subprocess.run(\n",
        "            [\"gsutil\", \"mb\", \"-l\", region, \"-p\", project_id, bucket_uri],\n",
        "            check=True,\n",
        "            stdout=subprocess.PIPE,\n",
        "            stderr=subprocess.PIPE,\n",
        "        )\n",
        "        print(f\"Bucket {bucket_uri} created successfully.\")\n",
        "    except subprocess.CalledProcessError as e:\n",
        "        error_msg = (\n",
        "            f\"Failed to create bucket {bucket_uri}. Error: {e.stderr.decode().strip()}\"\n",
        "        )\n",
        "        raise ValueError(error_msg) from e\n",
        "\n",
        "\n",
        "# --- Authentication (Colab/Workbench for Vertex AI) ---\n",
        "if not PROJECT_ID or PROJECT_ID == \"\":\n",
        "    try:\n",
        "        from google.colab import auth\n",
        "\n",
        "        auth.authenticate_user()\n",
        "\n",
        "        PROJECT_ID = (\n",
        "            subprocess.check_output([\"gcloud\", \"config\", \"get-value\", \"project\"])\n",
        "            .decode(\"utf-8\")\n",
        "            .strip()\n",
        "        )\n",
        "        print(f\"Retrieved Project ID: {PROJECT_ID}\")\n",
        "    except Exception as e:\n",
        "        print(\n",
        "            f\"Could not automatically retrieve Project ID. Please set it manually. Error: {e}\"\n",
        "        )\n",
        "\n",
        "# Ensure BUCKET_NAME is set, and attempt to create the bucket\n",
        "if not BUCKET_NAME or BUCKET_NAME == \"\":\n",
        "    if PROJECT_ID:\n",
        "        BUCKET_NAME = f\"{PROJECT_ID}-gemini-tuning-bucket\"\n",
        "        BUCKET_URI = f\"gs://{BUCKET_NAME}\"\n",
        "        print(f\"Bucket name not provided. Using default: {BUCKET_NAME}\")\n",
        "    else:\n",
        "        raise ValueError(\n",
        "            \"Please provide a valid GCS Bucket name or ensure PROJECT_ID is set for default bucket creation.\"\n",
        "        )\n",
        "\n",
        "ensure_gcs_bucket_exists(PROJECT_ID, REGION, BUCKET_URI)\n",
        "\n",
        "if PROJECT_ID:\n",
        "    print(\n",
        "        f\"Initializing Vertex AI for project: {PROJECT_ID} in {REGION} using bucket {BUCKET_URI}\"\n",
        "    )\n",
        "    # Initialize Vertex AI SDK (needed for launching the tuning job)\n",
        "    vertexai.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)\n",
        "    # Initialize the genai client specifically for Vertex AI operations (like tuning)\n",
        "    vertex_client = VertexClient(vertexai=True, project=PROJECT_ID, location=REGION)\n",
        "    print(\"Vertex AI SDK Initialized.\")\n",
        "else:\n",
        "    raise ValueError(\"PROJECT_ID must be set for Vertex AI operations.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fQ3ECMSARdQ6"
      },
      "source": [
        "### Imports and Global Configuration"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bV74HX1oRdQ6"
      },
      "outputs": [],
      "source": [
        "import io\n",
        "import json\n",
        "import random\n",
        "import time\n",
        "import warnings\n",
        "from typing import Any\n",
        "\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from PIL import Image, ImageDraw\n",
        "from google.cloud import storage\n",
        "from google.genai import types as genai_types\n",
        "\n",
        "# --- Global Settings ---\n",
        "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
        "warnings.filterwarnings(\"ignore\", category=FutureWarning)\n",
        "np.random.seed(42)\n",
        "random.seed(42)\n",
        "\n",
        "# --- Constants ---\n",
        "BASE_MODEL_ID = \"gemini-2.5-flash\"  # Tunable model ID on Vertex AI\n",
        "TUNED_MODEL_DISPLAY_NAME = f\"visual-defect-gemini-tuned-{int(time.time())}\"\n",
        "DATA_DIR_GCS = f\"{BUCKET_URI}/visual_defect_tuning_data\"\n",
        "IMAGE_DIR_GCS_PATH = \"visual_defect_tuning_data/images\"  # Relative path for client\n",
        "IMAGE_DIR_GCS_URI = f\"{DATA_DIR_GCS}/images\"\n",
        "TRAIN_JSONL_GCS_URI = f\"{DATA_DIR_GCS}/train_data.jsonl\"\n",
        "VALIDATION_JSONL_GCS_URI = f\"{DATA_DIR_GCS}/validation_data.jsonl\"\n",
        "TEST_JSONL_GCS_URI = f\"{DATA_DIR_GCS}/test_data.jsonl\"  # For qualitative eval later\n",
        "\n",
        "print(f\"Base model for tuning: {BASE_MODEL_ID}\")\n",
        "print(f\"Tuning data GCS path: {DATA_DIR_GCS}\")\n",
        "print(f\"Image GCS URI: {IMAGE_DIR_GCS_URI}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TEUFLjeYRdQ7"
      },
      "source": [
        "## Step 1: Generate Simulated Image Data\n",
        "\n",
        "Instead of loading data, we'll generate simulated product images (simple shapes) and upload them to GCS. We'll create 'Pass' images (clean) and 'Defect' images (with a visual flaw). We will return a Pandas DataFrame acting as a manifest file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8RqK0Td6RdQ7"
      },
      "outputs": [],
      "source": [
        "def _upload_image_blob(bucket, blob_name: str, image: Image.Image) -> bool:\n",
        "    \"\"\"Helper to safely upload a PIL image to GCS.\"\"\"\n",
        "    try:\n",
        "        blob = bucket.blob(blob_name)\n",
        "        with io.BytesIO() as output:\n",
        "            image.save(output, format=\"PNG\")\n",
        "            blob.upload_from_string(output.getvalue(), content_type=\"image/png\")\n",
        "        return True\n",
        "    except Exception as e:\n",
        "        print(f\"Failed to upload {blob_name}: {e}\")\n",
        "        return False\n",
        "\n",
        "\n",
        "def generate_and_upload_images(\n",
        "    bucket_name: str,\n",
        "    gcs_image_path: str,\n",
        "    num_images: int = 100,\n",
        "    defect_ratio: float = 0.4,\n",
        ") -> pd.DataFrame:\n",
        "    \"\"\"Generates simple images, uploads to GCS, and returns a manifest DataFrame.\"\"\"\n",
        "    print(f\"Generating {num_images} simulated images...\")\n",
        "    storage_client = storage.Client(project=PROJECT_ID)\n",
        "    bucket = storage_client.bucket(bucket_name)\n",
        "    manifest = []\n",
        "\n",
        "    for i in range(num_images):\n",
        "        img = Image.new(\"RGB\", (100, 100), color=\"#DDDDDD\")\n",
        "        draw = ImageDraw.Draw(img)\n",
        "        draw.rectangle((20, 20, 80, 80), fill=\"#5555AA\")  # Main product shape\n",
        "\n",
        "        is_defect = random.random() < defect_ratio\n",
        "        image_name = f\"product_image_{i}.png\"\n",
        "        gcs_blob_name = f\"{gcs_image_path}/{image_name}\"\n",
        "        gcs_uri = f\"gs://{bucket_name}/{gcs_blob_name}\"\n",
        "        defect_type = \"None\"\n",
        "        label = \"Status: Pass\"\n",
        "\n",
        "        if is_defect:\n",
        "            # Add a random defect\n",
        "            defect_type = random.choice([\"Scratch\", \"Crack\", \"Discoloration\"])\n",
        "            if defect_type == \"Scratch\":\n",
        "                draw.line((30, 30, 70, 70), fill=\"#FF3333\", width=2)\n",
        "                label = \"Status: Defect - Scratch detected on product surface.\"\n",
        "            elif defect_type == \"Crack\":\n",
        "                draw.line((30, 50, 70, 45), fill=\"#FFFFFF\", width=3)\n",
        "                label = \"Status: Defect - Crack identified in main body.\"\n",
        "            elif defect_type == \"Discoloration\":\n",
        "                draw.ellipse((55, 55, 75, 75), fill=\"#44AA44\")\n",
        "                label = \"Status: Defect - Discoloration spot found.\"\n",
        "\n",
        "        # Upload to GCS\n",
        "        if not _upload_image_blob(bucket, gcs_blob_name, img):\n",
        "            continue\n",
        "\n",
        "        manifest.append(\n",
        "            {\n",
        "                \"image_name\": image_name,\n",
        "                \"gcs_uri\": gcs_uri,\n",
        "                \"status\": \"Defect\" if is_defect else \"Pass\",\n",
        "                \"defect_type\": defect_type,\n",
        "                \"label\": label,\n",
        "            }\n",
        "        )\n",
        "\n",
        "        if (i + 1) % 20 == 0:\n",
        "            print(f\"  ...generated and uploaded {i + 1}/{num_images} images.\")\n",
        "\n",
        "    print(f\"Image generation complete. Uploaded {len(manifest)} images.\")\n",
        "    return pd.DataFrame(manifest)\n",
        "\n",
        "\n",
        "# Generate data (e.g., 200 samples for this demo)\n",
        "# For a real project, you'd need many more (100s or 1000s)\n",
        "image_manifest_df = generate_and_upload_images(\n",
        "    bucket_name=BUCKET_NAME,\n",
        "    gcs_image_path=IMAGE_DIR_GCS_PATH,\n",
        "    num_images=200,\n",
        "    defect_ratio=0.5,\n",
        ")\n",
        "\n",
        "print(\"\\n--- Image Manifest Sample ---\")\n",
        "print(image_manifest_df.head())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9jbGD8gFRdQ7"
      },
      "source": [
        "## Step 2: Prepare Tuning Data (JSONL)\n",
        "\n",
        "We convert the manifest DataFrame into the required JSON Lines format. Each line will contain a **multimodal prompt** (text + image) and the expected completion (the classification label)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w9pXnVJARdQ7"
      },
      "outputs": [],
      "source": [
        "def _build_tuning_example(prompt: str, image_uri: str, label: str) -> dict:\n",
        "    \"\"\"Constructs the dictionary for a single tuning example.\"\"\"\n",
        "    return {\n",
        "        \"contents\": [\n",
        "            {\n",
        "                \"role\": \"user\",\n",
        "                \"parts\": [\n",
        "                    {\"text\": prompt},\n",
        "                    {\"fileData\": {\"mimeType\": \"image/png\", \"fileUri\": image_uri}},\n",
        "                ],\n",
        "            },\n",
        "            {\"role\": \"model\", \"parts\": [{\"text\": label}]},\n",
        "        ]\n",
        "    }\n",
        "\n",
        "\n",
        "def create_tuning_jsonl_from_manifest(\n",
        "    manifest_df: pd.DataFrame,\n",
        ") -> list[dict[str, Any]]:\n",
        "    \"\"\"Creates JSONL data for Gemini multimodal supervised tuning.\"\"\"\n",
        "    print(\"\\n--- Preparing JSONL Tuning Data ---\")\n",
        "\n",
        "    jsonl_data = []\n",
        "    base_prompt = \"Analyze the following product image for manufacturing defects. Classify its status as 'Pass' or 'Defect' and provide a brief description if a defect is present.\"\n",
        "\n",
        "    for _, row in manifest_df.iterrows():\n",
        "        image_uri = row[\"gcs_uri\"]\n",
        "        target_label = row[\"label\"]\n",
        "\n",
        "        if not image_uri or pd.isna(image_uri):\n",
        "            print(\"Skipping row with missing image URI.\")\n",
        "            continue\n",
        "\n",
        "        # Format according to Gemini multimodal tuning requirements\n",
        "        # The 'user' role contains both the text prompt and the image file data.\n",
        "        instance = _build_tuning_example(base_prompt, image_uri, target_label)\n",
        "        jsonl_data.append(instance)\n",
        "\n",
        "    print(f\"Generated {len(jsonl_data)} JSONL instances.\")\n",
        "    return jsonl_data\n",
        "\n",
        "\n",
        "# Create JSONL data\n",
        "tuning_data_jsonl = create_tuning_jsonl_from_manifest(image_manifest_df)\n",
        "\n",
        "# Shuffle and Split data\n",
        "if tuning_data_jsonl:\n",
        "    random.shuffle(tuning_data_jsonl)\n",
        "    # Using 80% train, 10% validation, 10% test split\n",
        "    split_idx_val = int(len(tuning_data_jsonl) * 0.8)\n",
        "    split_idx_test = int(len(tuning_data_jsonl) * 0.9)\n",
        "\n",
        "    train_split = tuning_data_jsonl[:split_idx_val]\n",
        "    validation_split = tuning_data_jsonl[split_idx_val:split_idx_test]\n",
        "    test_split = tuning_data_jsonl[split_idx_test:]\n",
        "\n",
        "    print(\n",
        "        f\"Split sizes: Train={len(train_split)}, Validation={len(validation_split)}, Test={len(test_split)}\"\n",
        "    )\n",
        "\n",
        "    # Display a sample\n",
        "    print(\"\\n--- Sample JSONL Instance ---\")\n",
        "    print(json.dumps(train_split[0], indent=2))\n",
        "else:\n",
        "    print(\"\\nWarning: No tuning data generated.\")\n",
        "    # Initialize splits as empty lists to prevent errors later\n",
        "    train_split, validation_split, test_split = [], [], []"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-OARmnqbRdQ7"
      },
      "source": [
        "## Step 3: Upload Tuning Data to GCS\n",
        "\n",
        "The fine-tuning service reads data directly from Google Cloud Storage."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sb9OuMrjRdQ7"
      },
      "outputs": [],
      "source": [
        "import google.auth\n",
        "\n",
        "def save_jsonl_to_gcs(instances: list[dict[str, Any]], gcs_uri: str) -> None:\n",
        "    \"\"\"Saves a list of dictionaries as a JSONL file to GCS using Pandas.\"\"\"\n",
        "    if not instances:\n",
        "        print(f\"No instances to upload to {gcs_uri}. Skipping upload.\")\n",
        "        return\n",
        "\n",
        "    print(f\"Uploading {len(instances)} instances to {gcs_uri}...\")\n",
        "\n",
        "    try:\n",
        "        # Get the application default credentials\n",
        "        credentials, _ = google.auth.default()\n",
        "\n",
        "        # Convert list of dicts to DataFrame\n",
        "        df = pd.DataFrame(instances)\n",
        "\n",
        "        # Save DataFrame to GCS as JSONL\n",
        "        # We MUST pass the 'token' (credentials) to authenticate the request\n",
        "        storage_options = {\"project\": PROJECT_ID, \"token\": credentials}\n",
        "\n",
        "        df.to_json(\n",
        "            gcs_uri, orient=\"records\", lines=True, storage_options=storage_options\n",
        "        )\n",
        "\n",
        "        print(\"Upload complete.\")\n",
        "    except Exception as e:\n",
        "        print(f\"ERROR during GCS upload to {gcs_uri}: {e}\")\n",
        "        print(\n",
        "            \"Please ensure your GCS bucket is accessible and pandas has GCS permissions (installed via gcsfs).\"\n",
        "        )\n",
        "\n",
        "\n",
        "# Save splits to GCS\n",
        "save_jsonl_to_gcs(train_split, TRAIN_JSONL_GCS_URI)\n",
        "save_jsonl_to_gcs(validation_split, VALIDATION_JSONL_GCS_URI)\n",
        "save_jsonl_to_gcs(\n",
        "    test_split, TEST_JSONL_GCS_URI\n",
        ")  # Save test split for later evaluation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bPa8ylyURdQ8"
      },
      "source": [
        "## Step 4: Launch Fine-tuning Job\n",
        "\n",
        "We use the `google-genai` client **configured for Vertex AI** (`vertex_client`) to start the supervised tuning job, as fine-tuning management is a Vertex AI feature."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HF9z_1k3RdQ8"
      },
      "outputs": [],
      "source": [
        "TUNING_JOB_NAME = None  # Initialize\n",
        "if not train_split or not validation_split:\n",
        "    print(\"Skipping fine-tuning job launch as training or validation data is empty.\")\n",
        "else:\n",
        "    print(f\"Starting supervised fine-tuning job for model: {BASE_MODEL_ID}\")\n",
        "    print(f\"Tuned model display name: {TUNED_MODEL_DISPLAY_NAME}\")\n",
        "\n",
        "    training_dataset = {\n",
        "        \"gcs_uri\": TRAIN_JSONL_GCS_URI,\n",
        "    }\n",
        "\n",
        "    validation_dataset = genai_types.TuningValidationDataset(\n",
        "        gcs_uri=VALIDATION_JSONL_GCS_URI\n",
        "    )\n",
        "\n",
        "    try:\n",
        "        # Use the vertex_client configured specifically for Vertex AI operations\n",
        "        sft_tuning_job = vertex_client.tunings.tune(\n",
        "            base_model=BASE_MODEL_ID,\n",
        "            training_dataset=training_dataset,\n",
        "            config=genai_types.CreateTuningJobConfig(\n",
        "                adapter_size=\"ADAPTER_SIZE_FOUR\",  # Smaller adapter for faster tuning\n",
        "                epoch_count=5,  # Increased epochs for better specialization\n",
        "                tuned_model_display_name=TUNED_MODEL_DISPLAY_NAME,\n",
        "                validation_dataset=validation_dataset,\n",
        "            ),\n",
        "        )\n",
        "        print(\"\\nTuning job created:\")\n",
        "        print(sft_tuning_job)\n",
        "        TUNING_JOB_NAME = sft_tuning_job.name  # Save for monitoring\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"ERROR starting tuning job: {e}\")\n",
        "        # Attempt to list existing jobs with the same display name in case of interruption\n",
        "        try:\n",
        "            print(\n",
        "                f\"Checking for existing tuning jobs named '{TUNED_MODEL_DISPLAY_NAME}'...\"\n",
        "            )\n",
        "            existing_jobs = vertex_client.tunings.list(\n",
        "                page_size=100\n",
        "            )  # List might need pagination for many jobs\n",
        "            for job in existing_jobs:\n",
        "                # Check if config exists and has the attribute\n",
        "                job_config = getattr(job, \"config\", None)\n",
        "                if (\n",
        "                    job_config\n",
        "                    and getattr(job_config, \"tuned_model_display_name\", None)\n",
        "                    == TUNED_MODEL_DISPLAY_NAME\n",
        "                ):\n",
        "                    print(f\"Found existing job: {job.name} with state {job.state}\")\n",
        "                    TUNING_JOB_NAME = job.name  # Use the existing job name\n",
        "                    break\n",
        "        except Exception as list_e:\n",
        "            print(f\"Could not list existing tuning jobs: {list_e}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VV8iimM8RdQ8"
      },
      "source": [
        "**Note:** Fine-tuning can take a significant amount of time (potentially 30 minutes to several hours depending on the dataset size, base model, and adapter size)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vs9s8uLRRdQ8"
      },
      "source": [
        "## Step 5: Monitor Job"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9mVFYfL2RdQ8"
      },
      "outputs": [],
      "source": [
        "def monitor_tuning_job(job_name: str) -> Any:\n",
        "    \"\"\"Polls the tuning job until it reaches a terminal state.\"\"\"\n",
        "    print(f\"Monitoring tuning job: {job_name}\")\n",
        "    running_states = {\n",
        "        genai_types.JobState.JOB_STATE_PENDING,\n",
        "        genai_types.JobState.JOB_STATE_RUNNING,\n",
        "    }\n",
        "\n",
        "    while True:\n",
        "        try:\n",
        "            tuning_job = vertex_client.tunings.get(name=job_name)\n",
        "        except Exception as e:\n",
        "            print(f\"Error polling status: {e}. Retrying in 60s...\")\n",
        "            time.sleep(60)\n",
        "            continue\n",
        "\n",
        "        if tuning_job.state not in running_states:\n",
        "            return tuning_job\n",
        "\n",
        "        state_name = str(tuning_job.state).split(\".\")[-1]\n",
        "        print(f\"  Current state: {state_name}...\")\n",
        "        time.sleep(60)\n",
        "\n",
        "\n",
        "TUNED_MODEL_ENDPOINT = None\n",
        "\n",
        "if TUNING_JOB_NAME:\n",
        "    final_job = monitor_tuning_job(TUNING_JOB_NAME)\n",
        "\n",
        "    final_state_name = str(final_job.state).split(\".\")[-1]\n",
        "    print(f\"\\nTuning job finished with state: {final_state_name}\")\n",
        "\n",
        "    if final_job.state == genai_types.JobState.JOB_STATE_SUCCEEDED:\n",
        "        if hasattr(final_job, \"tuned_model\") and final_job.tuned_model.endpoint:\n",
        "            TUNED_MODEL_ENDPOINT = final_job.tuned_model.endpoint\n",
        "            print(f\"Tuned model endpoint ready: {TUNED_MODEL_ENDPOINT}\")\n",
        "        else:\n",
        "            print(\"Warning: Job succeeded but endpoint not found.\")\n",
        "    else:\n",
        "        error_msg = getattr(final_job, \"error\", \"Unknown error\")\n",
        "        print(f\"Tuning job failed. Error: {error_msg}\")\n",
        "else:\n",
        "    print(\"Skipping monitoring...\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JeiZwXsMRdQ8"
      },
      "source": [
        "## Step 6: Evaluate Tuned Model (Qualitative)\n",
        "\n",
        "We take a sample from our test set (which the model hasn't seen during tuning) and send the multimodal prompt (text + image) to the tuned endpoint to compare the prediction to the expected output."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0fG27PjrRdQ8"
      },
      "outputs": [],
      "source": [
        "import re  # <-- Import regex module\n",
        "\n",
        "from google.genai import types as genai_types  # Ensure genai_types is available\n",
        "\n",
        "def _extract_predicted_text(response: Any) -> str:\n",
        "    \"\"\"Safely extracts the predicted text using guard clauses.\"\"\"\n",
        "    if not response:\n",
        "        return \"(Response is None)\"\n",
        "\n",
        "    # 1. Check simplified text attribute\n",
        "    if hasattr(response, \"text\") and response.text:\n",
        "        return response.text.strip()\n",
        "\n",
        "    # 2. Validate candidates exist\n",
        "    if not hasattr(response, \"candidates\") or not response.candidates:\n",
        "        return \"(No candidates found)\"\n",
        "\n",
        "    first_candidate = response.candidates[0]\n",
        "\n",
        "    # 3. Validate finish reason\n",
        "    finish_reason = getattr(first_candidate, \"finish_reason\", None)\n",
        "    if finish_reason != genai_types.FinishReason.STOP:\n",
        "        return f\"(Generation stopped: {finish_reason})\"\n",
        "\n",
        "    # 4. Validate content parts\n",
        "    if not (hasattr(first_candidate, \"content\") and first_candidate.content.parts):\n",
        "        return \"(No content parts)\"\n",
        "\n",
        "    return first_candidate.content.parts[0].text.strip()\n",
        "\n",
        "\n",
        "def evaluate_qualitatively(\n",
        "    tuned_endpoint: str, test_data: list[dict[str, Any]], num_samples: int = 3\n",
        ") -> None:\n",
        "    \"\"\"Makes predictions with the tuned model and prints comparisons.\"\"\"\n",
        "    if not tuned_endpoint:\n",
        "        print(\"Tuned model endpoint not available. Skipping evaluation.\")\n",
        "        return\n",
        "\n",
        "    if not test_data:\n",
        "        print(\"No test data available for evaluation.\")\n",
        "        return\n",
        "\n",
        "    print(f\"\\n--- Qualitative Evaluation of Tuned Model ({tuned_endpoint}) ---\")\n",
        "\n",
        "    # Select random samples from the test set\n",
        "    samples = random.sample(test_data, min(num_samples, len(test_data)))\n",
        "\n",
        "    for i, sample in enumerate(samples):\n",
        "        print(f\"\\n--- Sample {i + 1} ---\")\n",
        "        # Ensure the sample structure is correct\n",
        "        try:\n",
        "            # Extract multimodal prompt parts\n",
        "            user_parts = sample[\"contents\"][0][\"parts\"]\n",
        "            expected_output = sample[\"contents\"][1][\"parts\"][0][\"text\"]\n",
        "\n",
        "            # Reconstruct the text and image parts for the prompt\n",
        "            prompt_text_part = user_parts[0][\"text\"]\n",
        "            image_file_part = user_parts[1][\"fileData\"]\n",
        "            image_uri = image_file_part[\"fileUri\"]\n",
        "            image_mime = image_file_part[\"mimeType\"]\n",
        "\n",
        "        except (KeyError, IndexError, TypeError) as e:\n",
        "            print(f\"Skipping sample due to unexpected format: {e}\")\n",
        "            print(f\"Problematic sample: {sample}\")\n",
        "            continue\n",
        "\n",
        "        print(f\"Input Prompt Text: {prompt_text_part}\")\n",
        "        print(f\"Input Image URI: {image_uri}\")\n",
        "        print(f\"\\nExpected Output: {expected_output}\")\n",
        "\n",
        "        try:\n",
        "            # Prepare contents for prediction (both text and image parts)\n",
        "            prediction_contents = [\n",
        "                {\n",
        "                    \"role\": \"user\",\n",
        "                    \"parts\": [\n",
        "                        {\"text\": prompt_text_part},\n",
        "                        {\"fileData\": {\"mimeType\": image_mime, \"fileUri\": image_uri}},\n",
        "                    ],\n",
        "                }\n",
        "            ]\n",
        "\n",
        "            # Use the vertex_client for predictions against the tuned endpoint\n",
        "            response = vertex_client.models.generate_content(\n",
        "                model=tuned_endpoint,\n",
        "                contents=prediction_contents,\n",
        "                config={\n",
        "                    \"temperature\": 0.1,  # Low temperature for deterministic classification\n",
        "                    \"max_output_tokens\": 2000,\n",
        "                },\n",
        "            )\n",
        "\n",
        "            # Safely access predicted text\n",
        "            predicted_output = _extract_predicted_text(response)\n",
        "\n",
        "            print(f\"Predicted Output: {predicted_output}\")\n",
        "\n",
        "            # --- NEW REGEX EVALUATION LOGIC ---\n",
        "            result_str = \"MISMATCH\"  # Default\n",
        "\n",
        "            if expected_output == \"Status: Pass\":\n",
        "                if predicted_output == \"Status: Pass\":\n",
        "                    result_str = \"MATCH\"\n",
        "            elif \"Defect\" in expected_output:\n",
        "                # 1. Check if prediction also says \"Defect\"\n",
        "                if \"Defect\" not in predicted_output:\n",
        "                    result_str = \"MISMATCH (Predicted 'Pass' for 'Defect')\"\n",
        "                else:\n",
        "                    # 2. Check if the key defect word is present\n",
        "                    key_defect_match = re.search(\n",
        "                        r\"(Scratch|Crack|Discoloration)\", expected_output\n",
        "                    )\n",
        "                    if key_defect_match:\n",
        "                        defect_type = key_defect_match.group(1)\n",
        "                        # Check if the predicted string contains the defect type (case-insensitive)\n",
        "                        if re.search(defect_type, predicted_output, re.IGNORECASE):\n",
        "                            result_str = \"MATCH (Regex)\"\n",
        "                        else:\n",
        "                            result_str = f\"MISMATCH (Missing key defect: {defect_type})\"\n",
        "                    else:\n",
        "                        # Fallback if it's a defect but not one of the known types\n",
        "                        result_str = \"MATCH (Regex - 'Defect' present)\"\n",
        "\n",
        "            print(f\"Result: {result_str}\")\n",
        "            # --- END OF NEW LOGIC ---\n",
        "\n",
        "        except Exception as e:\n",
        "            print(f\"ERROR during prediction for sample {i + 1}: {e}\")\n",
        "\n",
        "\n",
        "# Run qualitative evaluation (only if tuning succeeded and test data exists)\n",
        "evaluate_qualitatively(TUNED_MODEL_ENDPOINT, test_split)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jW2wMC1yRdQ8"
      },
      "source": [
        "## Step 7: Integrating Gemini for Reporting (Using Base Model)\n",
        "\n",
        "We can use a base Gemini model (accessed via Vertex AI) to summarize the fine-tuning job itself."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WZ5N1MAQRdQ8"
      },
      "outputs": [],
      "source": [
        "def generate_tuning_summary_with_gemini(tuning_job_details: Any) -> None:\n",
        "    \"\"\"Generates a summary of the tuning job using the Gemini API.\"\"\"\n",
        "    print(\"\\n--- Generating Tuning Job Summary with Gemini ---\")\n",
        "\n",
        "    if not tuning_job_details:\n",
        "        print(\"No tuning job details provided. Skipping summary.\")\n",
        "        return\n",
        "\n",
        "    # We will use the Vertex AI client, which is already initialized.\n",
        "    model_name_for_vertex_ai = \"gemini-2.5-flash\"  # Use a standard model for reporting\n",
        "    reporting_client = None\n",
        "\n",
        "    try:\n",
        "        # This uses the high-level vertexai SDK for base model generation\n",
        "        # Correctly import GenerativeModel from vertexai.preview.generative_models\n",
        "        from vertexai.preview.generative_models import GenerativeModel\n",
        "\n",
        "        reporting_client = GenerativeModel(model_name_for_vertex_ai)\n",
        "        print(f\"Using Vertex AI model ({model_name_for_vertex_ai}) for reporting.\")\n",
        "    except Exception as e:\n",
        "        print(\n",
        "            f\"Failed to initialize Vertex AI client for reporting with {model_name_for_vertex_ai}: {e}\"\n",
        "        )\n",
        "        print(\"Skipping summary generation.\")\n",
        "        return\n",
        "\n",
        "    try:\n",
        "        # Extract relevant details safely\n",
        "        job_name = getattr(tuning_job_details, \"name\", \"N/A\")\n",
        "        job_state_enum = getattr(\n",
        "            tuning_job_details, \"state\", genai_types.JobState.JOB_STATE_UNSPECIFIED\n",
        "        )  # Default to unspecified\n",
        "        job_state = str(job_state_enum).split(\".\")[\n",
        "            -1\n",
        "        ]  # Get 'SUCCEEDED', 'FAILED', etc.\n",
        "        base_model = getattr(tuning_job_details, \"base_model\", \"N/A\")\n",
        "        tuned_model_obj = getattr(tuning_job_details, \"tuned_model\", None)\n",
        "        tuned_endpoint = (\n",
        "            getattr(tuned_model_obj, \"endpoint\", \"N/A\") if tuned_model_obj else \"N/A\"\n",
        "        )\n",
        "        error_obj = getattr(tuning_job_details, \"error\", None)\n",
        "        error_message = str(error_obj) if error_obj else \"None\"\n",
        "        config_obj = getattr(tuning_job_details, \"config\", None)\n",
        "        display_name = (\n",
        "            getattr(config_obj, \"tuned_model_display_name\", \"N/A\")\n",
        "            if config_obj\n",
        "            else \"N/A\"\n",
        "        )\n",
        "\n",
        "        prompt = f\"\"\"Generate a brief status report for a Gemini model fine-tuning job for a 'Visual Defect Detection' use case.\n",
        "        Job Name: {job_name}\n",
        "        Base Model: {base_model}\n",
        "        Tuned Model Display Name: {display_name}\n",
        "        Final Status: {job_state}\n",
        "        Tuned Model Endpoint: {tuned_endpoint}\n",
        "        Error (if any): {error_message}\n",
        "\n",
        "        Summarize the outcome of this tuning job in 1-2 sentences, specifically mentioning its readiness for the manufacturing defect analysis task.\"\"\"\n",
        "\n",
        "        print(\"\\nSending request to Gemini...\")\n",
        "        # Use the selected reporting_client (Vertex AI based)\n",
        "        response = reporting_client.generate_content(prompt)\n",
        "\n",
        "        print(\"\\n--- Gemini Tuning Job Summary ---\")\n",
        "        # Handle potential response variations\n",
        "        response_text = _extract_predicted_text(response)\n",
        "        print(response_text)\n",
        "        print(\"---------------------------------\")\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"\\nERROR: Failed to generate Gemini summary: {e}\")\n",
        "        if (\n",
        "            \"permission denied\" in str(e).lower()\n",
        "            or \"consumer project\" in str(e).lower()\n",
        "        ):\n",
        "            print(\n",
        "                \"Please ensure the Vertex AI API is enabled in your project and the runtime environment has the correct permissions.\"\n",
        "            )\n",
        "        else:\n",
        "            print(\n",
        "                \"Please check your Vertex AI setup, model name, and network connection.\"\n",
        "            )\n",
        "\n",
        "\n",
        "# Get the final job details again using the vertex_client (which manages tuning)\n",
        "final_tuning_job = None\n",
        "if TUNING_JOB_NAME:\n",
        "    try:\n",
        "        # Use vertex_client to get the job status\n",
        "        final_tuning_job = vertex_client.tunings.get(name=TUNING_JOB_NAME)\n",
        "    except Exception as e:\n",
        "        print(f\"Error retrieving final tuning job details: {e}\")\n",
        "\n",
        "# Generate the summary using the Vertex Gemini client\n",
        "generate_tuning_summary_with_gemini(final_tuning_job)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "sft_gemini_visual_defect_detection.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
