{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1oLIdNNH4xTY"
      },
      "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": "dYFAt0Bs4xTZ"
      },
      "source": [
        "# Prepare High-Quality Preference Data for Gemini 2.5\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_data_prep_tuning.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_data_prep_tuning.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_data_prep_tuning.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_data_prep_tuning.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://storage.googleapis.com/github-repo/generative-ai/logos/GitHub_Invertocat_Dark.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_data_prep_tuning.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_data_prep_tuning.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_data_prep_tuning.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_data_prep_tuning.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_data_prep_tuning.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": "9pQonORh4xTa"
      },
      "source": [
        "| Authors |\n",
        "| --- |\n",
        "| Haichao Yu |\n",
        "| [Ivan Nardini](https://github.com/inardini) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cG8Zkh1u4xTa"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial teaches you how to prepare high-quality preference data for tuning Gemini 2.5. Before running Direct Preference Optimization (DPO), it's crucial to ensure your training data is high quality. We will use **Vertex AI Gen AI Evaluation SDK** to filter out low-quality preference pairs.\n",
        "\n",
        "### Why does data quality matter?\n",
        "\n",
        "Imagine you're teaching by showing examples. If half your \"good examples\" are actually mediocre, and half your \"bad examples\" are actually decent, you'll confuse the student. The same happens when training AI models.\n",
        "\n",
        "**Think of preference data like graded essays:**\n",
        "- Some essays marked \"A\" might actually be B-quality\n",
        "- Some essays marked \"F\" might actually deserve a C\n",
        "- If you don't filter these out, your model learns the wrong patterns\n",
        "\n",
        "### What you'll learn\n",
        "\n",
        "By the end of this tutorial, you will:\n",
        "1.  Load a preference dataset (UltraFeedback).\n",
        "2.  Use **Vertex AI Gen AI Evaluation SDK** to score responses.\n",
        "3.  Visualize quality scores to understand your data.\n",
        "4.  Apply filtering strategies (Absolute Threshold and Win Margin).\n",
        "5.  Create a Supervised Fine-Tuning (SFT) dataset from high-quality responses."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XCj5nYbJ4xTa"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m0QHzqS04xTb"
      },
      "source": [
        "### Install required packages\n",
        "\n",
        "We'll install the necessary packages for Google Cloud, data handling, and evaluation.\n",
        "\n",
        "**⚠️ Expected Behavior**: After running this cell, Colab will ask you to **restart the runtime**. This is normal and necessary."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XT-CQJGw4xTb"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet \"google-cloud-aiplatform[evaluation]\" datasets matplotlib"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KNH1nB2Q4xTc"
      },
      "source": [
        "### Authenticate (Colab only)\n",
        "\n",
        "If you're running this in Google Colab, you need to authenticate to access your Google Cloud project."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n2-GakFW4xTc"
      },
      "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": "i2RkKJ6m4xTc"
      },
      "source": [
        "### Set up your Google Cloud project\n",
        "\n",
        "Replace `[your-project-id]` below with your actual Google Cloud project ID."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tRILcSNa4xTc"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "from vertexai import Client, 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",
        "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",
        "LOCATION = \"us-central1\"\n",
        "BUCKET_NAME = f\"{PROJECT_ID}-data-prep\"\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}\"\n",
        "\n",
        "print(f\"📦 Creating GCS bucket: {BUCKET_NAME}...\")\n",
        "! gcloud storage buckets create --location {LOCATION} --project {PROJECT_ID} {BUCKET_URI} 2>/dev/null || echo \"(Bucket already exists, continuing...)\"\n",
        "\n",
        "# Initialize the Client\n",
        "client = Client(project=PROJECT_ID, location=LOCATION)\n",
        "\n",
        "print(f\"✅ Using project: {PROJECT_ID}\")\n",
        "print(f\"✅ Bucket: {BUCKET_URI}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yf9fb7rv4xTd"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lT-m_c4R4xTd"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from IPython.display import display\n",
        "from datasets import load_dataset\n",
        "\n",
        "print(\"✅ Libraries imported\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YL-i-tw64xTd"
      },
      "source": [
        "## Step 1: Load the training dataset\n",
        "\n",
        "We'll use the **UltraFeedback** dataset, same as in the DPO tutorial. This dataset contains prompts with human-rated responses."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CmsEnuRt4xTd"
      },
      "outputs": [],
      "source": [
        "print(\"📥 Loading UltraFeedback dataset...\")\n",
        "dataset = load_dataset(\"zhengr/ultrafeedback_binarized\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FuGVPOF_0MrI"
      },
      "outputs": [],
      "source": [
        "# Define the Filter Function\n",
        "def is_valid_example(example):\n",
        "    \"\"\"Returns True if the example has a valid prompt and valid responses.\"\"\"\n",
        "    # Check Prompt\n",
        "    if not example.get(\"prompt\"):\n",
        "        return False\n",
        "\n",
        "    # Check Chosen (Preferred) content\n",
        "    # Ensure list exists, is not empty, and last message has content\n",
        "    if not example.get(\"chosen\") or not example[\"chosen\"][-1].get(\"content\"):\n",
        "        return False\n",
        "\n",
        "    # Check Rejected (Dispreferred) content\n",
        "    if not example.get(\"rejected\") or not example[\"rejected\"][-1].get(\"content\"):\n",
        "        return False\n",
        "\n",
        "    return True\n",
        "\n",
        "\n",
        "# Apply Filter\n",
        "print(\"🧹 Filtering empty items from dataset...\")\n",
        "# This removes invalid rows from the entire dataset split efficiently\n",
        "clean_dataset = dataset[\"train_prefs\"].filter(is_valid_example)\n",
        "\n",
        "print(f\"Original size: {len(dataset['train_prefs'])}\")\n",
        "print(f\"Clean size:    {len(clean_dataset)}\")\n",
        "\n",
        "# Select Subset\n",
        "# Now we select 100 items knowing they are all valid\n",
        "# We use min() just in case the dataset is smaller than 100 after filtering\n",
        "n_examples = 100\n",
        "train_data = clean_dataset.select(range(min(n_examples, len(clean_dataset))))\n",
        "\n",
        "print(f\"✅ Loaded {len(train_data)} CLEAN examples for training.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cApaIy_F4xTd"
      },
      "source": [
        "## Step 2: Data Quality Evaluation\n",
        "\n",
        "We will use Vertex AI's Gen AI Evaluation SDK to score both preferred and dispreferred responses. This provides a streamlined way to run evaluations without manual API calls or parallelization."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tlq-Ou7O5rHN"
      },
      "source": [
        "### 2.1. Prepare Data for Evaluation\n",
        "\n",
        "We need to extract the prompt and the response for each example. We'll create two dataframes: one for preferred responses and one for dispreferred. The SDK expects columns named `prompt` and `response`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o9cRvIwR4xTd"
      },
      "outputs": [],
      "source": [
        "def prepare_eval_data(data, subset=\"preferred\"):\n",
        "    prompts = []\n",
        "    responses = []\n",
        "\n",
        "    print(f\"Processing subset: {subset}\")\n",
        "\n",
        "    for i, example in enumerate(data):\n",
        "        # Extract content first\n",
        "        current_prompt = example[\"prompt\"]\n",
        "\n",
        "        if subset == \"preferred\":\n",
        "            resp = example[\"chosen\"][-1][\"content\"]\n",
        "        else:\n",
        "            resp = example[\"rejected\"][-1][\"content\"]\n",
        "\n",
        "        # Append BOTH at the same time\n",
        "        # This guarantees prompts and responses have the exact same length\n",
        "        prompts.append(current_prompt)\n",
        "        responses.append(resp)\n",
        "\n",
        "    return pd.DataFrame({\"prompt\": prompts, \"response\": responses})\n",
        "\n",
        "\n",
        "df_preferred = prepare_eval_data(train_data, \"preferred\")\n",
        "df_dispreferred = prepare_eval_data(train_data, \"dispreferred\")\n",
        "\n",
        "print(\"✅ Data prepared for evaluation.\")\n",
        "display(df_preferred.head(2))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HeVLgsYu4xTd"
      },
      "source": [
        "### 2.2. Run Auto-Evaluation using Gen AI Evaluation SDK\n",
        "\n",
        "We'll use the `client.evals.evaluate` method. We'll use the `general_quality_v1` metric."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jF4BtXHN6rz9"
      },
      "outputs": [],
      "source": [
        "# Define the evaluation metric\n",
        "metric_to_use = types.RubricMetric.GENERAL_QUALITY\n",
        "metric_name_for_results = \"general_quality_v1\"\n",
        "\n",
        "# Helper Function to Extract Scores\n",
        "\n",
        "\n",
        "def extract_scores_from_result(result, metric_name):\n",
        "    \"\"\"Iterates through the EvaluationResult object to extract individual scores.\"\"\"\n",
        "    scores = []\n",
        "    # Ensure we process them in order of their index to match the DataFrame\n",
        "    sorted_cases = sorted(result.eval_case_results, key=lambda x: x.eval_case_index)\n",
        "\n",
        "    for case in sorted_cases:\n",
        "        try:\n",
        "            # Get the result for the first candidate (we only evaluate 1 response per row)\n",
        "            metric_result = case.response_candidate_results[0].metric_results[\n",
        "                metric_name\n",
        "            ]\n",
        "            scores.append(metric_result.score)\n",
        "        except (KeyError, IndexError, AttributeError):\n",
        "            scores.append(None)  # Handle cases where evaluation might have failed\n",
        "\n",
        "    return scores\n",
        "\n",
        "\n",
        "print(\"Running evaluation on preferred responses...\")\n",
        "result_preferred = client.evals.evaluate(\n",
        "    dataset=df_preferred,\n",
        "    metrics=[metric_to_use],\n",
        ")\n",
        "\n",
        "df_preferred[\"score\"] = extract_scores_from_result(\n",
        "    result_preferred, metric_name_for_results\n",
        ")\n",
        "\n",
        "print(\"Running evaluation on dispreferred responses...\")\n",
        "result_dispreferred = client.evals.evaluate(\n",
        "    dataset=df_dispreferred,\n",
        "    metrics=[metric_to_use],\n",
        ")\n",
        "df_dispreferred[\"score\"] = extract_scores_from_result(\n",
        "    result_dispreferred, metric_name_for_results\n",
        ")\n",
        "\n",
        "print(\"✅ Evaluation complete.\")\n",
        "print()\n",
        "display(df_preferred[[\"prompt\", \"score\"]].head())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FRUAaVAG4xTd"
      },
      "source": [
        "## Step 3: Visualize Quality Scores\n",
        "\n",
        "Now for the interesting part! Let's visualize the quality scores to see how well your \"preferred\" and \"dispreferred\" labels match actual quality.\n",
        "\n",
        "**What we're looking for:**\n",
        "- **Good data**: Preferred responses score higher than dispreferred ones (clear separation in the histogram)\n",
        "- **Problematic data**: Lots of overlap between the two distributions\n",
        "\n",
        "This visualization will help you decide where to set filtering thresholds in the next step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wr5kjp5u4xTd"
      },
      "outputs": [],
      "source": [
        "# Create histogram\n",
        "# Using 20 bins to get better granularity on the 0-1 scale\n",
        "bins = np.linspace(0, 1, 20)\n",
        "\n",
        "plt.figure(figsize=(10, 6))\n",
        "plt.hist(\n",
        "    df_preferred[\"score\"].dropna(),\n",
        "    bins,\n",
        "    alpha=0.5,\n",
        "    label=f\"Preferred (avg: {df_preferred['score'].mean():.2f})\",\n",
        "    color=\"green\",\n",
        ")\n",
        "plt.hist(\n",
        "    df_dispreferred[\"score\"].dropna(),\n",
        "    bins,\n",
        "    alpha=0.5,\n",
        "    label=f\"Dispreferred (avg: {df_dispreferred['score'].mean():.2f})\",\n",
        "    color=\"red\",\n",
        ")\n",
        "\n",
        "# Axis labels reflect the 0-1 scale\n",
        "plt.xlabel(\"Quality Score (0.0 = poor, 1.0 = perfect)\")\n",
        "plt.ylabel(\"Number of Responses\")\n",
        "plt.title(\"Distribution of Quality Scores: Preferred vs Dispreferred Responses\")\n",
        "plt.legend()\n",
        "plt.grid(True, alpha=0.3)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dH2JYRWT4xTe"
      },
      "source": [
        "## Step 4: Data Cleaning (Filtering)\n",
        "\n",
        "Now we'll filter the dataset based on the scores we just computed.\n",
        "\n",
        "### Filtering Strategies\n",
        "\n",
        "| Strategy | When to Use | How It Works |\n",
        "|----------|-------------|--------------|\n",
        "| **Absolute Threshold** | You want to set a minimum quality bar | Keep if preferred > threshold AND dispreferred < threshold |\n",
        "| **Win Margin** | You want clear winners | Keep if preferred score is at least X points higher than dispreferred |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DDeK5qPA5yEW"
      },
      "source": [
        "### Strategy 1: Absolute Threshold\n",
        "\n",
        "We keep pairs where:\n",
        "- Preferred score > 0.7 (out of 5)\n",
        "- Dispreferred score < 0.5"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XZjNaajT4xTe"
      },
      "outputs": [],
      "source": [
        "PREFERRED_THRESHOLD = 0.65\n",
        "DISPREFERRED_THRESHOLD = 0.95\n",
        "\n",
        "# Calculate how much data each rule removes.\n",
        "remove_mask_abs_preferred = df_preferred[\"score\"] <= PREFERRED_THRESHOLD\n",
        "remove_mask_abs_dispreferred = df_dispreferred[\"score\"] >= DISPREFERRED_THRESHOLD\n",
        "keep_mask_abs = (df_preferred[\"score\"] > PREFERRED_THRESHOLD) & (\n",
        "    df_dispreferred[\"score\"] < DISPREFERRED_THRESHOLD\n",
        ")\n",
        "\n",
        "print(\n",
        "    f\"Pairs removed with preferred response filtering only: {remove_mask_abs_preferred.sum()} out of {len(train_data)}\"\n",
        ")\n",
        "print(\n",
        "    f\"Pairs removed with dispreferred response filtering only: {remove_mask_abs_dispreferred.sum()} out of {len(train_data)}\"\n",
        ")\n",
        "print(\n",
        "    f\"Pairs kept with Absolute Threshold: {keep_mask_abs.sum()} out of {len(train_data)}\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hPNa27Sj4xTe"
      },
      "source": [
        "### Strategy 2: Win Margin\n",
        "\n",
        "We keep pairs where the preferred score is at least 1 point higher than the dispreferred score."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TIqEc5lb4xTe"
      },
      "outputs": [],
      "source": [
        "MARGIN = 0.05\n",
        "keep_mask_margin = (df_preferred[\"score\"] - df_dispreferred[\"score\"]) >= MARGIN\n",
        "\n",
        "print(f\"Pairs kept with Win Margin: {keep_mask_margin.sum()} out of {len(train_data)}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "McmR4HyJ4xTe"
      },
      "source": [
        "### Save Cleaned Data\n",
        "\n",
        "Let's save the cleaned data using the Win Margin strategy, transformed into Gemini format."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "spFOD9DI4xTe"
      },
      "outputs": [],
      "source": [
        "cleaned_data = []\n",
        "for i, keep in enumerate(keep_mask_margin):\n",
        "    if keep:\n",
        "        example = train_data[i]\n",
        "        gemini_format = {\n",
        "            \"contents\": [{\"role\": \"user\", \"parts\": [{\"text\": example[\"prompt\"]}]}],\n",
        "            \"completions\": [\n",
        "                {\n",
        "                    \"score\": 1.0,\n",
        "                    \"completion\": {\n",
        "                        \"role\": \"model\",\n",
        "                        \"parts\": [{\"text\": example[\"chosen\"][-1][\"content\"]}],\n",
        "                    },\n",
        "                },\n",
        "                {\n",
        "                    \"score\": 0.0,\n",
        "                    \"completion\": {\n",
        "                        \"role\": \"model\",\n",
        "                        \"parts\": [{\"text\": example[\"rejected\"][-1][\"content\"]}],\n",
        "                    },\n",
        "                },\n",
        "            ],\n",
        "        }\n",
        "        cleaned_data.append(gemini_format)\n",
        "\n",
        "output_file = \"cleaned_preference_data.jsonl\"\n",
        "with open(output_file, \"w\") as f:\n",
        "    for entry in cleaned_data:\n",
        "        f.write(json.dumps(entry) + \"\\n\")\n",
        "\n",
        "!gcloud storage cp {output_file} {BUCKET_URI}/{output_file}\n",
        "print(f\"✅ Cleaned data saved to {BUCKET_URI}/{output_file}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QT89EumD4xTe"
      },
      "source": [
        "## Step 5: Create SFT Dataset (Optional)\n",
        "\n",
        "**Bonus step!** You can extract just the preferred responses to create a supervised fine-tuning (SFT) dataset.\n",
        "\n",
        "### Why would you want this?\n",
        "\n",
        "The [preference optimization tutorial](https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/tuning/dpo_gemini.ipynb) recommends a **two-phase approach**:\n",
        "\n",
        "1. **Phase 1: SFT (Supervised Fine-Tuning)** - Train the model on just the good responses\n",
        "2. **Phase 2: DPO (Direct Preference Optimization)** - Show the model pairs to refine its judgment\n",
        "\n",
        "**Think of it like teaching:**\n",
        "- **SFT** = \"Here are examples of good essays\" (establishes baseline competence)\n",
        "- **DPO** = \"Now here's why this essay is better than that one\" (refines judgment)\n",
        "\n",
        "This two-phase approach often produces better results than jumping straight to preference optimization."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xHdzGzyK4xTe"
      },
      "outputs": [],
      "source": [
        "sft_data = []\n",
        "for entry in cleaned_data:\n",
        "    # Extract preferred response (score 1.0)\n",
        "    preferred_resp = next(\n",
        "        c[\"completion\"] for c in entry[\"completions\"] if c[\"score\"] == 1.0\n",
        "    )\n",
        "    sft_entry = {\"contents\": [entry[\"contents\"][0], preferred_resp]}\n",
        "    sft_data.append(sft_entry)\n",
        "\n",
        "sft_output_file = \"sft_data.jsonl\"\n",
        "with open(sft_output_file, \"w\") as f:\n",
        "    for entry in sft_data:\n",
        "        f.write(json.dumps(entry) + \"\\n\")\n",
        "\n",
        "!gcloud storage cp {sft_output_file} {BUCKET_URI}/{sft_output_file}\n",
        "print(f\"✅ SFT data saved to {BUCKET_URI}/{sft_output_file}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kSG0pCDR4xTe"
      },
      "source": [
        "## Conclusion\n",
        "\n",
        "You have now prepared a high-quality preference dataset and an SFT dataset using the Vertex AI Gen AI Evaluation SDK. These are ready to be used in the [Gemini DPO Tutorial](https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/tuning/dpo_gemini.ipynb).\n",
        "\n",
        "### Clean up\n",
        "\n",
        "To avoid charges, you can delete the GCS bucket created in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X6JD1n3Z4xTe"
      },
      "outputs": [],
      "source": [
        "# !gcloud storage rm --recursive {BUCKET_URI}"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "dpo_gemini_data_prep_tuning.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
