Generative AI Google Brain

embeddings-large-embs-generation-for-vvs

3 min read

Handling large-scale embedding generation for Agent Platform Vector Search

Source notebook

Repo path: embeddings/large-embs-generation-for-vvs.ipynb · Open on GitHub · intermediate

Generates large-scale text and image embeddings for Agent Platform Vector Search with throttling and checkpointing.

Summary

This notebook teaches how to generate text and multimodal embeddings for 10K, 100K, or 1M GBIF dataset items using the Google GenAI SDK against Vertex AI. It demonstrates quota-aware multithreaded embedding generation, queue-based error handling, periodic Cloud Storage checkpointing, and optional use of the JSONL output to create an Agent Platform Vector Search index.

Key code patterns

Initialize GenAI embedding client

from google import genai
 
embed_client = genai.Client(
    vertexai=True,
    project=PROJECT_ID,
    location=LOCATION,
)

Uses the Google GenAI SDK with Vertex AI project and region settings.

Text embeddings with task type

response = embed_client.models.embed_content(
    model="gemini-embedding-001",
    contents=names,
    config=EmbedContentConfig(
        task_type="SEMANTIC_SIMILARITY",
        output_dimensionality=768,
    ),
)

Shows how to batch text inputs and set task type and dimensionality.

Quota throttling loop

req_interval = 1.0 / (reqs_per_min_quota / 60)
for i in tqdm(range(0, len(items), items_per_req)):
    time.sleep(req_interval)
    items_slice = items[i : i + items_per_req]
    worker_thread = threading.Thread(
        target=run_worker_thread,
        args=(generation_func, items_slice, emb_queue, err_queue),
    )

Spaces requests by quota while running embedding calls concurrently.

Checkpoint embeddings to Cloud Storage

for _ in range(count):
    emb = emb_queue.get()
    embs += json.dumps(emb) + "\n"
gcs_file = storage_bucket.blob(f"{gcs_path}/{timestamp}_embs.json")
gcs_file.upload_from_string(embs, content_type="application/json")

Persists JSONL embedding batches so long jobs can retain progress.

Multimodal image embeddings

response = embed_client.models.embed_content(
    model="gemini-embedding-2-preview",
    contents=Part.from_uri(file_uri=gcsUri, mime_type="image/jpeg"),
    config=EmbedContentConfig(output_dimensionality=768),
)

Generates image embeddings from GCS URIs using the multimodal embedding model.

Models & APIs used

  • Models: gemini-embedding-001, gemini-embedding-2-preview
  • APIs / services: Vertex AI, Cloud Storage, Agent Platform Embeddings API, Agent Platform Vector Search
  • SDKs / libraries: google-genai, google-cloud-aiplatform, google.cloud.storage

When to use this

Use this pattern when generating large embedding datasets from a notebook for Agent Platform Vector Search development workflows.

Gotchas & caveats

  • Enable the Agent Platform API before running the notebook.
  • Colab requires explicit user authentication and may require manual PROJECT_ID entry.
  • The notebook installs packages and then restarts the runtime.
  • Embedding quotas vary by region and may change over time.
  • Exceeding request-per-minute quotas can cause ResourceExhausted errors.
  • Text embedding requests are limited to 250 input texts and 20000 tokens per request.
  • The notebook states this approach is for development, not production MLOps workflows.
  • gemini-embedding-2-preview has lower default quota and only takes one image per request in this notebook.

Best practices

  • Throttle API calls to stay within embedding quota limits.
  • Use multithreading to reduce latency impact and improve quota utilization.
  • Checkpoint generated embeddings periodically to Cloud Storage.
  • Record errors to a log file during long-running jobs.
  • Use moderate text batch sizes such as 20 instead of always using the maximum 250.
  • Convert long-running notebook jobs to a Python process with nohup.
  • Monitor embedding quota usage in the Google Cloud console.
  • Use Agent Platform Pipelines and Dataflow for production embedding generation workflows.

Graph View

Backlinks