Supervised Fine-tuning Gemini 2.5 Flash for Predictive Maintenance

Source notebook

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

Fine-tunes Gemini 2.5 Flash on simulated sensor data to classify equipment maintenance status.

Summary

The notebook demonstrates supervised fine-tuning of Gemini 2.5 Flash for predictive maintenance using simulated industrial sensor readings and failure logs. It generates labeled JSONL training data, uploads train and validation splits to Cloud Storage, launches and monitors a Vertex AI tuning job, then qualitatively tests the tuned endpoint. It also uses the base Gemini model to summarize the tuning job outcome.

Key code patterns

Initialize Vertex AI clients

vertexai.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)
vertex_client = VertexClient(
    vertexai=True,
    project=PROJECT_ID,
    location=REGION,
)

Configures both the Vertex AI SDK and Google GenAI client for Vertex AI tuning operations.

Format Gemini tuning JSONL

instance = {
    "contents": [
        {"role": "user", "parts": [{"text": prompt}]},
        {"role": "model", "parts": [{"text": target_status}]},
    ]
}

Creates supervised tuning examples in the chat-style contents format required by Gemini tuning.

Upload JSONL to GCS

credentials, _ = google.auth.default()
df = pd.DataFrame(instances)
df.to_json(
    gcs_uri,
    orient="records",
    lines=True,
    storage_options={"project": PROJECT_ID, "token": credentials},
)

Stores tuning splits in Cloud Storage so the Vertex AI fine-tuning service can read them.

Launch supervised tuning

sft_tuning_job = vertex_client.tunings.tune(
    base_model=BASE_MODEL_ID,
    training_dataset={"gcs_uri": TRAIN_JSONL_GCS_URI},
    config=genai_types.CreateTuningJobConfig(
        adapter_size="ADAPTER_SIZE_FOUR",
        epoch_count=3,
        tuned_model_display_name=TUNED_MODEL_DISPLAY_NAME,
        validation_dataset=validation_dataset,
    ),
)

Starts supervised fine-tuning with a base model, GCS training data, validation data, adapter size, and epoch count.

Predict with tuned endpoint

response = vertex_client.models.generate_content(
    model=tuned_endpoint,
    contents=[{"role": "user", "parts": [{"text": user_prompt}]}],
    config={"temperature": 0.1, "max_output_tokens": 50},
)

Uses the tuned model endpoint directly as the model argument for deterministic classification-style predictions.

Models & APIs used

  • Models: gemini-2.5-flash
  • APIs / services: Vertex AI, Cloud Storage
  • SDKs / libraries: google-genai, vertexai, pandas, numpy, google-cloud-storage, gcsfs, google-auth

When to use this

Use this pattern when adapting Gemini to a narrow classification task from structured or summarized operational data.

Gotchas & caveats

  • Restart the kernel after installing notebook packages.
  • PROJECT_ID, REGION, and BUCKET_NAME must be set or discoverable before Vertex AI operations run.
  • The Cloud Storage bucket must exist or be creatable with the configured project, region, and permissions.
  • Pandas writes to GCS through gcsfs and requires application default credentials passed as a token.
  • Fine-tuning can take 30 minutes to several hours depending on dataset size, base model, and adapter size.
  • The notebook skips tuning if the training or validation split is empty.
  • The tuned model endpoint may be missing even if the tuning job succeeds, requiring console inspection.
  • Prediction response text extraction must handle candidates, finish reasons, and safety-filter-like response variations.

Best practices

  • Use train, validation, and test splits so evaluation samples are not seen during tuning.
  • Use Cloud Storage URIs for tuning datasets consumed by Vertex AI.
  • Set low temperature and max output tokens for classification-style tuned model evaluation.
  • Poll tuning job state until it leaves pending or running states before using the endpoint.
  • Check for empty data and failed job creation before launching or monitoring tuning.
  • Use a small adapter size and low epoch count for a demonstration tuning run.
  • Seed numpy and random for reproducible simulated data and shuffled splits.