From 84cd8156b64308885b8c41ec89e01f8704a0e124 Mon Sep 17 00:00:00 2001
From: yishan-pu <107094052+yishan-pu@users.noreply.github.com>
Date: Wed, 7 May 2025 11:12:27 -0700
Subject: [PATCH] feat(generativeai): Create genai_sdk_supervised_example.py

---
 .../genai_sdk_supervised_example.py           | 78 +++++++++++++++++++
 1 file changed, 78 insertions(+)
 create mode 100644 generative_ai/model_tuning/genai_sdk_supervised_example.py

diff --git a/generative_ai/model_tuning/genai_sdk_supervised_example.py b/generative_ai/model_tuning/genai_sdk_supervised_example.py
new file mode 100644
index 00000000000..a1cc75f4911
--- /dev/null
+++ b/generative_ai/model_tuning/genai_sdk_supervised_example.py
@@ -0,0 +1,78 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+from google.genai import types
+
+PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
+
+
+def genai_sdk_gemini_tuning_basic() -> types.TuningJob:
+  # [START genaisdk_gemini_tuning_basic]
+  import time
+
+  from google import genai
+  from google.genai import types
+
+  # TODO(developer): Update and un-comment below lines
+  # PROJECT_ID = "your-project-id"
+  client = genai.Client(
+      vertexai=True,
+      project=PROJECT_ID,
+      location="us-central1",
+  )
+
+  tuning_job = client.tunings.tune(
+      base_model="gemini-2.0-flash-lite-001",
+      train_dataset="gs://cloud-samples-data/ai-platform/generative_ai/gemini-2_0/text/sft_train_data.jsonl",
+      config=types.CreateTuningJobConfig(
+          tuned_model_display_name="Example tuning job",
+          # Set to True to disable tuning intermediate checkpoints.
+          # export_last_checkpoint_only=True,
+      ),
+  )
+
+  running_states = set([
+      "JOB_STATE_PENDING",
+      "JOB_STATE_RUNNING",
+  ])
+
+  while tuning_job.state in running_states:
+    print(tuning_job.state)
+    tuning_job = client.tunings.get(name=tuning_job.name)
+    time.sleep(60)
+
+  print(tuning_job.tuned_model.model)
+  print(tuning_job.tuned_model.endpoint)
+  print(tuning_job.experiment)
+  # Example response:
+  # projects/123456789012/locations/us-central1/models/1234567890@1
+  # projects/123456789012/locations/us-central1/endpoints/123456789012345
+  # projects/123456789012/locations/us-central1/metadataStores/default/contexts/tuning-experiment-2025010112345678
+
+  if tuning_job.tuned_model.checkpoints:
+    for i in range(len(tuning_job.tuned_model.checkpoints)):
+      checkpoint = tuning_job.tuned_model.checkpoints[i]
+      print(f"Checkpoint {i + 1}: ", checkpoint)
+  # Example response:
+  # Checkpoint 1:  checkpoint_id='1' epoch=1 step=10 endpoint='projects/123456789012/locations/us-central1/endpoints/123456789000000'
+  # Checkpoint 2:  checkpoint_id='2' epoch=2 step=20 endpoint='projects/123456789012/locations/us-central1/endpoints/123456789012345'
+
+  # [END genaisdk_gemini_tuning_basic]
+  return tuning_job
+
+
+if __name__ == "__main__":
+  genai_sdk_gemini_tuning_basic()