chore: LLM - Renamed batch_predict function's source_uri to dataset.

This way we can support Pandas DataFrames and local datasets in the future.

PiperOrigin-RevId: 548256736

Author: Ark-kun

tests/system/aiplatform/test_language_models.py

@@ -168,7 +168,7 @@ def test_batch_prediction(self):
 
         model = TextGenerationModel.from_pretrained("text-bison@001")
         job = model.batch_predict(
-            source_uri=source_uri,
+            dataset=source_uri,
             destination_uri_prefix=destination_uri_prefix,
             model_parameters={"temperature": 0, "top_p": 1, "top_k": 5},
         )

tests/unit/aiplatform/test_language_models.py

@@ -1311,7 +1311,7 @@ def test_batch_prediction(self):
             attribute="create",
         ) as mock_create:
             model.batch_predict(
-                source_uri="gs://test-bucket/test_table.jsonl",
+                dataset="gs://test-bucket/test_table.jsonl",
                 destination_uri_prefix="gs://test-bucket/results/",
                 model_parameters={"temperature": 0.1},
             )

vertexai/language_models/_language_models.py

@@ -16,6 +16,7 @@
 
 import dataclasses
 from typing import Any, Dict, List, Optional, Sequence, Union
+import warnings
 
 from google.cloud import aiplatform
 from google.cloud.aiplatform import base
@@ -332,14 +333,14 @@ class _ModelWithBatchPredict(_LanguageModel):
     def batch_predict(
         self,
         *,
-        source_uri: Union[str, List[str]],
+        dataset: Union[str, List[str]],
         destination_uri_prefix: str,
         model_parameters: Optional[Dict] = None,
     ) -> aiplatform.BatchPredictionJob:
         """Starts a batch prediction job with the model.
 
         Args:
-            source_uri: The location of the dataset.
+            dataset: The location of the dataset.
                 `gs://` and `bq://` URIs are supported.
             destination_uri_prefix: The URI prefix for the prediction.
                 `gs://` and `bq://` URIs are supported.
@@ -351,22 +352,22 @@ def batch_predict(
             ValueError: When source or destination URI is not supported.
         """
         arguments = {}
-        first_source_uri = source_uri if isinstance(source_uri, str) else source_uri[0]
+        first_source_uri = dataset if isinstance(dataset, str) else dataset[0]
         if first_source_uri.startswith("gs://"):
-            if not isinstance(source_uri, str):
-                if not all(uri.startswith("gs://") for uri in source_uri):
+            if not isinstance(dataset, str):
+                if not all(uri.startswith("gs://") for uri in dataset):
                     raise ValueError(
-                        f"All URIs in the list must start with 'gs://': {source_uri}"
+                        f"All URIs in the list must start with 'gs://': {dataset}"
                     )
-            arguments["gcs_source"] = source_uri
+            arguments["gcs_source"] = dataset
         elif first_source_uri.startswith("bq://"):
-            if not isinstance(source_uri, str):
+            if not isinstance(dataset, str):
                 raise ValueError(
-                    f"Only single BigQuery source can be specified: {source_uri}"
+                    f"Only single BigQuery source can be specified: {dataset}"
                 )
-            arguments["bigquery_source"] = source_uri
+            arguments["bigquery_source"] = dataset
         else:
-            raise ValueError(f"Unsupported source_uri: {source_uri}")
+            raise ValueError(f"Unsupported source_uri: {dataset}")
 
         if destination_uri_prefix.startswith("gs://"):
             arguments["gcs_destination_prefix"] = destination_uri_prefix
@@ -391,8 +392,48 @@ def batch_predict(
         return job
 
 
+class _PreviewModelWithBatchPredict(_ModelWithBatchPredict):
+    """Model that supports batch prediction."""
+
+    def batch_predict(
+        self,
+        *,
+        destination_uri_prefix: str,
+        dataset: Optional[Union[str, List[str]]] = None,
+        model_parameters: Optional[Dict] = None,
+        **_kwargs: Optional[Dict[str, Any]],
+    ) -> aiplatform.BatchPredictionJob:
+        """Starts a batch prediction job with the model.
+
+        Args:
+            dataset: Required. The location of the dataset.
+                `gs://` and `bq://` URIs are supported.
+            destination_uri_prefix: The URI prefix for the prediction.
+                `gs://` and `bq://` URIs are supported.
+            model_parameters: Model-specific parameters to send to the model.
+            **_kwargs: Deprecated.
+
+        Returns:
+            A `BatchPredictionJob` object
+        Raises:
+            ValueError: When source or destination URI is not supported.
+        """
+        if "source_uri" in _kwargs:
+            warnings.warn("source_uri is deprecated, use dataset instead.")
+            if dataset:
+                raise ValueError("source_uri is deprecated, use dataset instead.")
+            dataset = _kwargs["source_uri"]
+        if not dataset:
+            raise ValueError("dataset must be specified")
+        return super().batch_predict(
+            dataset=dataset,
+            destination_uri_prefix=destination_uri_prefix,
+            model_parameters=model_parameters,
+        )
+
+
 class _PreviewTextGenerationModel(
-    TextGenerationModel, _TunableModelMixin, _ModelWithBatchPredict
+    TextGenerationModel, _TunableModelMixin, _PreviewModelWithBatchPredict
 ):
     """Preview text generation model."""