feat: LLM - Released the BatchPrediction to GA for TextGenerationModel

PiperOrigin-RevId: 548634713

Author: Ark-kun

tests/unit/aiplatform/test_language_models.py

@@ -1302,7 +1302,7 @@ def test_batch_prediction(self):
                 _TEXT_BISON_PUBLISHER_MODEL_DICT
             ),
         ):
-            model = preview_language_models.TextGenerationModel.from_pretrained(
+            model = language_models.TextGenerationModel.from_pretrained(
                 "text-bison@001"
             )
 

vertexai/language_models/_language_models.py

@@ -223,7 +223,7 @@ def __repr__(self):
         return self.text
 
 
-class TextGenerationModel(_LanguageModel):
+class _TextGenerationModel(_LanguageModel):
     """TextGenerationModel represents a general language model.
 
     Examples::
@@ -324,9 +324,6 @@ def _batch_predict(
         return results
 
 
-_TextGenerationModel = TextGenerationModel
-
-
 class _ModelWithBatchPredict(_LanguageModel):
     """Model that supports batch prediction."""
 
@@ -432,15 +429,19 @@ def batch_predict(
         )
 
 
+class TextGenerationModel(_TextGenerationModel, _ModelWithBatchPredict):
+    pass
+
+
 class _PreviewTextGenerationModel(
-    TextGenerationModel, _TunableModelMixin, _PreviewModelWithBatchPredict
+    _TextGenerationModel, _TunableModelMixin, _PreviewModelWithBatchPredict
 ):
     """Preview text generation model."""
 
     _LAUNCH_STAGE = _model_garden_models._SDK_PUBLIC_PREVIEW_LAUNCH_STAGE
 
 
-class _ChatModel(TextGenerationModel):
+class _ChatModel(_TextGenerationModel):
     """ChatModel represents a language model that is capable of chat.
 
     Examples::
@@ -457,10 +458,10 @@ class _ChatModel(TextGenerationModel):
 
     def start_chat(
         self,
-        max_output_tokens: int = TextGenerationModel._DEFAULT_MAX_OUTPUT_TOKENS,
-        temperature: float = TextGenerationModel._DEFAULT_TEMPERATURE,
-        top_k: int = TextGenerationModel._DEFAULT_TOP_K,
-        top_p: float = TextGenerationModel._DEFAULT_TOP_P,
+        max_output_tokens: int = _TextGenerationModel._DEFAULT_MAX_OUTPUT_TOKENS,
+        temperature: float = _TextGenerationModel._DEFAULT_TEMPERATURE,
+        top_k: int = _TextGenerationModel._DEFAULT_TOP_K,
+        top_p: float = _TextGenerationModel._DEFAULT_TOP_P,
     ) -> "_ChatSession":
         """Starts a chat session with the model.
 
@@ -491,10 +492,10 @@ class _ChatSession:
     def __init__(
         self,
         model: _ChatModel,
-        max_output_tokens: int = TextGenerationModel._DEFAULT_MAX_OUTPUT_TOKENS,
-        temperature: float = TextGenerationModel._DEFAULT_TEMPERATURE,
-        top_k: int = TextGenerationModel._DEFAULT_TOP_K,
-        top_p: float = TextGenerationModel._DEFAULT_TOP_P,
+        max_output_tokens: int = _TextGenerationModel._DEFAULT_MAX_OUTPUT_TOKENS,
+        temperature: float = _TextGenerationModel._DEFAULT_TEMPERATURE,
+        top_k: int = _TextGenerationModel._DEFAULT_TOP_K,
+        top_p: float = _TextGenerationModel._DEFAULT_TOP_P,
     ):
         self._model = model
         self._history = []
@@ -635,10 +636,10 @@ def start_chat(
         *,
         context: Optional[str] = None,
         examples: Optional[List[InputOutputTextPair]] = None,
-        max_output_tokens: int = TextGenerationModel._DEFAULT_MAX_OUTPUT_TOKENS,
-        temperature: float = TextGenerationModel._DEFAULT_TEMPERATURE,
-        top_k: int = TextGenerationModel._DEFAULT_TOP_K,
-        top_p: float = TextGenerationModel._DEFAULT_TOP_P,
+        max_output_tokens: int = _TextGenerationModel._DEFAULT_MAX_OUTPUT_TOKENS,
+        temperature: float = _TextGenerationModel._DEFAULT_TEMPERATURE,
+        top_k: int = _TextGenerationModel._DEFAULT_TOP_K,
+        top_p: float = _TextGenerationModel._DEFAULT_TOP_P,
         message_history: Optional[List[ChatMessage]] = None,
     ) -> "ChatSession":
         """Starts a chat session with the model.
@@ -754,10 +755,10 @@ def __init__(
         model: _ChatModelBase,
         context: Optional[str] = None,
         examples: Optional[List[InputOutputTextPair]] = None,
-        max_output_tokens: int = TextGenerationModel._DEFAULT_MAX_OUTPUT_TOKENS,
-        temperature: float = TextGenerationModel._DEFAULT_TEMPERATURE,
-        top_k: int = TextGenerationModel._DEFAULT_TOP_K,
-        top_p: float = TextGenerationModel._DEFAULT_TOP_P,
+        max_output_tokens: int = _TextGenerationModel._DEFAULT_MAX_OUTPUT_TOKENS,
+        temperature: float = _TextGenerationModel._DEFAULT_TEMPERATURE,
+        top_k: int = _TextGenerationModel._DEFAULT_TOP_K,
+        top_p: float = _TextGenerationModel._DEFAULT_TOP_P,
         is_code_chat_session: bool = False,
         message_history: Optional[List[ChatMessage]] = None,
     ):
@@ -885,10 +886,10 @@ def __init__(
         model: ChatModel,
         context: Optional[str] = None,
         examples: Optional[List[InputOutputTextPair]] = None,
-        max_output_tokens: int = TextGenerationModel._DEFAULT_MAX_OUTPUT_TOKENS,
-        temperature: float = TextGenerationModel._DEFAULT_TEMPERATURE,
-        top_k: int = TextGenerationModel._DEFAULT_TOP_K,
-        top_p: float = TextGenerationModel._DEFAULT_TOP_P,
+        max_output_tokens: int = _TextGenerationModel._DEFAULT_MAX_OUTPUT_TOKENS,
+        temperature: float = _TextGenerationModel._DEFAULT_TEMPERATURE,
+        top_k: int = _TextGenerationModel._DEFAULT_TOP_K,
+        top_p: float = _TextGenerationModel._DEFAULT_TOP_P,
         message_history: Optional[List[ChatMessage]] = None,
     ):
         super().__init__(