Skip to content

Commit

Permalink
feat(generative-ai): Update to gemini_batch_predict gcs specific (#12723
Browse files Browse the repository at this point in the history
)

* feat(generative-ai): Update to gemini_batch_predict gcs specific

* updating nox config to add back 3.11
  • Loading branch information
CadillacBurgess1 authored Oct 23, 2024
1 parent 2d8f95d commit ffbe8ed
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 7 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -23,11 +23,11 @@ def batch_predict_gemini_createjob(
) -> BatchPredictionJob:
"""Perform batch text prediction using a Gemini AI model.
Args:
input_uri (str): URI of the input file in BigQuery table or Google Cloud Storage.
Example: "gs://[BUCKET]/[DATASET].jsonl" OR "bq://[PROJECT].[DATASET].[TABLE]"
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 BigQuery table or Google Cloud Storage.
Example: "gs://[BUCKET]/[OUTPUT].jsonl" OR "bq://[PROJECT].[DATASET].[TABLE]"
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.
"""
Expand Down
6 changes: 3 additions & 3 deletions generative_ai/batch_predict/test_batch_predict_examples.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@

import batch_code_predict
import batch_text_predict
import gemini_batch_predict
import gemini_batch_predict_gcs


INPUT_BUCKET = "cloud-samples-data"
Expand Down Expand Up @@ -74,10 +74,10 @@ def test_batch_code_predict(output_folder: pytest.fixture()) -> None:
assert OUTPUT_PATH in job.output_info.gcs_output_directory


def test_batch_gemini_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"
job = _main_test(
test_func=lambda: gemini_batch_predict.batch_predict_gemini_createjob(
test_func=lambda: gemini_batch_predict_gcs.batch_predict_gemini_createjob(
input_uri, output_folder
)
)
Expand Down

0 comments on commit ffbe8ed

Please sign in to comment.