{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "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",
        "# 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": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Evaluate images with Gecko"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/evaluation/evaltask_approach/evaluate_images_with_gecko.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://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fevaluation%2Fevaltask_approach%2Fevaluate_images_with_gecko.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://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/evaluation/evaltask_approach/evaluate_images_with_gecko.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/generative-ai/blob/main/gemini/evaluation/evaltask_approach/evaluate_images_with_gecko.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/evaluation/evaltask_approach/evaluate_images_with_gecko.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/evaluation/evaltask_approach/evaluate_images_with_gecko.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/evaluation/evaltask_approach/evaluate_images_with_gecko.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/evaluation/evaltask_approach/evaluate_images_with_gecko.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/evaluation/evaltask_approach/evaluate_images_with_gecko.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UBAgFfnqvxy0"
      },
      "source": [
        " | | | | |\n",
        " |-|-|-|-|\n",
        " |Author(s): | [Greg Breard](https://github.com/gregbreard) | Anant Nawalgaria | Olivia Wiles |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This CoLAB shows how to leverage the Vertex AI evaluation service in order to run [Gecko](https://arxiv.org/abs/2404.16820).\n",
        "\n",
        "As with a more standard rubric approach, Gecko proceeds in two stages: a rubric generation step followed by a validator step. The key difference is that the rubric is generated based on the prompt.\n",
        "This allows for a more fine-grained metric that can be customized to prompts with differing challenges.\n",
        "\n",
        "In more detail, Gecko proceeds as follows, with two key steps: the QA generation step (ie the rubric generation step) and then the VQA step (ie the validator step).\n",
        "\n",
        "## The rubric generation step\n",
        "Given a prompt, such as `A teddy bear riding a skateboard`, we prompt the Gemini model to generate a set of questions, answer choices and corresponding ground truth (GT) answer. The question is also tagged with a question type. Depending on the prompt, these questions can either be `yes`/`no` questions or multiple choice ones.\n",
        "\n",
        "`A teddy bear riding a skateboard` -->\n",
        "\n",
        "- `Q1: Is there a teddy bear? Choices: [yes, no]. GT Answer: yes. Tag: Object.`\n",
        "- `Q2: Is there a skateboard? Choices: [yes, no]. GT Answer: yes. Tag: Object.`\n",
        "- `Q3: Is the teddy bear riding a skateboard? Choices: [yes, no]. GT Answer: yes. Tag: Action.`\n",
        "\n",
        "## The validator step\n",
        "Given a generated image and the questions above, we query the Gemini model for each question to give an answer. We then check if it matches the GT answer, with a result of 1 if it matches and 0 if it does not. We aggregate these results to give a final overall score, which can be broken down into scores per question. We can also aggregate scores based on tags.\n",
        "\n",
        "For example, imagine we have a generated image `<image1>` which includes a teddy bear but no skateboard, and Gemini outputs the following results:\n",
        "\n",
        "- `<image1> Is there a teddy bear? GT Answer: yes. Result: 1.`\n",
        "- `<image1> Is there a skateboard? GT Answer: no. Result: 0.`\n",
        "- `<image1> Is the teddy bear riding a skateboard? GT Answer: no. Result: 0.`\n",
        "\n",
        "The final score will be `0.33` with a score of `0.5` for the question tag and `0.0` for the action tag.\n",
        "\n",
        "## Further exploration\n",
        "We provide two prompts, engineered for video and image generation tasks. Below, we show how to run Gecko for the image modality on a set of generations.\n",
        "\n",
        "However, these prompts can be modified and changed as suits a developer's needs. The quality can be analysed by exploring what  questions are generated as well as the reliability of the validator step. Questions can also be manually added as desired for an application.\n",
        "\n",
        "## Steps\n",
        "\n",
        "1. Set up the environment.\n",
        "2. Define helper functions, prompt templates, and metric.\n",
        "3. Prepare the dataset for evaluation.\n",
        "4. Run the evaluation (including model inference).\n",
        "\n",
        "## Costs\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "- Vertex AI\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "# Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet google-cloud-aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5Xep4W9lq-Z"
      },
      "source": [
        "### Restart runtime (Colab only)\n",
        "\n",
        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SbmM4z7FOBpM"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"your-project-id\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
        "\n",
        "\n",
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EdvJRUWRNGHE"
      },
      "source": [
        "## Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R0peQPa48oTg"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "from vertexai.preview.evaluation import (\n",
        "    CustomOutputConfig,\n",
        "    EvalTask,\n",
        "    PointwiseMetric,\n",
        "    RubricBasedMetric,\n",
        "    RubricGenerationConfig,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e-TDamAG9mkM"
      },
      "source": [
        "# Set up eval metrics for Gecko"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SR2qRUm29otX"
      },
      "source": [
        "## Helper functions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1pjVIACIMYne"
      },
      "source": [
        "The outputs supported by Gecko are more sophisticated than the default outputs of predefined rubric based metrics. To handle this, custom parsing logic is required.\n",
        "\n",
        "The following code block defines 2 classes: `QARecord` and `QAResult`. The `QARecord` represents the questions created during rubric generation. The `QAResult` extends the `QARecord` with a result field that is populated after validation.\n",
        "\n",
        "There are also two parsing methods. The `parse_json_to_qa_records` method converts the text output of rubric generation to `QARecords` and the `parse_rubric_results` method extracts the answers from the validation step. These are passed into the metric definition and parsing is handled automatically during the generation and validation steps.\n",
        "\n",
        "Finally, the `compute_scores` method compares the `QARecord`s and rubric results to calculate a per row score and appends `QAResult`s and scores to the dataset.\n",
        "\n",
        "In addition, there are pretty printing methods provided to present the output in a human readable format."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GfLmFxkP-Kvr"
      },
      "outputs": [],
      "source": [
        "from collections.abc import Sequence\n",
        "from dataclasses import dataclass, field\n",
        "import json\n",
        "import re\n",
        "from typing import Any\n",
        "\n",
        "import numpy as np\n",
        "\n",
        "_QUESTION_REGEX = re.compile(r\"Question:(.*?)Verdict:\", re.DOTALL)\n",
        "_VERDICT_REGEX = re.compile(\"Verdict:(.*)\")\n",
        "_QUESTION_BLOCK_REGEX = re.compile(\"<question>(.*?)</question>\", re.DOTALL)\n",
        "_TABLE_STYLE = [\n",
        "    {\n",
        "        \"selector\": \"th\",\n",
        "        \"props\": [\n",
        "            (\"background-color\", \"#f2f2f2\"),\n",
        "            (\"border\", \"1px solid gray\"),\n",
        "            (\"color\", \"black\"),\n",
        "            (\"font-size\", \"11pt\"),\n",
        "            (\"text-align\", \"center\"),\n",
        "            (\"word-break\", \"break-all\"),\n",
        "        ],\n",
        "    },\n",
        "    {\"selector\": \"tr:nth-child(even)\", \"props\": [(\"background-color\", \"#f9f9f9\")]},\n",
        "    {\"selector\": \"tr:nth-child(odd)\", \"props\": [(\"background-color\", \"white\")]},\n",
        "    {\"selector\": \"tr:hover\", \"props\": [(\"background-color\", \"#94e6ff\")]},\n",
        "    {\"selector\": \"td:hover\", \"props\": [(\"background-color\", \"#ffffb3\")]},\n",
        "]\n",
        "\n",
        "\n",
        "@dataclass(kw_only=True, frozen=True)\n",
        "class QARecord:\n",
        "    \"\"\"A basic QA Record for storing question-answer pairs.\n",
        "\n",
        "    Attributes:\n",
        "        question: Question text.\n",
        "        question_type: Category of question.\n",
        "        gt_answer: Ground-truth answer to the question.\n",
        "        answer_choices: Possible answers for multiple choice questions.\n",
        "        justification: How the question relates to the prompt.\n",
        "    \"\"\"\n",
        "\n",
        "    question: str = \"\"\n",
        "    question_type: str = \"\"\n",
        "    gt_answer: str = \"\"\n",
        "    answer_choices: Sequence[str] = field(default_factory=list)\n",
        "    justification: str = \"\"\n",
        "\n",
        "\n",
        "class QAResult(QARecord):\n",
        "    \"\"\"A basic QA Result for storing question-answer results.\n",
        "\n",
        "    Attributes:\n",
        "        result: The result of answering the question.\n",
        "    \"\"\"\n",
        "\n",
        "    result: str = \"\"\n",
        "\n",
        "    def __init__(self, qa_record: QARecord, result: str):\n",
        "        super().__init__(\n",
        "            question=qa_record.question,\n",
        "            question_type=qa_record.question_type,\n",
        "            gt_answer=qa_record.gt_answer,\n",
        "            answer_choices=qa_record.answer_choices,\n",
        "            justification=qa_record.justification,\n",
        "        )\n",
        "        self.result = result\n",
        "\n",
        "\n",
        "def parse_json_to_qa_records(json_response: str) -> dict[str, Any]:\n",
        "    \"\"\"\n",
        "    Parse the JSON response and convert it to a questions and QARecords.\n",
        "\n",
        "    Args:\n",
        "        json_response: JSON string containing the QA data.\n",
        "\n",
        "    Returns:\n",
        "        Dict with keywords, questions, and QARecord objects.\n",
        "\n",
        "    Raises:\n",
        "        json.JSONDecodeError: If JSON parsing fails\n",
        "        KeyError: If expected keys are missing from the JSON structure\n",
        "    \"\"\"\n",
        "    json_response = re.sub(\n",
        "        r\"(.*```json|```.*)\",\n",
        "        \"\",\n",
        "        json_response.strip(),\n",
        "    )\n",
        "    try:\n",
        "        # Parse JSON string to Python object\n",
        "        data = json.loads(json_response)\n",
        "        qa_records = []\n",
        "\n",
        "        # Process each QA pair in the QAs array\n",
        "        rubrics = []\n",
        "        for qa in data[\"qas\"]:\n",
        "            record = QARecord(\n",
        "                question=qa[\"question\"],\n",
        "                question_type=qa[\"question_type\"],\n",
        "                gt_answer=qa[\"answer\"],\n",
        "                answer_choices=qa[\"choices\"],\n",
        "                justification=qa[\"justification\"],\n",
        "            )\n",
        "            qa_records.append(record)\n",
        "            rubrics.append(\n",
        "                f\"<question>{record.question}<choices>{','.join(record.answer_choices)}\"\n",
        "            )\n",
        "        return {\n",
        "            \"questions\": \"\\n\".join(rubrics),\n",
        "            \"keywords\": data[\"keywords\"],\n",
        "            \"qa_records\": qa_records,\n",
        "        }\n",
        "    except json.JSONDecodeError as e:\n",
        "        return {\n",
        "            \"questions\": f\"Error decoding JSON response: {str(e)}\",\n",
        "            \"keywords\": \"\",\n",
        "            \"qa_records\": json_response,\n",
        "        }\n",
        "    except KeyError as e:\n",
        "        return {\n",
        "            \"questions\": f\"Missing required key in JSON structure: {str(e)}\",\n",
        "            \"keywords\": \"\",\n",
        "            \"qa_records\": json_response,\n",
        "        }\n",
        "\n",
        "\n",
        "def parse_rubric_results(results: list[str]) -> dict[str, Any]:\n",
        "    \"\"\"Parses the rubric results from the rubric validator response.\"\"\"\n",
        "    rubric_results = {}\n",
        "    for result in results:\n",
        "        rubric_verdicts = _parse_question_blocks(result)\n",
        "        for rubric, verdict in rubric_verdicts:\n",
        "            rubric_results[rubric.lower()] = verdict.lower()\n",
        "    return {\"rubric_results\": rubric_results}\n",
        "\n",
        "\n",
        "def _parse_question_blocks(txt: str) -> list[tuple[str, bool]]:\n",
        "    \"\"\"Parses the question blocks from the rubric validator response.\"\"\"\n",
        "    responses = []\n",
        "    question_blocks = _QUESTION_BLOCK_REGEX.findall(txt)\n",
        "    if not question_blocks:\n",
        "        question_blocks = [txt]\n",
        "    for block in question_blocks:\n",
        "        q = _parse_question(block)\n",
        "        v = _parse_verdict(block)\n",
        "        if q is not None and v is not None:\n",
        "            responses.append((q, v))\n",
        "    return responses\n",
        "\n",
        "\n",
        "def _parse_question(txt: str):\n",
        "    \"\"\"Parses the question from the rubric validator response.\"\"\"\n",
        "    if not isinstance(txt, str) or not txt:\n",
        "        return None\n",
        "    try:\n",
        "        txt = txt.split(\"Verdict:\")[0]\n",
        "        if \"Question:\" in txt:\n",
        "            return txt.split(\"Question:\")[-1].strip()\n",
        "        if question := _QUESTION_REGEX.findall(txt):\n",
        "            return question[0].strip()\n",
        "    except Exception as e:\n",
        "        print(f\"Failed to parse question: {str(e)}\")\n",
        "        return None\n",
        "\n",
        "\n",
        "def _parse_verdict(txt: str):\n",
        "    \"\"\"Parses the verdict from the rubric validator response.\"\"\"\n",
        "    if not isinstance(txt, str) or not txt:\n",
        "        return None\n",
        "    try:\n",
        "        if verdict := _VERDICT_REGEX.findall(txt):\n",
        "            verdict = verdict[0].strip()\n",
        "            return verdict\n",
        "    except Exception as e:\n",
        "        print(f\"Failed to parse question: {str(e)}\")\n",
        "        return None\n",
        "\n",
        "\n",
        "def compute_scores(df: \"pd.DataFrame\") -> \"pd.DataFrame\":\n",
        "    \"\"\"Computes scores for each row based on QA results.\"\"\"\n",
        "    qa_results = []\n",
        "    final_scores = []\n",
        "    for idx, row in df.iterrows():\n",
        "        rubric_results = {}\n",
        "        for key in row.keys():\n",
        "            if \"rubric_results\" in key:\n",
        "                rubric_results = row[key]\n",
        "        scores = []\n",
        "        results = []\n",
        "        for qa in row[\"qa_records\"]:\n",
        "            q = qa.question.lower()\n",
        "            if q in rubric_results:\n",
        "                if qa.gt_answer.lower() in rubric_results[q]:\n",
        "                    results.append(QAResult(qa, f\"{qa.gt_answer} ✓\"))\n",
        "                    scores.append(1)\n",
        "                else:\n",
        "                    results.append(QAResult(qa, f\"{rubric_results[q]} 🗴\"))\n",
        "                    scores.append(0)\n",
        "            else:\n",
        "                results.append(QAResult(qa, \"no result\"))\n",
        "                scores.append(0)\n",
        "        qa_results.append(results)\n",
        "        final_scores.append(np.mean(scores))\n",
        "    df_with_score = df.assign(qa_results=qa_results, final_score=final_scores)\n",
        "    return df_with_score\n",
        "\n",
        "\n",
        "def pretty_print_qa_records_df(\n",
        "    df: \"pd.DataFrame\", hide_columns: list[str]\n",
        ") -> \"pd.Styler\":\n",
        "    \"\"\"Prints QA records data frame as stylized HTML table.\"\"\"\n",
        "    styled_df = df.copy()\n",
        "    for col in df.columns:\n",
        "        if (\n",
        "            isinstance(df[col][0], list)\n",
        "            and df[col][0]\n",
        "            and isinstance(df[col][0][0], QARecord)\n",
        "        ):\n",
        "            styled_df[col] = styled_df[col].apply(\n",
        "                lambda x: _qa_records_to_html_table(x)\n",
        "            )\n",
        "    styles = _TABLE_STYLE.copy()\n",
        "    styles.append(\n",
        "        {\n",
        "            \"selector\": \"td\",\n",
        "            \"props\": [\n",
        "                (\"border\", \"1px solid gray\"),\n",
        "                (\"color\", \"black\"),\n",
        "                (\"min-width\", \"100px\"),\n",
        "                (\"text-align\", \"center\"),\n",
        "            ],\n",
        "        }\n",
        "    )\n",
        "    return (\n",
        "        styled_df.style.hide(axis=\"index\")\n",
        "        .hide(subset=hide_columns, axis=1)\n",
        "        .set_table_styles(styles)\n",
        "    )\n",
        "\n",
        "\n",
        "def pretty_print_result_df(df: \"pd.DataFrame\", hide_columns: list[str]) -> \"pd.Styler\":\n",
        "    \"\"\"Prints results data frame as stylized HTML table.\"\"\"\n",
        "    styled_df = df.copy()\n",
        "    for col in df.columns:\n",
        "        if (\n",
        "            isinstance(df[col][0], list)\n",
        "            and df[col][0]\n",
        "            and isinstance(df[col][0][0], QARecord)\n",
        "        ):\n",
        "            styled_df[col] = styled_df[col].apply(\n",
        "                lambda x: _qa_records_to_html_table(x)\n",
        "            )\n",
        "    styles = _TABLE_STYLE.copy()\n",
        "    styles.append(\n",
        "        {\n",
        "            \"selector\": \"td\",\n",
        "            \"props\": [\n",
        "                (\"border\", \"1px solid gray\"),\n",
        "                (\"color\", \"black\"),\n",
        "                (\"min-width\", \"120px\"),\n",
        "                (\"text-align\", \"center\"),\n",
        "            ],\n",
        "        }\n",
        "    )\n",
        "    return (\n",
        "        styled_df.style.hide(axis=\"index\")\n",
        "        .hide(subset=hide_columns, axis=1)\n",
        "        .format({\"final_score\": \"{:,.1f}\"})\n",
        "        .set_table_styles(styles)\n",
        "    )\n",
        "\n",
        "\n",
        "def _qa_records_to_html_table(data: list[QARecord]) -> str:\n",
        "    \"\"\"Converts a list to an HTML table.\"\"\"\n",
        "    if not data:\n",
        "        return \"<i>No data to display.</i>\"\n",
        "    html_table = \"<table style='border-collapse: collapse'><thead><tr>\"\n",
        "    # Extract headers from the first element\n",
        "    keys = [\"question\", \"answer_choices\", \"gt_answer\"]\n",
        "    if isinstance(data[0], QAResult):\n",
        "        keys.append(\"result\")\n",
        "    for key in keys:\n",
        "        html_table += f\"<th>{key}</th>\"\n",
        "    html_table += \"</tr></thead><tbody>\"\n",
        "    # Add rows\n",
        "    for item in data:\n",
        "        html_table += \"<tr>\"\n",
        "        for key in keys:\n",
        "            html_table += f\"<td>{item.__dict__[key]}</td>\"\n",
        "        html_table += \"</tr>\"\n",
        "    html_table += \"</tbody></table>\"\n",
        "    return html_table"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tLn3FteP9ttL"
      },
      "source": [
        "## Prompt Templates"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k6bau3qOFInw"
      },
      "source": [
        "This cell defines the prompt templates that will be used for evaluation. The `RUBRIC_GENERATION_PROMPT` is used to generate questions relevant to the user input. The `RUBRIC_VALIDATOR_PROMPT` is then used to answer the questions for a generated image."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5--M9Kcz-LFt"
      },
      "outputs": [],
      "source": [
        "RUBRIC_GENERATION_PROMPT = \"\"\"\n",
        "In this task, you will help me measure generate question-answer pairs to verify\n",
        "an image description.\n",
        "\n",
        "You will first identify the key words to be validated, e.g. ignoring filler or\n",
        "redundant words.\n",
        "\n",
        "You will then, for each word, generate a question-answer pair for each word. The\n",
        "question should be simple and *cannot* be answered correctly based on common\n",
        "sense or without reading the description. You will also tag each question as\n",
        "having a type, which should be one off: object,\n",
        "human, animal, food, activity, attribute, counting, color, material, spatial,\n",
        "location, shape, other.\n",
        "\n",
        "**Important**: There should be one and only one question-answer pair per key word.\n",
        "\n",
        "\n",
        "Given a \"description\", your answer must have this format:\n",
        "{\n",
        "  \"keywords\": \"Your {1}[itemized] {2}[keywords]\",\n",
        "  \"qas\": [\n",
        "    The list of QAs in the format \"{\n",
        "      \"question_id\": i,\n",
        "      \"question\": \"the question\",\n",
        "      \"answer\": \"the answer: yes or no\",\n",
        "      \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"why is this about the keyword\",\n",
        "      \"question_type\": \"the question type. One of [object, human, animal, food, activity, attribute, counting, color, material, spatial, location, shape, other].\"\n",
        "      }\".,\n",
        "  ]\n",
        "}\n",
        "\n",
        "**Important**: There should be one and only one question-answer pair per key word.\n",
        "===\n",
        "Some examples are below.\n",
        "\n",
        "Description: A man posing for a selfie in a jacket and bow tie.\n",
        "Answer:\n",
        "{\n",
        "  \"keywords\": \"A {1}[man] {2}[posing] for a {3}[selfie] in a {4}[jacket] and a {5}[bow tie].\",\n",
        "  \"qas\": [\n",
        "    {\n",
        "      \"question_id\": 1, \"question\": \"is there a man in the image?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There is a man in the image.\", \"question_type\": \"human\"\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 2, \"question\": \"is the man posing for a selfie?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"The man is posing for a selfie.\", \"question_type\": \"activity\"\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 3, \"question\": \"Is the man taking a selfie?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"This is a selfie.\", \"question_type\": \"object\"\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 4, \"question\": \"Is the man wearing a jacket?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"The man is wearing a jacket.\", \"question_type\": \"object\"\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 5, \"question\": \"Is the man wearing a bow tie?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"The man is wearing a bow tie.\", \"question_type\": \"object\"\n",
        "    },\n",
        "  ]\n",
        "}\n",
        "\n",
        "Description: A horse and several cows feed on hay.\n",
        "Answer:\n",
        "{\n",
        "  \"keywords\": \"A {1}[horse] and {2}[several] {3}[cows] {4}[feed] on {5}[hay]\",\n",
        "  \"qas\": [\n",
        "    {\n",
        "      \"question_id\": 1, \"question\": \"is there a horse?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There is a horse in the image.\", \"question_type\": \"animal\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 2, \"question\": \"are there several cows?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There are several cows in the image.\", \"question_type\": \"counting\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 3, \"question\": \"are there cows?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There are cows in the image.\", \"question_type\": \"animal\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 4, \"question\": \"are the horse and cows feeding on hay?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"The horse and cows are feeding.\", \"question_type\": \"activity\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 5, \"question\": \"is there hay?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There is hay in the image.\", \"question_type\": \"object\",\n",
        "    },\n",
        "  ]\n",
        "}\n",
        "\n",
        "Description: A red colored dog.\n",
        "Answer:\n",
        "{\n",
        "  \"keywords\": \"A {1}[red colored] {2}[dog].\",\n",
        "  \"qas\": [\n",
        "    {\n",
        "      \"question_id\": 1, \"question\": \"is the dog red?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There is a red colored dog in the image.\", \"question_type\": \"color\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 2, \"question\": \"is there a dog?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There is a dog in the image.\", \"question_type\": \"animal\",\n",
        "    },\n",
        "  ]\n",
        "}\n",
        "\n",
        "Description: A busy intersection with an ice cream truck driving by.\n",
        "Answer:\n",
        "{\n",
        "  \"keywords\": \"A {1}[busy] {2}[intersection] with an {3}[ice cream truck] {4}[driving by].\",\n",
        "  \"qas\": [\n",
        "    {\n",
        "      \"question_id\": 1, \"question\": \"is this a busy intersection?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"The intersection is busy.\", \"question_type\": \"attribute\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 2, \"question\": \"is this an intersection?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There is a busy intersection.\", \"question_type\": \"object\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 3, \"question\": \"is there an ice cream truck?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There is an ice cream truck.\", \"question_type\": \"object\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 4, \"question\": \"is the ice cream truck driving by?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"The ice cream truck is driving by.\", \"question_type\": \"activity\",\n",
        "    },\n",
        "  ]\n",
        "}\n",
        "\n",
        "Description: Portrait of a gecko wearing a train conductor's hat and holding a flag that has a yin-yang symbol on it. Woodcut.\n",
        "Answer:\n",
        "{\n",
        "  \"keywords\": \"{1}[Portrait] of a {2}[gecko] {3}[wearing] a {4}[train conductor's hat] and {5}[holding] a {6}[flag] that has a {7}[yin-yang symbol] on it. {8}[Woodcut].\",\n",
        "  \"qas\": [\n",
        "    {\n",
        "      \"question_id\": 1, \"question\": \"is this a portrait?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"This is a portrait.\", \"question_type\": \"attribute\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 2, \"question\": \"is there a gecko?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There is a gecko.\", \"question_type\": \"animal\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 3, \"question\": \"is the gecko wearing a hat?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"The gecko is wearing a train conductor's hat.\", \"question_type\": \"activity\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 4, \"question\": \"is the gecko wearing a train conductor's hat?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"The gecko is wearing a train conductor's hat.\", \"question_type\": \"attribute\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 5, \"question\": \"is the gecko holding a flag\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"The gecko is holding a flag.\", \"question_type\": \"activity\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 6, \"question\": \"is there a flag?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There is a flag.\", \"question_type\": \"object\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 7, \"question\": \"does the flag have a yin-yang symbol?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"The flag has a yin-yang symbol on it.\", \"question_type\": \"attribute\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 8, \"question\": \"is this a woodcut?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"woodcut.\", \"question_type\": \"attribute\",\n",
        "    },\n",
        "  ]\n",
        "}\n",
        "\n",
        "Description: A woman is showing a watermelon slice to a woman on a scooter.\n",
        "Answer:\n",
        "{\n",
        "  \"keywords\": \"A {1}[woman] is {2}[showing] a {3}[watermelon slice] to a {4}[woman] {5}[on] a {6}[scooter].\",\n",
        "  \"qas\": [\n",
        "    {\n",
        "      \"question_id\": 1, \"question\": \"is there a woman?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"This is a woman.\", \"question_type\": \"human\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 2, \"question\": \"is one woman showing a watermelon slice to another woman?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"A woman is showing a watermelon slice to a woman.\", \"question_type\": \"activity\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 3, \"question\": \"is there a watermelon slice?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"The watermelon slice is a watermelon slice.\", \"question_type\": \"food\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 4, \"question\": \"are there two women?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"A woman is showing a watermelon slice to a woman.\", \"question_type\": \"human\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 5, \"question\": \"is one of the women on a scooter?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"A woman is on a scooter.\", \"question_type\": \"spatial\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 6, \"question\": \"is there a scooter?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There is a scooter.\", \"question_type\": \"object\",\n",
        "    },\n",
        "  ]\n",
        "}\n",
        "\n",
        "Description: A photo of three dogs.\n",
        "Answer:\n",
        "{\n",
        "  \"keywords\": \"A {1}[photo] of {2}[three] {3}[dogs].\",\n",
        "  \"qas\": [\n",
        "    {\n",
        "      \"question_id\": 1, \"question\": \"is this a photo?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"This is a photo of three dogs.\", \"question_type\": \"attribute\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 2, \"question\": \"are there three dogs?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"three dogs.\", \"question_type\": \"counting\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 3, \"question\": \"is there a dog?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There is a dog.\", \"question_type\": \"object\",\n",
        "    },\n",
        "  ]\n",
        "}\n",
        "\n",
        "Description: A white milk truck with a license plate that reads 'pro milk'.\n",
        "Answer:\n",
        "{\n",
        "  \"keywords\": \"A {1}[white] {2}[milk truck] with a {3}[license plate] that reads {4}['pro milk'].\",\n",
        "  \"qas\": [\n",
        "    {\n",
        "      \"question_id\": 1, \"question\": \"is this a white truck?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"This is a white milk truck.\", \"question_type\": \"attribute\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 2, \"question\": \"is there a milk truck?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"This is a white milk truck.\", \"question_type\": \"object\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 3, \"question\": \"is there a license plate on the vehicle?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"A white milk truck with a license plate.\", \"question_type\": \"object\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 4, \"question\": \"does the license plate read 'pro milk'?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"The license plate reads 'pro milk'.\", \"question_type\": \"attribute\",\n",
        "    },\n",
        "  ]\n",
        "}\n",
        "\n",
        "Description: A person sitting on a horse in air over gate in grass with people and trees in background.\n",
        "Answer:\n",
        "{\n",
        "  \"keywords\": \"A {1}[person] {2}[sitting] {3}[on] a {4}[horse] {5}[in air] {6}[over] {7}[gate] in {8}[grass] with {9}[people] and {10}[trees] in {11}[background].\",\n",
        "  \"qas\": [\n",
        "    {\n",
        "      \"question_id\": 1, \"question\": \"is there a person?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"This is a person.\", \"question_type\": \"human\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 2, \"question\": \"is the person sitting?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"A person is sitting on a horse.\", \"question_type\": \"activity\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 3, \"question\": \"is the person on a horse?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"A person is sitting on a horse.\", \"question_type\": \"spatial\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 4, \"question\": \"is there a horse?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"A person is sitting on a horse.\", \"question_type\": \"object\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 5, \"question\": \"is the horse in the air?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"The horse is in the air.\", \"question_type\": \"attribute\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 6, \"question\": \"is the horse over the gate?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"The horse is over the gate.\", \"question_type\": \"spatial\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 7, \"question\": \"is there a gate?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"The horse is over the gate.\", \"question_type\": \"object\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 8, \"question\": \"is there grass?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"The horse is in the grass.\", \"question_type\": \"object\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 9, \"question\": \"are there people?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There are people.\", \"question_type\": \"human\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 10, \"question\": \"are there trees?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There are trees.\", \"question_type\": \"object\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 11, \"question\": \"are there people and trees in the background?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There people and trees in the background.\", \"question_type\": \"human\",\n",
        "    },\n",
        "  ]\n",
        "}\n",
        "\n",
        "Description: a red blue and yellow train and some people on a platform\n",
        "Answer:\n",
        "{\n",
        "  \"keywords\": \"a {1}[red blue and yellow] {2}[train] and {3}[some] {4}[people] on a {5}[platform]\",\n",
        "  \"qas\": [\n",
        "    {\n",
        "      \"question_id\": 1, \"question\": \"is the train red blue and yellow?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"This is a red blue and yellow train.\", \"question_type\": \"color\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 2, \"question\": \"is there a train?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There is a train.\", \"question_type\": \"object\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 3, \"question\": \"Are there some people?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There are some people.\", \"question_type\": \"counting\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 4, \"question\": \"are the people on the platform?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There are some people on the platform.\", \"question_type\": \"human\",\n",
        "    },\n",
        "    {\n",
        "      \"question_id\": 5, \"question\": \"is there a platform?\", \"answer\": \"yes\", \"choices\": [\"yes\", \"no\"],\n",
        "      \"justification\": \"There is a platform.\", \"question_type\": \"object\",\n",
        "    },\n",
        "  ]\n",
        "}\n",
        "\n",
        "Description:\n",
        "{prompt}\n",
        "Answer:\n",
        "\"\"\"\n",
        "\n",
        "RUBRIC_VALIDATOR_PROMPT = \"\"\"\n",
        "# Instructions\n",
        "Look at the image carefully and answer each question with a yes or no:\n",
        "{rubrics}\n",
        "\n",
        "# Image\n",
        "{image}\n",
        "\n",
        "# Output Format\n",
        "<question>\n",
        "Question: repeat the original question\n",
        "Verdict: yes|no\n",
        "</question>\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Fdtcd-WF9ymb"
      },
      "source": [
        "## Define the metric"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xTyZIRi6L7CI"
      },
      "source": [
        "This cell configures the rubric generation and validator metric for rubric based evaluation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CeA0mrWi-LX1"
      },
      "outputs": [],
      "source": [
        "# Rubric Generation\n",
        "rubric_generation_config = RubricGenerationConfig(\n",
        "    prompt_template=RUBRIC_GENERATION_PROMPT,\n",
        "    parsing_fn=parse_json_to_qa_records,\n",
        ")\n",
        "\n",
        "# Rubric Validation\n",
        "pointwise_metric = PointwiseMetric(\n",
        "    metric=\"gecko_metric\",\n",
        "    metric_prompt_template=RUBRIC_VALIDATOR_PROMPT,\n",
        "    custom_output_config=CustomOutputConfig(\n",
        "        return_raw_output=True,\n",
        "        parsing_fn=parse_rubric_results,\n",
        "    ),\n",
        ")\n",
        "\n",
        "# Rubric Metric\n",
        "rubric_based_gecko = RubricBasedMetric(\n",
        "    generation_config=rubric_generation_config,\n",
        "    critique_metric=pointwise_metric,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nHjCRree99Px"
      },
      "source": [
        "# Prepare the dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nQQC5fuTbOqF"
      },
      "source": [
        "In the following dataset, two prompts are used for each generated image. The first is the prompt that corresponds to the generated content. The second is a counterexample that is similar but does not exactly match the generated content. This is done to demonstrate the difference in the Gecko evaluation for high quality and low quality responses."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U37FNLD8-Lp5"
      },
      "outputs": [],
      "source": [
        "prompts = [\n",
        "    \"steaming cup of coffee and a croissant on a table\",\n",
        "    \"steaming cup of coffee and toast in a cafe\",\n",
        "    \"sunset over a calm ocean\",\n",
        "    \"sunset over a tranquil forest\",\n",
        "    \"butterfly with colorful wings on a flower\",\n",
        "    \"butterfly fluttering over a leaf\",\n",
        "    \"musician playing guitar on a street corner\",\n",
        "    \"musician playing saxophone under lamp post\",\n",
        "    \"vintage camera with a worn leather strap\",\n",
        "    \"new camera with a power zoom lens\",\n",
        "    \"colorful abstract painting\",\n",
        "    \"black and white painting\",\n",
        "    \"baker decorating a cake with frosting\",\n",
        "    \"baker topping cupcakes with sprinkles\",\n",
        "    \"hot air balloon floating above a field of lavender\",\n",
        "    \"hot air balloon landing in a field of sunflowers\",\n",
        "]\n",
        "images = [\n",
        "    '{\"contents\": [{\"parts\": [{\"file_data\": {\"mime_type\": \"image/png\", \"file_uri\": \"gs://cloud-samples-data/generative-ai/evaluation/images/coffee.png\"}}]}]}',\n",
        "    '{\"contents\": [{\"parts\": [{\"file_data\": {\"mime_type\": \"image/png\", \"file_uri\": \"gs://cloud-samples-data/generative-ai/evaluation/images/coffee.png\"}}]}]}',\n",
        "    '{\"contents\": [{\"parts\": [{\"file_data\": {\"mime_type\": \"image/png\", \"file_uri\": \"gs://cloud-samples-data/generative-ai/evaluation/images/sunset.png\"}}]}]}',\n",
        "    '{\"contents\": [{\"parts\": [{\"file_data\": {\"mime_type\": \"image/png\", \"file_uri\": \"gs://cloud-samples-data/generative-ai/evaluation/images/sunset.png\"}}]}]}',\n",
        "    '{\"contents\": [{\"parts\": [{\"file_data\": {\"mime_type\": \"image/png\", \"file_uri\": \"gs://cloud-samples-data/generative-ai/evaluation/images/butterfly.png\"}}]}]}',\n",
        "    '{\"contents\": [{\"parts\": [{\"file_data\": {\"mime_type\": \"image/png\", \"file_uri\": \"gs://cloud-samples-data/generative-ai/evaluation/images/butterfly.png\"}}]}]}',\n",
        "    '{\"contents\": [{\"parts\": [{\"file_data\": {\"mime_type\": \"image/png\", \"file_uri\": \"gs://cloud-samples-data/generative-ai/evaluation/images/musician.png\"}}]}]}',\n",
        "    '{\"contents\": [{\"parts\": [{\"file_data\": {\"mime_type\": \"image/png\", \"file_uri\": \"gs://cloud-samples-data/generative-ai/evaluation/images/musician.png\"}}]}]}',\n",
        "    '{\"contents\": [{\"parts\": [{\"file_data\": {\"mime_type\": \"image/png\", \"file_uri\": \"gs://cloud-samples-data/generative-ai/evaluation/images/camera.png\"}}]}]}',\n",
        "    '{\"contents\": [{\"parts\": [{\"file_data\": {\"mime_type\": \"image/png\", \"file_uri\": \"gs://cloud-samples-data/generative-ai/evaluation/images/camera.png\"}}]}]}',\n",
        "    '{\"contents\": [{\"parts\": [{\"file_data\": {\"mime_type\": \"image/png\", \"file_uri\": \"gs://cloud-samples-data/generative-ai/evaluation/images/abstract.png\"}}]}]}',\n",
        "    '{\"contents\": [{\"parts\": [{\"file_data\": {\"mime_type\": \"image/png\", \"file_uri\": \"gs://cloud-samples-data/generative-ai/evaluation/images/abstract.png\"}}]}]}',\n",
        "    '{\"contents\": [{\"parts\": [{\"file_data\": {\"mime_type\": \"image/png\", \"file_uri\": \"gs://cloud-samples-data/generative-ai/evaluation/images/baker.png\"}}]}]}',\n",
        "    '{\"contents\": [{\"parts\": [{\"file_data\": {\"mime_type\": \"image/png\", \"file_uri\": \"gs://cloud-samples-data/generative-ai/evaluation/images/baker.png\"}}]}]}',\n",
        "    '{\"contents\": [{\"parts\": [{\"file_data\": {\"mime_type\": \"image/png\", \"file_uri\": \"gs://cloud-samples-data/generative-ai/evaluation/images/balloon.png\"}}]}]}',\n",
        "    '{\"contents\": [{\"parts\": [{\"file_data\": {\"mime_type\": \"image/png\", \"file_uri\": \"gs://cloud-samples-data/generative-ai/evaluation/images/balloon.png\"}}]}]}',\n",
        "]\n",
        "\n",
        "eval_dataset = pd.DataFrame(\n",
        "    {\n",
        "        \"prompt\": prompts,\n",
        "        \"image\": images,\n",
        "    }\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qm6mHoPc-DS0"
      },
      "source": [
        "# Run evaluation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BCdjWbaYCvIT"
      },
      "source": [
        "## Generate rubrics"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6DU7nuNPNoXQ"
      },
      "source": [
        "First we generate rubrics for the user prompts."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7EWsZkW-CzXZ"
      },
      "outputs": [],
      "source": [
        "dataset_with_rubrics = rubric_based_gecko.generate_rubrics(eval_dataset)\n",
        "pretty_print_qa_records_df(\n",
        "    dataset_with_rubrics, hide_columns=[\"prompt\", \"image\", \"rubrics\"]\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pq5uNh0hC3sz"
      },
      "source": [
        "## Evaluate with rubrics"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VcQwzMjYNsr7"
      },
      "source": [
        "Then we use the generated rubrics to evaluate the quality of the responses."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mMuDxiwy-MTt"
      },
      "outputs": [],
      "source": [
        "eval_task = EvalTask(\n",
        "    dataset=dataset_with_rubrics,\n",
        "    metrics=[rubric_based_gecko],\n",
        ")\n",
        "eval_result = eval_task.evaluate(response_column_name=\"image\")\n",
        "\n",
        "# Calculate overall score for metric.\n",
        "dataset_with_final_scores = compute_scores(eval_result.metrics_table)\n",
        "np.mean(dataset_with_final_scores[\"final_score\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zc6GvhGdK7D_"
      },
      "outputs": [],
      "source": [
        "pretty_print_result_df(\n",
        "    dataset_with_final_scores,\n",
        "    hide_columns=[\n",
        "        \"prompt\",\n",
        "        \"image\",\n",
        "        \"rubrics\",\n",
        "        \"qa_records\",\n",
        "        \"gecko_metric/rubric_results\",\n",
        "    ],\n",
        ")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "tLn3FteP9ttL",
        "Fdtcd-WF9ymb",
        "nHjCRree99Px"
      ],
      "name": "evaluate_images_with_gecko.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
