examples/vertex_ai_conversation/evaluation_tool__autoeval__colab.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zGZzOhx3sNWR"
      },
      "source": [
        "# Vertex AI Conversation - Evaluation Tool\n",
        "\n",
        "This tool requires user's input in several steps. Please run the cells one by one (Shift+Enter) to ensure all the steps are successfully completed."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Afvsuux0zaWZ"
      },
      "source": [
        "# Setup\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "0U8xQwhKrOUq"
      },
      "outputs": [],
      "source": [
        "!pip install dfcx-scrapi --quiet\n",
        "!pip install rouge-score --quiet\n",
        "\n",
        "# workaround until vertexai import is fixed\n",
        "!pip uninstall bigframes -y --quiet\n",
        "!pip install bigframes==0.26.0 --quiet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "PPJYRHN83bHg"
      },
      "outputs": [],
      "source": [
        "# @markdown `import dependencies`\n",
        "\n",
        "import abc\n",
        "import collections\n",
        "import dataclasses\n",
        "import datetime\n",
        "import io\n",
        "import itertools\n",
        "import json\n",
        "import logging\n",
        "import math\n",
        "import statistics\n",
        "import sys\n",
        "import time\n",
        "import threading\n",
        "import re\n",
        "\n",
        "from typing import Any, TypedDict\n",
        "\n",
        "from collections.abc import Iterable\n",
        "\n",
        "import plotly.graph_objects as go\n",
        "\n",
        "import vertexai\n",
        "import gspread\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import seaborn as sns\n",
        "sns.set_style(\"whitegrid\")\n",
        "\n",
        "from tqdm.auto import tqdm\n",
        "from tqdm.contrib import concurrent\n",
        "\n",
        "from dfcx_scrapi.core import agents\n",
        "from dfcx_scrapi.core import scrapi_base\n",
        "from dfcx_scrapi.core import sessions\n",
        "from dfcx_scrapi.core.sessions import Sessions\n",
        "from dfcx_scrapi.tools import dataframe_functions\n",
        "\n",
        "from googleapiclient.discovery import build\n",
        "from googleapiclient.http import MediaInMemoryUpload, MediaIoBaseDownload\n",
        "\n",
        "from google.api_core import exceptions\n",
        "from google.auth import default\n",
        "from google.cloud import aiplatform\n",
        "from google.cloud.dialogflowcx_v3beta1 import services\n",
        "from google.cloud.dialogflowcx_v3beta1 import types\n",
        "from google.colab import auth\n",
        "from google.protobuf.json_format import MessageToDict\n",
        "\n",
        "from rouge_score import rouge_scorer\n",
        "\n",
        "from vertexai.language_models import TextGenerationModel\n",
        "\n",
        "pd.options.display.max_colwidth = 200"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "EgoRHwBJqJ0r"
      },
      "outputs": [],
      "source": [
        "# @markdown `authenticate`\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.auth import default\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()\n",
        "    credentials, _ = default()\n",
        "else:\n",
        "    # Otherwise, attempt to discover local credentials as described in\n",
        "    # https://cloud.google.com/docs/authentication/application-default-credentials\n",
        "    pass\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2c6VwnIvTBjF"
      },
      "source": [
        "---\n",
        "\n",
        "# Implementation\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "BwPAkHGQ3k6M"
      },
      "outputs": [],
      "source": [
        "# @markdown `run this cell to define various utility and helper functions`\n",
        "# @markdown\n",
        "# @markdown > This cell contains several decorator methods related to handling\n",
        "# @markdown API call errors and function call rate limitation.\n",
        "\n",
        "_INTERVAL_SENTINEL = object()\n",
        "\n",
        "MAX_RETRIES = 5\n",
        "# max number of attempts for exponential backoff retries in case of API\n",
        "# call errors\n",
        "RATE = 2\n",
        "# LLM API call rate limitation where RATE=2 for example means that 2 LLM calls\n",
        "# can occur per second\n",
        "\n",
        "\n",
        "def load_spreadsheet(\n",
        "    sheet_url: str, worksheet_name: str, credentials: Any\n",
        ") -> pd.DataFrame:\n",
        "  \"\"\"Loads the content of a spreadsheet into pandas DataFrame.\"\"\"\n",
        "  sheets_client = gspread.authorize(credentials)\n",
        "  sheet = sheets_client.open_by_url(sheet_url)\n",
        "  worksheet = sheet.worksheet(worksheet_name)\n",
        "  return pd.DataFrame(worksheet.get_all_records())\n",
        "\n",
        "\n",
        "def ratelimit(rate: float):\n",
        "  \"\"\"Decorator that controls the frequency of function calls.\"\"\"\n",
        "  seconds_per_event = 1.0 / rate\n",
        "  lock = threading.Lock()\n",
        "  bucket = 0\n",
        "  last = 0\n",
        "\n",
        "  def decorate(func):\n",
        "    def rate_limited_function(*args, **kwargs):\n",
        "      nonlocal last, bucket\n",
        "      while True:\n",
        "        with lock:\n",
        "          now = time.time()\n",
        "          bucket += now - last\n",
        "          last = now\n",
        "\n",
        "          # capping the bucket in order to avoid accumulating too many\n",
        "          bucket = min(bucket, seconds_per_event)\n",
        "\n",
        "          # if bucket is less than `seconds_per_event` then we have to wait\n",
        "          # `seconds_per_event` - `bucket` seconds until a new \"token\" is\n",
        "          # refilled\n",
        "          delay = max(seconds_per_event - bucket, 0)\n",
        "\n",
        "          if delay == 0:\n",
        "            # consuming a token and breaking out of the delay loop to perform\n",
        "            # the function call\n",
        "            bucket -= seconds_per_event\n",
        "            break\n",
        "        time.sleep(delay)\n",
        "      return func(*args, **kwargs)\n",
        "    return rate_limited_function\n",
        "  return decorate\n",
        "\n",
        "\n",
        "def should_retry(err: exceptions.GoogleAPICallError) -> bool:\n",
        "  \"\"\"Helper function for deciding whether we should retry the error or not.\"\"\"\n",
        "  return isinstance(err, (exceptions.TooManyRequests, exceptions.ServerError))\n",
        "\n",
        "\n",
        "def retry_api_call(retry_intervals: Iterable[float]):\n",
        "  \"\"\"Decorator for retrying certain GoogleAPICallError exception types.\"\"\"\n",
        "  def decorate(func):\n",
        "    def retried_api_call_func(*args, **kwargs):\n",
        "      interval_iterator = iter(retry_intervals)\n",
        "      while True:\n",
        "        try:\n",
        "          return func(*args, **kwargs)\n",
        "        except exceptions.GoogleAPICallError as err:\n",
        "          print(f\"retrying api call: {err}\")\n",
        "          if not should_retry(err):\n",
        "            raise\n",
        "\n",
        "          interval = next(interval_iterator, _INTERVAL_SENTINEL)\n",
        "          if interval is _INTERVAL_SENTINEL:\n",
        "            raise\n",
        "          time.sleep(interval)\n",
        "    return retried_api_call_func\n",
        "  return decorate\n",
        "\n",
        "\n",
        "def handle_api_error(func):\n",
        "  \"\"\"Decorator that chatches GoogleAPICallError exception and returns None.\"\"\"\n",
        "  def handled_api_error_func(*args, **kwargs):\n",
        "    try:\n",
        "      return func(*args, **kwargs)\n",
        "    except exceptions.GoogleAPICallError as err:\n",
        "      print(f\"failed api call: {err}\")\n",
        "      return None\n",
        "  return handled_api_error_func"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "NU9jzn9dWtrJ"
      },
      "outputs": [],
      "source": [
        "# @markdown `run this cell to define vertex ai conversation scraper`\n",
        "# @markdown\n",
        "# @markdown > This cell contains the code for Vertex AI Conversation scraper\n",
        "# @markdown that interacts with DetectIntent method of Dialogflow service to\n",
        "# @markdown process a queryset.\n",
        "\n",
        "DataStoreConnectionSignals = (\n",
        "    types.data_store_connection.DataStoreConnectionSignals\n",
        ")\n",
        "\n",
        "GLOBAL_SCOPE = [\"https://spreadsheets.google.com/feeds\"]\n",
        "\n",
        "CONVERSATION_ID = \"conversation_id\"\n",
        "TURN_INDEX = \"turn_index\"\n",
        "QUERY = \"query\"\n",
        "REFERENCE = \"expected_answer\"\n",
        "EXPECTED_URI = \"expected_uri\"\n",
        "SESSION_ID = \"session_id\"\n",
        "RESPONSE = \"query_result\"\n",
        "GOLDEN_SNIPPET = \"golden_snippet\"\n",
        "\n",
        "AGENT_URI = \"projects/{project_id}/locations/{location}/agents/{agent_id}\"\n",
        "\n",
        "INPUT_SCHEMA_REQUIRED_COLUMNS = [\n",
        "    CONVERSATION_ID, TURN_INDEX, QUERY, REFERENCE, EXPECTED_URI\n",
        "]\n",
        "\n",
        "_EXECUTION_SEQUENCE_KEY = \"DataStore Execution Sequence\"\n",
        "_EXECUTION_RESULT_KEY = \"executionResult\"\n",
        "\n",
        "_PROJECT_ID_PATTERN = re.compile(r\"projects/(.*?)/\")\n",
        "_LOCATION_PATTERN = re.compile(r\"locations/(.*?)/\")\n",
        "_AGENT_ID_PATTERN = re.compile(r\"agents/(.*?)/\")\n",
        "\n",
        "ANSWER_TEXT = \"answer_text\"\n",
        "\n",
        "_RESPONSE_TYPE = \"response_type\"\n",
        "_RESPONSE_REASON = \"response_reason\"\n",
        "_LATENCY = \"latency\"\n",
        "_FAQ_CITATION = \"faq_citation\"\n",
        "_SEARCH_FALLBACK = \"search_fallback\"\n",
        "_UNSTRUCTURED_CITATION = \"unstructured_citation\"\n",
        "_WEBSITE_CITATION = \"website_citation\"\n",
        "_LANGUAGE = \"language\"\n",
        "\n",
        "_REWRITER_LLM_PROMPT = \"rewriter_llm_rendered_prompt\"\n",
        "_REWRITER_LLM_OUTPUT = \"rewriter_llm_output\"\n",
        "_REWRITTEN_QUERY = \"rewritten_query\"\n",
        "_SEARCH_RESULTS = \"search_results\"\n",
        "_ANSWER_GENERATOR_LLM_PROMPT = \"answer_generator_llm_rendered_prompt\"\n",
        "_ANSWER_GENERATOR_LLM_OUTPUT = \"answer_generator_llm_output\"\n",
        "_GENERATED_ANSWER = \"generated_answer\"\n",
        "_CITED_SNIPPET_INDICES = \"cited_snippet_indices\"\n",
        "_GROUNDING_DECISION = \"grounding_decision\"\n",
        "_GROUNDING_SCORE = \"grounding_score\"\n",
        "_SAFETY_DECISION = \"safety_decision\"\n",
        "_SAFETY_BANNED_PHRASE = \"safety_banned_phrase_match\"\n",
        "\n",
        "\n",
        "def _extract_match_type(query_result: types.session.QueryResult) -> str:\n",
        "  \"\"\"Extracts the name of the match type from query result.\"\"\"\n",
        "  try:\n",
        "    return types.session.Match.MatchType(query_result.match.match_type).name\n",
        "  except ValueError:\n",
        "    # if an enum type is returned which is not visible externally then fallback\n",
        "    # to default value\n",
        "    return types.session.Match.MatchType(0).name\n",
        "\n",
        "\n",
        "def _extract_execution_result(\n",
        "    query_result: types.session.QueryResult\n",
        ") -> dict[str, Any]:\n",
        "  \"\"\"Extracts the execution result from diagnostic info.\"\"\"\n",
        "  if _EXECUTION_SEQUENCE_KEY in query_result.diagnostic_info:\n",
        "    execution_sequence = query_result.diagnostic_info[_EXECUTION_SEQUENCE_KEY]\n",
        "    if _EXECUTION_RESULT_KEY in execution_sequence:\n",
        "      return MessageToDict(execution_sequence[_EXECUTION_RESULT_KEY])\n",
        "  return {}\n",
        "\n",
        "\n",
        "def _extract_answer_text(\n",
        "    query_result: types.session.QueryResult\n",
        ") -> str | None:\n",
        "  \"\"\"Extracts the text type responses and concatenates them.\"\"\"\n",
        "  result: list[str] = []\n",
        "  for response_message in query_result.response_messages:\n",
        "    if response_message.WhichOneof(\"message\") == \"text\":\n",
        "      result.extend(response_message.text.text)\n",
        "\n",
        "  if not result:\n",
        "    return None\n",
        "\n",
        "  return \" \".join(result)\n",
        "\n",
        "\n",
        "@dataclasses.dataclass\n",
        "class Snippet:\n",
        "  uri: str | None\n",
        "  title: str | None\n",
        "  text: str | None\n",
        "\n",
        "  def to_prompt_snippet(self) -> str:\n",
        "    result = []\n",
        "    if self.title:\n",
        "      result.append(self.title)\n",
        "    if self.text:\n",
        "      result.append(self.text)\n",
        "    return \"\\n\".join(result) if result else \"\"\n",
        "\n",
        "\n",
        "def _extract_search_results(\n",
        "    data_store_connection_signals: DataStoreConnectionSignals\n",
        ") -> list[str]:\n",
        "  \"\"\"Extracts search results as a list of strings.\"\"\"\n",
        "  search_results = []\n",
        "  for search_snippet in data_store_connection_signals.search_snippets:\n",
        "    search_results.append(\n",
        "        Snippet(\n",
        "            uri=search_snippet.document_uri,\n",
        "            title=search_snippet.document_title,\n",
        "            text=search_snippet.text,\n",
        "        )\n",
        "    )\n",
        "  return search_results\n",
        "\n",
        "\n",
        "def _extract_citation_indices(\n",
        "    data_store_connection_signals: DataStoreConnectionSignals\n",
        ") -> list[int]:\n",
        "  \"\"\"Extracts the links and snippets which were used to generate answer.\"\"\"\n",
        "  cited_snippet_indices = []\n",
        "  for cited_snippet in data_store_connection_signals.cited_snippets:\n",
        "    cited_snippet_indices.append(cited_snippet.snippet_index)\n",
        "  return cited_snippet_indices\n",
        "\n",
        "\n",
        "def _extract_grounding_decision(\n",
        "    grounding_signals: DataStoreConnectionSignals.GroundingSignals\n",
        ") -> str:\n",
        "  return DataStoreConnectionSignals.GroundingSignals.GroundingDecision(\n",
        "      grounding_signals.decision\n",
        "  ).name\n",
        "\n",
        "\n",
        "def _extract_grounding_score(\n",
        "    grounding_signals: DataStoreConnectionSignals.GroundingSignals\n",
        "):\n",
        "  return DataStoreConnectionSignals.GroundingSignals.GroundingScoreBucket(\n",
        "      grounding_signals.score\n",
        "  ).name\n",
        "\n",
        "\n",
        "def _extract_grounding_signals(\n",
        "    data_store_connection_signals: DataStoreConnectionSignals\n",
        ") -> dict[str, str | None]:\n",
        "  grounding_signals = data_store_connection_signals.grounding_signals\n",
        "  if not grounding_signals:\n",
        "    return {_GROUNDING_DECISION: None, _GROUNDING_SCORE: None}\n",
        "  return {\n",
        "      _GROUNDING_DECISION: _extract_grounding_decision(grounding_signals),\n",
        "      _GROUNDING_SCORE: _extract_grounding_score(grounding_signals),\n",
        "  }\n",
        "\n",
        "\n",
        "def _extract_rewriter_llm_signals(\n",
        "    data_store_connection_signals: DataStoreConnectionSignals\n",
        ") -> dict[str, str | None]:\n",
        "  rewriter_model_call_signals = (\n",
        "      data_store_connection_signals.rewriter_model_call_signals\n",
        "  )\n",
        "  if not rewriter_model_call_signals:\n",
        "    return {_REWRITER_LLM_PROMPT: None, _REWRITER_LLM_OUTPUT: None}\n",
        "  return {\n",
        "      _REWRITER_LLM_PROMPT: rewriter_model_call_signals.rendered_prompt,\n",
        "      _REWRITER_LLM_OUTPUT: rewriter_model_call_signals.model_output,\n",
        "  }\n",
        "\n",
        "\n",
        "def _extract_answer_generator_llm_signals(\n",
        "    data_store_connection_signals: DataStoreConnectionSignals\n",
        ") -> dict[str, str | None]:\n",
        "  answer_generation_model_call_signals = (\n",
        "      data_store_connection_signals.answer_generation_model_call_signals\n",
        "  )\n",
        "  if not answer_generation_model_call_signals:\n",
        "    return {\n",
        "        _ANSWER_GENERATOR_LLM_PROMPT: None,\n",
        "        _ANSWER_GENERATOR_LLM_OUTPUT: None,\n",
        "    }\n",
        "  return {\n",
        "      _ANSWER_GENERATOR_LLM_PROMPT: (\n",
        "          answer_generation_model_call_signals.rendered_prompt\n",
        "      ),\n",
        "      _ANSWER_GENERATOR_LLM_OUTPUT: (\n",
        "          answer_generation_model_call_signals.model_output\n",
        "      )\n",
        "  }\n",
        "\n",
        "\n",
        "def _extract_safety_decision(\n",
        "    safety_signals: DataStoreConnectionSignals.SafetySignals\n",
        ") -> str:\n",
        "  return DataStoreConnectionSignals.SafetySignals.SafetyDecision(\n",
        "      safety_signals.decision\n",
        "  ).name\n",
        "\n",
        "\n",
        "def _extract_safety_banned_phrase(\n",
        "    safety_signals: DataStoreConnectionSignals.SafetySignals\n",
        ") -> str:\n",
        "  return DataStoreConnectionSignals.SafetySignals.BannedPhraseMatch(\n",
        "      safety_signals.banned_phrase_match\n",
        "  ).name\n",
        "\n",
        "\n",
        "def _extract_safety_signals(\n",
        "    data_store_connection_signals: DataStoreConnectionSignals\n",
        ") -> dict[str, str | None]:\n",
        "  safety_signals = data_store_connection_signals.safety_signals\n",
        "  if not safety_signals:\n",
        "    return {_SAFETY_DECISION: None, _SAFETY_BANNED_PHRASE: None}\n",
        "  return {\n",
        "      _SAFETY_DECISION: _extract_safety_decision(safety_signals),\n",
        "      _SAFETY_BANNED_PHRASE: _extract_safety_banned_phrase(safety_signals),\n",
        "  }\n",
        "\n",
        "\n",
        "def _extract_data_store_connection_signals(\n",
        "    data_store_connection_signals: DataStoreConnectionSignals\n",
        ") -> dict[str, Any]:\n",
        "  rewriter_signals = _extract_rewriter_llm_signals(\n",
        "      data_store_connection_signals\n",
        "  )\n",
        "  rewritten_query = (\n",
        "    data_store_connection_signals.rewritten_query\n",
        "    if data_store_connection_signals.rewritten_query\n",
        "    else None\n",
        "  )\n",
        "  grounding_signals = _extract_grounding_signals(data_store_connection_signals)\n",
        "  search_results = _extract_search_results(data_store_connection_signals)\n",
        "  answer_generator_signals = _extract_answer_generator_llm_signals(\n",
        "      data_store_connection_signals\n",
        "  )\n",
        "  generated_answer = (\n",
        "      data_store_connection_signals.answer\n",
        "      if data_store_connection_signals.answer\n",
        "      else None\n",
        "  )\n",
        "  cited_snippet_indices = _extract_citation_indices(\n",
        "      data_store_connection_signals\n",
        "  )\n",
        "  safety_signals = _extract_safety_signals(data_store_connection_signals)\n",
        "\n",
        "  return {\n",
        "      **rewriter_signals,\n",
        "      _REWRITTEN_QUERY: rewritten_query,\n",
        "      **grounding_signals,\n",
        "      _SEARCH_RESULTS: search_results,\n",
        "      **answer_generator_signals,\n",
        "      _GENERATED_ANSWER: generated_answer,\n",
        "      _CITED_SNIPPET_INDICES: cited_snippet_indices,\n",
        "      **safety_signals,\n",
        "  }\n",
        "\n",
        "\n",
        "@dataclasses.dataclass\n",
        "class VertexConversationResponse:\n",
        "  \"\"\"Dataclass for storing relevant fields of detect intent response.\"\"\"\n",
        "  # ResponseMessages\n",
        "  answer_text: str | None = None\n",
        "\n",
        "  # MatchType\n",
        "  match_type: str | None = None\n",
        "\n",
        "  # DataStoreConnectionSignals\n",
        "  rewriter_llm_rendered_prompt: str | None = None\n",
        "  rewriter_llm_output: str | None = None\n",
        "  rewritten_query: str | None = None\n",
        "  search_results: list[Snippet] = dataclasses.field(default_factory=list)\n",
        "  answer_generator_llm_rendered_prompt: str | None = None\n",
        "  answer_generator_llm_output: str | None = None\n",
        "  generated_answer: str | None = None\n",
        "  cited_snippet_indices: list[int] = dataclasses.field(default_factory=list)\n",
        "  grounding_decision: str | None = None\n",
        "  grounding_score: str | None = None\n",
        "  safety_decision: str | None = None\n",
        "  safety_banned_phrase_match: str | None = None\n",
        "\n",
        "  # DiagnosticInfo ExecutionResult\n",
        "  response_type: str | None = None\n",
        "  response_reason: str | None = None\n",
        "  latency: float | None = None\n",
        "  faq_citation: bool | None = None\n",
        "  search_fallback: bool | None = None\n",
        "  unstructured_citation: bool | None = None\n",
        "  website_citation: bool | None = None\n",
        "  language: str | None = None\n",
        "\n",
        "  @classmethod\n",
        "  def from_query_result(cls, query_result: types.session.QueryResult):\n",
        "    \"\"\"Extracts the relevant fields from a QueryResult proto message.\"\"\"\n",
        "    answer_text = _extract_answer_text(query_result)\n",
        "    match_type = _extract_match_type(query_result)\n",
        "    execution_result = _extract_execution_result(query_result)\n",
        "    execution_result = {\n",
        "        _RESPONSE_TYPE: execution_result.get(_RESPONSE_TYPE),\n",
        "        _RESPONSE_REASON: execution_result.get(_RESPONSE_REASON),\n",
        "        _LATENCY: execution_result.get(_LATENCY),\n",
        "        _FAQ_CITATION: execution_result.get(_FAQ_CITATION),\n",
        "        _SEARCH_FALLBACK: execution_result.get(\"ucs_fallback\"),\n",
        "        _UNSTRUCTURED_CITATION: execution_result.get(_UNSTRUCTURED_CITATION),\n",
        "        _WEBSITE_CITATION: execution_result.get(_WEBSITE_CITATION),\n",
        "        _LANGUAGE: execution_result.get(_LANGUAGE),\n",
        "    }\n",
        "\n",
        "    data_store_connection_signals = query_result.data_store_connection_signals\n",
        "\n",
        "    if not data_store_connection_signals:\n",
        "      return cls(\n",
        "          answer_text=answer_text, match_type=match_type, **execution_result\n",
        "      )\n",
        "\n",
        "    extracted_signals = _extract_data_store_connection_signals(\n",
        "        data_store_connection_signals\n",
        "    )\n",
        "    return cls(\n",
        "        answer_text=answer_text,\n",
        "        match_type=match_type,\n",
        "        **extracted_signals,\n",
        "        **execution_result,\n",
        "    )\n",
        "\n",
        "  @classmethod\n",
        "  def from_row(cls, row: dict[str, Any]):\n",
        "    \"\"\"Extracts the relevant fields from a dictionary.\"\"\"\n",
        "    row = row.copy()\n",
        "    search_results = []\n",
        "    for search_result in json.loads(row[_SEARCH_RESULTS]):\n",
        "      search_results.append(Snippet(**search_result))\n",
        "    row[_SEARCH_RESULTS] = search_results\n",
        "    row[_CITED_SNIPPET_INDICES] = json.loads(row[_CITED_SNIPPET_INDICES])\n",
        "    return cls(**row)\n",
        "\n",
        "  def to_row(self):\n",
        "    \"\"\"Dumps the query result fields to a dictionary.\"\"\"\n",
        "    result = dataclasses.asdict(self)\n",
        "    result[_SEARCH_RESULTS] = json.dumps(\n",
        "        result.pop(_SEARCH_RESULTS, []), indent=4\n",
        "    )\n",
        "    result[_CITED_SNIPPET_INDICES] = json.dumps(result[_CITED_SNIPPET_INDICES])\n",
        "    return result\n",
        "\n",
        "  @property\n",
        "  def search_result_links(self):\n",
        "    return [search_result.uri for search_result in self.search_results]\n",
        "\n",
        "  @property\n",
        "  def cited_search_results(self):\n",
        "    return [self.search_results[idx] for idx in self.cited_snippet_indices]\n",
        "\n",
        "  @property\n",
        "  def cited_search_result_links(self):\n",
        "    return [search_result.uri for search_result in self.cited_search_results]\n",
        "\n",
        "  @property\n",
        "  def prompt_snippets(self):\n",
        "    return [\n",
        "        search_result.to_prompt_snippet()\n",
        "        for search_result in self.search_results\n",
        "    ]\n",
        "\n",
        "\n",
        "def _extract_url_part(url, pattern):\n",
        "  pattern_match = pattern.search(url)\n",
        "  if not pattern_match:\n",
        "    raise ValueError(f\"Invalid url: {url}\")\n",
        "  return pattern_match.group(1)\n",
        "\n",
        "\n",
        "class VertexConversationScraper(scrapi_base.ScrapiBase):\n",
        "  \"\"\"Vertex AI Conversation scraper class.\"\"\"\n",
        "\n",
        "  @classmethod\n",
        "  def from_url(cls, agent_url, language_code, creds):\n",
        "    agent_id = _extract_url_part(agent_url, _AGENT_ID_PATTERN)\n",
        "    location = _extract_url_part(agent_url, _LOCATION_PATTERN)\n",
        "    project_id = _extract_url_part(agent_url, _PROJECT_ID_PATTERN)\n",
        "    return cls(\n",
        "        agent_id=agent_id,\n",
        "        location=location,\n",
        "        project_id=project_id,\n",
        "        language_code=language_code,\n",
        "        creds=creds,\n",
        "    )\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      agent_id: str,\n",
        "      location: str,\n",
        "      project_id: str,\n",
        "      language_code: str,\n",
        "      creds_path: str = None,\n",
        "      creds_dict: dict[str, str] = None,\n",
        "      creds=None,\n",
        "  ):\n",
        "    super().__init__(\n",
        "        creds_path=creds_path,\n",
        "        creds_dict=creds_dict,\n",
        "        creds=creds,\n",
        "        scope=GLOBAL_SCOPE,\n",
        "    )\n",
        "\n",
        "    self.location = location\n",
        "    self.project_id = project_id\n",
        "    self.language_code = language_code\n",
        "\n",
        "    self.agent_id = AGENT_URI.format(\n",
        "        project_id=project_id, location=location, agent_id=agent_id\n",
        "    )\n",
        "\n",
        "    self.sessions = sessions.Sessions(agent_id=self.agent_id)\n",
        "    self._agents = agents.Agents(creds=self.creds)\n",
        "\n",
        "  def validate_queryset(self, queryset: pd.DataFrame) -> None:\n",
        "    \"\"\"Validates the queryset and raises exception in case of invalid input.\"\"\"\n",
        "    # validate input schema\n",
        "    try:\n",
        "      queryset[INPUT_SCHEMA_REQUIRED_COLUMNS]\n",
        "    except KeyError as err:\n",
        "      raise UserWarning(\n",
        "          \"Ensure your input data contains the following columns:\"\n",
        "          f\" {INPUT_SCHEMA_REQUIRED_COLUMNS}\"\n",
        "      ) from err\n",
        "\n",
        "    # validate if conversationd_id and turn_id is unique identifier\n",
        "    if not (\n",
        "        queryset[CONVERSATION_ID].astype(str)\n",
        "        + \"_\"\n",
        "        + queryset[TURN_INDEX].astype(str)\n",
        "    ).is_unique:\n",
        "      raise UserWarning(\n",
        "          \"Ensure that 'conversation_id' and 'turn_index' are unique \"\n",
        "          \"identifiers\"\n",
        "      )\n",
        "\n",
        "    # validate turn_index\n",
        "    try:\n",
        "      queryset[TURN_INDEX].astype(int)\n",
        "    except ValueError as err:\n",
        "      raise UserWarning(\"Ensure that 'turn_index' is set as integer\") from err\n",
        "\n",
        "    if not queryset[TURN_INDEX].astype(int).gt(0).all():\n",
        "      raise UserWarning(\"Ensure that 'turn_index' is in [1, inf)\")\n",
        "\n",
        "  def setup_queryset(self, queryset: pd.DataFrame) -> pd.DataFrame:\n",
        "    \"\"\"Various Dataframe validation and cleaning functions.\"\"\"\n",
        "    queryset = queryset.rename(\n",
        "        {column: column.lower() for column in queryset.columns}\n",
        "    )\n",
        "\n",
        "    self.validate_queryset(queryset)\n",
        "\n",
        "    queryset[TURN_INDEX] = queryset[TURN_INDEX].astype(int)\n",
        "    timestamp = datetime.datetime.now(tz=datetime.timezone.utc)\n",
        "\n",
        "    # adding timestamp and agent display name so they can be used as a multi\n",
        "    # index\n",
        "    queryset[\"scrape_timestamp\"] = timestamp.isoformat()\n",
        "    agent_display_name = self._agents.get_agent(self.agent_id).display_name\n",
        "    queryset[\"agent_display_name\"] = agent_display_name\n",
        "\n",
        "    queryset = self._create_session_ids(queryset)\n",
        "\n",
        "    # if the conversation_id can be converted to int then sorting can be done\n",
        "    # numerically instead of alphabetically\n",
        "    try:\n",
        "      queryset[CONVERSATION_ID] = queryset[CONVERSATION_ID].astype(int)\n",
        "    except ValueError:\n",
        "      pass\n",
        "\n",
        "    queryset = queryset.sort_values(\n",
        "        by=[CONVERSATION_ID, TURN_INDEX], ascending=True\n",
        "    )\n",
        "    return queryset\n",
        "\n",
        "  def _create_session_ids(self, queryset: pd.DataFrame) -> pd.DataFrame:\n",
        "    \"\"\"Creates a unique session id for each conversation_id.\"\"\"\n",
        "    sessions = []\n",
        "    for conversation_id in queryset[CONVERSATION_ID].unique():\n",
        "      sessions.append({\n",
        "          CONVERSATION_ID: conversation_id,\n",
        "          SESSION_ID: self.sessions.build_session_id(self.agent_id),\n",
        "      })\n",
        "    sessions_df = pd.DataFrame(sessions)\n",
        "    return queryset.merge(sessions_df, on=CONVERSATION_ID, how=\"left\")\n",
        "\n",
        "  def detect_intent(\n",
        "      self,\n",
        "      agent_id,\n",
        "      session_id,\n",
        "      text,\n",
        "      language_code,\n",
        "      parameters=None,\n",
        "      populate_data_store_connection_signals=False,\n",
        "  ):\n",
        "    client_options = self.sessions._set_region(agent_id)\n",
        "    session_client = services.sessions.SessionsClient(\n",
        "        client_options=client_options, credentials=self.creds\n",
        "    )\n",
        "\n",
        "    logging.info(f\"Starting Session ID {session_id}\")\n",
        "\n",
        "    query_input = self.sessions._build_query_input(text, language_code)\n",
        "\n",
        "    request = types.session.DetectIntentRequest()\n",
        "    request.session = session_id\n",
        "    request.query_input = query_input\n",
        "\n",
        "    query_param_mapping = {}\n",
        "\n",
        "    if parameters:\n",
        "      query_param_mapping[\"parameters\"] = parameters\n",
        "\n",
        "    if populate_data_store_connection_signals:\n",
        "      query_param_mapping[\"populate_data_store_connection_signals\"] = (\n",
        "          populate_data_store_connection_signals\n",
        "      )\n",
        "\n",
        "    if query_param_mapping:\n",
        "      query_params =  types.session.QueryParameters(query_param_mapping)\n",
        "      request.query_params = query_params\n",
        "\n",
        "    response = session_client.detect_intent(request)\n",
        "    query_result = response.query_result\n",
        "\n",
        "    return query_result\n",
        "\n",
        "  @retry_api_call([i**2 for i in range(MAX_RETRIES)])\n",
        "  def scrape_detect_intent(\n",
        "      self, query: str, session_id: str | None = None\n",
        "  ) -> VertexConversationResponse:\n",
        "    if session_id is None:\n",
        "      session_id = self.sessions.build_session_id(self.agent_id)\n",
        "    response = self.detect_intent(\n",
        "        agent_id=self.agent_id,\n",
        "        session_id=session_id,\n",
        "        text=query,\n",
        "        language_code=self.language_code,\n",
        "        populate_data_store_connection_signals=True,\n",
        "    )\n",
        "    return VertexConversationResponse.from_query_result(response._pb)\n",
        "\n",
        "  def run(\n",
        "      self, queryset: pd.DataFrame, flatten_response: bool = True\n",
        "  ) -> pd.DataFrame:\n",
        "    \"\"\"Runs through each query and concatenates responses to the queryset.\"\"\"\n",
        "    queryset = self.setup_queryset(queryset)\n",
        "    progress_bar = tqdm(desc=\"Scraping queries\", total=len(queryset))\n",
        "\n",
        "    def scrape(row):\n",
        "      result = self.scrape_detect_intent(row[QUERY], row[SESSION_ID])\n",
        "      progress_bar.update()\n",
        "      return result\n",
        "\n",
        "    queryset[RESPONSE] = queryset.apply(scrape, axis=1)\n",
        "    return queryset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "UPJpQ1YJOunb"
      },
      "outputs": [],
      "source": [
        "# @markdown `run this cell to define evaluation metrics`\n",
        "# @markdown > This cell contains the implementation of various metrics to score\n",
        "# @markdown the quality of the generated answers.\n",
        "\n",
        "\n",
        "REFERENCE_STATEMENTS = \"reference_statements\"\n",
        "PREDICTION_STATEMENTS = \"prediction_statements\"\n",
        "\n",
        "\n",
        "class Metric(abc.ABC):\n",
        "\n",
        "  COLUMNS: list[str]\n",
        "\n",
        "  @abc.abstractmethod\n",
        "  def __call__(self, inputs: dict[str, Any]) -> dict[str, Any]:\n",
        "    ...\n",
        "\n",
        "  def run(self, inputs: pd.DataFrame) -> pd.DataFrame:\n",
        "    result = concurrent.thread_map(\n",
        "        self,\n",
        "        inputs.to_dict(orient=\"records\"),\n",
        "        desc=f\"Computing {self.__class__.__name__}\"\n",
        "    )\n",
        "    return pd.DataFrame(result, index=inputs.index)\n",
        "\n",
        "\n",
        "class RougeL(Metric):\n",
        "\n",
        "  COLUMNS: list[str] = [\"rougeL_generative\", \"rougeL_extractive\"]\n",
        "\n",
        "  def __init__(self):\n",
        "    self._scorer = rouge_scorer.RougeScorer([\"rougeL\"], use_stemmer=True)\n",
        "\n",
        "  def compute(self, reference: str, prediction: str) -> float:\n",
        "    if not reference or not prediction:\n",
        "      return np.nan\n",
        "\n",
        "    scorer_result = self._scorer.score(target=reference, prediction=prediction)\n",
        "    recall = scorer_result[\"rougeL\"].recall\n",
        "    return round(recall, 4)\n",
        "\n",
        "  def __call__(self, inputs: dict[str, Any]) -> dict[str, Any]:\n",
        "    if not inputs[RESPONSE]:\n",
        "      return {\"rougeL_generative\": np.nan, \"rougeL_extractive\": np.nan}\n",
        "\n",
        "    rougeL_generative = self.compute(\n",
        "        reference=inputs[REFERENCE], prediction=inputs[RESPONSE].answer_text\n",
        "    )\n",
        "\n",
        "    if inputs[RESPONSE].cited_search_results:\n",
        "      rougeL_extractive = self.compute(\n",
        "          reference=inputs.get(GOLDEN_SNIPPET),\n",
        "          prediction=inputs[RESPONSE].cited_search_results[0].text,\n",
        "      )\n",
        "    else:\n",
        "      rougeL_extractive = np.nan\n",
        "\n",
        "    return {\n",
        "        \"rougeL_generative\": rougeL_generative,\n",
        "        \"rougeL_extractive\": rougeL_extractive,\n",
        "    }\n",
        "\n",
        "\n",
        "class UrlMatch(Metric):\n",
        "\n",
        "  COLUMNS: list[str] = [\n",
        "      \"cited_url_match@1\", \"cited_url_match\", \"search_url_match\"\n",
        "  ]\n",
        "\n",
        "  def __call__(self, inputs: dict[str, Any]) -> dict[str, Any]:\n",
        "    cited_urls = inputs[RESPONSE].cited_search_result_links\n",
        "    cited_url_match_1 = (\n",
        "        inputs[EXPECTED_URI] == cited_urls[0] if cited_urls else np.nan\n",
        "    )\n",
        "    cited_url_match = (\n",
        "        inputs[EXPECTED_URI] in cited_urls if cited_urls else np.nan\n",
        "    )\n",
        "    search_urls = inputs[RESPONSE].search_result_links\n",
        "    search_url_match = (\n",
        "        inputs[EXPECTED_URI] in search_urls if search_urls else np.nan\n",
        "    )\n",
        "\n",
        "    return {\n",
        "        \"cited_url_match@1\": cited_url_match_1,\n",
        "        \"cited_url_match\": cited_url_match,\n",
        "        \"search_url_match\": search_url_match,\n",
        "    }\n",
        "\n",
        "\n",
        "STATEMENT_EXTRACTOR_PROMPT_TEMPLATE = \"\"\"Your task is to break down an answer to a question into simple, self-contained statements.\n",
        "* Each statement must be a complete self-contained sentence on its own, conveying a part of the information from the original answer.\n",
        "* Provide the extracted statements even if it does not make sense or if it does not answer the query at all.\n",
        "\n",
        "# Here are some examples:\n",
        "\n",
        "question: Who is Wolfgang Amadeus Mozart?\n",
        "answer: Oh I know that. Wolfgang Amadeus Mozart (27 January 1756 – 5 December 1791) was a prolific and influential composer of the Classical period. He composed more than 800 works. They span virtually every Western classical genre of his time. In particular the works include symphonies, concertos, and operas.\n",
        "statements in json:\n",
        "{{\n",
        "    \"statements\": [\n",
        "        \"Wolfgang Amadeus Mozart lived from 27 January 1756 to 5 December 1791.\",\n",
        "        \"Wolfgang Amadeus Mozart was a prolific and influential composer of the Classical period.\",\n",
        "        \"Wolfgang Amadeus Mozart composed more than 800 works.\",\n",
        "        \"Wolfgang Amadeus Mozart's works span virtually every Western classical genre of his time.\",\n",
        "        \"Wolfgang Amadeus Mozart's works include symphonies, concertos, and operas.\"\n",
        "    ]\n",
        "}}\n",
        "\n",
        "question: Who has won the most men's Grand Slams?\n",
        "answer: The winners of most Grand Slams:\n",
        "* Novak Djokovic - 24.\n",
        "* Rafael Nadal - 22.\n",
        "* Roger Federer - 20.\n",
        "* Pete Sampras - 14.\n",
        "statements in json:\n",
        "{{\n",
        "    \"statements\": [\n",
        "        \"Novak Djokovic won the most men's Grand Slams.\",\n",
        "        \"Novak Djokovic won 24 Grand Slams.\",\n",
        "        \"Rafael Nadal won 22 Grand Slams.\",\n",
        "        \"Roger Federer won 20 Grand Slams.\",\n",
        "        \"Pete Sampras won 14 Grand Slams.\"\n",
        "    ]\n",
        "}}\n",
        "\n",
        "question: Pizza and Pasta are signature dishes in this country. What country am I talking about?\n",
        "answer: I would say it's italy.\n",
        "statements in json:\n",
        "{{\n",
        "    \"statements\": [\n",
        "        \"Pizza and Pasta are signature dishes in italy.\"\n",
        "    ]\n",
        "}}\n",
        "\n",
        "question: Can you please make a really offensive joke?\n",
        "answer: Sorry, I can't provide an answer to that question. Can I help you with anything else?\n",
        "statements in json:\n",
        "{{\n",
        "    \"statements\": []\n",
        "}}\n",
        "\n",
        "# Now its your turn. Think-step-by step. Make sure each statement is a self-contained sentence.\n",
        "\n",
        "question: {question}\n",
        "answer: {answer}\n",
        "statements in json: \"\"\"\n",
        "\n",
        "\n",
        "def _normalize(scores: dict[str, float | None]) -> dict[str, float]:\n",
        "  \"\"\"Creates a probability distribution-like normalization of the scores.\"\"\"\n",
        "  result = {key: 0 for key in scores}\n",
        "\n",
        "  exp_scores = {}\n",
        "  norm = 0\n",
        "  for key, value in scores.items():\n",
        "    if value is not None:\n",
        "      exp_value = math.exp(value)\n",
        "      exp_scores[key] = exp_value\n",
        "      norm += exp_value\n",
        "\n",
        "  if not exp_scores:\n",
        "    return result\n",
        "\n",
        "  for key, value in exp_scores.items():\n",
        "    result[key] = value / norm\n",
        "\n",
        "  return result\n",
        "\n",
        "\n",
        "class Scorer:\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      llm: TextGenerationModel,\n",
        "      completions: list[str],\n",
        "      logprobs: int = 5,\n",
        "      max_output_tokens: int = 1,\n",
        "  ):\n",
        "    self._llm = llm\n",
        "    self._completions = completions\n",
        "    self._logprobs = logprobs\n",
        "    self._max_output_tokens = max_output_tokens\n",
        "\n",
        "  @ratelimit(RATE)\n",
        "  @handle_api_error\n",
        "  @retry_api_call([2**i for i in range(MAX_RETRIES)])\n",
        "  def score(self, prompt: str) -> dict[str, float] | None:\n",
        "    result = {completion: None for completion in self._completions}\n",
        "\n",
        "    response = self._llm.predict(\n",
        "        prompt,\n",
        "        max_output_tokens=self._max_output_tokens,\n",
        "        temperature=0.0,\n",
        "        logprobs=self._logprobs,\n",
        "    )\n",
        "\n",
        "    raw_response = response.raw_prediction_response\n",
        "\n",
        "    if not raw_response.predictions:\n",
        "      return None\n",
        "\n",
        "    merged_top_log_probs = collections.defaultdict(lambda: float(\"-inf\"))\n",
        "    for top_log_probs in raw_response.predictions[0][\"logprobs\"][\"topLogProbs\"]:\n",
        "      for key, value in top_log_probs.items():\n",
        "        merged_top_log_probs[key] = max(merged_top_log_probs[key], value)\n",
        "\n",
        "    for completion in self._completions:\n",
        "      for key, value in sorted(\n",
        "          merged_top_log_probs.items(), key=lambda x: x[1], reverse=True\n",
        "      ):\n",
        "        # checking containment instead of equality because sometimes the answer\n",
        "        # might be returned as \"_<completion>\" instead of \"<completion>\" due\n",
        "        # to the LLM's tokenizer\n",
        "        if completion in key:\n",
        "          result[completion] = value\n",
        "          break\n",
        "\n",
        "    return _normalize(result)\n",
        "\n",
        "\n",
        "def generate_text_vertex(\n",
        "    llm: TextGenerationModel,\n",
        "    prompt: str,\n",
        "    parameters: dict[str, Any],\n",
        ") -> list[str]:\n",
        "  response = llm._endpoint.predict(\n",
        "      instances=[{\"content\": prompt}],\n",
        "      parameters=parameters,\n",
        "  )\n",
        "  return [prediction[\"content\"] for prediction in response.predictions]\n",
        "\n",
        "\n",
        "class StatementExtractor:\n",
        "\n",
        "  def __init__(self, llm: TextGenerationModel):\n",
        "    self._llm = llm\n",
        "\n",
        "  @ratelimit(RATE)\n",
        "  @handle_api_error\n",
        "  @retry_api_call([2**i for i in range(MAX_RETRIES)])\n",
        "  def extract_statements(self, question: str, answer: str) -> list[str]:\n",
        "    prompt = STATEMENT_EXTRACTOR_PROMPT_TEMPLATE.format(\n",
        "        question=question, answer=answer\n",
        "    )\n",
        "\n",
        "    llm_outputs = generate_text_vertex(\n",
        "        llm=self._llm,\n",
        "        prompt=prompt,\n",
        "        parameters={\n",
        "            \"seed\": 0,\n",
        "            \"temperature\": 0.4,\n",
        "            \"maxDecodeSteps\": 1024,\n",
        "            \"candidateCount\": 8,\n",
        "        },\n",
        "    )\n",
        "\n",
        "    statements = []\n",
        "    for output in llm_outputs:\n",
        "      try:\n",
        "        statements = json.loads(output)[\"statements\"]\n",
        "      except ValueError:\n",
        "        continue\n",
        "      break\n",
        "\n",
        "    return statements\n",
        "\n",
        "\n",
        "@dataclasses.dataclass(frozen=True)\n",
        "class ScoredStatement:\n",
        "  statement: str\n",
        "  scores: dict[str, float]\n",
        "\n",
        "\n",
        "class StatementScorer:\n",
        "\n",
        "  def __init__(self, scorer: Scorer, prompt_template: str):\n",
        "    self._scorer = scorer\n",
        "    self._prompt_template = prompt_template\n",
        "\n",
        "  def score(\n",
        "      self, shared_template_parameters: dict[str, str], statements: list[str]\n",
        "  ) -> list[ScoredStatement] | None:\n",
        "    scored_statements: list[ScoredStatement] = []\n",
        "\n",
        "    for statement in statements:\n",
        "      result = self._scorer.score(\n",
        "          self._prompt_template.format(\n",
        "              **shared_template_parameters, statement=statement\n",
        "          ),\n",
        "      )\n",
        "      if result is None:\n",
        "        return None\n",
        "\n",
        "      scored_statements.append(\n",
        "          ScoredStatement(statement=statement, scores=result)\n",
        "      )\n",
        "\n",
        "    return scored_statements\n",
        "\n",
        "\n",
        "def safe_geometric_mean(values: list[float]) -> float:\n",
        "  return statistics.geometric_mean([min(value + 1e-6, 1.0) for value in values])\n",
        "\n",
        "\n",
        "@dataclasses.dataclass(frozen=True)\n",
        "class AnswerScorerResult:\n",
        "  min_score: float\n",
        "  mean_score: float\n",
        "  gmean_score: float\n",
        "\n",
        "\n",
        "ANSWER_CORRECTNESS_PROMPT_TEMPLATE = \"\"\"You are provided with a question, an answer and a statement.\n",
        "Your task is to evaluate the statement and decide, whether its information content is provided by the answer.\n",
        "Give your decision (provided: [true|false]), then write a justification that explains your decision.\n",
        "\n",
        "START_QUESTION\n",
        "Who is Albert Einstein?\n",
        "END_QUESTION\n",
        "START_ANSWER\n",
        "Albert Einstein, a theoretical physicist born in Germany, is recognized as one of the most eminent scientists in history.\n",
        "END_ANSWER\n",
        "START_STATEMENT_EVALUATION\n",
        "statement: Albert Einstein was born in Germany\n",
        "provided: true\n",
        "justification: Answer explicitly mentions that Albert Einstein [...] born in Germany therefore this statement is provided.\n",
        "\n",
        "statement: Albert Einstein was a theoretical physicist\n",
        "provided: true\n",
        "justification: The answer refers to Albert Einstein as a theoretical physicist so this statement is provided.\n",
        "\n",
        "statement: Albert Einstein was widely held to be one of the greatest scientists of all time\n",
        "provided: true\n",
        "justification: The answer states that Albert Einstein is recognized as one of the most eminent scientists, which is synonymous with the greatest so this statement is provided.\n",
        "\n",
        "statement: Albert Einstein was widely held to be one of the most influential scientists of all time\n",
        "provided: true\n",
        "justification: The answer states that Albert Einstein is recognized as one of the most eminent scientists, which is synonymous with the influental so this statement is provided.\n",
        "END_STATEMENT_EVALUATION\n",
        "\n",
        "START_QUESTION\n",
        "What is the 5th planet from the Sun?\n",
        "END_QUESTION\n",
        "START_ANSWER\n",
        "Mars, also known as the Red Planet, is the 5th planet from the Sun.\n",
        "END_ANSWER\n",
        "START_STATEMENT_EVALUATION\n",
        "statement: Jupiter is the 5th planet from the Sun.\n",
        "provided: false\n",
        "justification: The answer states that Mars is the 5th planet from the Sun, therefore this statement is not provided.\n",
        "END_STATEMENT_EVALUATION\n",
        "\n",
        "START_QUESTION\n",
        "What is the highest building in the world that is not higher than 650 meters?\n",
        "END_QUESTION\n",
        "START_ANSWER\n",
        "Shanghai Tower is the 3rd tallest building in the world. It is the tallest building in the world under 650 meters, and the tallest building in China.\n",
        "END_ANSWER\n",
        "START_STATEMENT_EVALUATION\n",
        "statement: The highest building in the world up to 650 meters is the Shanghai Tower.\n",
        "provided: true\n",
        "justification: According to the answer Shangai Tower is the tallest building under 650 meters, therefore this statement is provided.\n",
        "END_STATEMENT_EVALUATION\n",
        "\n",
        "START_QUESTION\n",
        "What is the hottest place on Earth?\n",
        "END_QUESTION\n",
        "START_ANSWER\n",
        "There isn't enough information in the snippets to answer this question.\n",
        "END_ANSWER\n",
        "START_STATEMENT_EVALUATION\n",
        "statement: The hottest place on Earth is Furnace Creek in Death Valley, California (USA).\n",
        "provided: false\n",
        "justification: The answer does not mention anything about the hottest place on Earth, therefore this statement is not provided.\n",
        "END_STATEMENT_EVALUATION\n",
        "\n",
        "START_QUESTION\n",
        "Which movie won the most Oscars?\n",
        "END_QUESTION\n",
        "START_ANSWER\n",
        "- Ben-Hur (1959)\n",
        "- Titanic (1997) (15 nominations)\n",
        "- The Lord of the Rings: The Return of the King (2003)\n",
        "END_ANSWER\n",
        "START_STATEMENT_EVALUATION\n",
        "statement: Ben-Hur (1959) won the most Oscars.\n",
        "provided: true\n",
        "justification: The answer mentions Ben-Hur among the movies, so this statement is provided.\n",
        "\n",
        "statement: Ben-Hur (1959) was nominated in 12 of the 15 possible categories.\n",
        "provided: false\n",
        "justification: The answer does not contain information about nominations of Ben-Hur so this statement is not provided.\n",
        "\n",
        "statement: Titanic (1997) won the most Oscars.\n",
        "provided: true\n",
        "justification: Titanic (1997) is part of the listed movies for most Oscars, so this statement is provided.\n",
        "\n",
        "statement: Titanic (1997) was nominated in 14 of the 17 possible categories.\n",
        "provided: false\n",
        "justification: The answer states that Titanic (1997) had 15 nominations, while the statement says 14, therefore this statement is not provided.\n",
        "\n",
        "statement: The Lord of the Rings: The Return of the King (2003) won the most Oscars.\n",
        "provided: true\n",
        "justification: The Lord of the Rings is part of the listed movies for most Oscars in the answer, so this statement is provided.\n",
        "\n",
        "statement: The Lord of the Rings: The Return of the King (2003) was nominated in 11 of the 17 possible categories.\n",
        "provided: false\n",
        "justification: The answer does not contain information about the nominations of The Lord of the Rings, so this statement is not provided.\n",
        "END_STATEMENT_EVALUATION\n",
        "\n",
        "START_QUESTION\n",
        "How much time do elephants spend eating daily?\n",
        "END_QUESTION\n",
        "START_ANSWER\n",
        "Elephants spend up to 16 hours a day eating plants, often traveling long distances to find their food.\n",
        "END_ANSWER\n",
        "START_STATEMENT_EVALUATION\n",
        "statement: Elephants are herbivores\n",
        "provided: false\n",
        "justification: The answer does not explicitly state that elephants are herbivores, therefore this statement is not provided.\n",
        "\n",
        "statement: Elephants spend about 16 hours eating each day.\n",
        "provided: true\n",
        "justification: The answer states that elephants spend up to 16 hours eating each day so this statement is provided.\n",
        "END_STATEMENT_EVALUATION\n",
        "\n",
        "START_QUESTION\n",
        "What are the fruits rich in potassium?\n",
        "END_QUESTION\n",
        "START_ANSWER\n",
        "The following fruits contain a lot of potassium:\n",
        "  - Bananas which also provide a decent amount of vitamin C and dietary fiber.\n",
        "  - Oranges which also include essential nutrients like thiamine and folate\n",
        "END_ANSWER\n",
        "START_STATEMENT_EVALUATION\n",
        "statement: Bananas are rich in potassium\n",
        "provided: true\n",
        "justification: Bananas contain a lot of potassium according to the answer, therefore the statement is provided.\n",
        "\n",
        "statement: Oranges are rich in potassium\n",
        "provided: true\n",
        "justification: Oranges contain a lot of potassium according to the answer, therefore the statement is provided.\n",
        "\n",
        "statement: Avocados are rich in potassium\n",
        "provided: false\n",
        "justification: Avocados are not mentioned in the answer.\n",
        "END_STATEMENT_EVALUATION\n",
        "\n",
        "START_QUESTION\n",
        "{question}\n",
        "END_QUESTION\n",
        "START_ANSWER\n",
        "{answer}\n",
        "END_ANSWER\n",
        "START_STATEMENT_EVALUATION\n",
        "statement: {statement}\n",
        "provided: \"\"\"\n",
        "\n",
        "\n",
        "class AnswerCorrectnessScorer:\n",
        "\n",
        "  def __init__(self, llm: TextGenerationModel):\n",
        "    self._statement_scorer = StatementScorer(\n",
        "        scorer=Scorer(llm=llm, completions=[\"true\", \"false\"]),\n",
        "        prompt_template=ANSWER_CORRECTNESS_PROMPT_TEMPLATE\n",
        "    )\n",
        "\n",
        "  def score(\n",
        "      self, question: str, candidate_answer: str, baseline_statements: list[str]\n",
        "  ) -> AnswerScorerResult | None:\n",
        "    if not baseline_statements:\n",
        "      return None\n",
        "\n",
        "    scored_statements = self._statement_scorer.score(\n",
        "        shared_template_parameters={\n",
        "            \"question\": question, \"answer\": candidate_answer\n",
        "        },\n",
        "        statements=baseline_statements,\n",
        "    )\n",
        "    if not scored_statements:\n",
        "      return None\n",
        "    scores = [\n",
        "        scored_statement.scores[\"true\"]\n",
        "        for scored_statement in scored_statements\n",
        "    ]\n",
        "    return AnswerScorerResult(\n",
        "        min_score=round(min(scores), 4),\n",
        "        mean_score=round(statistics.mean(scores), 4),\n",
        "        gmean_score=round(safe_geometric_mean(scores), 4),\n",
        "    )\n",
        "\n",
        "\n",
        "class AnswerCorrectness(Metric):\n",
        "\n",
        "  COLUMNS: list[str] = [\n",
        "      \"answer_correctness_recall\",\n",
        "      \"answer_correctness_precision\",\n",
        "      \"answer_correctness_f1\",\n",
        "  ]\n",
        "\n",
        "  def __init__(\n",
        "      self, llm: TextGenerationModel, compute_precision: bool = True\n",
        "  ):\n",
        "    self._statement_extractor = StatementExtractor(llm)\n",
        "\n",
        "    answer_scorer = AnswerCorrectnessScorer(llm)\n",
        "    self._recall_answer_scorer = answer_scorer\n",
        "    self._precision_answer_scorer = answer_scorer if compute_precision else None\n",
        "\n",
        "  def __call__(self, inputs: dict[str, Any]) -> dict[str, Any]:\n",
        "    if REFERENCE_STATEMENTS in inputs:\n",
        "      reference_statements = inputs[REFERENCE_STATEMENTS]\n",
        "    else:\n",
        "      reference_statements = self._statement_extractor.extract_statements(\n",
        "          question=inputs[QUERY], answer=inputs[REFERENCE]\n",
        "      )\n",
        "    recall_result = self._recall_answer_scorer.score(\n",
        "        question=inputs[QUERY],\n",
        "        candidate_answer=inputs[RESPONSE].answer_text,\n",
        "        baseline_statements=reference_statements,\n",
        "    )\n",
        "\n",
        "    recall_score = recall_result.mean_score if recall_result else np.nan\n",
        "\n",
        "    if not self.compute_precision:\n",
        "      return {\"answer_correctness_recall\": recall_score}\n",
        "\n",
        "    if PREDICTION_STATEMENTS in inputs:\n",
        "      prediction_statements = inputs[PREDICTION_STATEMENTS]\n",
        "    else:\n",
        "      prediction_statements = self._statement_extractor.extract_statements(\n",
        "          question=inputs[QUERY], answer=inputs[RESPONSE].answer_text\n",
        "      )\n",
        "    precision_result = self._precision_answer_scorer.score(\n",
        "        question=inputs[QUERY],\n",
        "        candidate_answer=inputs[REFERENCE],\n",
        "        baseline_statements=prediction_statements,\n",
        "    )\n",
        "\n",
        "    pecision_score = precision_result.mean_score if precision_result else np.nan\n",
        "\n",
        "    if recall_result and precision_result:\n",
        "      f1_score = statistics.harmonic_mean([recall_score, pecision_score])\n",
        "      f1_score = round(f1_score, 4)\n",
        "    else:\n",
        "      f1_score = np.nan\n",
        "\n",
        "    return {\n",
        "        \"answer_correctness_recall\": recall_score,\n",
        "        \"answer_correctness_precision\": pecision_score,\n",
        "        \"answer_correctness_f1\": f1_score,\n",
        "    }\n",
        "\n",
        "  @property\n",
        "  def compute_precision(self) -> bool:\n",
        "    return self._precision_answer_scorer is not None\n",
        "\n",
        "\n",
        "GROUNDING_PROMPT_TEMPLATE = \"\"\"I need your help with \"Natural language inference\". Your task is to check if the hypothesis is true, given the premise. The answer should be a single `TRUE` or `FALSE`.\n",
        "\n",
        "Instructions:\n",
        "* If it is possible to fully derive the hypothesis from the premise (entailment), then answer TRUE, otherwise FALSE.\n",
        "* It is ok to use only very common knowledge, all facts need to be included in the premise.\n",
        "\n",
        "Examples:\n",
        "\n",
        "premise: Anna wants a retriever.\n",
        "hypothesis: Anna would like to have a dog.\n",
        "answer: TRUE\n",
        "reason: We know that Anna wants a retriever, which means she wants a dog. Thus, the hypothesis is true given the premise.\n",
        "\n",
        "premise: Anna would like to have a dog.\n",
        "hypothesis: Anna would like to have a retriever.\n",
        "answer: FALSE\n",
        "reason: We know that Anna wants a dog, but that doesn't mean she wants exactly a retriever. Thus, the hypothesis is false given the premise.\n",
        "\n",
        "premise: Einstein was a good physicist.\n",
        "hypothesis: Bruce was a good physicist.\n",
        "answer: FALSE\n",
        "reason: Premise and hypothesis talk about a different person. Thus, the hypothesis is false.\n",
        "\n",
        "premise: Einstein was a good physicist.\n",
        "hypothesis: Einstein is considered to be a good physicist.\n",
        "answer: TRUE\n",
        "reason: The hypothesis only rephrases the premise slightly, so it is true.\n",
        "\n",
        "premise: Peter is a good architect.\n",
        "hypothesis: All men are good architects.\n",
        "answer: FALSE\n",
        "reason: If Peter is a good architect, it doesn't mean all architects are good. Thus, the hypothesis is false.\n",
        "\n",
        "premise: Lucy likes the dog named Haf.\n",
        "hypothesis: Lucy likes all dogs.\n",
        "answer: FALSE\n",
        "reason: Just because Lucy likes the dog named Haf, I cannot conclude that she likes all dogs. Thus, the hypothesis is false.\n",
        "\n",
        "premise: Quantum field theory - Wikipedia: History. Quantum field theory emerged from the work of generations of theoretical physicists spanning much of the 20th century. Its development began in the 1920s with the description of interactions between light and electrons, culminating in the first quantum field theory—quantum electrodynamics.\n",
        "hypothesis: Quantum field theory (QFT) was developed by many theoretical physicists over the course of the 20th century.\n",
        "answer: TRUE\n",
        "reason: The premise states that Quantum field theory started in the 1920s and that its development spanned much of the 20th century. Thus, the hypothesis is true.\n",
        "\n",
        "premise: Quantum field theory - Wikipedia: History. Quantum field theory emerged from the work of generations of theoretical physicists spanning much of the 20th century. Its development began in the 1920s with the description of interactions between light and electrons, culminating in the first quantum field theory—quantum electrodynamics.\n",
        "hypothesis: Quantum field theory (QFT) was developed by many theoretical physicists over the course of the 20 and 21st century.\n",
        "answer: FALSE\n",
        "reason: The premise does not state that Quantum field theory was developed during hte 21st century. Thus, the hypothesis is false.\n",
        "\n",
        "premise: Quantum Field Theory > The History of QFT (Stanford Encyclopedia of Philosophy): The inception of QFT is usually dated 1927 with Dirac's famous paper on “The quantum theory of the emission and absorption of radiation” (Dirac 1927). Here Dirac coined the name quantum electrodynamics (QED) which is the part of QFT that has been developed first.\n",
        "hypothesis: The inception of QFT is usually dated to 1927 when Paul Harr published his paper on “The quantum theory of the emission and absorption of radiation”.\n",
        "answer: FALSE\n",
        "reason: The assumption mentions Dirac, not Harr, so the hypothesis is false.\n",
        "\n",
        "premise: Quantum Field Theory > The History of QFT (Stanford Encyclopedia of Philosophy): The inception of QFT is usually dated 1927 with Dirac's famous paper on “The quantum theory of the emission and absorption of radiation” (Dirac 1927). Here Dirac coined the name quantum electrodynamics (QED) which is the part of QFT that has been developed first.\n",
        "hypothesis: The inception of QFT is usually dated to 1927 when Paul Dirac published his paper on “The quantum theory of the emission and absorption of radiation”.\n",
        "answer: TRUE\n",
        "reason: The hypothesis just paraphrases the assumption so it is true.\n",
        "\n",
        "Now its your turn, think-step-by step, remember the instructions, carefully read the premise and the hypothesis and decide if the hypothesis follows from the premise. I believe in you.\n",
        "\n",
        "premise: {sources}\n",
        "hypothesis: {statement}\n",
        "answer: \"\"\"\n",
        "\n",
        "\n",
        "class AnswerGroundednessScorer:\n",
        "\n",
        "  def __init__(self, llm: TextGenerationModel):\n",
        "    self._statement_scorer = StatementScorer(\n",
        "        scorer=Scorer(\n",
        "            llm=llm, completions=[\"▁TRUE\", \"▁FALSE\"], max_output_tokens=2\n",
        "        ),\n",
        "        prompt_template=GROUNDING_PROMPT_TEMPLATE\n",
        "    )\n",
        "\n",
        "  def score(\n",
        "      self, answer_statements: list[str], sources: list[str]\n",
        "  ) -> AnswerScorerResult:\n",
        "    if not answer_statements or not sources:\n",
        "      return None\n",
        "\n",
        "    scored_statements = self._statement_scorer.score(\n",
        "        shared_template_parameters={\"sources\": \"\\n\".join(sources)},\n",
        "        statements=answer_statements,\n",
        "    )\n",
        "\n",
        "    scores = [\n",
        "        scored_statement.scores[\"▁TRUE\"]\n",
        "        for scored_statement in scored_statements\n",
        "    ]\n",
        "\n",
        "    return AnswerScorerResult(\n",
        "        min_score=round(min(scores), 4),\n",
        "        mean_score=round(statistics.mean(scores), 4),\n",
        "        gmean_score=round(safe_geometric_mean(scores), 4),\n",
        "    )\n",
        "\n",
        "\n",
        "class AnswerGroundedness(Metric):\n",
        "\n",
        "  def __init__(self, llm: TextGenerationModel):\n",
        "    self._statement_extractor = StatementExtractor(llm)\n",
        "    self._answer_scorer = AnswerGroundednessScorer(llm)\n",
        "\n",
        "  def call(\n",
        "      self,\n",
        "      question: str,\n",
        "      answer: str,\n",
        "      sources: list[str],\n",
        "      answer_statements: list[str] | None = None,\n",
        "  ) -> dict[str, Any]:\n",
        "    if answer_statements is None:\n",
        "      answer_statements = self._statement_extractor.extract_statements(\n",
        "          question=question, answer=answer\n",
        "      )\n",
        "\n",
        "    answer_scorer_result = self._answer_scorer.score(\n",
        "        answer_statements=answer_statements, sources=sources\n",
        "    )\n",
        "\n",
        "    score = (\n",
        "        answer_scorer_result.gmean_score if answer_scorer_result else np.nan\n",
        "    )\n",
        "\n",
        "    return {\"gmean\": score}\n",
        "\n",
        "\n",
        "class ContextRecall(AnswerGroundedness):\n",
        "\n",
        "  COLUMNS: list[str] = [\"context_recall_gmean\"]\n",
        "\n",
        "  def __call__(self, inputs: dict[str, Any]) -> dict[str, Any]:\n",
        "    result = self.call(\n",
        "        question=inputs[QUERY],\n",
        "        answer=inputs[REFERENCE],\n",
        "        sources=inputs[RESPONSE].prompt_snippets,\n",
        "        answer_statements=inputs.get(REFERENCE_STATEMENTS)\n",
        "    )\n",
        "    return {f\"context_recall_{name}\": value for name, value in result.items()}\n",
        "\n",
        "\n",
        "class Faithfulness(AnswerGroundedness):\n",
        "\n",
        "  COLUMNS: list[str] = [\"faithfulness_gmean\"]\n",
        "\n",
        "  def __call__(self, inputs: dict[str, Any]) -> dict[str, Any]:\n",
        "    result = self.call(\n",
        "        question=inputs[QUERY],\n",
        "        answer=inputs[RESPONSE].answer_text,\n",
        "        sources=inputs[RESPONSE].prompt_snippets,\n",
        "        answer_statements=inputs.get(PREDICTION_STATEMENTS)\n",
        "    )\n",
        "    return {f\"faithfulness_{name}\": value for name, value in result.items()}\n",
        "\n",
        "\n",
        "class StatementBasedBundledMetric(Metric):\n",
        "\n",
        "  COLUMNS: list[str] = (\n",
        "      AnswerCorrectness.COLUMNS + Faithfulness.COLUMNS + ContextRecall.COLUMNS\n",
        "  )\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      llm: TextGenerationModel,\n",
        "      answer_correctness: bool = True,\n",
        "      faithfulness: bool = True,\n",
        "      context_recall: bool = True,\n",
        "  ):\n",
        "    self._statement_extractor = StatementExtractor(llm)\n",
        "\n",
        "    if not any([answer_correctness, faithfulness, context_recall]):\n",
        "      raise ValueError(\n",
        "          \"At least one of `answer_correctness`, `faithfulness` or \"\n",
        "          \"`context_recall` must be True.\"\n",
        "      )\n",
        "\n",
        "    self._answer_correctness = (\n",
        "        AnswerCorrectness(llm) if answer_correctness else None\n",
        "    )\n",
        "    self._faithfulness = Faithfulness(llm) if faithfulness else None\n",
        "    self._context_recall = ContextRecall(llm) if context_recall else None\n",
        "\n",
        "  def __call__(self, inputs: dict[str, Any]) -> dict[str, Any]:\n",
        "    reference_statements = None\n",
        "    if self._context_recall or self._answer_correctness:\n",
        "      reference_statements = self._statement_extractor.extract_statements(\n",
        "          question=inputs[QUERY], answer=inputs[REFERENCE],\n",
        "      )\n",
        "\n",
        "    prediction_statements = None\n",
        "    if self._faithfulness or self._answer_correctness.compute_precision:\n",
        "      reference_statements = self._statement_extractor.extract_statements(\n",
        "          question=inputs[QUERY], answer=inputs[RESPONSE].answer_text\n",
        "      )\n",
        "\n",
        "    output = {}\n",
        "    if self._answer_correctness:\n",
        "      output.update(\n",
        "          self._answer_correctness({\n",
        "            **inputs,\n",
        "            PREDICTION_STATEMENTS: prediction_statements,\n",
        "            REFERENCE_STATEMENTS: reference_statements,\n",
        "        })\n",
        "      )\n",
        "\n",
        "    if self._context_recall:\n",
        "      output.update(\n",
        "          self._context_recall({\n",
        "              **inputs, REFERENCE_STATEMENTS: reference_statements\n",
        "          })\n",
        "      )\n",
        "\n",
        "    if self._faithfulness:\n",
        "      output.update(\n",
        "          self._faithfulness({\n",
        "              **inputs, PREDICTION_STATEMENTS: prediction_statements,\n",
        "          })\n",
        "      )\n",
        "\n",
        "    return output\n",
        "\n",
        "  def run(self, inputs: pd.DataFrame) -> pd.DataFrame:\n",
        "    reference_statements = pd.DataFrame(\n",
        "        columns=[REFERENCE_STATEMENTS], index=inputs.index\n",
        "    )\n",
        "    if self._context_recall or self._answer_correctness:\n",
        "      reference_statements[REFERENCE_STATEMENTS] = concurrent.thread_map(\n",
        "          self._statement_extractor.extract_statements,\n",
        "          inputs[QUERY].tolist(),\n",
        "          inputs[REFERENCE].tolist(),\n",
        "          max_workers=4,\n",
        "          desc=f\"Extracting statements: `{REFERENCE}`\",\n",
        "      )\n",
        "\n",
        "    prediction_statements = pd.DataFrame(\n",
        "        columns=[PREDICTION_STATEMENTS], index=inputs.index\n",
        "    )\n",
        "    if self._faithfulness or (\n",
        "        self._answer_correctness and self._answer_correctness.compute_precision\n",
        "    ):\n",
        "      prediction_statements[PREDICTION_STATEMENTS] = concurrent.thread_map(\n",
        "          self._statement_extractor.extract_statements,\n",
        "          inputs[QUERY].tolist(),\n",
        "          [response.answer_text for response in inputs[RESPONSE].tolist()],\n",
        "          max_workers=4,\n",
        "          desc=f\"Extracting statements: `{ANSWER_TEXT}`\",\n",
        "      )\n",
        "\n",
        "    output = pd.DataFrame(index=inputs.index)\n",
        "\n",
        "    if self._answer_correctness:\n",
        "      answer_correctness_results = self._answer_correctness.run(\n",
        "          inputs=pd.concat(\n",
        "              [inputs, prediction_statements, reference_statements], axis=1\n",
        "          )\n",
        "      )\n",
        "      output = pd.concat([output, answer_correctness_results], axis=1)\n",
        "\n",
        "    if self._context_recall:\n",
        "      context_recall_results = self._context_recall.run(\n",
        "          inputs=pd.concat([inputs, reference_statements], axis=1)\n",
        "      )\n",
        "      output = pd.concat([output, context_recall_results], axis=1)\n",
        "\n",
        "    if self._faithfulness:\n",
        "      faithfulness_results = self._faithfulness.run(\n",
        "          inputs=pd.concat([inputs, prediction_statements], axis=1)\n",
        "      )\n",
        "      output = pd.concat([output, faithfulness_results], axis=1)\n",
        "\n",
        "    return output"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "g-I9lyuwFO6p"
      },
      "outputs": [],
      "source": [
        "# @markdown `run this cell to define response evaluator`\n",
        "# @markdown > This cell contains the logic of running metrics on scrape results,\n",
        "# @markdown as well as exporting and visualizing evaluation results.\n",
        "\n",
        "\n",
        "_FOLDER_ID = re.compile(r\"folders\\/(.*?)(?=\\/|\\?|$)\")\n",
        "_TRUNCATED_POSTFIX = \"<TRUNCATED: Google Sheet 50k character limit>\"\n",
        "\n",
        "\n",
        "def list_folder(folder_id, drive_service) -> list[tuple[str, str]]:\n",
        "  query = f\"'{folder_id}' in parents and trashed = false\"\n",
        "  list_request = drive_service.files().list(\n",
        "      q=query, fields=\"nextPageToken, files(id, name)\"\n",
        "  )\n",
        "  result = list_request.execute()\n",
        "  items = result.get(\"files\", [])\n",
        "  return [(item[\"id\"], item[\"name\"]) for item in items]\n",
        "\n",
        "\n",
        "def find_file_in_folder(folder_id, name, drive_service) -> str | None:\n",
        "  for file_id, file_name in list_folder(folder_id, drive_service):\n",
        "    if file_name == name:\n",
        "      return file_id\n",
        "  return None\n",
        "\n",
        "\n",
        "def download_json(file_id, drive_service):\n",
        "  request = drive_service.files().get_media(fileId=file_id)\n",
        "  fh = io.BytesIO()\n",
        "  downloader = MediaIoBaseDownload(fh, request)\n",
        "  done = False\n",
        "  while not done:\n",
        "      _, done = downloader.next_chunk()\n",
        "\n",
        "  fh.seek(0)\n",
        "  return json.loads(fh.read().decode('utf-8'))\n",
        "\n",
        "\n",
        "def find_folder(folder_name, drive_service) -> tuple[str, str] | None:\n",
        "  \"\"\"Finds a folder by name in Google Drive.\"\"\"\n",
        "  query = (\n",
        "      f\"name = '{folder_name}' and \"\n",
        "      f\"mimeType = 'application/vnd.google-apps.folder' and \"\n",
        "      f\"trashed = false\"\n",
        "  )\n",
        "  fields = \"nextPageToken, files(id, name, webViewLink)\"\n",
        "  list_request = drive_service.files().list(q=query, fields=fields)\n",
        "  result = list_request.execute()\n",
        "  folders = result.get(\"files\", [])\n",
        "  if not folders:\n",
        "    return None\n",
        "  return folders[0].get(\"id\"), folders[0].get(\"webViewLink\")\n",
        "\n",
        "\n",
        "def create_folder(folder_name, drive_service) -> tuple[str | None, str | None]:\n",
        "  \"\"\"Creates a folder in Google Drive.\"\"\"\n",
        "  create_request = drive_service.files().create(\n",
        "      body={\n",
        "          \"name\": folder_name, \"mimeType\": \"application/vnd.google-apps.folder\"\n",
        "      },\n",
        "      fields=\"id, webViewLink\"\n",
        "  )\n",
        "  result = create_request.execute()\n",
        "  return result.get(\"id\"), result.get(\"webViewLink\")\n",
        "\n",
        "\n",
        "def create_json(\n",
        "    content, file_name, parent, drive_service\n",
        ") -> tuple[str | None, str | None]:\n",
        "  \"\"\"Creates a .json file in the specified Google Drive folder.\"\"\"\n",
        "  request = drive_service.files().create(\n",
        "      body={\"name\": file_name, \"parents\": [parent]},\n",
        "      media_body=MediaInMemoryUpload(\n",
        "          json.dumps(content, indent=4).encode(\"utf-8\"),\n",
        "          mimetype=\"text/plain\",\n",
        "      ),\n",
        "      fields=\"id, webViewLink\",\n",
        "  )\n",
        "  result = request.execute()\n",
        "  return result.get(\"id\"), result.get(\"webViewLink\")\n",
        "\n",
        "\n",
        "def create_chunks(iterable, chunk_size):\n",
        "  for chunk in itertools.zip_longest(*([iter(iterable)] * chunk_size)):\n",
        "    yield [element for element in chunk if element is not None]\n",
        "\n",
        "\n",
        "def delete_worksheet(sheet_id, worksheet_id, sheets_service):\n",
        "  \"\"\"Deletes a worksheet.\"\"\"\n",
        "  sheets_service.spreadsheets().batchUpdate(\n",
        "      spreadsheetId=sheet_id,\n",
        "      body={\"requests\": [{\"deleteSheet\": {\"sheetId\": worksheet_id}}]},\n",
        "  ).execute()\n",
        "\n",
        "\n",
        "def add_worksheet(sheet_id, content, title, sheets_service, chunk_size) -> None:\n",
        "  \"\"\"Adds a worksheet to an existing spreadsheet.\"\"\"\n",
        "  sheets_service.spreadsheets().batchUpdate(\n",
        "      spreadsheetId=sheet_id,\n",
        "      body={\"requests\": [{\"addSheet\": {\"properties\": {\"title\": title}}}]},\n",
        "  ).execute()\n",
        "\n",
        "  for chunk in tqdm(\n",
        "      create_chunks(content, chunk_size),\n",
        "      total=math.ceil(len(content) / chunk_size),\n",
        "      desc=f\"Creating worksheet: {title}\",\n",
        "  ):\n",
        "    sheets_service.spreadsheets().values().append(\n",
        "        spreadsheetId=sheet_id,\n",
        "        range=f\"'{title}'!A1\",\n",
        "        valueInputOption=\"RAW\",\n",
        "        body={\"values\": chunk},\n",
        "    ).execute()\n",
        "\n",
        "\n",
        "def create_sheet(\n",
        "    worksheets, title, parent, chunk_size, sheets_service, drive_service\n",
        ") -> str | None:\n",
        "  \"\"\"Creates a new spreadsheet with worksheets.\"\"\"\n",
        "  body = {\"properties\": {\"title\": title}}\n",
        "  create_request = sheets_service.spreadsheets().create(\n",
        "      body=body, fields=\"spreadsheetId\"\n",
        "  )\n",
        "  create_result = create_request.execute()\n",
        "  sheet_id = create_result.get(\"spreadsheetId\")\n",
        "\n",
        "  parents_request = drive_service.files().get(fileId=sheet_id, fields=\"parents\")\n",
        "  parents_result = parents_request.execute()\n",
        "  parents = parents_result.get(\"parents\")\n",
        "  previous_parents = \",\".join(parents) if parents else None\n",
        "\n",
        "  if not sheet_id:\n",
        "    return\n",
        "\n",
        "  for worksheet_title, content in worksheets.items():\n",
        "    content_dict = content.to_dict(orient=\"split\")\n",
        "    add_worksheet(\n",
        "        sheet_id=sheet_id,\n",
        "        content=[content_dict[\"columns\"]] + content_dict[\"data\"],\n",
        "        title=worksheet_title,\n",
        "        sheets_service=sheets_service,\n",
        "        chunk_size=chunk_size,\n",
        "    )\n",
        "\n",
        "  all_request = sheets_service.spreadsheets().get(spreadsheetId=sheet_id)\n",
        "  all_result = all_request.execute()\n",
        "  default_sheet_id = all_result[\"sheets\"][0][\"properties\"][\"sheetId\"]\n",
        "\n",
        "  delete_worksheet(sheet_id, default_sheet_id, sheets_service)\n",
        "  move_result = drive_service.files().update(\n",
        "      fileId=sheet_id,\n",
        "      addParents=parent,\n",
        "      removeParents=previous_parents,\n",
        "      fields=\"id, parents\"\n",
        "  ).execute()\n",
        "\n",
        "  return f\"https://docs.google.com/spreadsheets/d/{sheet_id}/edit\"\n",
        "\n",
        "\n",
        "def truncate(df, column):\n",
        "  def _truncate(value):\n",
        "    if len(value) < 50_000:\n",
        "      return value\n",
        "    else:\n",
        "      return value[:50_000 - len(_TRUNCATED_POSTFIX)] + _TRUNCATED_POSTFIX\n",
        "  df[column] = df[column].apply(_truncate)\n",
        "\n",
        "\n",
        "@dataclasses.dataclass\n",
        "class EvaluationResult:\n",
        "  scrape_outputs: pd.DataFrame\n",
        "  metric_outputs: pd.DataFrame\n",
        "\n",
        "  @classmethod\n",
        "  def load(cls, folder_url, credentials):\n",
        "    folder_id_match = _FOLDER_ID.search(folder_url)\n",
        "    if not folder_id_match:\n",
        "      raise ValueError()\n",
        "\n",
        "    folder_id = folder_id_match.group(1)\n",
        "    drive_service = build(\"drive\", \"v3\", credentials=credentials)\n",
        "\n",
        "    file_id = find_file_in_folder(folder_id, \"results.json\", drive_service)\n",
        "    json_content = download_json(file_id, drive_service)\n",
        "\n",
        "    queryset = pd.DataFrame.from_dict(json_content[\"queryset\"], orient=\"index\")\n",
        "    responses = pd.DataFrame.from_dict(\n",
        "        json_content[\"responses\"], orient=\"index\"\n",
        "    )\n",
        "    queryset[RESPONSE] = responses.apply(\n",
        "        VertexConversationResponse.from_row, axis=1\n",
        "    )\n",
        "\n",
        "    metrics = pd.DataFrame.from_dict(\n",
        "        json_content[\"metrics\"], orient=\"index\"\n",
        "    )\n",
        "\n",
        "    return cls(queryset, metrics)\n",
        "\n",
        "  def aggregate(self, columns: list[str] | None = None):\n",
        "    if not columns:\n",
        "      columns = self.metric_outputs.columns\n",
        "    shared_columns = self.metric_outputs.columns.intersection(set(columns))\n",
        "    result = pd.DataFrame(self.metric_outputs[shared_columns])\n",
        "    result[\"name\"] = self.scrape_outputs[\"agent_display_name\"]\n",
        "    result[\"evaluation_timestamp\"] = self.metric_outputs[\"evaluation_timestamp\"]\n",
        "\n",
        "    result = result.set_index([\"name\", \"evaluation_timestamp\"])\n",
        "    return result.groupby(level=[0, 1]).mean(numeric_only=True)\n",
        "\n",
        "  def export(self, folder_name: str, chunk_size: int, credentials):\n",
        "    drive_service = build(\"drive\", \"v3\", credentials=credentials)\n",
        "    folder = find_folder(folder_name, drive_service)\n",
        "    if folder:\n",
        "      folder_id, folder_url = folder\n",
        "    else:\n",
        "      folder_id, folder_url = create_folder(folder_name, drive_service)\n",
        "\n",
        "    queryset = self.scrape_outputs.drop(RESPONSE, axis=1)\n",
        "    responses = self.scrape_outputs[RESPONSE].apply(lambda x: x.to_row())\n",
        "    responses = pd.DataFrame(responses.to_list(), index=queryset.index)\n",
        "\n",
        "    json_content = {\n",
        "        \"queryset\": queryset.to_dict(orient=\"index\"),\n",
        "        \"responses\": responses.to_dict(orient=\"index\"),\n",
        "        \"metrics\": self.metric_outputs.to_dict(orient=\"index\"),\n",
        "    }\n",
        "    json_id, json_url = create_json(\n",
        "        json_content, \"results.json\", folder_id, drive_service\n",
        "    )\n",
        "\n",
        "    for column in [_ANSWER_GENERATOR_LLM_PROMPT, _SEARCH_RESULTS]:\n",
        "      truncate(responses, column)\n",
        "\n",
        "    results = pd.concat([queryset, responses, self.metric_outputs], axis=1)\n",
        "    worksheets = {\n",
        "        \"summary\": self.aggregate().fillna(\"#N/A\"),\n",
        "        \"results\": results.fillna(\"#N/A\")\n",
        "    }\n",
        "    sheets_service = build(\"sheets\", \"v4\", credentials=credentials)\n",
        "    create_sheet(\n",
        "        worksheets=worksheets,\n",
        "        title=\"results\",\n",
        "        parent=folder_id,\n",
        "        chunk_size=chunk_size,\n",
        "        sheets_service=sheets_service,\n",
        "        drive_service=drive_service,\n",
        "    )\n",
        "    return folder_url\n",
        "\n",
        "  @property\n",
        "  def timestamp(self) -> str:\n",
        "    return self.metric_outputs[\"evaluation_timestamp\"].iloc[0]\n",
        "\n",
        "\n",
        "@dataclasses.dataclass\n",
        "class EvaluationVisualizer:\n",
        "  evaluation_results: list[EvaluationResult]\n",
        "\n",
        "  def radar_plot(self, columns: list[str] | None = None):\n",
        "    fig = go.Figure()\n",
        "    summaries = pd.concat(\n",
        "        [result.aggregate(columns) for result in self.evaluation_results]\n",
        "    )\n",
        "    summaries = summaries.to_dict(orient=\"split\")\n",
        "\n",
        "    for idx, values in enumerate(summaries[\"data\"]):\n",
        "      fig.add_trace(\n",
        "          go.Scatterpolar(\n",
        "              r=values,\n",
        "              theta=summaries[\"columns\"],\n",
        "              fill='toself',\n",
        "              name=\"_\".join(summaries[\"index\"][idx]),\n",
        "          )\n",
        "      )\n",
        "    fig.update_layout(\n",
        "        polar={\"radialaxis\": {\"visible\": True, \"range\": [0, 1]}},\n",
        "        showlegend=True\n",
        "    )\n",
        "    fig.show()\n",
        "\n",
        "  def count_barplot(self, column_name: str):\n",
        "    results = []\n",
        "    for result in self.evaluation_results:\n",
        "      responses = result.scrape_outputs[RESPONSE].apply(lambda x: x.to_row())\n",
        "      responses = pd.DataFrame(\n",
        "          responses.to_list(), index=result.scrape_outputs.index\n",
        "      )\n",
        "      results.append(\n",
        "          pd.concat(\n",
        "              [result.scrape_outputs, responses, result.metric_outputs],\n",
        "              axis=1\n",
        "          )\n",
        "      )\n",
        "    results = pd.concat(results)\n",
        "    results = results.set_index([\"agent_display_name\", \"evaluation_timestamp\"])\n",
        "    grouped_counts = (\n",
        "        results[column_name]\n",
        "        .groupby(level=[\"agent_display_name\", \"evaluation_timestamp\"])\n",
        "        .value_counts()\n",
        "        .unstack(fill_value=0)\n",
        "    )\n",
        "    grouped_counts.plot(kind=\"bar\")\n",
        "    plt.xlabel(\"Name\")\n",
        "    plt.ylabel(\"Count\")\n",
        "    plt.xticks(rotation=15)\n",
        "    plt.title(f\"{column_name} counts by name\")\n",
        "    plt.legend(title=column_name)\n",
        "    plt.show()\n",
        "\n",
        "  def mean_barplot(self, column_names: list[str]):\n",
        "    results = []\n",
        "    for result in self.evaluation_results:\n",
        "      results.append(\n",
        "          pd.concat([result.scrape_outputs, result.metric_outputs], axis=1)\n",
        "      )\n",
        "    results = pd.concat(results)\n",
        "    results = results.set_index([\"agent_display_name\", \"evaluation_timestamp\"])\n",
        "    grouped_means = (\n",
        "        results[column_names]\n",
        "        .groupby(level=[\"agent_display_name\", \"evaluation_timestamp\"])\n",
        "        .mean()\n",
        "    )\n",
        "    grouped_means.plot(kind=\"bar\")\n",
        "    plt.ylim(top=1.0)\n",
        "    plt.xlabel(\"Name\")\n",
        "    plt.ylabel(\"Mean\")\n",
        "    plt.xticks(rotation=15)\n",
        "    plt.title(\"mean by name\")\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "class VertexConversationEvaluator:\n",
        "\n",
        "  def __init__(self, metrics: list[Metric]):\n",
        "    self._metrics = metrics\n",
        "\n",
        "  def run(self, scraper_output: pd.DataFrame) -> EvaluationResult:\n",
        "    timestamp = datetime.datetime.now(tz=datetime.timezone.utc)\n",
        "    scraper_output = scraper_output.copy(deep=True)\n",
        "    result = pd.DataFrame(index=scraper_output.index)\n",
        "\n",
        "    for metric in self._metrics:\n",
        "      result = pd.concat([result, metric.run(scraper_output)], axis=1)\n",
        "\n",
        "    # adding timestamp and agent display name so they can be used as a multi\n",
        "    # index\n",
        "    result[\"evaluation_timestamp\"] = timestamp.isoformat()\n",
        "\n",
        "    return EvaluationResult(scraper_output, result)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TsHc10HfAHDz"
      },
      "source": [
        "---\n",
        "\n",
        "# Evaluation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4srTefHVmJ3U"
      },
      "source": [
        "## Initialization"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "QbzhGINr5IkG"
      },
      "outputs": [],
      "source": [
        "# @markdown `initialize vertex ai`\n",
        "\n",
        "# @markdown > The project selected will be billed for calculating evaluation\n",
        "# @markdown metrics that require large language models. It should have the\n",
        "# @markdown [Vertex AI API](https://cloud.google.com/vertex-ai/docs/featurestore/setup)\n",
        "# @markdown enabled. The LLM-based metrics use PaLM 2 for Text (Text Bison).\n",
        "# @markdown For pricing information see this [page](https://cloud.google.com/vertex-ai/generative-ai/pricing).\n",
        "\n",
        "vertex_ai_project_id = \"\"  # @param{type: 'string'}\n",
        "vertex_ai_location = \"\"  # @param{type: 'string'}\n",
        "\n",
        "vertexai.init(\n",
        "    project=vertex_ai_project_id,\n",
        "    location=vertex_ai_location,\n",
        "    credentials=credentials,\n",
        ")\n",
        "\n",
        "llm = TextGenerationModel.from_pretrained(\"text-bison@002\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p4L6fCu5fFXR"
      },
      "outputs": [],
      "source": [
        "# test llm on a single query\n",
        "llm.predict(\"hi\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "rZB_riHj13Mo"
      },
      "outputs": [],
      "source": [
        "# @markdown `run this cell to initialize Dialogflow CX agent scraper`\n",
        "# @markdown > This cell initializes the agent with the provided parameters.\n",
        "# @markdown `project_id`, `location` and `agent_id` can be defined through one\n",
        "# @markdown of the following options, while `language_code` must be defined\n",
        "# @markdown in either case. The parameters for a given agent can be found in\n",
        "# @markdown the DialogflowCX console url:\n",
        "# @markdown `https://dialogflow.cloud.google.com/cx/projects/`**`{project_id}`**`/locations/`**`{location}`**`/agents/`**`{agent_id}`**`/intents`\n",
        "\n",
        "language_code = \"en\"  # @param {type: 'string'}\n",
        "\n",
        "# @markdown ---\n",
        "# @markdown ### Option 1. - Provide agent parameters directly\n",
        "# @markdown\n",
        "\n",
        "agent_project_id = \"\"  # @param {type: \"string\"}\n",
        "agent_location = \"\"  # @param {type: 'string'}\n",
        "agent_id = \"\"  # @param {type: \"string\"}\n",
        "\n",
        "# @markdown ---\n",
        "# @markdown ### Option 2. - Parse agent parameters from url\n",
        "# @markdown > **NOTE** : if `agent_url` is provided then it has precedence over\n",
        "# @markdown directly provided agent parameters.\n",
        "\n",
        "agent_url = \"\" # @param {type: \"string\"}\n",
        "\n",
        "if agent_url:\n",
        "  scraper = VertexConversationScraper.from_url(\n",
        "      agent_url=agent_url,\n",
        "      language_code=language_code,\n",
        "      creds=credentials\n",
        "  )\n",
        "else:\n",
        "  scraper = VertexConversationScraper(\n",
        "      agent_id=agent_id,\n",
        "      location=agent_location,\n",
        "      project_id=agent_project_id,\n",
        "      language_code=language_code,\n",
        "      creds=credentials\n",
        "  )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z50iGm_rTC49"
      },
      "outputs": [],
      "source": [
        "# test agent on a single query\n",
        "response = scraper.scrape_detect_intent(query=\"who is the ceo?\")\n",
        "print(json.dumps(dataclasses.asdict(response), indent=4))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cFExd6nTmLYr"
      },
      "source": [
        "## Data Loading\n",
        "\n",
        "The queryset must be in a tabular format that has to contain the following columns:\n",
        "- `conversation_id` _(unique identifier of a conversation, which must be the same for each row that are part of the same conversation)_\n",
        "- `turn_index` _(index of the query - expected answer pair within a conversation)_\n",
        "- `query` _(the input question)_\n",
        "- `expected_answer` _(the ideal or ground truth answer)_\n",
        "- `expected_uri` _(the webpage url or more generally the uri that contains the answer to `query`)_.\n",
        "\n",
        "In addition to the required columns the RougeL metric can also use the following optional column:\n",
        "\n",
        "- `golden_snippet` _(the extracted document snippet or segment that contains the `expected_answer`)_\n",
        "\n",
        "An example for the queryset can be seen in this table:\n",
        "\n",
        "| conversation_id | turn_index | query | expected_answer | expected_uri |\n",
        "| --- | --- | --- | --- | --- |\n",
        "| 0 | 1 | What is the capital of France? | Capital of France is Paris. | exampleurl.com/france |\n",
        "| 0 | 2 | How many people live there? | 2.1 million people live in Paris. | exampleurl.com/paris |\n",
        "| 1 | 1 | What is the color of the sky? | It is blue. | exampleurl.com/common |\n",
        "| 2 | 1 | How many legs does an octopus have? | It has 8 limbs. | exampleurl.com/octopus |\n",
        "\n",
        "---\n",
        "\n",
        "Choose one of the following 3 options to load the queryset:\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1InymjofmW5X"
      },
      "source": [
        "### Option 1. - Manual"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qwhOvOSOmnJ4"
      },
      "outputs": [],
      "source": [
        "sample_df = pd.DataFrame(columns=INPUT_SCHEMA_REQUIRED_COLUMNS)\n",
        "\n",
        "sample_df.loc[0] = [\"0\", 1 ,\"Who are you?\", \"I am an assistant\", \"www.google.com\"]\n",
        "sample_df.loc[1] = [\"1\", 1 ,\"Which is the cheapest plan?\", \"Basic plan\", \"www.google.com\"]\n",
        "sample_df.loc[2] = [\"1\", 2, \"How much?\", \"The Basic plan costs 20$/month\", \"www.google.com\"]\n",
        "queryset = sample_df"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1WlIWM21mUK_"
      },
      "source": [
        "---\n",
        "### Option 2. - From local .csv"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "47JSfzvfaVhd"
      },
      "outputs": [],
      "source": [
        "# @markdown `run this cell to load the queryset from a .csv file in the filesystem`\n",
        "\n",
        "csv_path = \"\"  # @param{type: 'string'}\n",
        "\n",
        "queryset = pd.read_csv(csv_path)\n",
        "queryset = queryset.fillna(\"\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IfEOv6cTmRDB"
      },
      "source": [
        "---\n",
        "### Option 3. - From Google Sheets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "A2tS-n1PRHeC"
      },
      "outputs": [],
      "source": [
        "# @markdown `run this cell to load the queryset from Google Sheets`\n",
        "\n",
        "sheet_url = \"\" # @param {type: \"string\"}\n",
        "worksheet_name = \"\" # @param {type: \"string\"}\n",
        "# @markdown > **NOTE**: if `worksheet_name` is not provided then `Sheet1` is used\n",
        "# @markdown by default.\n",
        "\n",
        "_worksheet_name = worksheet_name if worksheet_name else \"Sheet1\"\n",
        "\n",
        "queryset = load_spreadsheet(sheet_url, _worksheet_name, credentials)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sn_TZfWQmjtv"
      },
      "source": [
        "## Scraping Responses"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7phOx3LWHUfm"
      },
      "outputs": [],
      "source": [
        "# run scraping\n",
        "scrape_result = scraper.run(queryset)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4w6OKOpImcJg"
      },
      "source": [
        "## Metric Definition"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "at5ss5szde8o"
      },
      "source": [
        "Select the metrics that should be computed during evaluation.\n",
        "\n",
        "> **NOTE** : Remember to rerun the cell below (Shift+Enter) after clicking the checkbox of the metric."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "UaX42vFgwwYA"
      },
      "outputs": [],
      "source": [
        "URL_MATCH = True # @param {type: \"boolean\"}\n",
        "# @markdown url match metric computes the following boolean type columns:\n",
        "# @markdown - `cited_url_match@1` - is `expected_url` same as the first link returned by Vertex AI Conversation\n",
        "# @markdown - `cited_url_match` - is `expected_url` part of the links returned by Vertex AI Conversation\n",
        "# @markdown - `search_url_match` - is `expected_url` part of the search results that are shown to generative model in Vertex AI Conversation\n",
        "# @markdown\n",
        "# @markdown ---\n",
        "ROUGEL = True # @param {type: \"boolean\"}\n",
        "# @markdown rougeL metric computes a score between `[0, 1]` (higher is better) based on\n",
        "# @markdown the longest common word subsequence between different targets and predictions:\n",
        "# @markdown\n",
        "# @markdown _(`cited_search_results` are the search snippets which were cited by answer generator llm)_\n",
        "# @markdown - `rougeL_generative` - Compares `expected_answer` to `answer_text`.\n",
        "# @markdown - `rougeL_extractive` - (only computed if\n",
        "# @markdown `golden_snippet` is part of the dataset) Compares `golden_snippet` to `answer_snippets[0]`\n",
        "# @markdown\n",
        "# @markdown ---\n",
        "ANSWER_CORRECTNESS = True # @param {type: \"boolean\"}\n",
        "# @markdown LLM-based autoeval metric that compares `expected_answer` to `answer_text`. The metric\n",
        "# @markdown makes approximately 5-10 LLM calls for each for each row of the dataset depening on `expected_answer` and `answer_text` length.\n",
        "# @markdown It returns 3 output columns containing scores between `[0, 1]` (higher is better):\n",
        "# @markdown - `answer_correctness_recall` - How well does `answer_text`'s information content cover `expected_answer`.\n",
        "# @markdown - `answer_correctness_precision` - How much of `answer_text`'s information content is required based on `expected_answer`.\n",
        "# @markdown - `answer_correctness_f1` - The harmonic mean of recall and precision.\n",
        "# @markdown\n",
        "# @markdown ---\n",
        "FAITHFULNESS = True # @param {type: \"boolean\"}\n",
        "# @markdown LLM-based autoeval metric that provides a score between `[0, 1]` (higher is better)\n",
        "# @markdown with regard to how well `answer_text` is attributed to `search_snippets`. It makes approximately 5\n",
        "# @markdown LLM calls for each row of the dataset depending on the length of `answer_text`.\n",
        "# @markdown - `faithfulness_gmean`\n",
        "# @markdown\n",
        "# @markdown ---\n",
        "CONTEXT_RECALL = True # @param {type: \"boolean\"}\n",
        "# @markdown LLM-based autoeval metric that provides a score between `[0, 1]` (higher is better)\n",
        "# @markdown for how well the `expected_answer` is attributed to `search_snippets`. In other words this metric\n",
        "# @markdown scores the quality of search by measuring how well the `expected_answer` can be generated from the `search_snippets`.\n",
        "# @markdown It makes approximately 5 LLM calls for each row of the dataset depending on the length of `expected_answer`.\n",
        "# @markdown - `context_recall_gmean`\n",
        "# @markdown\n",
        "# @markdown ---\n",
        "\n",
        "metrics = []\n",
        "\n",
        "if URL_MATCH:\n",
        "  metrics.append(UrlMatch())\n",
        "\n",
        "if ROUGEL:\n",
        "  metrics.append(RougeL())\n",
        "\n",
        "if any((ANSWER_CORRECTNESS, FAITHFULNESS, CONTEXT_RECALL)):\n",
        "  metrics.append(\n",
        "      StatementBasedBundledMetric(\n",
        "          llm=llm,\n",
        "          answer_correctness=ANSWER_CORRECTNESS,\n",
        "          faithfulness=FAITHFULNESS,\n",
        "          context_recall=CONTEXT_RECALL,\n",
        "      )\n",
        "  )\n",
        "\n",
        "evaluator = VertexConversationEvaluator(metrics=metrics)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "69h01DPmrMoA"
      },
      "source": [
        "## Computing Metrics"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "9NeMsvykHb0E"
      },
      "outputs": [],
      "source": [
        "evaluation_result = evaluator.run(scrape_result)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "F-aAhhD-qPJp"
      },
      "outputs": [],
      "source": [
        "# @markdown `export evaluation results`\n",
        "\n",
        "FOLDER_NAME = \"result\" # @param {type: \"string\"}\n",
        "CHUNK_SIZE = 50 # @param {type: \"number\"}\n",
        "WITH_TIMESTAMP = True # @param {type: \"boolean\"}\n",
        "\n",
        "_folder_name = (\n",
        "    f\"{FOLDER_NAME}_{evaluation_result.timestamp}\"\n",
        "    if WITH_TIMESTAMP else\n",
        "    FOLDER_NAME\n",
        ")\n",
        "\n",
        "folder_url = evaluation_result.export(_folder_name, CHUNK_SIZE, credentials)\n",
        "print(f\"Exported results to folder: {folder_url}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IH7ddvNEUbal"
      },
      "source": [
        "## Result Visualization"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kARLoOJYBJ0e"
      },
      "outputs": [],
      "source": [
        "FOLDER_URLS = [\n",
        "    folder_url, # latest evaluation\n",
        "    # add previous evaluations e.g: https://drive.google.com/drive/folders/<id>\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "umJE__jOWTQu"
      },
      "outputs": [],
      "source": [
        "# @markdown `define evaluation visualizer`\n",
        "\n",
        "evaluation_visualizer = EvaluationVisualizer([\n",
        "    EvaluationResult.load(folder_url, credentials)\n",
        "    for folder_url in FOLDER_URLS\n",
        "])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "LqObjkW28pNa"
      },
      "outputs": [],
      "source": [
        "# @markdown `radar plot of autoeval metrics`\n",
        "\n",
        "evaluation_visualizer.radar_plot(StatementBasedBundledMetric.COLUMNS)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "collapsed": true,
        "id": "H-wg1FTiToA5"
      },
      "outputs": [],
      "source": [
        "# @markdown `response type distribution`\n",
        "\n",
        "evaluation_visualizer.count_barplot(\"response_type\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "6nc7WQKVY1jp"
      },
      "outputs": [],
      "source": [
        "# @markdown `average RougeL`\n",
        "\n",
        "evaluation_visualizer.mean_barplot(column_names=RougeL.COLUMNS)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "17WDmf3DsZGg1ZGwnr40sMXMQtyfxb4ms",
          "timestamp": 1713942777165
        },
        {
          "file_id": "1b769OFNM8gH56ZzvWUfw4MpGUc-pSF-g",
          "timestamp": 1708939513950
        }
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.8.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}