kubeflow - fine tuning gemini with text-based dataset

To fine tune gemini using kubeflow pipeline, we are using the following setup. There is no validation being configure here. You can easily added it tho. This will be a much more practical example of using kubeflow for :- 

from kfp import compiler, dsl

 

# ---------------------------- edit these ----------------------------

PROJECT_ID = "project-id"

LOCATION = "us-central1"

BUCKET = "my-vertex-bucket-ai"  # the bucket where you uploaded train.jsonl

 

TRAIN_URI = f"gs://{BUCKET}/text/train.jsonl"

VALIDATION_URI = f"gs://{BUCKET}/text/validation.jsonl"  # set to "" to skip validation

 

# Must be a Gemini model that supports supervised tuning. Check Google's

# supported-models list, because it changes over time.

SOURCE_MODEL = "gemini-2.5-flash"

TUNED_MODEL_NAME = "ticket-classifier"

 

PIPELINE_ROOT = f"gs://{BUCKET}/pipeline-root"  # where pipeline files are stored

# --------------------------------------------------------------------

@dsl.component(base_image="python:3.11", packages_to_install=["google-cloud-aiplatform"])

def tune_gemini(

    project_id: str,

    location: str,

    source_model: str,

    train_uri: str,

    tuned_model_name: str,

    validation_uri: str = "",

) -> str:

    """Start a supervised fine-tuning job, wait for it, return the tuned endpoint name."""

    import time

 

    import vertexai

    from vertexai.tuning import sft

 

    vertexai.init(project=project_id, location=location)

 

    kwargs = {}

    if validation_uri:

        kwargs["validation_dataset"] = validation_uri

 

    job = sft.train(

        source_model=source_model,

        train_dataset=train_uri,

        tuned_model_display_name=tuned_model_name,

        **kwargs,

    )

 

    while not job.has_ended:

        time.sleep(60)

        job.refresh()

 

    print("tuned model:   ", job.tuned_model_name)

    print("tuned endpoint:", job.tuned_model_endpoint_name)

    return job.tuned_model_endpoint_name

 

 

@dsl.pipeline(name="gemini-text-tuning", pipeline_root=PIPELINE_ROOT)

def pipeline(

    project_id: str,

    location: str,

    source_model: str,

    train_uri: str,

    tuned_model_name: str,

    validation_uri: str = "",

):

    tune = tune_gemini(

        project_id=project_id,

        location=location,

        source_model=source_model,

        train_uri=train_uri,

        tuned_model_name=tuned_model_name,

        validation_uri=validation_uri,

    )

    tune.set_caching_options(False)  # every run starts a new, billed tuning job

 

An then we are going to create our pipeline using the following command :-

compiler.Compiler().compile(pipeline_func=pipeline, package_path="gemini_tuning.yaml")

And finally to run our pipeline. 

job = aiplatform.PipelineJob(

    display_name="gemini-text-tuning",

    template_path="gemini_tuning.yaml",

    pipeline_root=PIPELINE_ROOT,

    parameter_values={

        "project_id": PROJECT_ID,

        "train_uri": "gs://my-vertex-bucket-ai/text/train.jsonl",

        "validation_uri": "gs://my-bucket/text/validation.jsonl",

    },

)

job.run()

And finally the sample train.jsonl would look like this :-

{"contents": [{"role": "user", "parts": [{"text": "Classify this support

ticket as billing, bug or account.\n\nTicket: My card was charged twice this month"}]}

, {"role": "model", "parts": [{"text": "billing"}]}]}

{"contents": [{"role": "user", "parts": [{"text": "Classify this support ticket

as billing, bug or account.\n\nTicket: I was billed after I cancelled my subscription"

}]}, {"role": "model", "parts": [{"text": "billing"}]}]}

The sample data format can be obtained here.

https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/tuning/text?hl=en

We can see the pipeline is running:-

And you can see that tuning is in progresss - under the Agentic Platform -> Tuning

You need to check this as pipeline can be successful but the fine tuning

does not take place

And you can also see some metric here

And then there dataset metrics that shows us the output per token and input token used.

And you can also see it in the endpoint sections:-

And if you wonder, how can you test it out? We can do it like this here. From the

Fine-Tuning goto Test.

And then do your prompt thingie.