Generative AI Google Brain

gemini-tuning-sft-gemini-summarization

4 min read

Supervised Fine Tuning with Gemini 2.5 Flash for Article Summarization

Source notebook

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

Fine-tunes Gemini 2.5 Flash on WikiLingua article summaries and evaluates ROUGE before and after tuning.

Summary

This notebook teaches supervised fine tuning of gemini-2.5-flash for article summarization using Vertex AI. It converts WikiLingua training data into the Gemini 2.5 tuning JSONL format, uploads train and validation files to Cloud Storage, runs a tuning job, and evaluates baseline versus tuned outputs with ROUGE metrics. It also retrieves tuning metrics from Vertex AI Experiment/Tensorboard resources and cleans up the experiment, endpoint, and bucket.

Key code patterns

Initialize clients

PROJECT_ID = "[your-project-id]"
REGION = "us-central1"
vertexai.init(project=PROJECT_ID, location=REGION)
client = genai.Client(vertexai=True, project=PROJECT_ID, location=REGION)

Sets the Google Cloud project, region, Vertex AI SDK, and Google GenAI SDK client for Vertex AI calls.

Convert tuning JSONL

def create_tuning_samples(file_path):
    instances = []
    with jsonlines.open(file_path) as reader:
        for obj in reader:
            contents = []
            for content in obj["messages"]:
                contents.append({"role": content["role"], "parts": [{"text": content["content"]}]})
            instances.append({"contents": contents})
    return instances

Transforms message records into the Gemini 2.5 tuning schema with contents, roles, and text parts.

Run supervised tuning

training_dataset = {"gcs_uri": f"{BUCKET_URI}/train/sft_train_samples.jsonl"}
validation_dataset = types.TuningValidationDataset(
    gcs_uri=f"{BUCKET_URI}/val/sft_val_samples.jsonl"
)
sft_tuning_job = client.tunings.tune(
    base_model=base_model,
    training_dataset=training_dataset,
    config=types.CreateTuningJobConfig(
        tuned_model_display_name=tuned_model_display_name,
        validation_dataset=validation_dataset,
    ),
)

Starts a Vertex AI supervised tuning job from GCS training and validation datasets.

Poll tuned endpoint

running_states = ["JOB_STATE_PENDING", "JOB_STATE_RUNNING"]
while tuning_job.state.name in running_states:
    tuning_job = client.tunings.get(name=tuning_job.name)
    time.sleep(10)
tuned_model = tuning_job.tuned_model.endpoint

Waits for the tuning job to finish and captures the deployed tuned model endpoint resource name.

ROUGE evaluation

scorer = rouge_scorer.RougeScorer(["rouge1", "rouge2", "rougeL"], use_stemmer=True)
response = client.models.generate_content(model=model, contents=document, config=config)
if not (response and response.candidates and response.candidates[0].content.parts):
    continue
scores = scorer.score(target=summary, prediction=response.text)

Compares generated summaries against reference summaries and skips blocked or empty responses.

Read tuning metrics

experiment = aiplatform.Experiment(experiment_name=experiment_name)
tensorboard_run = aiplatform.TensorboardRun(tensorboard_run_name)
metrics = tensorboard_run.read_time_series_data()
train_loss = get_metrics(metric="/train_total_loss")
eval_loss = get_metrics(metric="/eval_total_loss")

Extracts train and evaluation loss time series for post-tuning visualization.

Models & APIs used

  • Models: gemini-2.5-flash
  • APIs / services: Vertex AI, Cloud Storage
  • SDKs / libraries: google-genai, vertexai, google-cloud-aiplatform, jsonlines, pandas, plotly, rouge_score, tqdm

When to use this

Use this when you need Gemini 2.5 Flash to match a specific summarization style using labeled prompt and completion examples.

Gotchas & caveats

  • Requires a Google Cloud project, Vertex AI API enablement, and Colab or application-default authentication.
  • Training and validation data must be JSONL in Cloud Storage using the contents, role, parts, and text schema.
  • The notebook uses a Cloud Storage bucket for datasets and intermediate artifacts; same project is recommended.
  • Tuning automatically creates a Vertex AI endpoint and deploys the tuned model to it.
  • Vertex AI and Cloud Storage are billable components.
  • Colab users may need to restart the runtime after installing packages.
  • Evaluation and tuning take minutes on the provided batch and dataset.
  • Gemini 2.5 Flash adapter_size supports 1, 2, 4, and 8, with default value 4.

Best practices

  • Use high-quality, well-labeled, task-relevant training data.
  • Use a separate validation or evaluation dataset to measure tuned model performance.
  • Evaluate the base model before tuning and the tuned endpoint after tuning.
  • Choose task-appropriate metrics; the notebook uses ROUGE-L for summarization.
  • Experiment with generation parameters, prompt structure, epochs, and learning rate multiplier.
  • Monitor train and eval loss to inspect tuning behavior.
  • Handle blocked, empty, or errored model responses during evaluation.
  • Clean up experiments, endpoints, and Cloud Storage resources after the tutorial.

Graph View

Backlinks