Generative AI Google Brain

gemini-tuning-dpo-gemini-data-prep-tuning

3 min read

Prepare High-Quality Preference Data for Gemini 2.5

Source notebook

Repo path: gemini/tuning/dpo_gemini_data_prep_tuning.ipynb · Open on GitHub · intermediate

Prepares and filters Gemini preference data with Vertex AI Gen AI Evaluation SDK for DPO and SFT.

Summary

This notebook teaches how to load UltraFeedback preference pairs, validate rows, and score preferred and rejected responses with Vertex AI Gen AI Evaluation SDK. It visualizes score distributions, applies absolute-threshold and win-margin filtering, writes cleaned Gemini-format preference JSONL, and optionally creates an SFT JSONL dataset from preferred responses. The output files are copied to a Cloud Storage bucket for use in a Gemini DPO tutorial.

Key code patterns

Initialize Vertex AI client and bucket

PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT", "")
LOCATION = "us-central1"
BUCKET_URI = f"gs://{PROJECT_ID}-data-prep"
!gcloud storage buckets create --location {LOCATION} --project {PROJECT_ID} {BUCKET_URI}
client = Client(project=PROJECT_ID, location=LOCATION)

Sets the project, region, storage location, and Vertex AI client used by evaluation and data export.

Filter valid preference rows

def is_valid_example(example):
    if not example.get("prompt"):
        return False
    if not example.get("chosen") or not example["chosen"][-1].get("content"):
        return False
    if not example.get("rejected") or not example["rejected"][-1].get("content"):
        return False
    return True
 
clean_dataset = dataset["train_prefs"].filter(is_valid_example)

Prevents empty prompts or missing preferred/rejected responses from entering evaluation and tuning data.

Evaluate response quality

metric_to_use = types.RubricMetric.GENERAL_QUALITY
result_preferred = client.evals.evaluate(dataset=df_preferred, metrics=[metric_to_use])
result_dispreferred = client.evals.evaluate(dataset=df_dispreferred, metrics=[metric_to_use])
df_preferred["score"] = extract_scores_from_result(result_preferred, "general_quality_v1")
df_dispreferred["score"] = extract_scores_from_result(result_dispreferred, "general_quality_v1")

Uses the Gen AI Evaluation SDK to score each preferred and dispreferred response with the general quality metric.

Apply win-margin filtering

MARGIN = 0.05
keep_mask_margin = (df_preferred["score"] - df_dispreferred["score"]) >= MARGIN
for i, keep in enumerate(keep_mask_margin):
    if keep:
        example = train_data[i]
        cleaned_data.append({"contents": [...], "completions": [...]})

Keeps only pairs where the preferred response scores meaningfully higher than the rejected response.

Write Gemini JSONL datasets

with open("cleaned_preference_data.jsonl", "w") as f:
    for entry in cleaned_data:
        f.write(json.dumps(entry) + "\n")
!gcloud storage cp cleaned_preference_data.jsonl {BUCKET_URI}/cleaned_preference_data.jsonl

Exports cleaned preference data in Gemini tuning format and uploads it to Cloud Storage.

Models & APIs used

  • APIs / services: Vertex AI, Cloud Storage
  • SDKs / libraries: vertexai, google-cloud-aiplatform[evaluation], datasets, matplotlib, pandas, numpy

When to use this

Use this pattern before Gemini DPO tuning when preference pairs need quality scoring, filtering, and conversion to Gemini JSONL formats.

Gotchas & caveats

  • Colab authentication is required when running in Google Colab.
  • PROJECT_ID must be set manually or through GOOGLE_CLOUD_PROJECT.
  • The notebook uses LOCATION = “us-central1”.
  • The Cloud Storage bucket is created with gcloud and may already exist.
  • The install cell warns that Colab may require a runtime restart.
  • Evaluation extraction appends None when metric results are missing or malformed.
  • The text says preferred score > 0.7 and dispreferred score < 0.5, but the code uses 0.65 and 0.95.

Best practices

  • Filter examples with missing prompts, chosen content, or rejected content before evaluation.
  • Prepare separate prompt-response dataframes for preferred and dispreferred responses with columns named prompt and response.
  • Sort evaluation case results by eval_case_index before attaching scores back to the dataframe.
  • Visualize preferred and dispreferred score distributions before choosing filtering thresholds.
  • Use win-margin filtering to keep pairs with clear quality separation.
  • Create an SFT dataset from the preferred responses as part of the two-phase SFT then DPO approach.
  • Clean up the Cloud Storage bucket to avoid charges.

Graph View

Backlinks