Generative AI Google Brain

search-custom-embeddings-custom-embeddings

3 min read

Custom Embeddings with Vertex AI Search

Source notebook

Repo path: search/custom-embeddings/custom_embeddings.ipynb · Open on GitHub · intermediate

Builds a Vertex AI Search app using custom text-embedding-005 embeddings from Stack Overflow data.

Summary

The notebook teaches how to generate document embeddings with Vertex AI, reshape them into Vertex AI Search JSONL format, and import them into a custom search data store. It fetches Stack Overflow questions from BigQuery, embeds title and body text, uploads scraped HTML and JSONL metadata to Cloud Storage, configures Discovery Engine schema and ranking, then tests search summaries and snippets.

Key code patterns

Initialize Vertex AI

PROJECT_ID = "[your-project-id]"
LOCATION = "us-central1"
vertexai.init(project=PROJECT_ID, location=LOCATION)

Sets the project and region used by the Vertex AI embedding model.

Load embedding model

model = TextEmbeddingModel.from_pretrained("text-embedding-005")

Uses the exact Vertex AI text embedding model required for document vectors.

Batch document embeddings

def get_embeddings_wrapper(texts, batch_size=5):
    embs = []
    for i in tqdm(range(0, len(texts), batch_size)):
        result = model.get_embeddings([
            TextEmbeddingInput(text=text, task_type="RETRIEVAL_DOCUMENT")
            for text in texts[i:i + batch_size]
        ])
        embs.extend([e.values for e in result])
    return embs

Creates retrieval-optimized embeddings in small batches.

Vertex AI Search JSONL row

def format_row(row):
    return {
        "id": row["id"],
        "content": {"mimeType": "text/html", "uri": row["gcs_uri"]},
        "structData": {
            "embedding_vector": row["embedding"],
            "title": row["title"],
            "question_url": row["question_url"]
        },
    }

Maps document content and embedding vectors into the unstructured metadata format expected by Vertex AI Search.

requests.patch(
    url="https://discoveryengine.googleapis.com/v1alpha/.../servingConfigs/default_search",
    headers={"Authorization": f"Bearer {access_token}"},
    json={
        "embeddingConfig": {"fieldPath": "embedding_vector"},
        "ranking_expression": "0.5 * relevance_score",
    },
)

Sets the embedding field and ranking expression through REST because the notebook states this is not supported in client libraries.

Models & APIs used

  • Models: text-embedding-005
  • APIs / services: Vertex AI, BigQuery, Cloud Storage, Discovery Engine, Vertex AI Search
  • SDKs / libraries: google-cloud-aiplatform, google-cloud-discoveryengine, google-cloud-storage, google-cloud-bigquery, vertexai, pandas, requests, tqdm

When to use this

Use this pattern when you already have custom document embeddings and need Vertex AI Search to index and rank them with document content stored in Cloud Storage.

Gotchas & caveats

  • Colab requires google.colab auth.authenticate_user; Vertex AI Workbench does not.
  • The notebook tells users to restart the runtime after installing libraries.
  • The Stack Overflow BigQuery table has 23 million rows, so the tutorial limits the query to 500 rows.
  • Custom embedding schema uses dimension 768 for embedding_vector.
  • Embedding configuration and ranking expression are set with a REST PATCH because the notebook says client libraries do not support it.
  • The notebook shells out to gcloud auth print-access-token for the REST request.
  • DATA_STORE_LOCATION is global, while Vertex AI initialization uses us-central1.
  • Scraping Stack Overflow pages depends on requests returning status code 200 and non-empty content.

Best practices

  • Use TextEmbeddingInput with task_type=“RETRIEVAL_DOCUMENT” for document retrieval embeddings.
  • Batch embedding calls instead of embedding all texts at once.
  • Convert BigQuery integer IDs to strings before using them as document IDs.
  • Store HTML content in Cloud Storage and reference it from JSONL content.uri.
  • Define the embedding vector in the Discovery Engine schema before importing documents.
  • Use FULL reconciliation mode when importing the JSONL file for the tutorial data set.
  • Enable snippets, summaries with citations, query expansion, and spell correction when testing search.

Graph View

Backlinks