feat: LLM - Released CodeChatModel tuning to GA

PiperOrigin-RevId: 578949277
googleapis · Nov 2, 2023 · 621af52 · 621af52
1 parent 53d31b5
commit 621af52
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/tests/unit/aiplatform/test_language_models.py b/tests/unit/aiplatform/test_language_models.py
@@ -1915,9 +1915,7 @@ def test_tune_code_chat_model(
                 _CODECHAT_BISON_PUBLISHER_MODEL_DICT
             ),
         ):
-            model = preview_language_models.CodeChatModel.from_pretrained(
-                "codechat-bison@001"
-            )
+            model = language_models.CodeChatModel.from_pretrained("codechat-bison@001")
 
             # The tune_model call needs to be inside the PublisherModel mock
             # since it gets a new PublisherModel when tuning completes.

diff --git a/vertexai/language_models/_language_models.py b/vertexai/language_models/_language_models.py
@@ -1640,8 +1640,7 @@ def start_chat(
             stop_sequences=stop_sequences,
         )
 
-
-class CodeChatModel(_ChatModelBase):
+class CodeChatModel(_ChatModelBase, _TunableChatModelMixin):
     """CodeChatModel represents a model that is capable of completing code.
 
     Examples:
@@ -1693,7 +1692,7 @@ def start_chat(
         )
 
 
-class _PreviewCodeChatModel(CodeChatModel, _TunableChatModelMixin):
+class _PreviewCodeChatModel(CodeChatModel):
     __name__ = "CodeChatModel"
     __module__ = "vertexai.preview.language_models"