feat: GenAI - Added support for supervised fine-tuning

Ark-kun · copybara-github · commit 036d2d0306e5 · 2024-04-04T15:16:57.000-07:00
PiperOrigin-RevId: 621984253
diff --git a/tests/unit/vertexai/test_tuning.py b/tests/unit/vertexai/test_tuning.py
@@ -0,0 +1,189 @@
+# -*- coding: utf-8 -*-
+
+# 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
+#
+#     http://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.
+#
+"""Unit tests for generative model tuning."""
+# pylint: disable=protected-access,bad-continuation
+
+import copy
+import datetime
+from typing import Dict, Iterable
+from unittest import mock
+import uuid
+
+import vertexai
+from google.cloud.aiplatform import compat
+from google.cloud.aiplatform import initializer
+from google.cloud.aiplatform import utils as aiplatform_utils
+from google.cloud.aiplatform_v1.services import gen_ai_tuning_service
+from google.cloud.aiplatform_v1.types import job_state
+from google.cloud.aiplatform_v1.types import tuning_job as gca_tuning_job
+from vertexai.preview import tuning
+from vertexai.preview.tuning import sft as supervised_tuning
+
+import pytest
+
+from google.rpc import status_pb2
+
+
+_TEST_PROJECT = "test-project"
+_TEST_LOCATION = "us-central1"
+
+
+_global_tuning_jobs: Dict[str, gca_tuning_job.TuningJob] = {}
+
+
+class MockGenAiTuningServiceClient(gen_ai_tuning_service.GenAiTuningServiceClient):
+    @property
+    def _tuning_jobs(self) -> Dict[str, gca_tuning_job.TuningJob]:
+        return _global_tuning_jobs
+
+    def create_tuning_job(
+        self,
+        *,
+        parent: str,
+        tuning_job: gca_tuning_job.TuningJob,
+        **_,
+    ) -> gca_tuning_job.TuningJob:
+        tuning_job = copy.deepcopy(tuning_job)
+        resource_id = uuid.uuid4().hex
+        resource_name = f"{parent}/tuningJobs/{resource_id}"
+        tuning_job.name = resource_name
+        current_time = datetime.datetime.now(datetime.timezone.utc)
+        tuning_job.create_time = current_time
+        tuning_job.update_time = current_time
+        tuning_job.state = job_state.JobState.JOB_STATE_PENDING
+        self._tuning_jobs[resource_name] = tuning_job
+        return tuning_job
+
+    def _progress_tuning_job(self, name: str):
+        tuning_job: gca_tuning_job.TuningJob = self._tuning_jobs[name]
+        current_time = datetime.datetime.now(datetime.timezone.utc)
+        if tuning_job.state == job_state.JobState.JOB_STATE_PENDING:
+            if (
+                "invalid_dataset"
+                in tuning_job.supervised_tuning_spec.training_dataset_uri
+            ):
+                tuning_job.state = job_state.JobState.JOB_STATE_FAILED
+                tuning_job.error = status_pb2.Status(
+                    code=400, message="Invalid dataset."
+                )
+            else:
+                tuning_job.state = job_state.JobState.JOB_STATE_RUNNING
+            tuning_job.update_time = current_time
+        elif tuning_job.state == job_state.JobState.JOB_STATE_RUNNING:
+            parent = tuning_job.name.partition("/tuningJobs/")[0]
+            tuning_job.state = job_state.JobState.JOB_STATE_SUCCEEDED
+            experiment_id = uuid.uuid4().hex
+            tuned_model_id = uuid.uuid4().hex
+            tuned_model_endpoint_id = uuid.uuid4().hex
+            tuning_job.experiment = (
+                f"{parent}/metadataStores/default/contexts/{experiment_id}"
+            )
+            tuning_job.tuned_model = gca_tuning_job.TunedModel(
+                model=f"{parent}/models/{tuned_model_id}",
+                endpoint=f"{parent}/endpoints/{tuned_model_endpoint_id}",
+            )
+            tuning_job.end_time = current_time
+            tuning_job.update_time = current_time
+        else:
+            pass
+
+    def get_tuning_job(self, *, name: str, **_) -> gca_tuning_job.TuningJob:
+        tuning_job = self._tuning_jobs[name]
+        tuning_job = copy.deepcopy(tuning_job)
+        self._progress_tuning_job(name)
+
+        return tuning_job
+
+    def list_tuning_jobs(
+        self, *, parent: str, **_
+    ) -> Iterable[gca_tuning_job.TuningJob]:
+        return [
+            tuning_job
+            for name, tuning_job in self._tuning_jobs.items()
+            if name.startswith(parent + "/")
+        ]
+
+    def cancel_tuning_job(self, *, name: str, **_) -> None:
+        tuning_job = self._tuning_jobs[name]
+        assert tuning_job.state in (
+            job_state.JobState.JOB_STATE_RUNNING,
+            job_state.JobState.JOB_STATE_PENDING,
+        )
+        tuning_job.state = job_state.JobState.JOB_STATE_CANCELLED
+
+
+class MockTuningJobClientWithOverride(aiplatform_utils.ClientWithOverride):
+    _is_temporary = False
+    _default_version = compat.V1
+    _version_map = (
+        (compat.V1, MockGenAiTuningServiceClient),
+        # v1beta1 version does not exist
+        # (compat.V1BETA1, gen_ai_tuning_service_v1beta1.client.JobServiceClient),
+    )
+
+
+@pytest.mark.usefixtures("google_auth_mock")
+class TestgenerativeModelTuning:
+    """Unit tests for generative model tuning."""
+
+    def setup_method(self):
+        vertexai.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+        )
+
+    def teardown_method(self):
+        initializer.global_pool.shutdown(wait=True)
+
+    @mock.patch.object(
+        target=tuning.TuningJob,
+        attribute="client_class",
+        new=MockTuningJobClientWithOverride,
+    )
+    def test_genai_tuning_service_supervised_tuning_tune_model(self):
+        sft_tuning_job = supervised_tuning.train(
+            source_model="gemini-1.0-pro-001",
+            train_dataset="gs://some-bucket/some_dataset.jsonl",
+            # Optional:
+            validation_dataset="gs://some-bucket/some_dataset.jsonl",
+            epochs=300,
+            learning_rate_multiplier=1.0,
+        )
+        assert sft_tuning_job.state == job_state.JobState.JOB_STATE_PENDING
+        assert not sft_tuning_job.has_ended
+        assert not sft_tuning_job.has_succeeded
+
+        # Refreshing the job
+        sft_tuning_job.refresh()
+        assert sft_tuning_job.state == job_state.JobState.JOB_STATE_PENDING
+        assert not sft_tuning_job.has_ended
+        assert not sft_tuning_job.has_succeeded
+
+        # Refreshing the job
+        sft_tuning_job.refresh()
+        assert sft_tuning_job.state == job_state.JobState.JOB_STATE_RUNNING
+        assert not sft_tuning_job.has_ended
+        assert not sft_tuning_job.has_succeeded
+
+        # Refreshing the job
+        sft_tuning_job.refresh()
+        assert sft_tuning_job.state == job_state.JobState.JOB_STATE_SUCCEEDED
+        assert sft_tuning_job.has_ended
+        assert sft_tuning_job.has_succeeded
+        assert sft_tuning_job._experiment_name
+        assert sft_tuning_job.tuned_model_name
+        assert sft_tuning_job.tuned_model_endpoint_name
diff --git a/vertexai/generative_models/_generative_models.py b/vertexai/generative_models/_generative_models.py
@@ -145,13 +145,16 @@ def __init__(
 
         Args:
             model_name: Model Garden model resource name.
+                Alternatively, a tuned model endpoint resource name can be provided.
             generation_config: Default generation config to use in generate_content.
             safety_settings: Default safety settings to use in generate_content.
             tools: Default tools to use in generate_content.
             system_instruction: Default system instruction to use in generate_content.
                 Note: Only text should be used in parts.
                 Content of each part will become a separate paragraph.
         """
+        if not model_name:
+            raise ValueError("model_name must not be empty")
         if "/" not in model_name:
             model_name = "publishers/google/models/" + model_name
         if model_name.startswith("models/"):
@@ -160,10 +163,13 @@ def __init__(
         project = aiplatform_initializer.global_config.project
         location = aiplatform_initializer.global_config.location
 
+        if model_name.startswith("publishers/"):
+            prediction_resource_name = f"projects/{project}/locations/{location}/{model_name}"
+        else:
+            prediction_resource_name = model_name
+
         self._model_name = model_name
-        self._prediction_resource_name = (
-            f"projects/{project}/locations/{location}/{model_name}"
-        )
+        self._prediction_resource_name = prediction_resource_name
         self._generation_config = generation_config
         self._safety_settings = safety_settings
         self._tools = tools
diff --git a/vertexai/preview/tuning/__init__.py b/vertexai/preview/tuning/__init__.py
@@ -0,0 +1,23 @@
+# 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
+#
+#     http://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.
+#
+"""Classes for tuning models."""
+
+# We just want to re-export certain classes
+# pylint: disable=g-multiple-import,g-importing-member
+from vertexai.tuning._tuning import TuningJob
+
+__all__ = [
+    "TuningJob",
+]
diff --git a/vertexai/preview/tuning/sft.py b/vertexai/preview/tuning/sft.py
@@ -0,0 +1,27 @@
+# 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
+#
+#     http://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.
+#
+"""Classes for supervised tuning."""
+
+# We just want to re-export certain classes
+# pylint: disable=g-multiple-import,g-importing-member
+from vertexai.tuning._supervised_tuning import (
+    train,
+    SupervisedTuningJob,
+)
+
+__all__ = [
+    "train",
+    "SupervisedTuningJob",
+]
diff --git a/vertexai/tuning/_supervised_tuning.py b/vertexai/tuning/_supervised_tuning.py
@@ -0,0 +1,71 @@
+# 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
+#
+#     http://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.
+#
+
+from typing import Optional, Union
+
+from google.cloud.aiplatform_v1.types import tuning_job as gca_tuning_job_types
+
+from vertexai import generative_models
+from vertexai.tuning import _tuning
+
+
+def train(
+    *,
+    source_model: Union[str, generative_models.GenerativeModel],
+    train_dataset: str,
+    validation_dataset: Optional[str] = None,
+    tuned_model_display_name: Optional[str] = None,
+    epochs: Optional[int] = None,
+    learning_rate_multiplier: Optional[float] = None,
+) -> "SupervisedTuningJob":
+    """Tunes a model using supervised training.
+
+    Args:
+        source_model (str):
+            Model name for tuning, e.g., "gemini-1.0-pro" or "gemini-1.0-pro-001".
+        train_dataset: Cloud Storage path to file containing training dataset for tuning.
+            The dataset should be in JSONL format.
+        validation_dataset: Cloud Storage path to file containing validation dataset for tuning.
+            The dataset should be in JSONL format.
+        tuned_model_display_name: The display name of the
+            [TunedModel][google.cloud.aiplatform.v1.Model]. The name can
+            be up to 128 characters long and can consist of any UTF-8 characters.
+        epochs: Number of training epoches for this tuning job.
+        learning_rate_multiplier: Learning rate multiplier for tuning.
+
+    Returns:
+        A `TuningJob` object.
+    """
+    supervised_tuning_spec = gca_tuning_job_types.SupervisedTuningSpec(
+        training_dataset_uri=train_dataset,
+        validation_dataset_uri=validation_dataset,
+        hyper_parameters=gca_tuning_job_types.SupervisedHyperParameters(
+            epoch_count=epochs,
+            learning_rate_multiplier=learning_rate_multiplier,
+        ),
+    )
+
+    if isinstance(source_model, generative_models.GenerativeModel):
+        source_model = source_model._prediction_resource_name.rpartition('/')[-1]
+
+    return SupervisedTuningJob._create(  # pylint: disable=protected-access
+        base_model=source_model,
+        tuning_spec=supervised_tuning_spec,
+        tuned_model_display_name=tuned_model_display_name,
+    )
+
+
+class SupervisedTuningJob(_tuning.TuningJob):
+    pass
diff --git a/vertexai/tuning/_tuning.py b/vertexai/tuning/_tuning.py
