feat: LLM - Support tuning for the code-bison model (preview)

PiperOrigin-RevId: 550751052
googleapis · Jul 25, 2023 · e4b23a2 · e4b23a2
1 parent 75eb777
commit e4b23a2
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git a/vertexai/_model_garden/_model_garden_models.py b/vertexai/_model_garden/_model_garden_models.py
@@ -31,7 +31,8 @@
 _SUPPORTED_PUBLISHERS = ["google"]
 
 _SHORT_MODEL_ID_TO_TUNING_PIPELINE_MAP = {
-    "text-bison": "https://us-kfp.pkg.dev/ml-pipeline/large-language-model-pipelines/tune-large-model/v2.0.0"
+    "text-bison": "https://us-kfp.pkg.dev/ml-pipeline/large-language-model-pipelines/tune-large-model/v2.0.0",
+    "code-bison": "https://us-kfp.pkg.dev/ml-pipeline/large-language-model-pipelines/tune-large-model/v3.0.0",
 }
 
 _SDK_PRIVATE_PREVIEW_LAUNCH_STAGE = frozenset(

diff --git a/vertexai/language_models/_language_models.py b/vertexai/language_models/_language_models.py
@@ -54,7 +54,9 @@ def _get_model_id_from_tuning_model_id(tuning_model_id: str) -> str:
         return tuning_model_id.replace(
             "text-bison-", "publishers/google/models/text-bison@"
         )
-    raise ValueError(f"Unsupported tuning model ID {tuning_model_id}")
+    if "/" not in tuning_model_id:
+        return "publishers/google/models/" + tuning_model_id
+    return tuning_model_id
 
 
 class _LanguageModel(_model_garden_models._ModelGardenModel):
@@ -1007,6 +1009,10 @@ def predict(
         )
 
 
+class _PreviewCodeGenerationModel(CodeGenerationModel, _TunableModelMixin):
+    _LAUNCH_STAGE = _model_garden_models._SDK_PUBLIC_PREVIEW_LAUNCH_STAGE
+
+
 ###### Model tuning
 # Currently, tuning can only work in this location
 _TUNING_LOCATIONS = ("europe-west4", "us-central1")

diff --git a/vertexai/preview/language_models.py b/vertexai/preview/language_models.py
@@ -16,20 +16,21 @@
 
 from vertexai.language_models._language_models import (
     _PreviewChatModel,
+    _PreviewCodeGenerationModel,
     _PreviewTextEmbeddingModel,
     _PreviewTextGenerationModel,
     ChatMessage,
     ChatModel,
     ChatSession,
     CodeChatModel,
     CodeChatSession,
-    CodeGenerationModel,
     InputOutputTextPair,
     TextEmbedding,
     TextGenerationResponse,
 )
 
 ChatModel = _PreviewChatModel
+CodeGenerationModel = _PreviewCodeGenerationModel
 TextGenerationModel = _PreviewTextGenerationModel
 TextEmbeddingModel = _PreviewTextEmbeddingModel