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

Ark-kun · copybara-github · commit e4b23a254aad · 2023-07-24T20:17:18.000-07:00
PiperOrigin-RevId: 550751052
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
 

Original file line number	Diff line number	Diff line change
`@@ -31,7 +31,8 @@`
`31`	`31`	`_SUPPORTED_PUBLISHERS = ["google"]`
`32`	`32`
`33`	`33`	`_SHORT_MODEL_ID_TO_TUNING_PIPELINE_MAP = {`
`34`		`- "text-bison": "https://us-kfp.pkg.dev/ml-pipeline/large-language-model-pipelines/tune-large-model/v2.0.0"`
	`34`	`+ "text-bison": "https://us-kfp.pkg.dev/ml-pipeline/large-language-model-pipelines/tune-large-model/v2.0.0",`
	`35`	`+ "code-bison": "https://us-kfp.pkg.dev/ml-pipeline/large-language-model-pipelines/tune-large-model/v3.0.0",`
`35`	`36`	`}`
`36`	`37`
`37`	`38`	`_SDK_PRIVATE_PREVIEW_LAUNCH_STAGE = frozenset(`