Supervised Fine-tuning Gemini for Predictive Maintenance

[aniketagrawal]

Forging Specialist AI: Coder’s Guide to Tuning

Foundation models like Gemini are the raw material of the AI revolution—powerful, generalist intellects. But raw material doesn’t solve specific business problems. Real-world value is unlocked by forging these generalists into specialists: a legal AI that understands contract law, a financial AI that detects sophisticated fraud, a manufacturing AI that can hear the subtle signs of impending machine failure.

Related Assets:

1. Code: sft_gemini_predictive_maintenance.ipynb
2. Blog: Image tuning Gemini as GenAI quality control inspector

How do you build this AI factory? How do you create a process that is not just a one-off science experiment, but a repeatable, governed, and scalable engineering discipline?

This is the Vertex AI Blueprint.

This guide is a definitive, architectural deep-dive into using Google’s unified AI platform to systematically forge specialist models. We will use a real-world, high-value problem—building a “Digital Technician” for predictive maintenance—as our illustrative thread. We will dissect each Vertex AI capability, providing heavily annotated code directly from our working notebook and explaining the strategic engineering decisions that enable speed, governance, and scale.

The Customer Journey: From MLOps Chaos to a Unified AI Factory

Imagine Alex, a Lead ML Platform Engineer.

The “Before” State (The Pain): Alex’s world is a fragmented landscape of operational friction.

• Data Silos: Sensor data is in a data historian, logs are in Elasticsearch, and training data ends up as random CSVs in a GCS bucket.
• Isolated Development: Data scientists work in local Jupyter notebooks. Their code is not versioned, and their environments are not reproducible.
• Manual Deployment Hell: To deploy a model, Alex must manually containerize it with Docker, write a custom Flask API, configure a Kubernetes cluster with autoscaling, and set up a load balancer. It’s a two-week process for every new model version.
• No Governance: When a production model’s performance degrades, there is no audit trail. Which dataset was it trained on? What were the hyperparameters? No one knows for sure.

The “After” State (The Vision): Alex will architect a new workflow entirely on Vertex AI, creating a true AI factory.

• Unified Data Source: All data is staged in a versioned, secure GCS bucket, acting as the single source of truth.
• Serverless Specialization: A foundation model is efficiently fine-tuned using a managed, serverless training service. No GPUs to manage.
• Automated Governance: Every artifact, metric, and action is automatically tracked, versioned, and connected in a queryable lineage graph.
• One-Click Deployment: The specialized model is deployed to a secure, auto-scaling endpoint with a single command.

Section 0: Prerequisites and Environment Setup

A production-grade system requires a solid foundation. This section mirrors the initial setup cells of our notebook.

0.1. Project and API Configuration

1. Google Cloud Project: You need a Google Cloud project with billing enabled.
2. Enable APIs: Ensure the following APIs are enabled: aiplatform.googleapis.com (Vertex AI) and storage.googleapis.com (Cloud Storage).
3. Permissions (IAM): Your user account or service account needs Vertex AI User, Storage Admin, and Service Account User roles.

0.2. Development Environment and Authentication

A Vertex AI Workbench instance is highly recommended. The following code, taken directly from our notebook, handles authentication and configuration.

import os
import google.generativeai as genai
import vertexai
from google.genai import Client as VertexGenaiLowLevelClient

# --- Vertex AI Configuration ---
PROJECT_ID = ""  # @param {type: "string"}
REGION = ""      # @param {type: "string"}
BUCKET_NAME = "" # @param {type: "string"}
BUCKET_URI = f"gs://{BUCKET_NAME}"

# --- Authentication ---
# This block handles authentication, which is often seamless in a Vertex AI Workbench.
# It attempts to get the project ID from the gcloud environment if not set.
if not PROJECT_ID:
    from google.colab import auth
    auth.authenticate_user()
    import subprocess
    PROJECT_ID = (
        subprocess.check_output(["gcloud", "config", "get-value", "project"])
        .decode("utf-8")
        .strip()
    )

# --- GCS Bucket Creation ---
# This command ensures the GCS bucket exists, creating it if necessary.
# It's an idempotent operation, safe to run multiple times.
creation_command = f"gsutil ls {BUCKET_URI} > /dev/null 2>&1 || gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
os.system(creation_command)

# --- SDK Initialization ---
# This is a critical step. It configures all subsequent SDK calls to use your
# specific project, region, and staging bucket for any temporary artifacts.
vertexai.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)

# The google-genai library has two modes. We explicitly create a client
# configured for the Vertex AI backend for platform operations like tuning.
vertex_client = VertexGenaiLowLevelClient(
    vertexai=True, project=PROJECT_ID, location=REGION
)

0.3. Global Constants

Defining constants at the start of the notebook ensures consistency and avoids magic strings.

# The specific, tunable version of the Gemini model we will use as our base.
BASE_MODEL_ID = "gemini-2.0-flash-001"

# A dynamic name for our tuned model to ensure uniqueness in the Model Registry.
TUNED_MODEL_DISPLAY_NAME = f"pred-maint-gemini-tuned-{int(time.time())}"

# GCS paths for our datasets, derived from the bucket URI.
DATA_DIR_GCS = f"{BUCKET_URI}/pred_maint_tuning_data"
TRAIN_JSONL_GCS_URI = f"{DATA_DIR_GCS}/train_data.jsonl"
VALIDATION_JSONL_GCS_URI = f"{DATA_DIR_GCS}/validation_data.jsonl"
TEST_JSONL_GCS_URI = f"{DATA_DIR_GCS}/test_data.jsonl"

---

Part 1: The Data Foundation - Native Integration with Cloud Storage

Vertex AI is designed around a data-centric workflow with Google Cloud Storage (GCS) at its core.

1.1. Capability Spotlight: GCS as an AI Datalake

• Direct Integration: Vertex AI services read directly from GCS URIs (gs://...), eliminating complex data loaders.
• Performance: GCS provides high-throughput access, crucial for feeding GPUs without I/O bottlenecks.
• Versioning: GCS object versioning creates an immutable history of your datasets, a cornerstone of reproducibility.

1.2. The JSONL Format: The Lingua Franca of Tuning

We must prepare our data in the JSON Lines (JSONL) format. Unlike a single massive JSON array, this format is streamable. The tuning service can read and process the file line by line, allowing it to handle terabyte-scale datasets without loading the entire file into memory.

1.3. Code in Action: Data Transformation and Upload

The notebook’s create_tuning_jsonl function performs the vital translation from raw sensor data into the structured prompt-response format. The following function then uploads this data.

def save_jsonl_to_gcs(instances: list[dict], gcs_uri: str):
    """
    Serializes a list of Python dictionaries to a JSONL string and uploads it to GCS.
    This function bridges the gap between our in-memory data and the cloud data source.
    """
    if not instances:
        print(f"No instances to upload to {gcs_uri}. Skipping upload.")
        return

    # 1. Parse the GCS URI to get the bucket and the path (blob name) for the file.
    parts = gcs_uri[5:].split("/", 1)
    bucket_name = parts[0]
    blob_name = parts[1]

    # 2. Convert each Python dictionary into a JSON string and join them with newlines.
    jsonl_content = "\n".join([json.dumps(inst) for inst in instances])

    # 3. Instantiate the GCS client and upload the content.
    storage_client = storage.Client(project=PROJECT_ID)
    bucket = storage_client.bucket(bucket_name)
    blob = bucket.blob(blob_name)
    blob.upload_from_string(jsonl_content, content_type="application/jsonl")
    print(f"Uploading {len(instances)} instances to {gcs_uri}...Upload complete.")

# The main script calls this for each data split.
# save_jsonl_to_gcs(train_split, TRAIN_JSONL_GCS_URI)
# save_jsonl_to_gcs(validation_split, VALIDATION_JSONL_GCS_URI)
# save_jsonl_to_gcs(test_split, TEST_JSONL_GCS_URI)

---

Part 2: The Tuning Engine - Serverless, Efficient Specialization

The Vertex AI Training service is our AI factory’s core engine.

2.1. The Serverless Paradigm: Abstracting Infrastructure Away

When you launch a tuning job, you do not manage VMs, GPUs, networks, or clusters. You submit a job definition, and Vertex AI’s fleet management system handles provisioning, orchestration, fault tolerance (restarting from checkpoints if a node fails), and tear-down. This transforms a complex infrastructure problem into a simple API call.

2.2. The “Why” of PEFT/LoRA: A Deep-Dive into Efficient Specialization

Full fine-tuning is slow, expensive, and risks “catastrophic forgetting.” PEFT (Parameter-Efficient Fine-Tuning) via LoRA (Low-Rank Adaptation) is the solution. It freezes the base model’s weights and injects small, trainable “adapter” matrices. We only train these adapters, which represent the delta for the new task. On Vertex AI, this state-of-the-art technique is enabled with a single parameter.

2.3. Code in Action: Launching the Tuning Job

The Vertex AI SDK provides a declarative interface. We define what we want, and the platform figures out how.

# This code is from the notebook's "Step 4: Launch Fine-tuning Job"
TUNING_JOB_NAME = None
if not train_split or not validation_split:
    print("Skipping fine-tuning job launch as data is empty.")
else:
    # 1. Define the datasets using the GCS URIs.
    training_dataset = {"gcs_uri": TRAIN_JSONL_GCS_URI}
    validation_dataset = genai_types.TuningValidationDataset(gcs_uri=VALIDATION_JSONL_GCS_URI)

    # 2. Define the tuning configuration.
    tuning_config = genai_types.CreateTuningJobConfig(
        # This is the key to enabling PEFT/LoRA. 'ADAPTER_SIZE_FOUR' is a hyperparameter
        # controlling the capacity of the adapter. It's a simple knob for a complex technique.
        adapter_size="ADAPTER_SIZE_FOUR",
        
        # An epoch is one full pass over the training data.
        epoch_count=3,
        
        # This name registers the resulting model in the Vertex AI Model Registry.
        tuned_model_display_name=TUNED_MODEL_DISPLAY_NAME,
        
        # The validation dataset provides an unbiased measure of performance.
        validation_dataset=validation_dataset,
    )

    # 3. Launch the serverless tuning job. This call is asynchronous.
    sft_tuning_job = vertex_client.tunings.tune(
        base_model=BASE_MODEL_ID,
        training_dataset=training_dataset,
        config=tuning_config,
    )
    
    TUNING_JOB_NAME = sft_tuning_job.name
    print(f"\nTuning job created: {sft_tuning_job}")

---

Part 3: The Governance Layer - Automatic Tracking and Auditing

A model without a history is a liability. Vertex AI automatically provides governance through three interconnected services.

Capability Spotlight: Experiments, Lineage, and Model Registry

1. Vertex AI Experiments: Every tuning job is automatically logged as an experiment run, capturing hyperparameters, data sources, and metrics.
2. Vertex AI Model Registry: The tuned model is automatically versioned and registered here, becoming the canonical, deployable artifact.
3. Vertex AI Lineage: This is the crown jewel of MLOps. It automatically creates a graph connecting the data source, the training job, and the resulting model artifact, providing an unbreakable and automated audit trail.

---

Part 4: The Deployment Engine - Scalable, Secure Endpoints

Once our model is in the Model Registry, deployment is a simple, robust process.

Capability Spotlight: Vertex AI Endpoints

• One-Click Deployment: Deploy your model to a secure HTTPS endpoint with a single SDK call.
• Serverless & Scalable: The endpoint automatically scales based on traffic, from zero to thousands of RPS.
• Integrated Monitoring: Endpoints come with built-in monitoring for traffic, error rates, and latency.
• Traffic Splitting: Safely roll out new versions using A/B testing and canary deployments.

Code in Action: Evaluating the Deployed Specialist

The evaluate_qualitatively function in the notebook demonstrates how to query the new endpoint.

# This code is from "Step 6: Evaluate Tuned Model"
def evaluate_qualitatively(tuned_endpoint: str, test_data: list[dict]):
    """Makes predictions with the tuned model and prints comparisons."""
    if not tuned_endpoint:
        print("Tuned model endpoint not available. Skipping evaluation.")
        return
    
    # 1. Select a random, unseen sample from our test set.
    sample = random.choice(test_data)
    user_prompt = sample["contents"][0]["parts"][0]["text"]
    expected_output = sample["contents"][1]["parts"][0]["text"]

    # 2. Prepare the prediction request. It only contains the user prompt.
    prediction_contents = [{"role": "user", "parts": [{"text": user_prompt}]}]

    # 3. Call the endpoint using the Vertex AI client.
    # Note: The 'model' argument is the full resource name of the deployed endpoint.
    response = vertex_client.models.generate_content(
        model=tuned_endpoint,
        contents=prediction_contents,
        config={
            # Temperature controls randomness. For classification, we want the most
            # likely, factual answer, so we set it to a very low value.
            "temperature": 0.1,
            "max_output_tokens": 50,
        },
    )
    
    predicted_output = response.text.strip()
    print(f"Input Prompt:\n{user_prompt}")
    print(f"\nExpected Output: {expected_output}")
    print(f"Predicted Output: {predicted_output}")
    print(f"Result: {'MATCH' if predicted_output == expected_output else 'MISMATCH'}")

# The main script waits for the tuning job to finish and then calls this function.
# TUNED_MODEL_ENDPOINT = sft_tuning_job.tuned_model.endpoint
# evaluate_qualitatively(TUNED_MODEL_ENDPOINT, test_split)

Conclusion: Vertex AI is the Specialist Factory

We have executed the Vertex AI Blueprint. We transformed a generalist Gemini model into a high-performing “Digital Technician” with 97.7% accuracy. But the true victory is the establishment of a repeatable, governed, and scalable process.

• We used Cloud Storage as our integrated data foundation.
• We used the serverless Training service with built-in PEFT as our efficient tuning engine.
• We used the Model Registry, Experiments, and Lineage as our automated governance layer.
• We used Endpoints as our scalable, secure deployment engine.

This blueprint is the definitive guide to moving beyond AI experimentation and into the realm of industrial-scale AI production.

---

Key Technical References and Further Reading

1. Vertex AI Documentation: The central hub for all Vertex AI services.
2. Introduction to Vertex AI Lineage: The official documentation for automated artifact and execution tracking.
3. Deploy models to an endpoint: The definitive guide to deploying models for real-time predictions.
4. LoRA: Low-Rank Adaptation of Large Language Models (arXiv): The foundational research paper on the PEFT technique that makes efficient tuning possible.
5. Code: sft_gemini_predictive_maintenance.ipynb
6. Blog: Image tuning Gemini as GenAI quality control inspector

---

Let’s keep the conversation going! Share your thoughts, questions, and ideas in the comments.

Note: Should you have any concerns or queries about this post or my implementation, please feel free to connect with me on LinkedIn! Thanks!