Vertex AI批量预测（Gemini Flash API）安全配置设置及报错问题

报错原因 #

BatchPredictionJob.submit()方法不支持model_parameters参数，Gemini模型的批量预测中，模型相关参数（包括安全设置）需要通过prediction_parameters参数传递。

正确配置方法 #

将原代码中的model_parameters替换为prediction_parameters，并确保安全设置的格式符合Vertex AI要求：

• 安全设置的每个条目需包含category（危害类别）、threshold（拦截阈值），无需指定method字段（Gemini模型默认使用SEVERITY方法）
• 支持的危害类别包括：HARM_CATEGORY_DANGEROUS_CONTENT、HARM_CATEGORY_HATE_SPEECH、HARM_CATEGORY_HARASSMENT、HARM_CATEGORY_SEXUALLY_EXPLICIT
• 阈值可选值：BLOCK_NONE、BLOCK_ONLY_HIGH、BLOCK_MEDIUM_AND_ABOVE、BLOCK_LOW_AND_ABOVE

修改后的完整代码 #

# File Format
# {"request":{"contents": [{"role": "user", "parts": [{"text": "Translate this sentence to French: Hello, how are you today?"}]}]}}

import os
import vertexai
from vertexai.preview.batch_prediction import BatchPredictionJob

PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")


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 BigQuery table or Google Cloud Storage.
            Example: "gs://[BUCKET]/[DATASET].jsonl" OR "bq://[PROJECT].[DATASET].[TABLE]"

        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]"
    Returns:
        batch_prediction_job: The batch prediction job object containing details of the job.
    """

    import time

    input_uri ="gs://ge/in/vertex_batch_input.jsonl"
    output_uri ="gs://ge/out/"

    # Initialize vertexai
    vertexai.init(project=PROJECT_ID, location="us-central1")

    # 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,
        prediction_parameters={
            "safety_settings": [
                {
                    "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
                    "threshold": "BLOCK_ONLY_HIGH"
                },
                {
                    "category": "HARM_CATEGORY_HATE_SPEECH",
                    "threshold": "BLOCK_ONLY_HIGH"
                },
                {
                    "category": "HARM_CATEGORY_HARASSMENT",
                    "threshold": "BLOCK_ONLY_HIGH"
                },
                {
                    "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
                    "threshold": "BLOCK_ONLY_HIGH"
                }
            ]
        }
    )
    
    # 可选：等待任务完成并打印状态
    print(f"Job started: {batch_prediction_job.resource_name}")
    while batch_prediction_job.state != "SUCCEEDED":
        time.sleep(60)
        batch_prediction_job.refresh()
        print(f"Current job state: {batch_prediction_job.state}")
    
    return batch_prediction_job

额外说明 #

• 确保使用的Vertex AI Python SDK版本为>=1.42.0，旧版本可能不支持通过prediction_parameters传递安全设置
• 若需验证安全设置是否生效，可在输入数据中包含对应类别的测试内容，检查输出是否被拦截

内容的提问来源于stack exchange，提问作者JohnDotOwl