anthropic batch predict samples

Author: CadillacBurgess1

generative_ai/model_garden/claude_3_batch_prediciton_bq.py

@@ -0,0 +1,71 @@
+# 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_anthropic_gemini_createjob_bigquery]
+    import time
+    import vertexai
+
+    from vertexai.batch_prediction import BatchPredictionJob
+
+    # TODO(developer): Update and un-comment below line
+    # PROJECT_ID = "your-project-id"
+
+    # Initialize vertexai
+    vertexai.init(project=PROJECT_ID, location="us-east5")
+
+    input_uri = "bq://kellysun-test-project.bp_llm_input.claude_50_requests"
+
+    # Submit a batch prediction job with Gemini model
+    batch_prediction_job = BatchPredictionJob.submit(
+        source_model="publishers/anthropic/models/claude-3-5-haiku",
+        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_anthropic_gemini_createjob_bigquery]
+    return batch_prediction_job
+
+
+if __name__ == "__main__":
+    batch_predict_gemini_createjob(output_uri)

generative_ai/model_garden/claude_3_batch_prediction_gcs.py

@@ -0,0 +1,74 @@
+
+# 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 = "gs://python-docs-samples-tests"
+
+
+def batch_predict_createjob(output_uri: str) -> str:
+    "Perform batch text prediction using a Gemini AI model and returns the output location"
+
+    # [START generativeaionvertexai_batch_predict_anthropic_gemini_createjob]
+    import time
+    import vertexai
+
+    from vertexai.batch_prediction import BatchPredictionJob
+
+    # TODO(developer): Update and un-comment below line
+    # PROJECT_ID = "your-project-id"
+
+    # Initialize vertexai
+    vertexai.init(project=PROJECT_ID, location="us-east5")
+
+    input_uri = "gs://kellysun-test-project-europe-west1/input/claude_varied_input.jsonl"
+
+    # Submit a batch prediction job with Gemini model
+    batch_prediction_job = BatchPredictionJob.submit(
+        source_model="publishers/anthropic/models/claude-3-5-haiku",
+        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: gs://your-bucket/gen-ai-batch-prediction/prediction-model-year-month-day-hour:minute:second.12345
+
+    # [END generativeaionvertexai_batch_predict_anthropic_gemini_createjob]
+    return batch_prediction_job
+
+
+if __name__ == "__main__":
+    batch_predict_createjob(output_uri)

generative_ai/model_garden/noxfile_config.py

@@ -22,7 +22,7 @@
 
 TEST_CONFIG_OVERRIDE = {
     # You can opt out from the test for specific Python versions.
-    "ignored_versions": ["2.7", "3.7", "3.8", "3.10", "3.11", "3.13"],
+    "ignored_versions": ["2.7", "3.7", "3.8", "3.10", "3.11", "3.12"],
     # Old samples are opted out of enforcing Python type hints
     # All new samples should feature them
     "enforce_type_hints": True,

generative_ai/model_garden/test_model_garden_examples.py

@@ -16,6 +16,8 @@
 
 from google.api_core.exceptions import ResourceExhausted
 
+import claude_3_batch_prediciton_bq
+import claude_3_batch_prediction_gcs
 import claude_3_streaming_example
 import claude_3_tool_example
 import claude_3_unary_example
@@ -40,3 +42,15 @@ def test_tool_use() -> None:
 def test_generate_text() -> None:
     responses = claude_3_unary_example.generate_text()
     assert "bread" in responses.model_dump_json(indent=2)
+
+
+@backoff.on_exception(backoff.expo, ResourceExhausted, max_time=10)
+def test_generate_text_gcs() -> None:
+    responses = claude_3_batch_prediction_gcs.generate_text()
+    assert "bread" in responses.model_dump_json(indent=2)
+
+
+@backoff.on_exception(backoff.expo, ResourceExhausted, max_time=10)
+def test_generate_text_bq() -> None:
+    responses = claude_3_batch_prediciton_bq.generate_text()
+    assert "bread" in responses.model_dump_json(indent=2)