From 62ff30daa718ac7869714c68e55d6955d6355945 Mon Sep 17 00:00:00 2001
From: Alexey Volkov <avolkov@google.com>
Date: Sun, 20 Aug 2023 19:49:30 -0700
Subject: [PATCH] feat: LLM - Released `TextGenerationModel` tuning to GA

Changes from Preview:

* `tune_model()` no longer blocks while the model is being tuned.
* `tune_model()` no longer updates the model in-place once tuning has fnished. Instead, it returns a job that can be used to get the newly tuned model.
* `tune_model()` now returns a tuning job object. The `tuning_job.tuned_model` property can be used to get the tuned model, waiting for the tuning to finish if needed. (This is also working in Preview)
* The `learning_rate` parameter has been removed. Use `learning_rate_multiplier` instead.
* The default value for `train_steps` has changed from 1000 to the tuning pipeline default value (usually 300).

PiperOrigin-RevId: 558650122
---
 .../system/aiplatform/test_language_models.py |  2 +-
 tests/unit/aiplatform/test_language_models.py | 80 +++++++++++++++++++
 vertexai/language_models/_language_models.py  |  4 +-
 3 files changed, 84 insertions(+), 2 deletions(-)

diff --git a/tests/system/aiplatform/test_language_models.py b/tests/system/aiplatform/test_language_models.py
index b0ea461712..607d523d24 100644
--- a/tests/system/aiplatform/test_language_models.py
+++ b/tests/system/aiplatform/test_language_models.py
@@ -165,7 +165,7 @@ def test_tuning(self, shared_state):
         """Test tuning, listing and loading models."""
         aiplatform.init(project=e2e_base._PROJECT, location=e2e_base._LOCATION)
 
-        model = TextGenerationModel.from_pretrained("google/text-bison@001")
+        model = language_models.TextGenerationModel.from_pretrained("text-bison@001")
 
         import pandas
 
diff --git a/tests/unit/aiplatform/test_language_models.py b/tests/unit/aiplatform/test_language_models.py
index d8a612c9e9..c3836f88cb 100644
--- a/tests/unit/aiplatform/test_language_models.py
+++ b/tests/unit/aiplatform/test_language_models.py
@@ -1384,6 +1384,86 @@ def test_tune_text_generation_model(
                 == test_constants.EndpointConstants._TEST_ENDPOINT_NAME
             )
 
+    @pytest.mark.parametrize(
+        "job_spec",
+        [_TEST_PIPELINE_SPEC_JSON, _TEST_PIPELINE_JOB],
+    )
+    @pytest.mark.parametrize(
+        "mock_request_urlopen",
+        ["https://us-central1-kfp.pkg.dev/proj/repo/pack/latest"],
+        indirect=True,
+    )
+    def test_tune_text_generation_model_ga(
+        self,
+        mock_pipeline_service_create,
+        mock_pipeline_job_get,
+        mock_pipeline_bucket_exists,
+        job_spec,
+        mock_load_yaml_and_json,
+        mock_gcs_from_string,
+        mock_gcs_upload,
+        mock_request_urlopen,
+        mock_get_tuned_model,
+    ):
+        """Tests tuning the text generation model."""
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+            encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME,
+        )
+        with mock.patch.object(
+            target=model_garden_service_client.ModelGardenServiceClient,
+            attribute="get_publisher_model",
+            return_value=gca_publisher_model.PublisherModel(
+                _TEXT_BISON_PUBLISHER_MODEL_DICT
+            ),
+        ):
+            model = language_models.TextGenerationModel.from_pretrained(
+                "text-bison@001"
+            )
+
+            tuning_job_location = "europe-west4"
+            evaluation_data_uri = "gs://bucket/eval.jsonl"
+            evaluation_interval = 37
+            enable_early_stopping = True
+            tensorboard_name = f"projects/{_TEST_PROJECT}/locations/{tuning_job_location}/tensorboards/123"
+
+            tuning_job = model.tune_model(
+                training_data=_TEST_TEXT_BISON_TRAINING_DF,
+                tuning_job_location=tuning_job_location,
+                tuned_model_location="us-central1",
+                learning_rate_multiplier=2.0,
+                train_steps=10,
+                tuning_evaluation_spec=preview_language_models.TuningEvaluationSpec(
+                    evaluation_data=evaluation_data_uri,
+                    evaluation_interval=evaluation_interval,
+                    enable_early_stopping=enable_early_stopping,
+                    tensorboard=tensorboard_name,
+                ),
+            )
+            call_kwargs = mock_pipeline_service_create.call_args[1]
+            pipeline_arguments = call_kwargs[
+                "pipeline_job"
+            ].runtime_config.parameter_values
+            assert pipeline_arguments["learning_rate_multiplier"] == 2.0
+            assert pipeline_arguments["train_steps"] == 10
+            assert pipeline_arguments["evaluation_data_uri"] == evaluation_data_uri
+            assert pipeline_arguments["evaluation_interval"] == evaluation_interval
+            assert pipeline_arguments["enable_early_stopping"] == enable_early_stopping
+            assert pipeline_arguments["tensorboard_resource_id"] == tensorboard_name
+            assert pipeline_arguments["large_model_reference"] == "text-bison@001"
+            assert (
+                call_kwargs["pipeline_job"].encryption_spec.kms_key_name
+                == _TEST_ENCRYPTION_KEY_NAME
+            )
+
+            # Testing the tuned model
+            tuned_model = tuning_job.get_tuned_model()
+            assert (
+                tuned_model._endpoint_name
+                == test_constants.EndpointConstants._TEST_ENDPOINT_NAME
+            )
+
     @pytest.mark.parametrize(
         "job_spec",
         [_TEST_PIPELINE_SPEC_JSON],
diff --git a/vertexai/language_models/_language_models.py b/vertexai/language_models/_language_models.py
index 7ece93d930..cfa484d9dc 100644
--- a/vertexai/language_models/_language_models.py
+++ b/vertexai/language_models/_language_models.py
@@ -894,7 +894,9 @@ def batch_predict(
         )
 
 
-class TextGenerationModel(_TextGenerationModel, _ModelWithBatchPredict):
+class TextGenerationModel(
+    _TextGenerationModel, _TunableTextModelMixin, _ModelWithBatchPredict
+):
     pass