feat: New text Generation Samples (GoogleCloudPlatform#12543)

* Added 2 Samples for Batch Predict

Author: Thoughtseize1

generative_ai/batch_predict/batch_code_predict.py

@@ -0,0 +1,60 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from google.cloud.aiplatform import BatchPredictionJob
+
+
+def batch_code_prediction(
+    input_uri: str = None, output_uri: str = None
+) -> BatchPredictionJob:
+    """Perform batch code prediction using a pre-trained code generation model.
+    Args:
+        input_uri (str, optional): URI of the input dataset. Could be a BigQuery table or a Google Cloud Storage file.
+            E.g. "gs://[BUCKET]/[DATASET].jsonl" OR "bq://[PROJECT].[DATASET].[TABLE]"
+        output_uri (str, optional): URI where the output will be stored.
+            Could be a BigQuery table or a Google Cloud Storage file.
+            E.g. "gs://[BUCKET]/[OUTPUT].jsonl" OR "bq://[PROJECT].[DATASET].[TABLE]"
+    Returns:
+        batch_prediction_job: The batch prediction job object containing details of the job.
+    """
+
+    # [START aiplatform_batch_code_predict]
+    from vertexai.preview.language_models import CodeGenerationModel
+
+    # Example of using Google Cloud Storage bucket as the input and output data source
+    # TODO (Developer): Replace the input_uri and output_uri with your own GCS paths
+    # input_uri = "gs://cloud-samples-data/batch/prompt_for_batch_code_predict.jsonl"
+    # output_uri = "gs://your-bucket-name/batch_code_predict_output"
+
+    code_model = CodeGenerationModel.from_pretrained("code-bison")
+
+    batch_prediction_job = code_model.batch_predict(
+        dataset=input_uri,
+        destination_uri_prefix=output_uri,
+        # Optional:
+        model_parameters={
+            "maxOutputTokens": "200",
+            "temperature": "0.2",
+        },
+    )
+    print(batch_prediction_job.display_name)
+    print(batch_prediction_job.resource_name)
+    print(batch_prediction_job.state)
+
+    # [END aiplatform_batch_code_predict]
+
+    return batch_prediction_job
+
+
+if __name__ == "__main__":
+    batch_code_prediction()

generative_ai/batch_predict/batch_text_predict.py

@@ -0,0 +1,62 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from google.cloud.aiplatform import BatchPredictionJob
+
+
+def batch_text_prediction(
+    input_uri: str = None, output_uri: str = None
+) -> BatchPredictionJob:
+    """Perform batch text prediction using a pre-trained text generation model.
+    Args:
+        input_uri (str, optional): URI of the input dataset. Could be a BigQuery table or a Google Cloud Storage file.
+            E.g. "gs://[BUCKET]/[DATASET].jsonl" OR "bq://[PROJECT].[DATASET].[TABLE]"
+        output_uri (str, optional): URI where the output will be stored.
+            Could be a BigQuery table or a Google Cloud Storage file.
+            E.g. "gs://[BUCKET]/[OUTPUT].jsonl" OR "bq://[PROJECT].[DATASET].[TABLE]"
+    Returns:
+        batch_prediction_job: The batch prediction job object containing details of the job.
+    """
+
+    # [START aiplatform_batch_text_predict]
+    from vertexai.preview.language_models import TextGenerationModel
+
+    # Example of using Google Cloud Storage bucket as the input and output data source
+    # TODO (Developer): Replace the input_uri and output_uri with your own GCS paths
+    # input_uri = "gs://cloud-samples-data/batch/prompt_for_batch_text_predict.jsonl"
+    # output_uri = "gs://your-bucket-name/batch_text_predict_output"
+
+    # Initialize the text generation model from a pre-trained model named "text-bison"
+    text_model = TextGenerationModel.from_pretrained("text-bison")
+
+    batch_prediction_job = text_model.batch_predict(
+        dataset=input_uri,
+        destination_uri_prefix=output_uri,
+        # Optional:
+        model_parameters={
+            "maxOutputTokens": "200",
+            "temperature": "0.2",
+            "topP": "0.95",
+            "topK": "40",
+        },
+    )
+    print(batch_prediction_job.display_name)
+    print(batch_prediction_job.resource_name)
+    print(batch_prediction_job.state)
+
+    # [END aiplatform_batch_text_predict]
+    return batch_prediction_job
+
+
+if __name__ == "__main__":
+    batch_text_prediction()

generative_ai/batch_predict/test_batch_predict.py

@@ -0,0 +1,72 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Callable
+
+import batch_code_predict
+import batch_text_predict
+
+from google.cloud import storage
+from google.cloud.aiplatform import BatchPredictionJob
+from google.cloud.aiplatform_v1 import JobState
+
+import pytest
+
+INPUT_BUCKET = "cloud-samples-data"
+OUTPUT_BUCKET = "python-docs-samples-tests"
+OUTPUT_PATH = "batch/batch_text_predict_output"
+
+
+def _clean_resources() -> None:
+    storage_client = storage.Client()
+    bucket = storage_client.get_bucket(OUTPUT_BUCKET)
+    blobs = bucket.list_blobs(prefix=OUTPUT_PATH)
+    for blob in blobs:
+        blob.delete()
+
+
+@pytest.fixture(scope="session")
+def output_folder() -> str:
+    yield f"gs://{OUTPUT_BUCKET}/{OUTPUT_PATH}"
+    _clean_resources()
+
+
+def _main_test(test_func: Callable) -> BatchPredictionJob:
+    job = None
+    try:
+        job = test_func()
+        assert job.state == JobState.JOB_STATE_SUCCEEDED
+        return job
+    finally:
+        if job is not None:
+            job.delete()
+
+
+def test_batch_text_predict(output_folder: pytest.fixture()) -> None:
+    input_uri = f"gs://{INPUT_BUCKET}/batch/prompt_for_batch_text_predict.jsonl"
+    job = _main_test(
+        test_func=lambda: batch_text_predict.batch_text_prediction(
+            input_uri, output_folder
+        )
+    )
+    assert OUTPUT_PATH in job.output_info.gcs_output_directory
+
+
+def test_batch_code_predict(output_folder: pytest.fixture()) -> None:
+    input_uri = f"gs://{INPUT_BUCKET}/batch/prompt_for_batch_code_predict.jsonl"
+    job = _main_test(
+        test_func=lambda: batch_code_predict.batch_code_prediction(
+            input_uri, output_folder
+        )
+    )
+    assert OUTPUT_PATH in job.output_info.gcs_output_directory