Gemini 批量生成

MODEL ZOO Nov 19, 2024

LLM 非常适合按需生成内容，但如果不加以控制，你最终可能会收到一大笔账单。在我的“使用上下文缓存控制 LLM 成本”帖子中，我谈到了如何使用上下文缓存来限制成本。批量生成是另一种可用于以折扣价节省时间的技术。

1、什么是批量生成？

Gemini 中的批量生成允许你批量发送多个生成 AI 请求，而不是逐个发送，并在 Cloud Storage 存储桶或 BigQuery 表中异步获取响应。这不仅简化了大型数据集的处理，而且还节省了时间和金钱，因为批量请求是并行处理的，并且比标准请求便宜 50%。

考虑一家拥有数千本书的在线书店。Gemini 批量生成可以并行生成所有描述，而不是逐个生成每本书的描述，这会很耗时。这以折扣价减少了总体处理时间。

让我们看看如何使用批量生成。

2、创建 Cloud Storage 存储桶

你可以使用 Cloud Storage 或 BigQuery 来准备和保存批处理作业结果。在这里，我们将使用 Cloud Storage。

首先，让我们创建一个存储桶来保存批处理请求输入文件：

PROJECT_ID=your-project-id
INPUT_BUCKET_URI=gs://$PROJECT_ID-batch-processing-input
gsutil mb $INPUT_BUCKET_URI

你还需要一个存储桶来保存批处理请求结果：

PROJECT_ID=your-project-id
OUTPUT_BUCKET_URI=gs://$PROJECT_ID-batch-processing-output
gsutil mb $OUTPUT_BUCKET_URI

3、准备批处理生成输入文件

接下来，你需要在 jsonl 文件中准备批处理预测输入文件。

例如，查看 batch_request_text_input.jsonl，其中包含用于生成不同蛋糕食谱的文本提示：

{
  "request": {
    "contents": [
      {
        "parts": {
          "text": "Give me a recipe for banana bread."
        },
        "role": "user"
      }
    ]
  }
}{
  "request": {
    "contents": [
      {
        "parts": {
          "text": "Give me a recipe for chocolate cake."
        },
        "role": "user"
      }
    ]
  }
}
...
{
  "request": {
    "contents": [
      {
        "parts": {
          "text": "Give me a recipe for pound cake."
        },
        "role": "user"
      }
    ]
  }
}

你还可以使用带有文本、图像和视频的多模式提示来批量生成内容，如 batch_request_multimodal_input.jsonl 中所示：

{
  "request": {
    "contents": [
      {
        "role": "user",
        "parts": [
          {
            "text": "List objects in this image."
          },
          {
            "file_data": {
              "file_uri": "gs://cloud-samples-data/generative-ai/image/office-desk.jpeg",
              "mime_type": "image/jpeg"
            }
          }
        ]
      }
    ]
  }
}{
  "request": {
    "contents": [
      {
        "role": "user",
        "parts": [
          {
            "text": "List objects in this image."
          },
          {
            "file_data": {
              "file_uri": "gs://cloud-samples-data/generative-ai/image/gardening-tools.jpeg",
              "mime_type": "image/jpeg"
            }
          }
        ]
      }
    ]
  }
}{
  "request": {
    "contents": [
      {
        "role": "user",
        "parts": [
          {
            "text": "What is the relation between the following video and image samples?"
          },
          {
            "fileData": {
              "fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4",
              "mimeType": "video/mp4"
            }
          },
          {
            "fileData": {
              "fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg",
              "mimeType": "image/jpeg"
            }
          }
        ]
      }
    ]
  }
}

将两个文件上传到输入存储桶：

gsutil cp batch_request_text_input.jsonl $INPUT_BUCKET_URI
gsutil cp batch_request_multimodal_input.jsonl $INPUT_BUCKET_URI

4、运行批量生成

要运行批量生成作业，你需要提交一个包含输入文件和输出的 BatchPredictionJob bucket：

vertexai.init(project=args.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=args.input_dataset_uri,
    output_uri_prefix=args.output_bucket_uri,
)

然后，你需要等待批量生成完成：

while not batch_prediction_job.has_ended:
    print(f"Job state: {batch_prediction_job.state.name}")
    time.sleep(10)
    batch_prediction_job.refresh()

运行批量生成文本提示：

python main.py --project_id $PROJECT_ID \
  --input_dataset_uri $INPUT_BUCKET_URI/batch_request_text_input.jsonl \
  --output_bucket_uri $OUTPUT_BUCKET_URI

运行时，你可以在 Cloud Console 上查看其状态：

运行多模式提示的批量生成：

python main.py --project_id $PROJECT_ID \
  --input_dataset_uri $INPUT_BUCKET_URI/batch_request_multimodal_input.jsonl \
  --output_bucket_uri $OUTPUT_BUCKET_URI

最后，你将看到两个批处理作业都已完成：