feat(genai): add Gemini batch prediction samples for BQ and GCS (GoogleCloudPlatform#12731)

* feat(genai): add Gemini batch prediction samples for BQ and GCS

* cleaning up test file removing extra comments

* updating import and removing sample data

* move comments in region tag

* removing white space

* update values

* update values

* removing space

* update values

* update imprt

* update projectid

* updating testing file

* update assert value

* changing input value for gcs testing. Tests are failing due to permission erros

* update output path

Author: CadillacBurgess1

generative_ai/batch_predict/gemini_batch_predict_bigquery.py

@@ -0,0 +1,70 @@
+# 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.
+import os
+
+PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
+
+output_uri = "bq://storage-samples.generative_ai.gen_ai_batch_prediction.predictions"
+
+
+def batch_predict_gemini_createjob(output_uri: str) -> str:
+    """Perform batch text prediction using a Gemini AI model and returns the output location"""
+
+    # [START generativeaionvertexai_batch_predict_gemini_createjob_bigquery]
+    import time
+    import vertexai
+
+    from vertexai.batch_prediction import BatchPredictionJob
+
+    # TODO(developer): Update and un-comment below lines
+
+    # Initialize vertexai
+    vertexai.init(project=PROJECT_ID, location="us-central1")
+
+    input_uri = "bq://storage-samples.generative_ai.batch_requests_for_multimodal_input"
+
+    # Submit a batch prediction job with Gemini model
+    batch_prediction_job = BatchPredictionJob.submit(
+        source_model="gemini-1.5-flash-002",
+        input_dataset=input_uri,
+        output_uri_prefix=output_uri,
+    )
+
+    # Check job status
+    print(f"Job resource name: {batch_prediction_job.resource_name}")
+    print(f"Model resource name with the job: {batch_prediction_job.model_name}")
+    print(f"Job state: {batch_prediction_job.state.name}")
+
+    # Refresh the job until complete
+    while not batch_prediction_job.has_ended:
+        time.sleep(5)
+        batch_prediction_job.refresh()
+
+    # Check if the job succeeds
+    if batch_prediction_job.has_succeeded:
+        print("Job succeeded!")
+    else:
+        print(f"Job failed: {batch_prediction_job.error}")
+
+    # Check the location of the output
+    print(f"Job output location: {batch_prediction_job.output_location}")
+
+    # Example response:
+    #  Job output location: bq://Project-ID/gen-ai-batch-prediction/predictions-model-year-month-day-hour:minute:second.12345
+    # [END generativeaionvertexai_batch_predict_gemini_createjob_bigquery]
+    return batch_prediction_job
+
+
+if __name__ == "__main__":
+    batch_predict_gemini_createjob(output_uri)

generative_ai/batch_predict/gemini_batch_predict_gcs.py

@@ -13,24 +13,14 @@
 # limitations under the License.
 import os
 
-from vertexai.batch_prediction import BatchPredictionJob
 
 PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
 
+output_uri = "gs://python-docs-samples-tests"
 
-def batch_predict_gemini_createjob(
-    input_uri: str, output_uri: str
-) -> BatchPredictionJob:
-    """Perform batch text prediction using a Gemini AI model.
-    Args:
-        input_uri (str): URI of the input file in Google Cloud Storage.
-            Example: "gs://[BUCKET]/[DATASET].jsonl"
 
-        output_uri (str): URI of the output folder in Google Cloud Storage.
-            Example: "gs://[BUCKET]/[OUTPUT].jsonl"
-    Returns:
-        batch_prediction_job: The batch prediction job object containing details of the job.
-    """
+def batch_predict_gemini_createjob(output_uri: str) -> str:
+    "Perform batch text prediction using a Gemini AI model and returns the output location"
 
     # [START generativeaionvertexai_batch_predict_gemini_createjob]
     import time
@@ -39,12 +29,12 @@ def batch_predict_gemini_createjob(
     from vertexai.batch_prediction import BatchPredictionJob
 
     # TODO(developer): Update and un-comment below lines
-    # input_uri ="gs://[BUCKET]/[OUTPUT].jsonl" # Example
-    # output_uri ="gs://[BUCKET]"
 
     # Initialize vertexai
     vertexai.init(project=PROJECT_ID, location="us-central1")
 
+    input_uri = "gs://cloud-samples-data/batch/prompt_for_batch_gemini_predict.jsonl"
+
     # Submit a batch prediction job with Gemini model
     batch_prediction_job = BatchPredictionJob.submit(
         source_model="gemini-1.5-flash-002",
@@ -74,16 +64,9 @@ def batch_predict_gemini_createjob(
     # Example response:
     #  Job output location: gs://your-bucket/gen-ai-batch-prediction/prediction-model-year-month-day-hour:minute:second.12345
 
-    # https://storage.googleapis.com/cloud-samples-data/batch/prompt_for_batch_gemini_predict.jsonl
-
-    return batch_prediction_job
     # [END generativeaionvertexai_batch_predict_gemini_createjob]
+    return batch_prediction_job
 
 
 if __name__ == "__main__":
-    # TODO(developer): Update your Cloud Storage bucket and uri file paths
-    GCS_BUCKET = "gs://your-bucket"
-    batch_predict_gemini_createjob(
-        input_uri=f"gs://{GCS_BUCKET}/batch_data/sample_input_file.jsonl",
-        output_uri=f"gs://{GCS_BUCKET}/batch_predictions/sample_output/",
-    )
+    batch_predict_gemini_createjob(output_uri)

generative_ai/batch_predict/test_batch_predict_examples.py

@@ -11,22 +11,32 @@
 # 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.
+import os
+
 from typing import Callable
 
+
 from google.cloud import storage
 from google.cloud.aiplatform import BatchPredictionJob
 from google.cloud.aiplatform_v1 import JobState
 
+
 import pytest
 
+
 import batch_code_predict
 import batch_text_predict
+import gemini_batch_predict_bigquery
 import gemini_batch_predict_gcs
 
+PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
+
 
 INPUT_BUCKET = "cloud-samples-data"
 OUTPUT_BUCKET = "python-docs-samples-tests"
 OUTPUT_PATH = "batch/batch_text_predict_output"
+GCS_OUTPUT_PATH = "gs://python-docs-samples-tests/"
+OUTPUT_TABLE = f"bq://{PROJECT_ID}.gen_ai_batch_prediction.predictions"
 
 
 def _clean_resources() -> None:
@@ -75,10 +85,20 @@ def test_batch_code_predict(output_folder: pytest.fixture()) -> None:
 
 
 def test_batch_gemini_predict_gcs(output_folder: pytest.fixture()) -> None:
-    input_uri = f"gs://{INPUT_BUCKET}/batch/prompt_for_batch_gemini_predict.jsonl"
+    output_uri = "gs://python-docs-samples-tests"
     job = _main_test(
         test_func=lambda: gemini_batch_predict_gcs.batch_predict_gemini_createjob(
-            input_uri, output_folder
+            output_uri
         )
     )
-    assert OUTPUT_PATH in job.output_location
+    assert GCS_OUTPUT_PATH in job.output_location
+
+
+def test_batch_gemini_predict_bigquery(output_folder: pytest.fixture()) -> None:
+    output_uri = f"bq://{PROJECT_ID}.gen_ai_batch_prediction.predictions"
+    job = _main_test(
+        test_func=lambda: gemini_batch_predict_bigquery.batch_predict_gemini_createjob(
+            output_uri
+            )
+    )
+    assert OUTPUT_TABLE in job.output_location