Skip to content

Commit a163a42

Browse files
yishan-puYishan Pumsampathkumar
authored
feat(genai): Add tuning examples (#13360)
* feat(genai): Add tuning examples * fix(genai): fix tuning examples test issues - Updates tuning file names to fix cicd - word `test` is reserved for test files - Update region tag and file names. - add word `job` to filenames & region tags * feat: update test file - rename methods and function according to prior changes --------- Co-authored-by: Yishan Pu <[email protected]> Co-authored-by: Sampath Kumar <[email protected]>
1 parent 87e1fc6 commit a163a42

8 files changed

+335
-0
lines changed

‎genai/tuning/noxfile_config.py

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
# Copyright 2021 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# Default TEST_CONFIG_OVERRIDE for python repos.
16+
17+
# You can copy this file into your directory, then it will be imported from
18+
# the noxfile.py.
19+
20+
# The source of truth:
21+
# https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/noxfile_config.py
22+
23+
TEST_CONFIG_OVERRIDE = {
24+
# You can opt out from the test for specific Python versions.
25+
"ignored_versions": ["2.7", "3.7", "3.8", "3.10", "3.11", "3.12"],
26+
# Old samples are opted out of enforcing Python type hints
27+
# All new samples should feature them
28+
"enforce_type_hints": True,
29+
# An envvar key for determining the project id to use. Change it
30+
# to 'BUILD_SPECIFIC_GCLOUD_PROJECT' if you want to opt in using a
31+
# build specific Cloud project. You can also use your own string
32+
# to use your own Cloud project.
33+
"gcloud_project_env": "GOOGLE_CLOUD_PROJECT",
34+
# 'gcloud_project_env': 'BUILD_SPECIFIC_GCLOUD_PROJECT',
35+
# If you need to use a specific version of pip,
36+
# change pip_version_override to the string representation
37+
# of the version number, for example, "20.2.4"
38+
"pip_version_override": None,
39+
# A dictionary you want to inject into your test. Don't put any
40+
# secrets here. These values will override predefined values.
41+
"envs": {},
42+
}

‎genai/tuning/requirements-test.txt

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
google-api-core==2.24.0
2+
google-cloud-storage==2.19.0
3+
pytest==8.2.0

‎genai/tuning/requirements.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
google-genai==1.7.0

‎genai/tuning/test_tuning_examples.py

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from unittest.mock import MagicMock, patch
16+
17+
from google.genai import types
18+
19+
import tuning_job_create
20+
import tuning_job_get
21+
import tuning_job_list
22+
import tuning_textgen_with_txt
23+
24+
25+
@patch("google.genai.Client")
26+
def test_tuning_job_create(mock_genai_client: MagicMock) -> None:
27+
# Mock the API response
28+
mock_tuning_job = types.TuningJob(
29+
name="test-tuning-job",
30+
experiment="test-experiment",
31+
tuned_model=types.TunedModel(
32+
model="test-model",
33+
endpoint="test-endpoint"
34+
)
35+
)
36+
mock_genai_client.return_value.tunings.tune.return_value = mock_tuning_job
37+
38+
response = tuning_job_create.create_tuning_job()
39+
40+
mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1"))
41+
mock_genai_client.return_value.tunings.tune.assert_called_once()
42+
assert response == "test-tuning-job"
43+
44+
45+
@patch("google.genai.Client")
46+
def test_tuning_job_get(mock_genai_client: MagicMock) -> None:
47+
# Mock the API response
48+
mock_tuning_job = types.TuningJob(
49+
name="test-tuning-job",
50+
experiment="test-experiment",
51+
tuned_model=types.TunedModel(
52+
model="test-model",
53+
endpoint="test-endpoint"
54+
)
55+
)
56+
mock_genai_client.return_value.tunings.get.return_value = mock_tuning_job
57+
58+
response = tuning_job_get.get_tuning_job("test-tuning-job")
59+
60+
mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1"))
61+
mock_genai_client.return_value.tunings.get.assert_called_once()
62+
assert response == "test-tuning-job"
63+
64+
65+
@patch("google.genai.Client")
66+
def test_tuning_job_list(mock_genai_client: MagicMock) -> None:
67+
# Mock the API response
68+
mock_tuning_job = types.TuningJob(
69+
name="test-tuning-job",
70+
experiment="test-experiment",
71+
tuned_model=types.TunedModel(
72+
model="test-model",
73+
endpoint="test-endpoint"
74+
)
75+
)
76+
mock_genai_client.return_value.tunings.list.return_value = [mock_tuning_job]
77+
78+
tuning_job_list.list_tuning_jobs()
79+
80+
mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1"))
81+
mock_genai_client.return_value.tunings.list.assert_called_once()
82+
83+
84+
@patch("google.genai.Client")
85+
def test_tuning_textgen_with_txt(mock_genai_client: MagicMock) -> None:
86+
# Mock the API response
87+
mock_tuning_job = types.TuningJob(
88+
name="test-tuning-job",
89+
experiment="test-experiment",
90+
tuned_model=types.TunedModel(
91+
model="test-model",
92+
endpoint="test-endpoint"
93+
)
94+
)
95+
mock_response = types.GenerateContentResponse._from_response( # pylint: disable=protected-access
96+
response={
97+
"candidates": [
98+
{
99+
"content": {
100+
"parts": [{"text": "This is a mocked answer."}]
101+
}
102+
}
103+
]
104+
},
105+
kwargs={},
106+
)
107+
108+
mock_genai_client.return_value.tunings.get.return_value = mock_tuning_job
109+
mock_genai_client.return_value.models.generate_content.return_value = mock_response
110+
111+
tuning_textgen_with_txt.test_tuned_endpoint("test-tuning-job")
112+
113+
mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1"))
114+
mock_genai_client.return_value.tunings.get.assert_called_once()
115+
mock_genai_client.return_value.models.generate_content.assert_called_once()

‎genai/tuning/tuning_job_create.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
16+
def create_tuning_job() -> str:
17+
# [START googlegenaisdk_tuning_job_create]
18+
import time
19+
20+
from google import genai
21+
from google.genai.types import HttpOptions, CreateTuningJobConfig
22+
23+
client = genai.Client(http_options=HttpOptions(api_version="v1"))
24+
25+
tuning_job = client.tunings.tune(
26+
base_model="gemini-2.0-flash-lite-001",
27+
training_dataset="gs://cloud-samples-data/ai-platform/generative_ai/gemini-2_0/text/sft_train_data.jsonl",
28+
config=CreateTuningJobConfig(
29+
tuned_model_display_name="Example tuning job",
30+
),
31+
)
32+
33+
running_states = set([
34+
"JOB_STATE_PENDING",
35+
"JOB_STATE_RUNNING",
36+
])
37+
38+
while tuning_job.state in running_states:
39+
print(tuning_job.state)
40+
tuning_job = client.tunings.get(name=tuning_job.name)
41+
time.sleep(60)
42+
43+
print(tuning_job.tuned_model.model)
44+
print(tuning_job.tuned_model.endpoint)
45+
print(tuning_job.experiment)
46+
# Example response:
47+
# projects/123456789012/locations/us-central1/models/1234567890@1
48+
# projects/123456789012/locations/us-central1/endpoints/123456789012345
49+
# projects/123456789012/locations/us-central1/metadataStores/default/contexts/tuning-experiment-2025010112345678
50+
51+
# [END googlegenaisdk_tuning_job_create]
52+
return tuning_job.name
53+
54+
55+
if __name__ == "__main__":
56+
create_tuning_job()

‎genai/tuning/tuning_job_get.py

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
16+
def get_tuning_job(name: str) -> str:
17+
# [START googlegenaisdk_tuning_job_get]
18+
from google import genai
19+
from google.genai.types import HttpOptions
20+
21+
client = genai.Client(http_options=HttpOptions(api_version="v1"))
22+
23+
# Get the tuning job and the tuned model.
24+
# Eg. name = "projects/123456789012/locations/us-central1/tuningJobs/123456789012345"
25+
tuning_job = client.tunings.get(name=name)
26+
27+
print(tuning_job.tuned_model.model)
28+
print(tuning_job.tuned_model.endpoint)
29+
print(tuning_job.experiment)
30+
# Example response:
31+
# projects/123456789012/locations/us-central1/models/1234567890@1
32+
# projects/123456789012/locations/us-central1/endpoints/123456789012345
33+
# projects/123456789012/locations/us-central1/metadataStores/default/contexts/tuning-experiment-2025010112345678
34+
35+
# [END googlegenaisdk_tuning_job_get]
36+
return tuning_job.name
37+
38+
39+
if __name__ == "__main__":
40+
tuning_job_name = input("Tuning job name: ")
41+
get_tuning_job(tuning_job_name)

‎genai/tuning/tuning_job_list.py

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
16+
def list_tuning_jobs() -> None:
17+
# [START googlegenaisdk_tuning_job_list]
18+
from google import genai
19+
from google.genai.types import HttpOptions
20+
21+
client = genai.Client(http_options=HttpOptions(api_version="v1"))
22+
23+
responses = client.tunings.list()
24+
for response in responses:
25+
print(response.name)
26+
# Example response:
27+
# projects/123456789012/locations/us-central1/tuningJobs/123456789012345
28+
29+
# [END googlegenaisdk_tuning_job_list]
30+
return
31+
32+
33+
if __name__ == "__main__":
34+
tuning_job_name = input("Tuning job name: ")
35+
list_tuning_jobs()

‎genai/tuning/tuning_textgen_with_txt.py

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
# Copyright 2024 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
16+
def test_tuned_endpoint(name: str) -> str:
17+
# [START googlegenaisdk_tuning_textgen_with_txt]
18+
from google import genai
19+
from google.genai.types import HttpOptions
20+
21+
client = genai.Client(http_options=HttpOptions(api_version="v1"))
22+
23+
# Get the tuning job and the tuned model.
24+
# Eg. name = "projects/123456789012/locations/us-central1/tuningJobs/123456789012345"
25+
tuning_job = client.tunings.get(name=name)
26+
27+
contents = "Why is the sky blue?"
28+
29+
# Tests the default checkpoint
30+
response = client.models.generate_content(
31+
model=tuning_job.tuned_model.endpoint,
32+
contents=contents,
33+
)
34+
print(response.text)
35+
36+
# [END googlegenaisdk_tuning_textgen_with_txt]
37+
return response.text
38+
39+
40+
if __name__ == "__main__":
41+
tuning_job_name = input("Tuning job name: ")
42+
test_tuned_endpoint(tuning_job_name)

0 commit comments

Comments
 (0)