feat: LLM - Released the Codey models to GA

PiperOrigin-RevId: 544218292
googleapis · Jun 29, 2023 · 89609c9 · 89609c9
1 parent 16299d1
commit 89609c9
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 15 deletions.
diff --git a/tests/unit/aiplatform/test_language_models.py b/tests/unit/aiplatform/test_language_models.py
@@ -101,7 +101,7 @@
     "name": "publishers/google/models/codechat-bison",
     "version_id": "001",
     "open_source_category": "PROPRIETARY",
-    "launch_stage": gca_publisher_model.PublisherModel.LaunchStage.PUBLIC_PREVIEW,
+    "launch_stage": gca_publisher_model.PublisherModel.LaunchStage.GA,
     "publisher_model_template": "projects/{user-project}/locations/{location}/publishers/google/models/codechat-bison@001",
     "predict_schemata": {
         "instance_schema_uri": "gs://google-cloud-aiplatform/schema/predict/instance/codechat_generation_1.0.0.yaml",
@@ -114,7 +114,7 @@
     "name": "publishers/google/models/code-bison",
     "version_id": "001",
     "open_source_category": "PROPRIETARY",
-    "launch_stage": gca_publisher_model.PublisherModel.LaunchStage.PUBLIC_PREVIEW,
+    "launch_stage": gca_publisher_model.PublisherModel.LaunchStage.GA,
     "publisher_model_template": "projects/{user-project}/locations/{location}/publishers/google/models/code-bison@001",
     "predict_schemata": {
         "instance_schema_uri": "gs://google-cloud-aiplatform/schema/predict/instance/code_generation_1.0.0.yaml",
@@ -127,7 +127,7 @@
     "name": "publishers/google/models/code-gecko",
     "version_id": "001",
     "open_source_category": "PROPRIETARY",
-    "launch_stage": gca_publisher_model.PublisherModel.LaunchStage.PUBLIC_PREVIEW,
+    "launch_stage": gca_publisher_model.PublisherModel.LaunchStage.GA,
     "publisher_model_template": "projects/{user-project}/locations/{location}/publishers/google/models/code-gecko@001",
     "predict_schemata": {
         "instance_schema_uri": "gs://google-cloud-aiplatform/schema/predict/instance/code_generation_1.0.0.yaml",
@@ -853,7 +853,7 @@ def test_code_chat(self):
                 _CODECHAT_BISON_PUBLISHER_MODEL_DICT
             ),
         ) as mock_get_publisher_model:
-            model = preview_language_models.CodeChatModel.from_pretrained(
+            model = language_models.CodeChatModel.from_pretrained(
                 "google/codechat-bison@001"
             )
 
@@ -949,7 +949,7 @@ def test_code_generation(self):
                 _CODE_GENERATION_BISON_PUBLISHER_MODEL_DICT
             ),
         ) as mock_get_publisher_model:
-            model = preview_language_models.CodeGenerationModel.from_pretrained(
+            model = language_models.CodeGenerationModel.from_pretrained(
                 "google/code-bison@001"
             )
 
@@ -976,11 +976,9 @@ def test_code_generation(self):
         # Validating the parameters
         predict_temperature = 0.1
         predict_max_output_tokens = 100
-        default_temperature = (
-            preview_language_models.CodeGenerationModel._DEFAULT_TEMPERATURE
-        )
+        default_temperature = language_models.CodeGenerationModel._DEFAULT_TEMPERATURE
         default_max_output_tokens = (
-            preview_language_models.CodeGenerationModel._DEFAULT_MAX_OUTPUT_TOKENS
+            language_models.CodeGenerationModel._DEFAULT_MAX_OUTPUT_TOKENS
         )
 
         with mock.patch.object(
@@ -1017,7 +1015,7 @@ def test_code_completion(self):
                 _CODE_COMPLETION_BISON_PUBLISHER_MODEL_DICT
             ),
         ) as mock_get_publisher_model:
-            model = preview_language_models.CodeGenerationModel.from_pretrained(
+            model = language_models.CodeGenerationModel.from_pretrained(
                 "google/code-gecko@001"
             )
 
@@ -1044,11 +1042,9 @@ def test_code_completion(self):
         # Validating the parameters
         predict_temperature = 0.1
         predict_max_output_tokens = 100
-        default_temperature = (
-            preview_language_models.CodeGenerationModel._DEFAULT_TEMPERATURE
-        )
+        default_temperature = language_models.CodeGenerationModel._DEFAULT_TEMPERATURE
         default_max_output_tokens = (
-            preview_language_models.CodeGenerationModel._DEFAULT_MAX_OUTPUT_TOKENS
+            language_models.CodeGenerationModel._DEFAULT_MAX_OUTPUT_TOKENS
         )
 
         with mock.patch.object(

diff --git a/vertexai/language_models/__init__.py b/vertexai/language_models/__init__.py
@@ -15,6 +15,9 @@
 """Classes for working with language models."""
 
 from vertexai.language_models._language_models import (
+    CodeChatModel,
+    CodeChatSession,
+    CodeGenerationModel,
     InputOutputTextPair,
     TextEmbedding,
     TextEmbeddingModel,
@@ -23,6 +26,9 @@
 )
 
 __all__ = [
+    "CodeChatModel",
+    "CodeChatSession",
+    "CodeGenerationModel",
     "InputOutputTextPair",
     "TextEmbedding",
     "TextEmbeddingModel",

diff --git a/vertexai/language_models/_language_models.py b/vertexai/language_models/_language_models.py
@@ -668,6 +668,7 @@ class CodeChatModel(_ChatModelBase):
     """
 
     _INSTANCE_SCHEMA_URI = "gs://google-cloud-aiplatform/schema/predict/instance/codechat_generation_1.0.0.yaml"
+    _LAUNCH_STAGE = _model_garden_models._SDK_GA_LAUNCH_STAGE
 
     _DEFAULT_MAX_OUTPUT_TOKENS = 128
     _DEFAULT_TEMPERATURE = 0.5
@@ -917,7 +918,7 @@ class CodeGenerationModel(_LanguageModel):
 
     _INSTANCE_SCHEMA_URI = "gs://google-cloud-aiplatform/schema/predict/instance/code_generation_1.0.0.yaml"
 
-    _LAUNCH_STAGE = _model_garden_models._SDK_PUBLIC_PREVIEW_LAUNCH_STAGE
+    _LAUNCH_STAGE = _model_garden_models._SDK_GA_LAUNCH_STAGE
     _DEFAULT_TEMPERATURE = 0.0
     _DEFAULT_MAX_OUTPUT_TOKENS = 128