Generative AI Google Brain

gemini-rag-engine-rag-engine-evaluation

3 min read

Advanced RAG Techniques - Vertex RAG Engine Retrieval Quality Evaluation and Hyperparameters Tuning

Source notebook

Repo path: gemini/rag-engine/rag_engine_evaluation.ipynb · Open on GitHub · advanced

Evaluates Vertex AI RAG Engine retrieval quality and tunes chunking, top-k, threshold, and embedding settings.

Summary

The notebook shows how to evaluate retrieval quality for a Vertex AI RAG Engine corpus using BEIR-FiQA or another BEIR dataset. It creates or reuses a RagCorpus, imports documents from Cloud Storage with configurable chunk size and overlap, runs retrieval queries, and computes recall@k, precision@k, and nDCG@k. It then explains how to tune retrieval hyperparameters and clean up the RagCorpus.

Key code patterns

Initialize Vertex AI

import vertexai
 
PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT")
LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")
vertexai.init(project=PROJECT_ID, location=LOCATION)

Sets the Google Cloud project and region before using Vertex AI RAG Engine.

Create RAG corpus

from vertexai.preview import rag
 
embedding_model_config = rag.EmbeddingModelConfig(
    publisher_model="publishers/google/models/text-embedding-005"
)
rag_corpus = rag.create_corpus(
    display_name="test-corpus",
    description="A test corpus where we import the BEIR-FiQA-2018 dataset",
    embedding_model_config=embedding_model_config,
)

Creates a RagCorpus with an explicit embedding model for retrieval evaluation.

Import files with chunking

paths = [CURRENT_BUCKET_PATH + f"/{i}/" for i in range(num_subdirectories)]
 
import_rag_files_from_gcs(
    paths=paths,
    chunk_size=512,
    chunk_overlap=102,
    corpus_name=rag_corpus.name,
)

Ingests Cloud Storage files into the corpus with tunable chunk size and overlap.

Retrieve contexts

contexts = rag.retrieval_query(
    rag_resources=[rag.RagResource(rag_corpus=corpus_name)],
    text=query,
    similarity_top_k=top_k,
    vector_distance_threshold=0.5,
).contexts.contexts

Runs RAG Engine retrieval with configurable top-k and vector distance threshold.

Evaluate retrieval metrics

calculate_document_level_metrics(
    queries,
    qrels,
    [5, 10, 100],
    corpus_name=rag_corpus.name,
)

Computes retrieval quality metrics across multiple k values.

Models & APIs used

  • Models: publishers/google/models/text-embedding-005
  • APIs / services: Vertex AI, Cloud Storage
  • SDKs / libraries: google-cloud-aiplatform, beir, google-cloud-storage, vertexai

When to use this

Use this pattern when tuning a Vertex AI RAG Engine corpus and measuring retrieval quality before evaluating generated answers.

Gotchas & caveats

  • Requires an existing Google Cloud project with the Vertex AI API enabled.
  • Colab requires authenticate_user before using Google Cloud resources.
  • The notebook installs packages and then restarts the runtime.
  • Application Default Credentials and quota project are configured with gcloud commands.
  • Option A uses only the public BEIR-FiQA dataset; other BEIR datasets require Option B upload and ingestion.
  • Do not run all corpus setup options together.
  • Import uses max_embedding_requests_per_min=1400, implying embedding request rate limits matter.
  • The notebook uses vertexai.preview.rag, indicating preview API surface.

Best practices

  • Evaluate retrieval quality because poor retrieval can lead to irrelevant, incomplete, or hallucinated output.
  • Tune chunk size, chunk overlap, top-k, vector distance threshold, and embedding model based on recall, precision, and nDCG.
  • Use recall@k, precision@k, and nDCG@k to evaluate retrieval from different perspectives.
  • Reduce chunk size, increase chunk overlap, or increase top-k when recall is too low.
  • Reduce top-k, reduce chunk overlap, or increase chunk size when precision is too low.
  • Consider changing the embedding model when nDCG is low.
  • Clean up the RagCorpus after evaluation to free up resources.

Graph View

Backlinks