Generative AI Google Brain

embeddings-intro-embeddings-tuning

3 min read

Get started with embeddings tuning on Agent Platform

Source notebook

Repo path: embeddings/intro_embeddings_tuning.ipynb · Open on GitHub · advanced

Tunes text-embedding-005 for retrieval using synthetic Gemini queries and Document AI PDF chunks.

Summary

This notebook teaches how to prepare a supervised text embedding tuning dataset from a PDF using Document AI OCR, LangChain chunking, and Gemini-generated queries. It runs an Agent Platform pipeline to tune text-embedding-005, evaluates NDCG metrics, deploys the tuned model to an endpoint, and retrieves similar corpus chunks with dot-product similarity.

Key code patterns

Create OCR processor

client_options = ClientOptions(api_endpoint=f"{location}-documentai.googleapis.com")
client = documentai.DocumentProcessorServiceClient(client_options=client_options)
processor = client.create_processor(
    parent=client.common_location_path(project_id, location),
    processor=documentai.Processor(display_name=PROCESSOR_ID, type_="OCR_PROCESSOR"),
)

Creates the Document AI processor used to extract PDF text before chunking.

Generate synthetic queries

model = GenerativeModel("gemini-2.0-flash")
response = model.generate_content(
    [prompt_template.format(chunk=chunk.page_content)],
    generation_config=GenerationConfig(max_output_tokens=2048, temperature=0.9, top_p=1),
    safety_settings=safety_settings,
).text
query = Document(page_content=response, metadata={"page": chunk.metadata["page"]})

Uses Gemini to create query-document pairs for supervised embedding tuning.

Build tuning files

corpus_df.to_json(corpus_path, orient="records", lines=True)
query_df.to_json(query_path, orient="records", lines=True)
train_df.to_csv(train_path, sep="\t", header=True, index=False)
test_df.to_csv(test_path, sep="\t", header=True, index=False)

Writes the required corpus, query, train labels, and test labels files to Cloud Storage paths.

Run tuning pipeline

params = {
    "base_model_version_id": "text-embedding-005",
    "queries_path": query_path,
    "corpus_path": corpus_path,
    "train_label_path": train_path,
    "test_label_path": test_path,
}
job = aiplatform.PipelineJob(template_path=template_uri, parameter_values=params)
job.run()

Submits the Agent Platform text embedding tuning pipeline with dataset paths and base model.

Deploy tuned model

endpoint = aiplatform.Endpoint.create(display_name="tuned_custom_embedding_endpoint")
model = get_uploaded_model(job)
endpoint.deploy(model, accelerator_type="NVIDIA_TESLA_A100", accelerator_count=1, machine_type="a2-highgpu-1g")

Creates an endpoint and deploys the uploaded tuned embedding model for prediction.

Retrieve top-k chunks

response = endpoint.predict(instances=instances)
query_embedding = np.asarray(response.predictions)
similarity = corpus_embeddings.dot(query_embedding.T)
topk_index = pd.DataFrame({c: v.nlargest(n=k).index for c, v in similarity.items()})

Embeds queries with the tuned endpoint and ranks corpus embeddings by dot-product similarity.

Models & APIs used

  • Models: gemini-2.0-flash, text-embedding-005
  • APIs / services: Document AI, Agent Platform, Vertex AI, Cloud Storage, Gemini API
  • SDKs / libraries: google-cloud-aiplatform, google-cloud-documentai, google-cloud-documentai-toolbox, vertexai, langchain, langchain-core, langchain-text-splitters, langchain-google-community, gcsfs, etils

When to use this

Use this pattern when domain-specific retrieval quality needs improvement through supervised tuning of text embeddings.

Gotchas & caveats

  • Requires billing and enabling aiplatform.googleapis.com and documentai.googleapis.com.
  • Requires a Cloud Storage bucket for datasets, OCR output, pipeline root, and training artifacts.
  • Document AI processing and Gemini query generation can take minutes depending on chunk count.
  • Resource names use a timestamp and random suffix to avoid collisions in shared projects.
  • Training and prediction machine types and accelerators are explicitly configured and may affect cost and availability.
  • The notebook sets safety thresholds to BLOCK_NONE for synthetic query generation.

Best practices

  • Prepare corpus, query, and labels files in the format required by embedding tuning.
  • Split labels into train and test sets before running the tuning pipeline.
  • Use a timestamped Cloud Storage path for tuning artifacts.
  • Evaluate tuned embeddings with pipeline-produced NDCG metrics.
  • Create cleanup flags for endpoint, model, job, bucket objects, and local tutorial files.

Graph View

Backlinks