RAG Embeddings recipe (#184)

* updating readme and adding rag embeddings recipe

* updating readme

* resolving gemini suggestions

* resolving gemini typo suggestions

Author: jaivaldesai

genai-on-vertex-ai/gemini/model_upgrades/README.md

@@ -20,4 +20,4 @@ The goal is to accelerate the process of upgrading to the latest version of Gemi
 | Multi-turn Chat | [view](./multiturn_chat/vertex_colab/multiturn_chat_eval.ipynb) | [view](./multiturn_chat/vertex_script/README.md) | [view](./multiturn_chat/promptfoo/README.md) |
 | Instruction Following | [view](./instruction_following/vertex_colab/instruction_following_eval.ipynb) | [view](./instruction_following/vertex_script/README.md) | [view](./instruction_following/promptfoo/README.md) |
 | Image-Prompt Alignment | [view](./image_prompt_alignment/vertex_colab/image_prompt_alignment_eval.ipynb) | [view](./image_prompt_alignment/vertex_script/README.md) | [view](./image_prompt_alignment/promptfoo/README.md) |
-
+| RAG Embeddings | [view](./rag_embeddings/vertex_colab/rag_embeddings_eval.ipynb) | [view](./rag_embeddings/vertex_script/README.md) |  |

genai-on-vertex-ai/gemini/model_upgrades/rag_embeddings/vertex_colab/baseline_dataset.jsonl

@@ -0,0 +1,26 @@
+You are a professional RAG evaluator. 
+
+You will be assessing context retrieval quality. It measures how accurately does the retrieved context cover all the key information present in the reference answer.
+
+You will assign the writing response a score on a scale of 1-5, following the INDIVIDUAL RATING RUBRIC and EVALUATION STEPS.
+
+CRITERIA DEFINITIONS:
+Context Retrieval quality: The retrieved context should contain all the key information that is present in the reference. The context can contain additional information that is not found in the reference, but it should not miss any key information found in the reference.
+
+INDIVIDUAL RATING RUBRIC:
+1 : The retrieved context does not contain any of the key information from the reference.
+2 : The retrieved context covers just a little of the key information from the reference.
+3 : The retrieved context covers about half of the key information from the reference.
+4 : The retrieved context covers majority of the key information from the reference.
+5 : The retrieved context covers all of the key information from the reference.
+
+EVALUATION STEPS:
+STEP 1: Assess the context retrieval quality based on the criteria.
+STEP 2: Score based on the rubrics.
+
+Give step by step explanations for your scoring, and only choose scores from the individual rating rubric above.
+
+
+CONTEXT: {retrieved_context}
+
+REFERENCE: {reference}

genai-on-vertex-ai/gemini/model_upgrades/rag_embeddings/vertex_colab/candidate_dataset.jsonl

@@ -0,0 +1,250 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Copyright 2025 Google LLC\n",
+    "#\n",
+    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+    "# you may not use this file except in compliance with the License.\n",
+    "# You may obtain a copy of the License at\n",
+    "#\n",
+    "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+    "#\n",
+    "# Unless required by applicable law or agreed to in writing, software\n",
+    "\n",
+    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+    "# See the License for the specific language governing permissions and\n",
+    "# limitations under the License."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## **RAG Embeddings Retrieval Eval Recipe**\n",
+    "\n",
+    "This Eval Recipe demonstrates how to compare performance of two embedding models on a RAG dataset using [Vertex AI Evaluation Service](https://cloud.google.com/vertex-ai/generative-ai/docs/models/evaluation-overview).\n",
+    "\n",
+    "We will be looking at `text-embedding-004` as our baseline model and `text-embedding-005` as our candidate model. Please follow the [documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/get-text-embeddings) here to get an understanding of the various text embedding models. "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "<table align=\"left\">\n",
+    "  <td style=\"text-align: center\">\n",
+    "    <a href=\"https://art-analytics.appspot.com/r.html?uaid=G-FHXEFWTT4E&utm_source=aRT-rag_retrieval&utm_medium=aRT-clicks&utm_campaign=rag_retrieval&destination=rag_retrieval&url=https%3A%2F%2Fcolab.research.google.com%2Fgithub%2FGoogleCloudPlatform%2Fapplied-ai-engineering-samples%2Fblob%2Fmain%2Fgenai-on-vertex-ai%2Fgemini%2Fmodel_upgrades%2Frag_embeddings%2Fvertex_colab%2Frag_embeddings_eval.ipynb\">\n",
+    "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
+    "    </a>\n",
+    "  </td>\n",
+    "  <td style=\"text-align: center\">\n",
+    "    <a href=\"https://art-analytics.appspot.com/r.html?uaid=G-FHXEFWTT4E&utm_source=aRT-rag_retrieval&utm_medium=aRT-clicks&utm_campaign=rag_retrieval&destination=rag_retrieval&url=https%3A%2F%2Fconsole.cloud.google.com%2Fvertex-ai%2Fcolab%2Fimport%2Fhttps%3A%252F%252Fraw.githubusercontent.com%252FGoogleCloudPlatform%252Fapplied-ai-engineering-samples%252Fmain%252Fgenai-on-vertex-ai%252Fgemini%252Fmodel_upgrades%252Frag_embeddings%252Fvertex_colab%252Frag_embeddings_eval.ipynb\">\n",
+    "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
+    "    </a>\n",
+    "  </td>\n",
+    "  <td style=\"text-align: center\">\n",
+    "    <a href=\"https://art-analytics.appspot.com/r.html?uaid=G-FHXEFWTT4E&utm_source=aRT-rag_retrieval&utm_medium=aRT-clicks&utm_campaign=rag_retrieval&destination=rag_retrieval&url=https%3A%2F%2Fconsole.cloud.google.com%2Fvertex-ai%2Fworkbench%2Fdeploy-notebook%3Fdownload_url%3Dhttps%3A%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fapplied-ai-engineering-samples%2Fmain%2Fgenai-on-vertex-ai%2Fgemini%2Fmodel_upgrades%2Frag_embeddings%2Fvertex_colab%2Frag_embeddings_eval.ipynb\">\n",
+    "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
+    "    </a>\n",
+    "  </td>\n",
+    "  <td style=\"text-align: center\">\n",
+    "    <a href=\"https://github.com/GoogleCloudPlatform/applied-ai-engineering-samples/blob/main/genai-on-vertex-ai/gemini/model_upgrades/rag_embeddings/vertex_colab/rag_embeddings_eval.ipynb\">\n",
+    "      <img width=\"32px\" src=\"https://upload.wikimedia.org/wikipedia/commons/9/91/Octicons-mark-github.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
+    "    </a>\n",
+    "  </td>\n",
+    "</table>"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "- Use case: RAG retrieval\n",
+    "\n",
+    "- Metric: This eval uses a Pointwise Retrieval quality template to evaluate the responses and pick an embedding model as the winner. We will define `retrieval quality` as the metric here. It checks whether the `retrieved_context` contains all the key information present in `reference`.\n",
+    "\n",
+    "- Evaluation Datasets are based on [RAG Dataset](https://www.kaggle.com/datasets/samuelmatsuoharris/single-topic-rag-evaluation-dataset) in compliance with the following [license](https://www.mit.edu/~amini/LICENSE.md). They include 8 randomly sampled prompts in JSONL files `baseline_dataset.jsonl` and `candidate_dataset.jsonl` with the following structure:\n",
+    "    - `question`: User inputted question \n",
+    "    - `reference`: The golden truth answer for the question\n",
+    "    - `retrieved_context`: The context retrieved from the model \n",
+    "\n",
+    "\n",
+    "- Prompt Template is a zero-shot prompt located in [`prompt_template.txt`](./prompt_template.txt) with two prompt variables ( `reference` and `retrieved_context`) that are automatically populated from our dataset.\n",
+    "\n",
+    "- This eval recipe uses an LLM judge model(gemini-2.0-flash) to evaluate the retrieval quality of the embedding models. "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## **Prerequisite**\n",
+    "\n",
+    "This recipe assumes that the user has already created datasets for the baseline embedding model and the candidate embedding model. The user needs to generate the datasets for the baseline(text-embedding-004) and candidate(text-embedding-005) embedding models. Please refer to [RAG Engine generation notebook](https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_eval_service_sdk.ipynb) to create two separate RAG engines and set up corresponding datasets. The `retrieved_context` column in the dataset is the context retrieved from the respective RAG engine for each one of the questions."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Configure Eval Settings"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%writefile .env\n",
+    "PROJECT_ID=your-project-id        # Google Cloud Project ID\n",
+    "LOCATION=us-central1                  # Region for all required Google Cloud services\n",
+    "EXPERIMENT_NAME=rag-embeddings-eval-recipe-demo      # Creates Vertex AI Experiment to track the eval runs\n",
+    "BASELINE_EMBEDDING_MODEL=text-embedding-004\n",
+    "CANDIDATE_EMBEDDING_MODEL=text-embedding-005\n",
+    "MODEL=gemini-2.0-flash # This model will be the judge for performing evaluations\n",
+    "BASELINE_DATASET_URI=\"gs://gemini_assets/rag_embeddings/baseline_dataset.jsonl\"  # Baseline embedding model dataset in Google Cloud Storage\n",
+    "CANDIDATE_DATASET_URI=\"gs://gemini_assets/rag_embeddings/candidate_dataset.jsonl\"  # Candidate embedding model dataset in Google Cloud Storage\n",
+    "PROMPT_TEMPLATE_URI=\"gs://gemini_assets/rag_embeddings/prompt_template.txt\"  # Text file in Google Cloud Storage\n",
+    "METRIC_NAME=\"retrieval_quality\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Install Python Libraries"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install --upgrade --quiet google-cloud-aiplatform[evaluation] python-dotenv\n",
+    "# The error \"session crashed\" is expected. Please ignore it and proceed to the next cell.\n",
+    "import IPython\n",
+    "IPython.Application.instance().kernel.do_shutdown(True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import json\n",
+    "import pandas as pd\n",
+    "import sys\n",
+    "import vertexai\n",
+    "from dotenv import load_dotenv\n",
+    "from google.cloud import storage\n",
+    "\n",
+    "from datetime import datetime\n",
+    "from IPython.display import clear_output\n",
+    "from vertexai.evaluation import EvalTask, EvalResult, PointwiseMetric"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Authenticate to Google Cloud (requires permission to open a popup window)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "load_dotenv(override=True)\n",
+    "if os.getenv(\"PROJECT_ID\") == \"your-project-id\":\n",
+    "    raise ValueError(\"Please configure your Google Cloud Project ID in the first cell.\")\n",
+    "if \"google.colab\" in sys.modules:  \n",
+    "    from google.colab import auth  \n",
+    "    auth.authenticate_user()\n",
+    "vertexai.init(project=os.getenv('PROJECT_ID'), location=os.getenv('LOCATION'))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Run the eval on both models on the Pairwise Autorater"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def load_file(gcs_uri: str) -> str:\n",
+    "    blob = storage.Blob.from_string(gcs_uri, storage.Client())\n",
+    "    return blob.download_as_string().decode('utf-8')\n",
+    "\n",
+    "def load_dataset(dataset_uri: str):\n",
+    "    jsonl = load_file(dataset_uri)\n",
+    "    samples = [json.loads(line) for line in jsonl.splitlines() if line.strip()]\n",
+    "    df = pd.DataFrame(samples)\n",
+    "    return df\n",
+    "\n",
+    "def load_prompt_template() -> str:\n",
+    "    blob = storage.Blob.from_string(os.getenv(\"PROMPT_TEMPLATE_URI\"), storage.Client())\n",
+    "    return blob.download_as_string().decode('utf-8')\n",
+    "\n",
+    "def run_eval(model: str, embedding_model: str, dataset_uri: str) -> EvalResult:\n",
+    "  timestamp = f\"{datetime.now().strftime('%b-%d-%H-%M-%S')}\".lower()\n",
+    "  return EvalTask(\n",
+    "      dataset=dataset_uri,\n",
+    "      metrics=[PointwiseMetric(\n",
+    "               metric=os.getenv('METRIC_NAME'),\n",
+    "               metric_prompt_template= load_prompt_template()\n",
+    "               )   \n",
+    "               ],\n",
+    "      experiment=os.getenv('EXPERIMENT_NAME')\n",
+    "  ).evaluate(\n",
+    "      response_column_name= 'retrieved_context',\n",
+    "      experiment_run_name=f\"{timestamp}-{embedding_model}-{model.replace('.', '-')}\"\n",
+    "  )\n",
+    "\n",
+    "\n",
+    "baseline_metrics = run_eval(os.getenv(\"MODEL\"), os.getenv(\"BASELINE_EMBEDDING_MODEL\"), os.getenv(\"BASELINE_DATASET_URI\"))\n",
+    "candidate_metrics = run_eval(os.getenv(\"MODEL\"), os.getenv(\"CANDIDATE_EMBEDDING_MODEL\"), os.getenv(\"CANDIDATE_DATASET_URI\"))\n",
+    "clear_output()\n",
+    "print(\"Average score for baseline model retrieval quality:\", round(baseline_metrics.summary_metrics[f'{os.getenv(\"METRIC_NAME\")}/mean'],3))\n",
+    "print(\"Average score for candidate model retrieval quality:\", round(candidate_metrics.summary_metrics[f'{os.getenv(\"METRIC_NAME\")}/mean'],3))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

genai-on-vertex-ai/gemini/model_upgrades/rag_embeddings/vertex_colab/prompt_template.txt

@@ -0,0 +1,72 @@
+# RAG Retrieval
+
+### _Eval Recipe for model migration_
+
+This Eval Recipe demonstrates how to compare performance of two embedding models on a RAG dataset using [Vertex AI Evaluation Service](https://cloud.google.com/vertex-ai/generative-ai/docs/models/evaluation-overview).
+
+We will be looking at `text-embedding-004` as our baseline model and `text-embedding-005` as our candidate model. Please follow the [documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/get-text-embeddings) here to get an understanding of the various text embedding models. 
+
+- Use case: RAG retrieval
+
+- Metric: This eval uses a Pointwise Retrieval quality template to evaluate the responses and pick a model as the winner. We will define `retrieval quality` as the metric here. It checks whether the `retrieved_context` contains all the key information present in `reference`.
+
+- Evaluation Datasets are based on [RAG Dataset](https://www.kaggle.com/datasets/samuelmatsuoharris/single-topic-rag-evaluation-dataset) in compliance with the following [license](https://www.mit.edu/~amini/LICENSE.md). They include 8 randomly sampled prompts in JSONL files `baseline_dataset.jsonl` and `candidate_dataset.jsonl` with the following structure:
+    - `question`: User inputted question 
+    - `reference`: The golden truth answer for the question
+    - `retrieved_context`: The context retrieved from the model.
+
+- Prompt Template is a zero-shot prompt located in [`prompt_template.txt`](./prompt_template.txt) with two prompt variables ( `reference` and `retrieved_context`) that are automatically populated from our dataset.
+
+- This eval recipe uses an LLM judge model(gemini-2.0-flash) to evaluate the retrieval quality of the embedding models. 
+
+### _Prerequisite_
+
+This recipe assumes that the user has already created datasets for the baseline embedding model and the candidate embedding model. The user needs to generate the datasets for the baseline(text-embedding-004) and candidate(text-embedding-005) embedding models. Please refer to [RAG Engine generation notebook](https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_eval_service_sdk.ipynb) to create two separate RAG engines and set up corresponding datasets. The `retrieved_context` column in the dataset is the context retrieved from the respective RAG engine for each one of the questions.
+
+- Python script [`eval.py`](./eval.py) configures the evaluation:
+    - `run_eval`: configures the evaluation task, runs it on the 2 models and prints the results.
+    - `load_dataset`: loads the dataset including the contents of all documents.
+
+- Shell script [`run.sh`](./run.sh) installs the required Python libraries and runs `eval.py` 
+
+- Google Cloud Shell is the easiest option as it automatically clones our Github repo:
+
+    <a href="https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/GoogleCloudPlatform/applied-ai-engineering-samples&cloudshell_git_branch=main&cloudshell_workspace=genai-on-vertex-ai/gemini/model_upgrades">
+        <img alt="Open in Cloud Shell" src="http://gstatic.com/cloudssh/images/open-btn.png">
+    </a>
+
+- Alternatively, you can use the following command to clone this repo to any Linux environment with configured [Google Cloud Environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment):
+
+    ``` bash
+    git clone --filter=blob:none --sparse https://github.com/GoogleCloudPlatform/applied-ai-engineering-samples.git && \
+    cd applied-ai-engineering-samples && \
+    git sparse-checkout init && \
+    git sparse-checkout set genai-on-vertex-ai/gemini/model_upgrades && \
+    git pull origin main
+    cd genai-on-vertex-ai/gemini/model_upgrades
+    ```
+
+1. Navigate to the Eval Recipe directory in terminal, set your Google Cloud Project ID and run the shell script `run.sh`.
+
+    ``` bash
+    cd rag_embeddings/vertex_script
+    export PROJECT_ID="[your-project-id]"
+    ./run.sh
+    ```
+
+1. The resulting metrics will be displayed in the script output. 
+
+1. You can use [Vertex AI Experiments](https://console.cloud.google.com/vertex-ai/experiments) to view the history of evaluations for each experiment, including the final metrics scores.
+
+## How to customize this Eval Recipe:
+
+We will have two runs, one for the baseline model and the candidate model
+
+1. Edit the Python script `eval.py`:
+    - set the `project` parameter of vertexai.init to your Google Cloud Project ID.
+    - set the parameter `model`  in the run_eval calls (e.g., 'gemini-2.0-flash') to the LLM you want to use for performing the evaluation task.
+    - set the parameter `embedding_model` to the model that you want to run the evaluation for
+    - configure a unique `experiment_name` for tracking purposes
+    - set the parameter `dataset_local_path` to the file you are running the evaluations for 
+1. Replace the contents of `prompt_template.txt` with your custom prompt template. Make sure that prompt template variables map to the dataset attributes.
+1. Please refer to our [documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/models/determine-eval) if you want to further customize your evaluation. Vertex AI Evaluation Service has a lot of features that are not included in this recipe.