gemini-evaluation-evaltask-approach-evaluate-rag-batch-pipeline

Evaluate generated answers from Retrieval-Augmented Generation (RAG) using Rapid Evaluation and Dataflow ML with Vertex AI pipelines

Source notebook

Repo path: gemini/evaluation/evaltask_approach/evaluate_rag_batch_pipeline.ipynb · Open on GitHub · advanced

Builds a Vertex AI Pipeline to batch-evaluate RAG Q&A outputs with Rapid Eval API and Dataflow ML.

Summary

This notebook teaches how to evaluate generated RAG question-answering responses at batch scale using Vertex AI Pipelines, Dataflow ML, and the Vertex AI Rapid Eval API. It prepares a JSONL evaluation dataset, validates required columns, runs an Apache Beam Dataflow job to call Rapid Eval metrics, aggregates score statistics, and retrieves row-level plus summary results from the completed pipeline.

Key code patterns

Initialize Vertex AI

from google.cloud import aiplatform
 
aiplatform.init(
    project=PROJECT_ID,
    location=REGION,
    staging_bucket=BUCKET_URI,
)

Configures the project, region, and staging bucket used by Vertex AI Pipelines.

Validate JSONL eval data

eval_df = pd.read_json(input_dataset_path, lines=True)
invalid_indices = validate_dataframe(eval_df, num_processes=max_processes)
valid_eval_df = eval_df.drop(index=invalid_indices)
valid_eval_df.to_json(valid_eval_dataset_file_path, orient="records", lines=True)

Ensures required RAG fields are present and typed before remote evaluation runs.

Run Dataflow from pipeline

dataflow_python_op = DataflowPythonJobOp(
    project=project_id,
    location=location,
    python_module_path=python_file_path,
    temp_location=temp_location,
    requirements_file_path=requirements_file_path,
    args=prepare_args_op.outputs["args"],
)

Submits the Apache Beam evaluation module as a Dataflow job inside Vertex AI Pipelines.

Compile and run pipeline

compiler.Compiler().compile(
    pipeline_func=pipeline,
    package_path=str(PIPELINE_PATH) + "/eval_pipeline.json",
)
pipeline_job = aiplatform.PipelineJob(
    display_name="evaluate_rag_batch_eval",
    template_path=str(PIPELINE_PATH / "eval_pipeline.json"),
    parameter_values=pipeline_params,
    pipeline_root=str(PIPELINE_ROOT_URI),
    enable_caching=False,
)
pipeline_job.run()

Turns the KFP DSL workflow into a Vertex AI Pipeline job and executes it.

Aggregate Rapid Eval metrics

score_columns = [col for col in eval_result_df.columns if col.endswith("_score")]
metrics[f"{col}_mean"] = round(float(mean), 3)
metrics[f"{col}_std"] = round(float(std_dev), 3)
output_eval_summary_metrics.log_metric(key, value)

Computes mean and standard deviation for Rapid Eval score columns and logs pipeline metrics.

Models & APIs used

• APIs / services: Vertex AI, Dataflow, Vertex AI Rapid Eval API, Vertex AI Pipelines, Cloud Storage
• SDKs / libraries: google-cloud-aiplatform, google-cloud-pipeline-components, kfp, apache-beam, pandas, plotly, etils

When to use this

Use this pattern when you need a repeatable Vertex AI Pipeline for batch evaluation of RAG Q&A predictions with Rapid Eval metrics.

Gotchas & caveats

• Vertex AI API must be enabled for the Google Cloud project.
• A Cloud Storage bucket is required for datasets, source files, pipeline artifacts, and temporary files.
• The service account needs Vertex AI User, Storage Object Admin, Dataflow Worker, and Dataflow Developer roles.
• Colab requires authentication and runtime restart after installing packages.
• The evaluation dataset must be JSONL with instruction, context, and prediction columns.
• The pipeline run is stated to require about 15 minutes.

Best practices

• Validate the evaluation dataset before starting remote evaluation.
• Store pipeline data, source modules, requirements, outputs, and temporary files in Cloud Storage paths.
• Package the Apache Beam module with requirements.txt and setup.py for Dataflow workers.
• Use WaitGcpResourcesOp after DataflowPythonJobOp before reading output artifacts.
• Retrieve both row-level metrics and aggregated summary metrics after pipeline completion.
• Provide cleanup flags for deleting the bucket, pipeline jobs, and local tutorial directory.

Related

• Concepts: Evaluation · RAG & Grounding · MLOps & Deployment
• Entities: Vertex AI · Vertex AI SDK · Cloud Storage · Grounding
• Area: Gemini Notebooks
• Best practices: Evaluation - Best Practices · RAG & Grounding - Best Practices · MLOps & Deployment - Best Practices