Generative AI Google Brain

gemini-tuning-sft-gemini-qa

3 min read

Supervised fine-tuning with Gemini 2.0 Flash for Q&A using the Google Gen AI SDK

Source notebook

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

Fine-tunes Gemini 2.0 Flash on SQuAD Q&A with Vertex AI supervised tuning and evaluates EM/F1 gains.

Summary

This notebook teaches how to prepare SQuAD question-answering data for Gemini supervised fine-tuning using the Google Gen AI SDK on Vertex AI. It establishes a baseline with gemini-2.0-flash-001, converts training and validation data to JSONL with system instructions, uploads it to Cloud Storage, launches a tuning job, reads tuning metrics from Vertex AI Tensorboard, and re-evaluates the tuned model with EM and F1 scores.

Key code patterns

Initialize Gen AI SDK for Vertex AI

from google import genai
from google.genai import types
 
PROJECT_ID = str(os.environ.get("GOOGLE_CLOUD_PROJECT"))
LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")
client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)

Connects the Google Gen AI SDK to Vertex AI using project and region settings.

Baseline Gemini prediction

response = client.models.generate_content(
    model=base_model,
    contents=prompt,
    config={"system_instruction": systemInstruct, "temperature": 0.3},
)
return response.text

Uses the base Gemini model with a system instruction to measure pre-tuning Q&A behavior.

Build supervised tuning JSONL

jsonl_obj = {
  "systemInstruction": {"parts": [{"text": systemInstruct}]},
  "contents": [
    {"role": "user", "parts": [{"text": row.input_question}]},
    {"role": "model", "parts": [{"text": row.answers}]},
  ],
}

Formats each training example as user input plus target model output for Gemini supervised tuning.

Launch tuning job

sft_tuning_job = client.tunings.tune(
    base_model=base_model,
    training_dataset={"gcs_uri": train_dataset},
    config=types.CreateTuningJobConfig(
        adapter_size="ADAPTER_SIZE_EIGHT",
        epoch_count=1,
        tuned_model_display_name="gemini-flash-1.5-qa",
    ),
)

Starts Vertex AI supervised tuning from a GCS-hosted JSONL dataset.

Evaluate EM and F1

em = np.mean(y_true.combine(y_pred, exact_match_score))
f1 = np.mean(y_true.combine(y_pred, f1_score_squad))
return em, f1

Compares generated answers to ground truth using QA-specific exact match and token-overlap metrics.

Models & APIs used

  • Models: gemini-2.0-flash-001
  • APIs / services: Vertex AI, Cloud Storage
  • SDKs / libraries: google-genai, google-cloud-aiplatform, vertexai, pandas, numpy, plotly

When to use this

Use this pattern when you need a Gemini model to produce task-specific short-form Q&A answers from labeled examples.

Gotchas & caveats

  • Requires an existing Google Cloud project and Vertex AI API enabled.
  • Requires Colab authentication when running in Google Colab.
  • Training and validation data must be JSONL and stored at a Cloud Storage URI.
  • The notebook states Gemini tuning requires at least 100 examples.
  • Tuning uses billable Vertex AI and Cloud Storage resources.
  • The example sets epoch_count=1 to keep time and cost low.
  • The notebook warns the provided tuning job may take about 30 minutes.
  • Metrics visualizations require the tuning job to complete and validation metrics require a validation dataset.

Best practices

  • Establish a baseline with the default model before fine-tuning.
  • Use high-quality, well-labeled, task-relevant training data.
  • Keep test data formatted like training data to prevent training and serving skew.
  • Use system instructions to define desired model behavior and response style.
  • Use a separate evaluation set to assess model performance.
  • Choose evaluation metrics that reflect the QA task, such as EM and F1.
  • Start with default hyperparameters because the notebook states they are recommended for initial use.

Graph View

Backlinks