Generative AI Google Brain

embeddings-intro-textemb-vectorsearch

4 min read

Getting Started with Text Embeddings + Agent Platform Vector Search

Source notebook

Repo path: embeddings/intro-textemb-vectorsearch.ipynb · Open on GitHub · intermediate

Builds text embeddings from Stack Overflow titles and serves semantic search with Vector Search.

Summary

This notebook teaches embeddings, dot-product similarity, and approximate nearest neighbor search using Google Cloud services. It loads Stack Overflow questions from BigQuery, embeds titles with gemini-embedding-001 through the Google Gen AI SDK, writes embeddings to Cloud Storage as JSONL, builds and deploys a Vector Search index, then queries nearest neighbors. It also explains cleanup, quota pacing, endpoint security, and production tradeoffs versus RAG and hybrid search.

Key code patterns

Initialize Gen AI client

client = genai.Client(
    vertexai=True,
    project=PROJECT_ID,
    location=LOCATION,
)

Uses the Google Gen AI SDK against Vertex AI-hosted Agent Platform models.

Batch embedding wrapper

def get_embeddings_wrapper(texts: list[str]) -> list[list[float]]:
    embeddings = []
    for i in tqdm.tqdm(range(0, len(texts), BATCH_SIZE)):
        time.sleep(1)
        response = client.models.embed_content(
            model=TEXT_EMBEDDING_MODEL_ID,
            contents=texts[i : i + BATCH_SIZE],
            config=genai.types.EmbedContentConfig(output_dimensionality=768),
        )
        embeddings += [e.values for e in response.embeddings]
    return embeddings

Batches up to 5 texts and sleeps between calls to avoid request quota errors.

Local dot-product similarity

embs = np.array(df.embedding.to_list())
similarities = np.dot(embs[key], embs.T)
sorted_questions = sorted(
    zip(df.title, similarities), key=lambda x: x[1], reverse=True
)[:20]

Demonstrates the distance metric used for this embedding model before using Vector Search.

Export embeddings to JSONL

jsonl_string = df[["id", "embedding"]].to_json(
    orient="records", lines=True
)
with open("questions.json", "w") as f:
    f.write(jsonl_string)
! gsutil mb -l $LOCATION -p {PROJECT_ID} {BUCKET_URI}
! gsutil cp questions.json {BUCKET_URI}

Prepares the id and embedding columns in the JSONL format required by Vector Search.

Create Vector Search index

my_index = aiplatform.MatchingEngineIndex.create_tree_ah_index(
    display_name=f"embvs-tutorial-index-{UID}",
    contents_delta_uri=BUCKET_URI,
    dimensions=768,
    approximate_neighbors_count=20,
    distance_measure_type="DOT_PRODUCT_DISTANCE",
)

Builds a Tree-AH index using the embedding dimensionality and dot-product distance.

Deploy and query index

my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint.create(
    display_name=f"embvs-tutorial-index-endpoint-{UID}",
    public_endpoint_enabled=True,
)
my_index_endpoint.deploy_index(index=my_index, deployed_index_id=DEPLOYED_INDEX_ID)
response = my_index_endpoint.find_neighbors(
    deployed_index_id=DEPLOYED_INDEX_ID,
    queries=test_embeddings,
    num_neighbors=20,
)

Creates a secured public endpoint, deploys the index, and retrieves nearest neighbors.

Models & APIs used

  • Models: gemini-embedding-001
  • APIs / services: Agent Platform Embeddings for Text, Agent Platform Vector Search, BigQuery, Cloud Storage, Compute Engine, Agent Platform Workbench
  • SDKs / libraries: google-genai, google-cloud-aiplatform, google-cloud-storage, google-cloud-bigquery, numpy, tqdm

When to use this

Use this pattern when you need low-latency semantic search over business text using managed embeddings and Vector Search.

Gotchas & caveats

  • A Google Cloud project linked to billing is required outside Qwiklab.
  • The notebook enables compute.googleapis.com, aiplatform.googleapis.com, storage.googleapis.com, and bigquery.googleapis.com.
  • The default compute service account needs Agent Platform User, BigQuery User, and Storage Admin roles.
  • Colab requires explicit user authentication; Agent Platform Workbench is pre-authenticated.
  • The embeddings API accepts up to 5 texts per call in this notebook and may hit request-per-minute quotas without pacing.
  • The notebook uses 768-dimensional embeddings and DOT_PRODUCT_DISTANCE.
  • Index creation can take minutes or about 50 minutes or more depending on dataset size.
  • First index deployment can take around 25 minutes, and Colab runtimes may disconnect during long waits.
  • Indexes, index endpoints, Cloud Storage buckets, and Workbench instances can incur costs if not deleted.

Best practices

  • Limit tutorial data size from the 23 million row Stack Overflow table before loading into memory.
  • Batch embedding requests and add delay to avoid quota errors.
  • Use dot product for similarity with the Google embedding model shown.
  • Store Vector Search input as JSONL with id and embedding fields.
  • Match index dimensions to the embedding output dimensionality.
  • Use public endpoints unless VPC access is specifically required; access remains protected by IAM.
  • Reuse existing indexes and endpoints after runtime loss instead of recreating them.
  • Delete index endpoints, indexes, and Cloud Storage buckets after the tutorial to avoid unexpected costs.

Graph View

Backlinks