googleapis / python-aiplatform

A Python SDK for Vertex AI, a fully managed, end-to-end platform for data science and machine learning.
Apache License 2.0
589 stars 318 forks source link

Get batch predictions for Gemini #3871

Closed LindaLawton closed 1 month ago

LindaLawton commented 1 month ago

I have been chasing my tail all day trying to get Get batch predictions for Gemini to work.

I have been diging around in all the samples and documentation that i can find and have not been able to get this to work. From what i can tell the library supports this yet. I am faced with this error.

grpc.aio._call.AioRpcError: <AioRpcError of RPC that terminated with: status = StatusCode.UNIMPLEMENTED details = "Received http2 header with status: 404" debug_error_string = "UNKNOWN:Error received from peer {grpc_message:"Received http2 header with status: 404", grpc_status:12, created_time:"2024-06-05T17:00:46.9053534+00:00"}"

Any help would be greatly apricated

async def sample_create_batch_prediction_job():
    input_config = aiplatform_v1.types.BatchPredictionJob.InputConfig(
        bigquery_source=aiplatform_v1.types.BigQuerySource(
            input_uri="bq://sample.text_input"
        )
    )
    output_config = aiplatform_v1.types.BatchPredictionJob.OutputConfig(
        bigquery_destination=aiplatform_v1.types.BigQueryDestination(
            output_uri="bq://sample.llm_dataset.embedding_out_BP_sample_publisher_BQ_20230712_134650"
        )
    )
    # https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/batch-prediction-api
    batch_prediction_job_1 = BatchPredictionJob(
        name=os.getenv("PROJECT_ID"),
        display_name='my-batch-prediction-job',
        model=f"projects/{os.getenv("PROJECT_ID")}/locations/us-central1/models/gemini-1.0-pro-001",
        input_config=input_config,
        output_config=output_config,
        )

    # Create a client
    client = aiplatform_v1.JobServiceAsyncClient()

    request = aiplatform_v1.CreateBatchPredictionJobRequest(
        parent=f"projects/{os.getenv("PROJECT_ID")}/locations/us-central1",
        batch_prediction_job=batch_prediction_job_1,
    )

    # Make the request
    response = await client.create_batch_prediction_job(request=request)

    # Handle the response
    print(response)

I thought it might be the big query URLs i tried adding the table ID for my own tables but there is no change. The documentation doesn't actually state how or what to pass for these big query URLs. Nor what the data model should look like.

jaycee-li commented 1 month ago

Hi @LindaLawton , batch prediction for Gemini has just launch in preview in google-cloud-aiplatform==1.52.0 and our tech writer hasn't published the documentation yet. Here is a code sample:

import time
import vertexai
from vertexai.preview.batch_prediction import BatchPredictionJob

vertexai.init(project="your-project", location="your-location")

job = BatchPredictionJob.submit(
    "gemini-1.0-pro-001",   # source_model 
    "bq://projectId.datasetId.tableId",  # input_dataset
    output_uri_prefix = "bq://projectId.datasetId"  # Optional, output_uri_prefix
)

# Check job status
print(f"Job resouce name: {job.resource_name}")
print(f"Model resource name with the job: {job.model_name}")
print(f"Job state: {job.state.name}")

# Refresh the job until complete
while not job.has_ended:
  time.sleep(5)
  job.refresh()

# Check if the job succeeds
if job.has_succeeded:
  print("Job succeeded!")
else:
  print(f"Job failed: {job.error}")

# Check the location of the output
print(f"Job output location: {job.output_location}")

# List all the GenAI batch prediction jobs under the project
for bpj in BatchPredictionJob.list():
  print(f"Job ID: '{bpj.name}', Job state: {bpj.state.name}, Job model: {bpj.model_name}")

I took a look at your code and looks like you used the wrong model name. It should be publishers/google/models/gemini-1.0-pro-001. If you still prefer the gapic library, you can change the name and try again.

Thanks!

LindaLawton commented 1 month ago

You are my hero its working thank you:

BTW I was trying to do this with flash

The docs state

image

Yet I am getting

google.api_core.exceptions.InvalidArgument: 400 Do not support publisher model gemini-1.5-flash-001

1.5 pro doesn't appear to work either

google.api_core.exceptions.InvalidArgument: 400 Do not support publisher model gemini-1.5-pro-001

Will this change in the near future?

jaycee-li commented 1 month ago

@LindaLawton Glad to hear it's working.

gemini-1.5-flash-001 and gemini-1.5-pro-001 are currently in staging and should be publicly available by this week or early next week.

giancarlo-metitieri commented 1 month ago

@LindaLawton Glad to hear it's working.

gemini-1.5-flash-001 and gemini-1.5-pro-001 are currently in staging and should be publicly available by this week or early next week.

Is it possible that the bug that I posted may be related to this? https://issuetracker.google.com/issues/345948283

While the us-central1 throws the error that @LindaLawton posted, it does not seem to be the case for the europe-west6 region, which starts a job, then fails.

LindaLawton commented 1 month ago

Its working with Flash now so i will close this.