fix: update tuning code samples to improve readability (#13384)

Co-authored-by: Yishan Pu <[email protected]>

Author: yishan-pu

genai/tuning/test_tuning_examples.py

@@ -113,7 +113,7 @@ def test_tuning_textgen_with_txt(mock_genai_client: MagicMock) -> None:
     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_textgen_with_txt.test_tuned_endpoint("test-tuning-job")
+    tuning_textgen_with_txt.predict_with_tuned_endpoint("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()
@@ -277,7 +277,7 @@ def test_tuning_with_checkpoints_textgen_with_txt(mock_genai_client: MagicMock)
     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")
+    tuning_with_checkpoints_textgen_with_txt.predict_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()

genai/tuning/tuning_job_get.py

@@ -13,16 +13,16 @@
 # limitations under the License.
 
 
-def get_tuning_job(name: str) -> str:
+def get_tuning_job(tuning_job_name: str) -> str:
     # [START googlegenaisdk_tuning_job_get]
     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)
+    # Eg. tuning_job_name = "projects/123456789012/locations/us-central1/tuningJobs/123456789012345"
+    tuning_job = client.tunings.get(name=tuning_job_name)
 
     print(tuning_job.tuned_model.model)
     print(tuning_job.tuned_model.endpoint)
@@ -37,5 +37,5 @@ def get_tuning_job(name: str) -> str:
 
 
 if __name__ == "__main__":
-    tuning_job_name = input("Tuning job name: ")
-    get_tuning_job(tuning_job_name)
+    input_tuning_job_name = input("Tuning job name: ")
+    get_tuning_job(input_tuning_job_name)

genai/tuning/tuning_textgen_with_txt.py

@@ -13,30 +13,32 @@
 # limitations under the License.
 
 
-def test_tuned_endpoint(name: str) -> str:
+def predict_with_tuned_endpoint(tuning_job_name: str) -> str:
     # [START googlegenaisdk_tuning_textgen_with_txt]
     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)
+    # Eg. tuning_job_name = "projects/123456789012/locations/us-central1/tuningJobs/123456789012345"
+    tuning_job = client.tunings.get(name=tuning_job_name)
 
     contents = "Why is the sky blue?"
 
-    # Tests the default checkpoint
+    # Predicts with the tuned endpoint.
     response = client.models.generate_content(
         model=tuning_job.tuned_model.endpoint,
         contents=contents,
     )
     print(response.text)
+    # Example response:
+    # The sky is blue because ...
 
     # [END googlegenaisdk_tuning_textgen_with_txt]
     return response.text
 
 
 if __name__ == "__main__":
-    tuning_job_name = input("Tuning job name: ")
-    test_tuned_endpoint(tuning_job_name)
+    input_tuning_job_name = input("Tuning job name: ")
+    predict_with_tuned_endpoint(input_tuning_job_name)

genai/tuning/tuning_with_checkpoints_get_model.py

@@ -13,16 +13,16 @@
 # limitations under the License.
 
 
-def get_tuned_model_with_checkpoints(name: str) -> str:
+def get_tuned_model_with_checkpoints(tuning_job_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)
+    # Eg. tuning_job_name = "projects/123456789012/locations/us-central1/tuningJobs/123456789012345"
+    tuning_job = client.tunings.get(name=tuning_job_name)
     tuned_model = client.models.get(model=tuning_job.tuned_model.model)
     print(tuned_model)
     # Example response:
@@ -44,5 +44,5 @@ def get_tuned_model_with_checkpoints(name: str) -> str:
 
 
 if __name__ == "__main__":
-    tuning_job_name = input("Tuning job name: ")
-    get_tuned_model_with_checkpoints(tuning_job_name)
+    input_tuning_job_name = input("Tuning job name: ")
+    get_tuned_model_with_checkpoints(input_tuning_job_name)

genai/tuning/tuning_with_checkpoints_list_checkpoints.py

@@ -13,16 +13,16 @@
 # limitations under the License.
 
 
-def list_checkpoints(name: str) -> str:
+def list_checkpoints(tuning_job_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)
+    # Eg. tuning_job_name = "projects/123456789012/locations/us-central1/tuningJobs/123456789012345"
+    tuning_job = client.tunings.get(name=tuning_job_name)
 
     if tuning_job.tuned_model.checkpoints:
         for i, checkpoint in enumerate(tuning_job.tuned_model.checkpoints):
@@ -36,5 +36,5 @@ def list_checkpoints(name: str) -> str:
 
 
 if __name__ == "__main__":
-    tuning_job_name = input("Tuning job name: ")
-    list_checkpoints(tuning_job_name)
+    input_tuning_job_name = input("Tuning job name: ")
+    list_checkpoints(input_tuning_job_name)

genai/tuning/tuning_with_checkpoints_set_default_checkpoint.py

@@ -13,16 +13,16 @@
 # limitations under the License.
 
 
-def set_default_checkpoint(name: str, checkpoint_id: str) -> str:
+def set_default_checkpoint(tuning_job_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)
+    # Eg. tuning_job_name = "projects/123456789012/locations/us-central1/tuningJobs/123456789012345"
+    tuning_job = client.tunings.get(name=tuning_job_name)
     tuned_model = client.models.get(model=tuning_job.tuned_model.model)
 
     print(f"Default checkpoint: {tuned_model.default_checkpoint_id}")
@@ -49,6 +49,6 @@ def set_default_checkpoint(name: str, checkpoint_id: str) -> str:
 
 
 if __name__ == "__main__":
-    tuning_job_name = input("Tuning job name: ")
+    input_tuning_job_name = input("Tuning job name: ")
     default_checkpoint_id = input("Default checkpoint id: ")
-    set_default_checkpoint(tuning_job_name, default_checkpoint_id)
+    set_default_checkpoint(input_tuning_job_name, default_checkpoint_id)

genai/tuning/tuning_with_checkpoints_textgen_with_txt.py

@@ -13,44 +13,50 @@
 # limitations under the License.
 
 
-def test_checkpoint(name: str) -> str:
+def predict_with_checkpoints(tuning_job_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)
+    # Eg. tuning_job_name = "projects/123456789012/locations/us-central1/tuningJobs/123456789012345"
+    tuning_job = client.tunings.get(name=tuning_job_name)
 
     contents = "Why is the sky blue?"
 
-    # Tests the default checkpoint
+    # Predicts with the default checkpoint.
     response = client.models.generate_content(
         model=tuning_job.tuned_model.endpoint,
         contents=contents,
     )
     print(response.text)
+    # Example response:
+    # The sky is blue because ...
 
-    # Tests Checkpoint 1
+    # Predicts with Checkpoint 1.
     checkpoint1_response = client.models.generate_content(
         model=tuning_job.tuned_model.checkpoints[0].endpoint,
         contents=contents,
     )
     print(checkpoint1_response.text)
+    # Example response:
+    # The sky is blue because ...
 
-    # Tests Checkpoint 2
+    # Predicts with Checkpoint 2.
     checkpoint2_response = client.models.generate_content(
         model=tuning_job.tuned_model.checkpoints[1].endpoint,
         contents=contents,
     )
     print(checkpoint2_response.text)
+    # Example response:
+    # The sky is blue because ...
 
     # [END googlegenaisdk_tuning_with_checkpoints_test]
     return response.text
 
 
 if __name__ == "__main__":
-    tuning_job_name = input("Tuning job name: ")
-    test_checkpoint(tuning_job_name)
+    input_tuning_job_name = input("Tuning job name: ")
+    predict_with_checkpoints(input_tuning_job_name)