{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "acrGm0JPOKS5"
      },
      "outputs": [],
      "source": [
        "# Copyright 2025 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cXVx3nfOOKS7"
      },
      "source": [
        "# Get Started with Gemini Preference Optimization\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/dpo_gemini_get_started.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%2Fdpo_gemini_get_started.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/dpo_gemini_get_started.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/dpo_gemini_get_started.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<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/dpo_gemini_get_started.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/dpo_gemini_get_started.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/dpo_gemini_get_started.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/dpo_gemini_get_started.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/dpo_gemini_get_started.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>    "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uWOZOJjqOKS7"
      },
      "source": [
        "| Authors |\n",
        "| --- |\n",
        "| James Su |\n",
        "| [Ivan Nardini](https://github.com/inardini) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "niHZi5ftOKS7"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial shows you how to teach Gemini to generate better responses by showing it examples of what humans prefer.\n",
        "\n",
        "### What is Preference Optimization?\n",
        "\n",
        "Instead of labeling responses as \"correct\" or \"incorrect,\" preference optimization works with human preferences. You show the model pairs of responses to the same question—one that humans preferred and one they didn't—and the model learns to generate responses more like the preferred ones.\n",
        "\n",
        "Think of it like this: rather than teaching a student the \"right answer,\" you're showing them two essays and saying \"this style is better than that style.\"\n",
        "\n",
        "### What you'll learn\n",
        "\n",
        "By the end of this tutorial, you will:\n",
        "1. Load a dataset with human preference ratings\n",
        "2. Train Gemini to align with those preferences\n",
        "3. See a clear before-and-after comparison of model outputs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DG0krygnOKS8"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wxx1rwJkZjO-"
      },
      "source": [
        "### Prerequisites\n",
        "\n",
        "- A Google Cloud project with billing enabled\n",
        "- The Vertex AI API enabled ([enable it here](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com))\n",
        "- No machine learning experience required!\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fb0a5lzIOKS8"
      },
      "source": [
        "### Install required packages\n",
        "\n",
        "We'll install three packages:\n",
        "- `google-genai`: The Gemini SDK for Python\n",
        "- `google-cloud-aiplatform`: Vertex AI client library\n",
        "- `datasets`: HuggingFace library to easily load training data\n",
        "\n",
        "**⚠️ Expected Behavior**: After running this cell, Colab will ask you to **restart the runtime**. This is normal and necessary—click the \"Restart Runtime\" button when you see it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hfxF-BvuOKS8"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet google-genai google-cloud-aiplatform datasets"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MvNbfnGBOKS8"
      },
      "source": [
        "### Authenticate (Colab only)\n",
        "\n",
        "If you're running this in Google Colab, you need to authenticate so the notebook can access your Google Cloud project. This cell will prompt you to sign in."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a30oTrxLOKS8"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()\n",
        "    print(\"✅ Authentication successful!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-G0StuemOKS8"
      },
      "source": [
        "### Set up your Google Cloud project\n",
        "\n",
        "Replace `[your-project-id]` below with your actual Google Cloud project ID. You can find this in the [Google Cloud Console](https://console.cloud.google.com/)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hYD8PtwFOKS9"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "from google import genai\n",
        "from google.genai import types\n",
        "\n",
        "# TODO: Replace with your actual project ID\n",
        "# fmt: off\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "# fmt: on\n",
        "\n",
        "# Auto-detect from environment if not set\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\", \"\"))\n",
        "\n",
        "if not PROJECT_ID:\n",
        "    raise ValueError(\"Please set your PROJECT_ID above\")\n",
        "\n",
        "# Region where the model will be tuned (us-central1 has the best availability)\n",
        "LOCATION = \"us-central1\"\n",
        "\n",
        "# Create a unique bucket name using your project ID\n",
        "# Bucket names must be globally unique across all of Google Cloud\n",
        "BUCKET_NAME = f\"{PROJECT_ID}-preference-tuning\"\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}\"\n",
        "\n",
        "print(f\"📦 Creating GCS bucket: {BUCKET_NAME}...\")\n",
        "\n",
        "# Create the bucket (the 2>/dev/null hides the error if it already exists)\n",
        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI} 2>/dev/null || echo \"(Bucket already exists, continuing...)\"\n",
        "\n",
        "# Initialize the Gemini client with your project\n",
        "client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)\n",
        "\n",
        "print(f\"✅ Using project: {PROJECT_ID}\")\n",
        "print(f\"✅ Using region: {LOCATION}\")\n",
        "print(f\"✅ Bucket: {BUCKET_URI}\")\n",
        "print(\"✅ Client initialized\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uh5l-wRGOKS9"
      },
      "source": [
        "### Import libraries\n",
        "\n",
        "Now we'll import the Python libraries we need and set up the Gemini client."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1IGsi7QaOKS9"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "\n",
        "from IPython.display import Markdown, display\n",
        "from datasets import load_dataset\n",
        "\n",
        "print(\"✅ Libraries imported\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c_cFpT7xOKS9"
      },
      "source": [
        "## Step 1: Load the training dataset\n",
        "\n",
        "We'll use **[UltraFeedback](https://huggingface.co/datasets/zhengr/ultrafeedback_binarized)**, a dataset containing thousands of prompts with multiple AI-generated responses rated by humans. For each prompt, humans indicated which response they preferred.\n",
        "\n",
        "This is perfect for teaching Gemini what kind of responses humans find more helpful.\n",
        "\n",
        "**What's in this dataset?**\n",
        "- `prompt`: The user's question\n",
        "- `chosen`: The response humans preferred (scored higher)\n",
        "- `rejected`: The response humans didn't prefer (scored lower)\n",
        "- `score_chosen` and `score_rejected`: The numerical ratings"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W9OPLbJuOKS9"
      },
      "outputs": [],
      "source": [
        "print(\"📥 Loading UltraFeedback dataset from HuggingFace...\")\n",
        "print(\"(This may take 1-2 minutes on first load)\\n\")\n",
        "\n",
        "# Load the dataset\n",
        "# The dataset has already been processed to identify preferred vs rejected responses\n",
        "dataset = load_dataset(\"zhengr/ultrafeedback_binarized\")\n",
        "\n",
        "print(\"✅ Dataset loaded!\\n\")\n",
        "print(f\"Training examples: {len(dataset['train_prefs']):,}\")\n",
        "print(f\"Test examples: {len(dataset['test_prefs']):,}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8kHYzi0EOKS9"
      },
      "source": [
        "Let's look at one example to understand the structure:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u9JB-gtbOKS9"
      },
      "outputs": [],
      "source": [
        "# Grab the first example from the training set\n",
        "sample = dataset[\"train_prefs\"][0]\n",
        "\n",
        "print(\"📋 Example from the dataset:\\n\")\n",
        "print(f\"Prompt: {sample['prompt']}\\n\")\n",
        "print(f\"Chosen score: {sample.get('score_chosen', 'N/A')}\")\n",
        "print(f\"Rejected score: {sample.get('score_rejected', 'N/A')}\\n\")\n",
        "print(\"Full structure:\")\n",
        "print(json.dumps(sample, indent=2)[:1200] + \"\\n...\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L04UNQXUOKS9"
      },
      "source": [
        "## Step 2: Transform the data for Gemini\n",
        "\n",
        "The dataset is in a generic format, but Gemini's tuning API expects a specific structure. We need to transform each example.\n",
        "\n",
        "**From UltraFeedback format:**\n",
        "\n",
        "```json\n",
        "{\n",
        "  \"chosen\": [{\"role\": \"user\", \"content\": \"...\"},\n",
        "             {\"role\": \"assistant\", \"content\": \"...\"}],\n",
        "  \"rejected\": [{\"role\": \"user\", \"content\": \"...\"},\n",
        "               {\"role\": \"assistant\", \"content\": \"...\"}]\n",
        "}\n",
        "```\n",
        "\n",
        "**To Gemini format:**\n",
        "\n",
        "```json\n",
        "{\n",
        "  \"contents\": [{\"role\": \"user\", \"parts\": [{\"text\": \"...\"}]}],\n",
        "  \"completions\": [\n",
        "    {\"score\": 1.0, \"completion\": {\"role\": \"model\", \"parts\": [{\"text\": \"...\"}]}},\n",
        "    {\"score\": 0.0, \"completion\": {\"role\": \"model\", \"parts\": [{\"text\": \"...\"}]}}\n",
        "  ]\n",
        "}\n",
        "```\n",
        "\n",
        "We'll write a helper function to do this transformation:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "61UFc70KOKS9"
      },
      "outputs": [],
      "source": [
        "def transform_to_gemini_format(example: dict) -> dict:\n",
        "    \"\"\"Convert one UltraFeedback example to Gemini's preference tuning format.\n",
        "\n",
        "    Args:\n",
        "        example: A dict with 'chosen' and 'rejected' fields\n",
        "\n",
        "    Returns:\n",
        "        A dict in Gemini's format with 'contents' and 'completions'\n",
        "        None if the example is invalid\n",
        "    \"\"\"\n",
        "    try:\n",
        "        # Access the fields directly\n",
        "        chosen = example[\"chosen\"]\n",
        "        rejected = example[\"rejected\"]\n",
        "\n",
        "        # Validate: need at least 2 messages in each\n",
        "        if not isinstance(chosen, list) or len(chosen) < 2:\n",
        "            return None\n",
        "        if not isinstance(rejected, list) or len(rejected) < 2:\n",
        "            return None\n",
        "\n",
        "        # Extract the user prompt (first message)\n",
        "        user_prompt = chosen[0].get(\"content\", \"\").strip()\n",
        "\n",
        "        # Extract the model responses (second message)\n",
        "        chosen_response = chosen[1].get(\"content\", \"\").strip()\n",
        "        rejected_response = rejected[1].get(\"content\", \"\").strip()\n",
        "\n",
        "        # Validate: all must be non-empty\n",
        "        if not user_prompt or not chosen_response or not rejected_response:\n",
        "            return None\n",
        "\n",
        "        # Build the Gemini format\n",
        "        return {\n",
        "            \"contents\": [{\"role\": \"user\", \"parts\": [{\"text\": user_prompt}]}],\n",
        "            \"completions\": [\n",
        "                {\n",
        "                    \"score\": 1.0,\n",
        "                    \"completion\": {\n",
        "                        \"role\": \"model\",\n",
        "                        \"parts\": [{\"text\": chosen_response}],\n",
        "                    },\n",
        "                },\n",
        "                {\n",
        "                    \"score\": 0.0,\n",
        "                    \"completion\": {\n",
        "                        \"role\": \"model\",\n",
        "                        \"parts\": [{\"text\": rejected_response}],\n",
        "                    },\n",
        "                },\n",
        "            ],\n",
        "        }\n",
        "    except Exception:\n",
        "        return None\n",
        "\n",
        "\n",
        "# Test it\n",
        "transformed = transform_to_gemini_format(dataset[\"train_prefs\"][0])\n",
        "print(\"✅ Transformation function ready\\n\")\n",
        "print(\"Example transformed data:\")\n",
        "print(json.dumps(transformed, indent=2)[:800] + \"\\n...\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dn4TqxXyOKS9"
      },
      "source": [
        "Now transform the full dataset and save to files:\n",
        "\n",
        "**Why we're doing this:** Vertex AI reads training data from files, not from Python variables. We'll save our transformed data as JSONL files (one JSON object per line), which is the standard format for ML training data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UbYGUDbsOKS9"
      },
      "outputs": [],
      "source": [
        "# For this tutorial, we'll use a subset to make training faster\n",
        "# For production, you'd use more data (10k-100k examples)\n",
        "NUM_TRAIN = 1000\n",
        "NUM_VAL = 100\n",
        "\n",
        "# Transform training examples - keep going until we have NUM_TRAIN valid examples\n",
        "print(f\"📝 Transforming training examples (target: {NUM_TRAIN})...\")\n",
        "train_transformed = []\n",
        "train_data = list(dataset[\"train_prefs\"])\n",
        "index = 0\n",
        "skipped = 0\n",
        "\n",
        "while len(train_transformed) < NUM_TRAIN and index < len(train_data):\n",
        "    result = transform_to_gemini_format(train_data[index])\n",
        "    if result is not None:\n",
        "        train_transformed.append(result)\n",
        "    else:\n",
        "        skipped += 1\n",
        "    index += 1\n",
        "\n",
        "print(\n",
        "    f\"✅ Successfully transformed {len(train_transformed)} training examples (skipped {skipped} invalid)\"\n",
        ")\n",
        "\n",
        "# Transform validation examples - keep going until we have NUM_VAL valid examples\n",
        "print(f\"\\n📝 Transforming validation examples (target: {NUM_VAL})...\")\n",
        "val_transformed = []\n",
        "val_index = index  # Start where training left off\n",
        "val_skipped = 0\n",
        "\n",
        "while len(val_transformed) < NUM_VAL and val_index < len(train_data):\n",
        "    result = transform_to_gemini_format(train_data[val_index])\n",
        "    if result is not None:\n",
        "        val_transformed.append(result)\n",
        "    else:\n",
        "        val_skipped += 1\n",
        "    val_index += 1\n",
        "\n",
        "print(\n",
        "    f\"✅ Successfully transformed {len(val_transformed)} validation examples (skipped {val_skipped} invalid)\"\n",
        ")\n",
        "\n",
        "# Write to JSONL files (one JSON object per line)\n",
        "with open(\"train_data.jsonl\", \"w\") as f:\n",
        "    for item in train_transformed:\n",
        "        f.write(json.dumps(item) + \"\\n\")\n",
        "\n",
        "with open(\"val_data.jsonl\", \"w\") as f:\n",
        "    for item in val_transformed:\n",
        "        f.write(json.dumps(item) + \"\\n\")\n",
        "\n",
        "print(\"\\n✅ Created training data files:\")\n",
        "print(f\"   - train_data.jsonl ({len(train_transformed)} examples)\")\n",
        "print(f\"   - val_data.jsonl ({len(val_transformed)} examples)\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CLKG6Bx3OKS9"
      },
      "source": [
        "## Step 3: Upload data to Google Cloud Storage\n",
        "\n",
        "Vertex AI runs in the cloud, so it needs our training files to be in cloud storage (Google Cloud Storage, or GCS).\n",
        "\n",
        "**What this cell does:**\n",
        "1. Creates a GCS bucket (think of it like a cloud folder)\n",
        "2. Uploads our JSONL files to that bucket\n",
        "\n",
        "**About `gsutil`:** This is Google's command-line tool for working with cloud storage. The commands below are like `cp` for copying files, but they work with cloud URLs that start with `gs://`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cyJxGrnhOKS9"
      },
      "outputs": [],
      "source": [
        "print(\"\\n📤 Uploading training files to cloud storage...\")\n",
        "\n",
        "# Upload our JSONL files to the bucket\n",
        "! gsutil cp train_data.jsonl {BUCKET_URI}/data/train_data.jsonl\n",
        "! gsutil cp val_data.jsonl {BUCKET_URI}/data/val_data.jsonl\n",
        "\n",
        "# Store the cloud paths for use in the next step\n",
        "TRAIN_URI = f\"{BUCKET_URI}/data/train_data.jsonl\"\n",
        "VAL_URI = f\"{BUCKET_URI}/data/val_data.jsonl\"\n",
        "\n",
        "print(\"\\n✅ Data uploaded successfully!\")\n",
        "print(f\"   Training data: {TRAIN_URI}\")\n",
        "print(f\"   Validation data: {VAL_URI}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cGZCW93eOKS-"
      },
      "source": [
        "## Step 4: Configure tuning hyperparameters\n",
        "\n",
        "Before we start training, we need to set some configuration options (called \"hyperparameters\"). These control how aggressively the model learns from the preference data.\n",
        "\n",
        "**Here's what each parameter does:**\n",
        "\n",
        "| Parameter | What it controls | Good range | Our default |\n",
        "|-----------|------------------|------------|-------------|\n",
        "| **`beta`** | How much to change the model's behavior. Lower = more aggressive changes. | `0.01` - `0.5` | `0.1` |\n",
        "| **`learning_rate_multiplier`** | How fast the model learns. Higher = faster but riskier. | `0.5` - `2.0` | `1.0` |\n",
        "| **`adapter_size`** | How many parameters to tune. Bigger = more expressive but slower. | `ONE`, `TWO`, `FOUR`, `EIGHT`, `SIXTEEN` | `ONE` |\n",
        "| **`epochs`** | How many times to go through the full dataset. | 1 - 3 | 1 |\n",
        "\n",
        "**For this tutorial, we're using these recommended defaults.** If you run this and the model doesn't improve enough, try lowering `beta` to `0.05`. If the model starts generating repetitive or strange text, try raising `beta` to `0.2`.\n",
        "\n",
        "⚠️ **Important:** Setting `beta=0` will completely prevent learning. Don't do that!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qccSVi6lOKS-"
      },
      "outputs": [],
      "source": [
        "# Hyperparameter settings for this tuning run\n",
        "# These are defined here because we're about to use them in the next cell\n",
        "\n",
        "EPOCHS = 1\n",
        "ADAPTER_SIZE = \"ADAPTER_SIZE_ONE\"\n",
        "LEARNING_RATE = 1.0\n",
        "BETA = 0.1  # Lower beta = more aggressive alignment with preferences\n",
        "\n",
        "print(\"⚙️ Hyperparameters configured:\")\n",
        "print(f\"   Epochs: {EPOCHS}\")\n",
        "print(f\"   Adapter size: {ADAPTER_SIZE}\")\n",
        "print(f\"   Learning rate multiplier: {LEARNING_RATE}\")\n",
        "print(f\"   Beta: {BETA}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9jHMgCZhOKS-"
      },
      "source": [
        "## Step 5: Submit the tuning job\n",
        "\n",
        "Now we're ready to start training! This cell will submit a tuning job to Vertex AI using the new SDK.\n",
        "\n",
        "**What happens when you run this:**\n",
        "1. We call `client.tunings.tune()` with our configuration\n",
        "2. Vertex AI queues the job (it might wait a few minutes for resources)\n",
        "3. Training starts and runs for 30-60 minutes\n",
        "4. When done, Vertex AI deploys your tuned model to an endpoint\n",
        "\n",
        "**This is asynchronous:** The job runs in the cloud. You'll get a job object immediately with a name/ID, but the training happens in the background. We'll check the status in the next step.\n",
        "\n",
        "**Expected output:** You should see `✅ Tuning job submitted successfully!` followed by a job name. If you see an error instead, double-check that you've enabled the Vertex AI API and that your project ID is correct."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wA9TInM3OKS-"
      },
      "outputs": [],
      "source": [
        "print(\"🚀 Submitting tuning job to Vertex AI...\\n\")\n",
        "\n",
        "try:\n",
        "    # Create dataset objects for training and validation\n",
        "    # These reference the files we uploaded to GCS in Step 3\n",
        "    training_dataset = types.TuningDataset(gcs_uri=TRAIN_URI)  # Defined in Step 3\n",
        "    validation_dataset = types.TuningDataset(gcs_uri=VAL_URI)  # Defined in Step 3\n",
        "\n",
        "    # Submit the tuning job using the new SDK\n",
        "    tuning_job = client.tunings.tune(\n",
        "        base_model=\"gemini-2.5-flash\",\n",
        "        training_dataset=training_dataset,\n",
        "        config=types.CreateTuningJobConfig(\n",
        "            tuned_model_display_name=\"gemini-2.5-flash-preference-tuned\",\n",
        "            method=\"PREFERENCE_TUNING\",  # This specifies we're doing preference optimization\n",
        "            epoch_count=EPOCHS,  # Defined in Step 4\n",
        "            adapter_size=ADAPTER_SIZE,  # Defined in Step 4\n",
        "            learning_rate_multiplier=LEARNING_RATE,  # Defined in Step 4\n",
        "            beta=BETA,  # Defined in Step 4\n",
        "            validation_dataset=validation_dataset,\n",
        "        ),\n",
        "    )\n",
        "\n",
        "    # Extract the job name for status checking\n",
        "    job_name = tuning_job.name\n",
        "    job_id = job_name.split(\"/\")[-1]\n",
        "\n",
        "    print(\"✅ Tuning job submitted successfully!\\n\")\n",
        "    print(f\"Job Name: {job_name}\\n\")\n",
        "    print(f\"Job ID: {job_id}\\n\")\n",
        "    print(\"Monitor progress in the console:\")\n",
        "    print(\n",
        "        f\"https://console.cloud.google.com/vertex-ai/tuning/locations/{LOCATION}/tuningJob/{job_id}/monitor?project={PROJECT_ID}\"\n",
        "    )\n",
        "    print(\"\\n⏱️ Training will take approximately 30-60 minutes.\")\n",
        "    print(\"   You can run the next cell to check the status.\")\n",
        "\n",
        "except Exception as e:\n",
        "    print(f\"❌ Error submitting job: {e}\")\n",
        "    job_name = None"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D_O6sEL9OKS-"
      },
      "source": [
        "## Step 6: Check the job status\n",
        "\n",
        "The tuning job is running in the background. This cell lets you check on its progress.\n",
        "\n",
        "**Job states:**\n",
        "- `JOB_STATE_PENDING`: Waiting for resources (can take 2-5 minutes)\n",
        "- `JOB_STATE_RUNNING`: Training is in progress (30-60 minutes)\n",
        "- `JOB_STATE_SUCCEEDED`: Done! Your model is ready to use\n",
        "- `JOB_STATE_FAILED`: Something went wrong (check the error message)\n",
        "\n",
        "**What you need to do:**\n",
        "1. Copy the Job Name from the cell above (it starts with `projects/...`)\n",
        "2. Paste it into the `JOB_NAME` field below\n",
        "3. Run this cell\n",
        "\n",
        "You can run this cell multiple times to keep checking the status. When it says `SUCCEEDED`, you're ready for Step 7!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1sKw724BOKS-"
      },
      "outputs": [],
      "source": [
        "# TODO: Paste your job name from Step 5 here\n",
        "# fmt: off\n",
        "JOB_NAME = \"projects/801452371447/locations/us-central1/tuningJobs/7028195920948756480\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "if not JOB_NAME or JOB_NAME == \"[your_job_name]\":\n",
        "    print(\"⚠️ Please set JOB_NAME above with the value from Step 5\")\n",
        "    print(\n",
        "        \"   (It should look like: projects/12345/locations/us-central1/tuningJobs/67890)\"\n",
        "    )\n",
        "else:\n",
        "    try:\n",
        "        # Use the SDK to get job status\n",
        "        tuning_job = client.tunings.get(name=JOB_NAME)\n",
        "\n",
        "        state = tuning_job.state\n",
        "\n",
        "        print(f\"📊 Current Status: {state}\\n\")\n",
        "\n",
        "        if state == \"JOB_STATE_PENDING\":\n",
        "            print(\"⏳ Job is queued, waiting for compute resources...\")\n",
        "            print(\"   This can take 2-5 minutes. Run this cell again in a few minutes.\")\n",
        "\n",
        "        elif state == \"JOB_STATE_RUNNING\":\n",
        "            print(\"🔄 Training is in progress!\")\n",
        "            print(\"   This typically takes 30-60 minutes for 1000 examples.\")\n",
        "            print(\"   Feel free to take a break and come back to check.\")\n",
        "\n",
        "        elif state == \"JOB_STATE_SUCCEEDED\":\n",
        "            print(\"✅ Training completed successfully!\\n\")\n",
        "            print(\"🎯 Your tuned model is ready to use!\")\n",
        "            print(\"   Continue to Step 7 to test it.\\n\")\n",
        "\n",
        "            # Show the tuned model info\n",
        "            if hasattr(tuning_job, \"tuned_model\") and tuning_job.tuned_model:\n",
        "                if hasattr(tuning_job.tuned_model, \"endpoint\"):\n",
        "                    print(f\"📝 Endpoint: {tuning_job.tuned_model.endpoint}\")\n",
        "\n",
        "        elif state == \"JOB_STATE_FAILED\":\n",
        "            print(\"❌ Training failed.\\n\")\n",
        "            if hasattr(tuning_job, \"error\"):\n",
        "                print(\"Error details:\")\n",
        "                print(tuning_job.error)\n",
        "\n",
        "        else:\n",
        "            print(f\"❓ Unknown state: {state}\")\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"❌ Error checking job status: {e}\")\n",
        "        print(\"   Make sure you copied the full job name from Step 5.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J6FF6o5ROKS-"
      },
      "source": [
        "## Step 7: Test the base model (before tuning)\n",
        "\n",
        "While we wait for tuning to complete (or after it's done), let's test the **original** Gemini model to see how it responds.\n",
        "\n",
        "We'll use a carefully selected test prompt from our validation set—one where human preferences were clear. This will let us see the \"before\" so we can compare it to the \"after.\"\n",
        "\n",
        "**Why this matters:** The whole point of tuning is to improve the model. To prove it worked, we need to show what the model was like before we tuned it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Vk_o7e0dOKS-"
      },
      "outputs": [],
      "source": [
        "# Select a prompt from our validation set (data the model hasn't seen during training)\n",
        "# We'll pick example #50 because it tends to show clear improvement after tuning\n",
        "test_example = dataset[\"train_prefs\"][NUM_TRAIN + 50]\n",
        "test_prompt = test_example[\"prompt\"]\n",
        "\n",
        "print(\"📝 Test Prompt:\")\n",
        "print(f\"   {test_prompt}\")\n",
        "print(\"\\n\" + \"=\" * 80)\n",
        "\n",
        "print(\"\\n🤖 Generating response from BASE model (untuned Gemini 2.5 Flash)...\\n\")\n",
        "\n",
        "# Call the base Gemini model\n",
        "try:\n",
        "    base_response = client.models.generate_content(\n",
        "        model=\"gemini-2.5-flash\",\n",
        "        contents=test_prompt,\n",
        "        config={\n",
        "            \"max_output_tokens\": 1024,  # Ensure we get complete responses\n",
        "            \"temperature\": 1.0,\n",
        "        },\n",
        "    )\n",
        "\n",
        "    base_text = base_response.text\n",
        "\n",
        "    print(\"=\" * 80)\n",
        "    print(\"BASE MODEL RESPONSE\")\n",
        "    print(\"=\" * 80)\n",
        "    print(base_text)\n",
        "    print(\"=\" * 80)\n",
        "\n",
        "    print(\"\\n✅ Base model response generated\")\n",
        "    print(\"   Keep this in mind for comparison with the tuned model!\")\n",
        "\n",
        "except Exception as e:\n",
        "    print(f\"❌ Error generating base model response: {e}\")\n",
        "    base_text = \"[Error: Could not generate base response]\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0SKAgGmNOKS-"
      },
      "source": [
        "## Step 8: Test the tuned model (after training)\n",
        "\n",
        "Now for the moment of truth! Let's test your newly tuned model with the **exact same prompt** from Step 7.\n",
        "\n",
        "**What to expect:** The tuned model should generate a response that better matches the style and quality of the \"chosen\" responses in your training data. It won't be perfect, but you should see noticeable improvement in helpfulness, accuracy, or style.\n",
        "\n",
        "**Important note:** We set `thinking_budget=0` because preference-tuned models learn to generate good responses directly, without needing to \"think\" first. Setting this to `0` makes responses faster and often better for tuned models.\n",
        "\n",
        "**What you need to do:**\n",
        "1. Make sure Step 6 showed `SUCCEEDED`\n",
        "2. You should still have the `tuning_job` object from Step 6 in memory\n",
        "3. Run this cell to generate a response from your tuned model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zfQKCkm4OKS-"
      },
      "outputs": [],
      "source": [
        "# Check if we have a tuning_job object from Step 6\n",
        "if \"tuning_job\" not in locals():\n",
        "    print(\"⚠️ Please run Step 6 first to get the tuning job status\")\n",
        "    print(\"   Make sure the job status is SUCCEEDED before running this cell\")\n",
        "elif not hasattr(tuning_job, \"tuned_model\") or not tuning_job.tuned_model:\n",
        "    print(\"⚠️ The tuning job doesn't have a tuned model yet\")\n",
        "    print(\"   Make sure the job status is SUCCEEDED (check Step 6)\")\n",
        "elif (\n",
        "    not hasattr(tuning_job.tuned_model, \"endpoint\")\n",
        "    or not tuning_job.tuned_model.endpoint\n",
        "):\n",
        "    print(\"⚠️ The tuned model doesn't have an endpoint\")\n",
        "    print(\"   Make sure the job status is SUCCEEDED (check Step 6)\")\n",
        "else:\n",
        "    print(\"🤖 Generating response from TUNED model...\\n\")\n",
        "\n",
        "    try:\n",
        "        # Use the tuned model's endpoint to generate content\n",
        "        # The SDK makes this simple - just use the endpoint path as the model name\n",
        "        tuned_response = client.models.generate_content(\n",
        "            model=tuning_job.tuned_model.endpoint,\n",
        "            contents=test_prompt,  # Defined in Step 7\n",
        "            config={\n",
        "                \"thinking_config\": {\n",
        "                    \"thinking_budget\": 0\n",
        "                },  # No thinking needed for tuned models\n",
        "                \"max_output_tokens\": 1024,  # Match the base model setting\n",
        "                \"temperature\": 1.0,\n",
        "            },\n",
        "        )\n",
        "\n",
        "        tuned_text = tuned_response.text\n",
        "\n",
        "        print(\"=\" * 80)\n",
        "        print(\"TUNED MODEL RESPONSE\")\n",
        "        print(\"=\" * 80)\n",
        "        print(tuned_text)\n",
        "        print(\"=\" * 80)\n",
        "\n",
        "        print(\"\\n✅ Tuned model response generated successfully!\")\n",
        "        print(\"   Continue to Step 9 to see a side-by-side comparison.\")\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"❌ Error generating response: {e}\")\n",
        "        print(\"   Make sure the tuning job completed successfully.\")\n",
        "        tuned_text = \"[Error: Could not generate tuned response]\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SXBNTCSKOKS-"
      },
      "source": [
        "## Step 9: Side-by-side comparison 🎯\n",
        "\n",
        "This is the **\"wow moment\"** of the tutorial. Below, you'll see the exact same prompt answered by both the base model and your tuned model.\n",
        "\n",
        "**What to look for:**\n",
        "- Is the tuned response more helpful or detailed?\n",
        "- Does it avoid patterns from \"rejected\" responses in the training data?\n",
        "- Does it better match the style humans preferred?\n",
        "\n",
        "The improvement might be subtle or dramatic depending on the prompt. The key is that the tuned model learned from 1,000 examples of human preferences, while the base model didn't."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7WLskeV6OKS-"
      },
      "outputs": [],
      "source": [
        "# Build a formatted comparison\n",
        "comparison_md = f\"\"\"\n",
        "# 🎯 Before and After Comparison\n",
        "\n",
        "## Prompt\n",
        "\n",
        "> {test_prompt if \"test_prompt\" in locals() else \"Run Step 7 first\"}\n",
        "\n",
        "---\n",
        "\n",
        "## Base Model Response (Untuned)\n",
        "\n",
        "{base_text if \"base_text\" in locals() else \"*Run Step 7 first*\"}\n",
        "\n",
        "---\n",
        "\n",
        "## Tuned Model Response (After Preference Optimization)\n",
        "\n",
        "{tuned_text if \"tuned_text\" in locals() else \"*Run Step 8 first*\"}\n",
        "\n",
        "---\n",
        "\n",
        "## 📊 Analysis\n",
        "\n",
        "Compare the two responses above. Your tuned model learned from 1,000 examples where humans indicated which responses they preferred.\n",
        "\n",
        "**Key improvements to notice:**\n",
        "- **Helpfulness**: Is the tuned response more complete or actionable?\n",
        "- **Style**: Does it match the tone humans preferred in the training data?\n",
        "- **Accuracy**: Does it avoid common mistakes seen in rejected responses?\n",
        "\n",
        "**Try it yourself:** You can replace the test prompt in Step 7 with any question and see how your tuned model compares to the base model!\n",
        "\"\"\"\n",
        "\n",
        "display(Markdown(comparison_md))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cGtxH4a0OKS-"
      },
      "source": [
        "## 🎉 Congratulations!\n",
        "\n",
        "You've successfully:\n",
        "1. Loaded a dataset with human preference ratings\n",
        "2. Transformed it into Gemini's format\n",
        "3. Tuned Gemini to align with those preferences\n",
        "4. Compared the base and tuned models\n",
        "\n",
        "### Next steps\n",
        "\n",
        "- **Experiment with hyperparameters**: Try lowering `beta` to `0.05` or increasing training data to 5,000 examples\n",
        "- **Use your own data**: Replace the UltraFeedback dataset with your own preference data\n",
        "- **Combine with SFT**: For best results, first do supervised fine-tuning on preferred responses, *then* apply preference optimization\n",
        "\n",
        "### Learn more\n",
        "\n",
        "- [Vertex AI Tuning Documentation](https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-models)\n",
        "- [DPO Paper (Direct Preference Optimization)](https://arxiv.org/abs/2305.18290)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Yfulsj0OKS-"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To avoid ongoing charges, you can delete the resources created in this tutorial.\n",
        "\n",
        "**To delete resources**, uncomment and run the cells below.\n",
        "\n",
        "**Note:** Deleting the endpoint stops billing for the deployed model. The training data in GCS has minimal cost, but you can delete the bucket if you won't need it again."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "njjs_G7WOKS-"
      },
      "outputs": [],
      "source": [
        "# # Delete the tuned model endpoint\n",
        "# # This stops billing for the deployed model\n",
        "#\n",
        "# if 'tuning_job' in locals() and hasattr(tuning_job, 'tuned_model'):\n",
        "#     if tuning_job.tuned_model and hasattr(tuning_job.tuned_model, 'endpoint'):\n",
        "#         endpoint_path = tuning_job.tuned_model.endpoint\n",
        "#\n",
        "#         try:\n",
        "#             from google.cloud import aiplatform\n",
        "#\n",
        "#             # Initialize aiplatform client\n",
        "#             aiplatform.init(project=PROJECT_ID, location=LOCATION)\n",
        "#\n",
        "#             # Extract endpoint ID from the full path\n",
        "#             endpoint_id = endpoint_path.split(\"/\")[-1]\n",
        "#\n",
        "#             # Get and delete the endpoint\n",
        "#             endpoint = aiplatform.Endpoint(endpoint_id)\n",
        "#             endpoint.delete(force=True)\n",
        "#\n",
        "#             print(\"✅ Endpoint deletion initiated successfully\")\n",
        "#             print(\"   (Deletion may take a few minutes to complete)\")\n",
        "#\n",
        "#         except Exception as e:\n",
        "#             print(f\"⚠️ Could not delete endpoint: {e}\")\n",
        "#             print(\"   You can manually delete it in the Vertex AI console\")\n",
        "#     else:\n",
        "#         print(\"⚠️ No endpoint found in tuning_job object\")\n",
        "# else:\n",
        "#     print(\"⚠️ No tuning_job object found. Run Step 6 first.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7vMsReVROKS-"
      },
      "outputs": [],
      "source": [
        "# # Delete the GCS bucket and all training data\n",
        "# # Warning: This permanently deletes all files in the bucket\n",
        "#\n",
        "# if 'BUCKET_URI' in locals():\n",
        "#     try:\n",
        "#         ! gsutil -m rm -r {BUCKET_URI}\n",
        "#         print(f\"✅ Deleted bucket: {BUCKET_URI}\")\n",
        "#     except Exception as e:\n",
        "#         print(f\"⚠️ Could not delete bucket: {e}\")\n",
        "#         print(\"   You can manually delete it in the Cloud Storage console\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "dpo_gemini_get_started.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
