From b8718be2d6f937be1ab9da0bcf2910e0be491281 Mon Sep 17 00:00:00 2001 From: Yishan Pu Date: Tue, 13 May 2025 16:44:28 -0700 Subject: [PATCH 1/2] feat: add code samples for tuning with intermediate checkpoints --- genai/tuning/tuning_job_create.py | 7 ++ .../tuning/tuning_with_checkpoints_create.py | 65 +++++++++++++++++++ .../tuning_with_checkpoints_get_model.py | 48 ++++++++++++++ ...uning_with_checkpoints_list_checkpoints.py | 40 ++++++++++++ ...with_checkpoints_set_default_checkpoint.py | 54 +++++++++++++++ ...uning_with_checkpoints_textgen_with_txt.py | 56 ++++++++++++++++ 6 files changed, 270 insertions(+) create mode 100644 genai/tuning/tuning_with_checkpoints_create.py create mode 100644 genai/tuning/tuning_with_checkpoints_get_model.py create mode 100644 genai/tuning/tuning_with_checkpoints_list_checkpoints.py create mode 100644 genai/tuning/tuning_with_checkpoints_set_default_checkpoint.py create mode 100644 genai/tuning/tuning_with_checkpoints_textgen_with_txt.py diff --git a/genai/tuning/tuning_job_create.py b/genai/tuning/tuning_job_create.py index 706b0fa80c..e411027451 100644 --- a/genai/tuning/tuning_job_create.py +++ b/genai/tuning/tuning_job_create.py @@ -48,6 +48,13 @@ def create_tuning_job() -> str: # 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, checkpoint in enumerate(tuning_job.tuned_model.checkpoints): + 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 googlegenaisdk_tuning_job_create] return tuning_job.name diff --git a/genai/tuning/tuning_with_checkpoints_create.py b/genai/tuning/tuning_with_checkpoints_create.py new file mode 100644 index 0000000000..5427f2fa57 --- /dev/null +++ b/genai/tuning/tuning_with_checkpoints_create.py @@ -0,0 +1,65 @@ +# Copyright 2025 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. + + +def create_with_checkpoints() -> str: + # [START googlegenaisdk_tuning_with_checkpoints_create] + import time + + from google import genai + from google.genai.types import HttpOptions, CreateTuningJobConfig + + client = genai.Client(http_options=HttpOptions(api_version="v1")) + + tuning_job = client.tunings.tune( + base_model="gemini-2.0-flash-lite-001", + training_dataset="gs://cloud-samples-data/ai-platform/generative_ai/gemini-2_0/text/sft_train_data.jsonl", + config=CreateTuningJobConfig( + tuned_model_display_name="Example tuning job", + # Set to True to disable tuning intermediate checkpoints. Default is False. + export_last_checkpoint_only=False, + ), + ) + + 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, checkpoint in enumerate(tuning_job.tuned_model.checkpoints): + 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 googlegenaisdk_tuning_with_checkpoints_create] + return tuning_job.name + + +if __name__ == "__main__": + create_with_checkpoints() diff --git a/genai/tuning/tuning_with_checkpoints_get_model.py b/genai/tuning/tuning_with_checkpoints_get_model.py new file mode 100644 index 0000000000..c093ab2b96 --- /dev/null +++ b/genai/tuning/tuning_with_checkpoints_get_model.py @@ -0,0 +1,48 @@ +# Copyright 2025 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. + + +def get_tuned_model_with_checkpoints(name: str) -> str: + # [START googlegenaisdk_tuning_with_checkpoints_get_model] + from google import genai + from google.genai.types import HttpOptions + + client = genai.Client(http_options=HttpOptions(api_version="v1")) + + # Get the tuning job and the tuned model. + # Eg. name = "projects/123456789012/locations/us-central1/tuningJobs/123456789012345" + tuning_job = client.tunings.get(name=name) + tuned_model = client.models.get(model=tuning_job.tuned_model.model) + print(tuned_model) + # Example response: + # Model(name='projects/123456789012/locations/us-central1/models/1234567890@1', ...) + + print(f"Default checkpoint: {tuned_model.default_checkpoint_id}") + # Example response: + # Default checkpoint: 2 + + if tuned_model.checkpoints: + for _, checkpoint in enumerate(tuned_model.checkpoints): + print(f"Checkpoint {checkpoint.checkpoint_id}: ", checkpoint) + # Example response: + # Checkpoint 1: checkpoint_id='1' epoch=1 step=10 + # Checkpoint 2: checkpoint_id='2' epoch=2 step=20 + + # [END googlegenaisdk_tuning_with_checkpoints_get_model] + return tuned_model.name + + +if __name__ == "__main__": + tuning_job_name = input("Tuning job name: ") + get_tuned_model_with_checkpoints(tuning_job_name) diff --git a/genai/tuning/tuning_with_checkpoints_list_checkpoints.py b/genai/tuning/tuning_with_checkpoints_list_checkpoints.py new file mode 100644 index 0000000000..be688b6a5d --- /dev/null +++ b/genai/tuning/tuning_with_checkpoints_list_checkpoints.py @@ -0,0 +1,40 @@ +# Copyright 2025 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. + + +def list_checkpoints(name: str) -> str: + # [START googlegenaisdk_tuning_with_checkpoints_list_checkpoints] + from google import genai + from google.genai.types import HttpOptions + + client = genai.Client(http_options=HttpOptions(api_version="v1")) + + # Get the tuning job and the tuned model. + # Eg. name = "projects/123456789012/locations/us-central1/tuningJobs/123456789012345" + tuning_job = client.tunings.get(name=name) + + if tuning_job.tuned_model.checkpoints: + for i, checkpoint in enumerate(tuning_job.tuned_model.checkpoints): + 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 googlegenaisdk_tuning_with_checkpoints_list_checkpoints] + return tuning_job.name + + +if __name__ == "__main__": + tuning_job_name = input("Tuning job name: ") + list_checkpoints(tuning_job_name) diff --git a/genai/tuning/tuning_with_checkpoints_set_default_checkpoint.py b/genai/tuning/tuning_with_checkpoints_set_default_checkpoint.py new file mode 100644 index 0000000000..a1997e7fd1 --- /dev/null +++ b/genai/tuning/tuning_with_checkpoints_set_default_checkpoint.py @@ -0,0 +1,54 @@ +# Copyright 2025 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. + + +def set_default_checkpoint(name: str, checkpoint_id: str) -> str: + # [START googlegenaisdk_tuning_with_checkpoints_set_default] + from google import genai + from google.genai.types import HttpOptions, UpdateModelConfig + + client = genai.Client(http_options=HttpOptions(api_version="v1")) + + # Get the tuning job and the tuned model. + # Eg. name = "projects/123456789012/locations/us-central1/tuningJobs/123456789012345" + tuning_job = client.tunings.get(name=name) + tuned_model = client.models.get(model=tuning_job.tuned_model.model) + + print(f"Default checkpoint: {tuned_model.default_checkpoint_id}") + print(f"Tuned model endpoint: {tuning_job.tuned_model.endpoint}") + # Example response: + # Default checkpoint: 2 + # projects/123456789012/locations/us-central1/endpoints/123456789012345 + + # Set a new default checkpoint. + # Eg. checkpoint_id = "1" + tuned_model = client.models.update( + model=tuned_model.name, + config=UpdateModelConfig(default_checkpoint_id=checkpoint_id), + ) + + print(f"Default checkpoint: {tuned_model.default_checkpoint_id}") + print(f"Tuned model endpoint: {tuning_job.tuned_model.endpoint}") + # Example response: + # Default checkpoint: 1 + # projects/123456789012/locations/us-central1/endpoints/123456789000000 + + # [END googlegenaisdk_tuning_with_checkpoints_set_default] + return tuned_model.default_checkpoint_id + + +if __name__ == "__main__": + tuning_job_name = input("Tuning job name: ") + default_checkpoint_id = input("Default checkpoint id: ") + set_default_checkpoint(tuning_job_name, default_checkpoint_id) diff --git a/genai/tuning/tuning_with_checkpoints_textgen_with_txt.py b/genai/tuning/tuning_with_checkpoints_textgen_with_txt.py new file mode 100644 index 0000000000..68ae7a2dbc --- /dev/null +++ b/genai/tuning/tuning_with_checkpoints_textgen_with_txt.py @@ -0,0 +1,56 @@ +# Copyright 2025 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. + + +def test_checkpoint(name: str) -> str: + # [START googlegenaisdk_tuning_with_checkpoints_test] + from google import genai + from google.genai.types import HttpOptions + + client = genai.Client(http_options=HttpOptions(api_version="v1")) + + # Get the tuning job and the tuned model. + # Eg. name = "projects/123456789012/locations/us-central1/tuningJobs/123456789012345" + tuning_job = client.tunings.get(name=name) + + contents = "Why is the sky blue?" + + # Tests the default checkpoint + response = client.models.generate_content( + model=tuning_job.tuned_model.endpoint, + contents=contents, + ) + print(response.text) + + # Tests Checkpoint 1 + checkpoint1_response = client.models.generate_content( + model=tuning_job.tuned_model.checkpoints[0].endpoint, + contents=contents, + ) + print(checkpoint1_response.text) + + # Tests Checkpoint 2 + checkpoint2_response = client.models.generate_content( + model=tuning_job.tuned_model.checkpoints[1].endpoint, + contents=contents, + ) + print(checkpoint2_response.text) + + # [END googlegenaisdk_tuning_with_checkpoints_test] + return response.text + + +if __name__ == "__main__": + tuning_job_name = input("Tuning job name: ") + test_checkpoint(tuning_job_name) From bd6e7c8ce7162c3a6feaa7169efbbcc118236037 Mon Sep 17 00:00:00 2001 From: Yishan Pu Date: Tue, 13 May 2025 16:44:28 -0700 Subject: [PATCH 2/2] feat: add code samples for tuning with intermediate checkpoints --- genai/tuning/requirements.txt | 2 +- genai/tuning/test_tuning_examples.py | 175 +++++++++++++++++- genai/tuning/tuning_job_create.py | 7 + .../tuning/tuning_with_checkpoints_create.py | 65 +++++++ .../tuning_with_checkpoints_get_model.py | 48 +++++ ...uning_with_checkpoints_list_checkpoints.py | 40 ++++ ...with_checkpoints_set_default_checkpoint.py | 54 ++++++ ...uning_with_checkpoints_textgen_with_txt.py | 56 ++++++ 8 files changed, 445 insertions(+), 2 deletions(-) create mode 100644 genai/tuning/tuning_with_checkpoints_create.py create mode 100644 genai/tuning/tuning_with_checkpoints_get_model.py create mode 100644 genai/tuning/tuning_with_checkpoints_list_checkpoints.py create mode 100644 genai/tuning/tuning_with_checkpoints_set_default_checkpoint.py create mode 100644 genai/tuning/tuning_with_checkpoints_textgen_with_txt.py diff --git a/genai/tuning/requirements.txt b/genai/tuning/requirements.txt index 73d0828cb4..54b85900f5 100644 --- a/genai/tuning/requirements.txt +++ b/genai/tuning/requirements.txt @@ -1 +1 @@ -google-genai==1.7.0 +google-genai==1.15.0 diff --git a/genai/tuning/test_tuning_examples.py b/genai/tuning/test_tuning_examples.py index f8d8553680..f028f7eda1 100644 --- a/genai/tuning/test_tuning_examples.py +++ b/genai/tuning/test_tuning_examples.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest.mock import MagicMock, patch +from unittest.mock import call, MagicMock, patch from google.genai import types @@ -20,6 +20,11 @@ import tuning_job_get import tuning_job_list import tuning_textgen_with_txt +import tuning_with_checkpoints_create +import tuning_with_checkpoints_get_model +import tuning_with_checkpoints_list_checkpoints +import tuning_with_checkpoints_set_default_checkpoint +import tuning_with_checkpoints_textgen_with_txt @patch("google.genai.Client") @@ -113,3 +118,171 @@ def test_tuning_textgen_with_txt(mock_genai_client: MagicMock) -> None: mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1")) mock_genai_client.return_value.tunings.get.assert_called_once() mock_genai_client.return_value.models.generate_content.assert_called_once() + + +@patch("google.genai.Client") +def test_tuning_job_create_with_checkpoints(mock_genai_client: MagicMock) -> None: + # Mock the API response + mock_tuning_job = types.TuningJob( + name="test-tuning-job", + experiment="test-experiment", + tuned_model=types.TunedModel( + model="test-model", + endpoint="test-endpoint-2", + checkpoints=[ + types.TunedModelCheckpoint(checkpoint_id="1", epoch=1, step=10, endpoint="test-endpoint-1"), + types.TunedModelCheckpoint(checkpoint_id="2", epoch=2, step=20, endpoint="test-endpoint-2"), + ] + ) + ) + mock_genai_client.return_value.tunings.tune.return_value = mock_tuning_job + + response = tuning_with_checkpoints_create.create_with_checkpoints() + + mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1")) + mock_genai_client.return_value.tunings.tune.assert_called_once() + assert response == "test-tuning-job" + + +@patch("google.genai.Client") +def test_tuning_with_checkpoints_get_model(mock_genai_client: MagicMock) -> None: + # Mock the API response + mock_tuning_job = types.TuningJob( + name="test-tuning-job", + experiment="test-experiment", + tuned_model=types.TunedModel( + model="test-model", + endpoint="test-endpoint-2", + checkpoints=[ + types.TunedModelCheckpoint(checkpoint_id="1", epoch=1, step=10, endpoint="test-endpoint-1"), + types.TunedModelCheckpoint(checkpoint_id="2", epoch=2, step=20, endpoint="test-endpoint-2"), + ] + ) + ) + mock_model = types.Model( + name="test-model", + default_checkpoint_id="2", + checkpoints=[ + types.Checkpoint(checkpoint_id="1", epoch=1, step=10), + types.Checkpoint(checkpoint_id="2", epoch=2, step=20), + ] + ) + mock_genai_client.return_value.tunings.get.return_value = mock_tuning_job + mock_genai_client.return_value.models.get.return_value = mock_model + + response = tuning_with_checkpoints_get_model.get_tuned_model_with_checkpoints("test-tuning-job") + + mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1")) + mock_genai_client.return_value.tunings.get.assert_called_once_with(name="test-tuning-job") + mock_genai_client.return_value.models.get.assert_called_once_with(model="test-model") + assert response == "test-model" + + +@patch("google.genai.Client") +def test_tuning_with_checkpoints_list_checkpoints(mock_genai_client: MagicMock) -> None: + # Mock the API response + mock_tuning_job = types.TuningJob( + name="test-tuning-job", + experiment="test-experiment", + tuned_model=types.TunedModel( + model="test-model", + endpoint="test-endpoint-2", + checkpoints=[ + types.TunedModelCheckpoint(checkpoint_id="1", epoch=1, step=10, endpoint="test-endpoint-1"), + types.TunedModelCheckpoint(checkpoint_id="2", epoch=2, step=20, endpoint="test-endpoint-2"), + ] + ) + ) + mock_genai_client.return_value.tunings.get.return_value = mock_tuning_job + + response = tuning_with_checkpoints_list_checkpoints.list_checkpoints("test-tuning-job") + + mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1")) + mock_genai_client.return_value.tunings.get.assert_called_once_with(name="test-tuning-job") + assert response == "test-tuning-job" + + +@patch("google.genai.Client") +def test_tuning_with_checkpoints_set_default_checkpoint(mock_genai_client: MagicMock) -> None: + # Mock the API response + mock_tuning_job = types.TuningJob( + name="test-tuning-job", + experiment="test-experiment", + tuned_model=types.TunedModel( + model="test-model", + endpoint="test-endpoint-2", + checkpoints=[ + types.TunedModelCheckpoint(checkpoint_id="1", epoch=1, step=10, endpoint="test-endpoint-1"), + types.TunedModelCheckpoint(checkpoint_id="2", epoch=2, step=20, endpoint="test-endpoint-2"), + ] + ) + ) + mock_model = types.Model( + name="test-model", + default_checkpoint_id="2", + checkpoints=[ + types.Checkpoint(checkpoint_id="1", epoch=1, step=10), + types.Checkpoint(checkpoint_id="2", epoch=2, step=20), + ] + ) + mock_updated_model = types.Model( + name="test-model", + default_checkpoint_id="1", + checkpoints=[ + types.Checkpoint(checkpoint_id="1", epoch=1, step=10), + types.Checkpoint(checkpoint_id="2", epoch=2, step=20), + ] + ) + mock_genai_client.return_value.tunings.get.return_value = mock_tuning_job + mock_genai_client.return_value.models.get.return_value = mock_model + mock_genai_client.return_value.models.update.return_value = mock_updated_model + + response = tuning_with_checkpoints_set_default_checkpoint.set_default_checkpoint("test-tuning-job", "1") + + mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1")) + mock_genai_client.return_value.tunings.get.assert_called_once_with(name="test-tuning-job") + mock_genai_client.return_value.models.get.assert_called_once_with(model="test-model") + mock_genai_client.return_value.models.update.assert_called_once() + assert response == "1" + + +@patch("google.genai.Client") +def test_tuning_with_checkpoints_textgen_with_txt(mock_genai_client: MagicMock) -> None: + # Mock the API response + mock_tuning_job = types.TuningJob( + name="test-tuning-job", + experiment="test-experiment", + tuned_model=types.TunedModel( + model="test-model", + endpoint="test-endpoint-2", + checkpoints=[ + types.TunedModelCheckpoint(checkpoint_id="1", epoch=1, step=10, endpoint="test-endpoint-1"), + types.TunedModelCheckpoint(checkpoint_id="2", epoch=2, step=20, endpoint="test-endpoint-2"), + ] + ) + ) + mock_response = types.GenerateContentResponse._from_response( # pylint: disable=protected-access + response={ + "candidates": [ + { + "content": { + "parts": [{"text": "This is a mocked answer."}] + } + } + ] + }, + kwargs={}, + ) + + mock_genai_client.return_value.tunings.get.return_value = mock_tuning_job + mock_genai_client.return_value.models.generate_content.return_value = mock_response + + tuning_with_checkpoints_textgen_with_txt.test_checkpoint("test-tuning-job") + + mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1")) + mock_genai_client.return_value.tunings.get.assert_called_once() + assert mock_genai_client.return_value.models.generate_content.call_args_list == [ + call(model="test-endpoint-2", contents="Why is the sky blue?"), + call(model="test-endpoint-1", contents="Why is the sky blue?"), + call(model="test-endpoint-2", contents="Why is the sky blue?"), + ] diff --git a/genai/tuning/tuning_job_create.py b/genai/tuning/tuning_job_create.py index 706b0fa80c..e411027451 100644 --- a/genai/tuning/tuning_job_create.py +++ b/genai/tuning/tuning_job_create.py @@ -48,6 +48,13 @@ def create_tuning_job() -> str: # 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, checkpoint in enumerate(tuning_job.tuned_model.checkpoints): + 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 googlegenaisdk_tuning_job_create] return tuning_job.name diff --git a/genai/tuning/tuning_with_checkpoints_create.py b/genai/tuning/tuning_with_checkpoints_create.py new file mode 100644 index 0000000000..5427f2fa57 --- /dev/null +++ b/genai/tuning/tuning_with_checkpoints_create.py @@ -0,0 +1,65 @@ +# Copyright 2025 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. + + +def create_with_checkpoints() -> str: + # [START googlegenaisdk_tuning_with_checkpoints_create] + import time + + from google import genai + from google.genai.types import HttpOptions, CreateTuningJobConfig + + client = genai.Client(http_options=HttpOptions(api_version="v1")) + + tuning_job = client.tunings.tune( + base_model="gemini-2.0-flash-lite-001", + training_dataset="gs://cloud-samples-data/ai-platform/generative_ai/gemini-2_0/text/sft_train_data.jsonl", + config=CreateTuningJobConfig( + tuned_model_display_name="Example tuning job", + # Set to True to disable tuning intermediate checkpoints. Default is False. + export_last_checkpoint_only=False, + ), + ) + + 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, checkpoint in enumerate(tuning_job.tuned_model.checkpoints): + 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 googlegenaisdk_tuning_with_checkpoints_create] + return tuning_job.name + + +if __name__ == "__main__": + create_with_checkpoints() diff --git a/genai/tuning/tuning_with_checkpoints_get_model.py b/genai/tuning/tuning_with_checkpoints_get_model.py new file mode 100644 index 0000000000..c093ab2b96 --- /dev/null +++ b/genai/tuning/tuning_with_checkpoints_get_model.py @@ -0,0 +1,48 @@ +# Copyright 2025 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. + + +def get_tuned_model_with_checkpoints(name: str) -> str: + # [START googlegenaisdk_tuning_with_checkpoints_get_model] + from google import genai + from google.genai.types import HttpOptions + + client = genai.Client(http_options=HttpOptions(api_version="v1")) + + # Get the tuning job and the tuned model. + # Eg. name = "projects/123456789012/locations/us-central1/tuningJobs/123456789012345" + tuning_job = client.tunings.get(name=name) + tuned_model = client.models.get(model=tuning_job.tuned_model.model) + print(tuned_model) + # Example response: + # Model(name='projects/123456789012/locations/us-central1/models/1234567890@1', ...) + + print(f"Default checkpoint: {tuned_model.default_checkpoint_id}") + # Example response: + # Default checkpoint: 2 + + if tuned_model.checkpoints: + for _, checkpoint in enumerate(tuned_model.checkpoints): + print(f"Checkpoint {checkpoint.checkpoint_id}: ", checkpoint) + # Example response: + # Checkpoint 1: checkpoint_id='1' epoch=1 step=10 + # Checkpoint 2: checkpoint_id='2' epoch=2 step=20 + + # [END googlegenaisdk_tuning_with_checkpoints_get_model] + return tuned_model.name + + +if __name__ == "__main__": + tuning_job_name = input("Tuning job name: ") + get_tuned_model_with_checkpoints(tuning_job_name) diff --git a/genai/tuning/tuning_with_checkpoints_list_checkpoints.py b/genai/tuning/tuning_with_checkpoints_list_checkpoints.py new file mode 100644 index 0000000000..be688b6a5d --- /dev/null +++ b/genai/tuning/tuning_with_checkpoints_list_checkpoints.py @@ -0,0 +1,40 @@ +# Copyright 2025 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. + + +def list_checkpoints(name: str) -> str: + # [START googlegenaisdk_tuning_with_checkpoints_list_checkpoints] + from google import genai + from google.genai.types import HttpOptions + + client = genai.Client(http_options=HttpOptions(api_version="v1")) + + # Get the tuning job and the tuned model. + # Eg. name = "projects/123456789012/locations/us-central1/tuningJobs/123456789012345" + tuning_job = client.tunings.get(name=name) + + if tuning_job.tuned_model.checkpoints: + for i, checkpoint in enumerate(tuning_job.tuned_model.checkpoints): + 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 googlegenaisdk_tuning_with_checkpoints_list_checkpoints] + return tuning_job.name + + +if __name__ == "__main__": + tuning_job_name = input("Tuning job name: ") + list_checkpoints(tuning_job_name) diff --git a/genai/tuning/tuning_with_checkpoints_set_default_checkpoint.py b/genai/tuning/tuning_with_checkpoints_set_default_checkpoint.py new file mode 100644 index 0000000000..a1997e7fd1 --- /dev/null +++ b/genai/tuning/tuning_with_checkpoints_set_default_checkpoint.py @@ -0,0 +1,54 @@ +# Copyright 2025 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. + + +def set_default_checkpoint(name: str, checkpoint_id: str) -> str: + # [START googlegenaisdk_tuning_with_checkpoints_set_default] + from google import genai + from google.genai.types import HttpOptions, UpdateModelConfig + + client = genai.Client(http_options=HttpOptions(api_version="v1")) + + # Get the tuning job and the tuned model. + # Eg. name = "projects/123456789012/locations/us-central1/tuningJobs/123456789012345" + tuning_job = client.tunings.get(name=name) + tuned_model = client.models.get(model=tuning_job.tuned_model.model) + + print(f"Default checkpoint: {tuned_model.default_checkpoint_id}") + print(f"Tuned model endpoint: {tuning_job.tuned_model.endpoint}") + # Example response: + # Default checkpoint: 2 + # projects/123456789012/locations/us-central1/endpoints/123456789012345 + + # Set a new default checkpoint. + # Eg. checkpoint_id = "1" + tuned_model = client.models.update( + model=tuned_model.name, + config=UpdateModelConfig(default_checkpoint_id=checkpoint_id), + ) + + print(f"Default checkpoint: {tuned_model.default_checkpoint_id}") + print(f"Tuned model endpoint: {tuning_job.tuned_model.endpoint}") + # Example response: + # Default checkpoint: 1 + # projects/123456789012/locations/us-central1/endpoints/123456789000000 + + # [END googlegenaisdk_tuning_with_checkpoints_set_default] + return tuned_model.default_checkpoint_id + + +if __name__ == "__main__": + tuning_job_name = input("Tuning job name: ") + default_checkpoint_id = input("Default checkpoint id: ") + set_default_checkpoint(tuning_job_name, default_checkpoint_id) diff --git a/genai/tuning/tuning_with_checkpoints_textgen_with_txt.py b/genai/tuning/tuning_with_checkpoints_textgen_with_txt.py new file mode 100644 index 0000000000..68ae7a2dbc --- /dev/null +++ b/genai/tuning/tuning_with_checkpoints_textgen_with_txt.py @@ -0,0 +1,56 @@ +# Copyright 2025 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. + + +def test_checkpoint(name: str) -> str: + # [START googlegenaisdk_tuning_with_checkpoints_test] + from google import genai + from google.genai.types import HttpOptions + + client = genai.Client(http_options=HttpOptions(api_version="v1")) + + # Get the tuning job and the tuned model. + # Eg. name = "projects/123456789012/locations/us-central1/tuningJobs/123456789012345" + tuning_job = client.tunings.get(name=name) + + contents = "Why is the sky blue?" + + # Tests the default checkpoint + response = client.models.generate_content( + model=tuning_job.tuned_model.endpoint, + contents=contents, + ) + print(response.text) + + # Tests Checkpoint 1 + checkpoint1_response = client.models.generate_content( + model=tuning_job.tuned_model.checkpoints[0].endpoint, + contents=contents, + ) + print(checkpoint1_response.text) + + # Tests Checkpoint 2 + checkpoint2_response = client.models.generate_content( + model=tuning_job.tuned_model.checkpoints[1].endpoint, + contents=contents, + ) + print(checkpoint2_response.text) + + # [END googlegenaisdk_tuning_with_checkpoints_test] + return response.text + + +if __name__ == "__main__": + tuning_job_name = input("Tuning job name: ") + test_checkpoint(tuning_job_name)