wehub-resource-sync
762 lines
35 KiB
Plaintext
762 lines
35 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 1,
|
|
"metadata": {
|
|
"id": "V0QBQDRJ4kOG"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Copyright 2025 Google LLC\n",
|
|
"#\n",
|
|
"# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
|
|
"# you may not use this file except in compliance with the License.\n",
|
|
"# You may obtain a copy of the License at\n",
|
|
"#\n",
|
|
"# https://www.apache.org/licenses/LICENSE-2.0\n",
|
|
"#\n",
|
|
"# Unless required by applicable law or agreed to in writing, software\n",
|
|
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
|
|
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
|
|
"# See the License for the specific language governing permissions and\n",
|
|
"# limitations under the License."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "XRUrOsMa42ar"
|
|
},
|
|
"source": [
|
|
"<table align=\"left\">\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/search/ranking-api/ranking_api_beir_evaluation.ipynb\">\n",
|
|
" <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fsearch%2Franking-api%2Franking_api_beir_evaluation.ipynb\">\n",
|
|
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/search/ranking-api/ranking_api_beir_evaluation.ipynb\">\n",
|
|
" <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/search/ranking-api/ranking_api_beir_evaluation.ipynb\">\n",
|
|
" <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
"</table>\n",
|
|
"\n",
|
|
"<div style=\"clear: both;\"></div>\n",
|
|
"\n",
|
|
"<b>Share to:</b>\n",
|
|
"\n",
|
|
"<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/search/ranking-api/ranking_api_beir_evaluation.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/search/ranking-api/ranking_api_beir_evaluation.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/search/ranking-api/ranking_api_beir_evaluation.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/search/ranking-api/ranking_api_beir_evaluation.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/search/ranking-api/ranking_api_beir_evaluation.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
|
|
"</a>"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "VPVofWuB5BFe"
|
|
},
|
|
"source": [
|
|
"| Author |\n",
|
|
"| --- |\n",
|
|
"| [Jannis Grönberg](https://github.com/puebster) |"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "alpsCA9D5e5m"
|
|
},
|
|
"source": [
|
|
"## Overview\n",
|
|
"### Ranking API BEIR Evaluation\n",
|
|
"\n",
|
|
"This notebook is designed to evaluate the performance of the ranking model `googles/semantic-ranker-default-004` against established information retrieval benchmarks. It uses the BEIR (Benchmarking IR) dataset, which provides diverse collections and tasks for evaluating ranking models. The ground truth for evaluation consist of human-annotated relevance judgments.\n",
|
|
"\n",
|
|
"**Context:**\n",
|
|
"\n",
|
|
"This notebook performs the evaluation in several key stages:\n",
|
|
"\n",
|
|
"1. **Data Loading:** It begins by accessing a public Google Cloud Storage bucket (`ranking_datasets`). From this bucket, it loads humanly annotated data corresponding to the selected BEIR subdatasets: `nfcorpus`, `scidocs`, `trec-covid`, `webis-touche2020`, `dbpedia-entity`, and `msmarco`. This data includes the queries, documents, and their associated relevance judgments.\n",
|
|
"\n",
|
|
"2. **Score Computation:** Using the loaded data, the notebook proceeds to compute ranking scores. It performs the scoring operation (using the target model, `googles/semantic-ranker-default-004`) necessary for ranking. These computed scores are then stored back into the *same* designated Google Cloud Storage bucket, making them available for later use or inspection.\n",
|
|
"\n",
|
|
"3. **Recomputation Option:** For users who wish to re-run the scoring process or store the results in a different location, the notebook allows for configuration of the target Google Cloud Storage bucket to one within their own Google Cloud Platform project.\n",
|
|
"\n",
|
|
"4. **Score Loading for Evaluation:** The notebook loads these computed scores from the specified bucket.\n",
|
|
"\n",
|
|
"5. **Metric Calculation:** Finally, utilizing the `pytrec_eval` library, the notebook calculates standard ranking performance metrics. Specifically, it computes the Normalized Discounted Cumulative Gain (NDCG) at various cutoffs: NDCG@1, NDCG@3, NDCG@5, and NDCG@10. These values provide a quantitative measure of how well the model ranks relevant documents for the given queries, based on the loaded ground truth relevance judgments and the computed scores.\n",
|
|
"\n",
|
|
"**Important Notes on Evaluation:**\n",
|
|
"\n",
|
|
"* **Labeled Data Only:** We have restricted our evaluation to only the labeled datapoints within each BEIR subdataset. This decision stems from the observation that many datasets contain examples that are clearly relevant but lack explicit labels. Including these unlabeled but relevant examples would lead to an underestimation of the model's performance, making a fair evaluation challenging.\n",
|
|
"* **Datasets with Multiple Labels:** We have further filtered the BEIR datasets to include only those that have more than one unique label assigned. This is because computing meaningful Normalized Discounted Cumulative Gain (NDCG) scores on datasets where each query has only a single relevant document is problematic and can lead to misleading results.\n",
|
|
"* **Enable the API** User need to first enable the discoveryengine api by visiting [Enable Discoveryengine API](https://console.developers.google.com/apis/api/discoveryengine.googleapis.com/)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "yd2T58Vf5cOM"
|
|
},
|
|
"source": [
|
|
"### Install Google Cloud SDKs and other required packages\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "X99z1Ml65c1V"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"%pip install --upgrade --quiet google-cloud-discoveryengine pytrec_eval numpy pandas scikit-learn tqdm gcsfs"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "VA_Ko1_66t6q"
|
|
},
|
|
"source": [
|
|
"# Setup"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "mQVTqw4H5wH2"
|
|
},
|
|
"source": [
|
|
"### Authenticate your notebook environment (Colab only)\n",
|
|
"\n",
|
|
"If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 1,
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "zHlYMqp-5v4Q"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# @title Colab authentification\n",
|
|
"import sys\n",
|
|
"\n",
|
|
"if \"google.colab\" in sys.modules:\n",
|
|
" from google.colab import auth\n",
|
|
"\n",
|
|
" auth.authenticate_user()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "sZiEMnBLy84D"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# @title Import Libraries\n",
|
|
"\n",
|
|
"# Standard library imports\n",
|
|
"import json\n",
|
|
"import os\n",
|
|
"import time\n",
|
|
"\n",
|
|
"# Third-party library imports\n",
|
|
"import numpy as np\n",
|
|
"import pandas as pd\n",
|
|
"import pytrec_eval\n",
|
|
"from google.cloud import discoveryengine_v1 as discoveryengine\n",
|
|
"from google.cloud import storage\n",
|
|
"from sklearn.metrics import roc_auc_score\n",
|
|
"from tqdm.auto import tqdm"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 3,
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "A61SFkLs5ojJ"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# @title Helper functions\n",
|
|
"\n",
|
|
"\n",
|
|
"def call_model_in_batches(\n",
|
|
" df: pd.DataFrame, model_name: str, project_id: str, location: str\n",
|
|
"):\n",
|
|
" \"\"\"Generates ranking scores for document IDs based on provided queries using\n",
|
|
"\n",
|
|
" the Google Cloud Discovery Engine Rank API.\n",
|
|
"\n",
|
|
" The function processes a DataFrame of queries and associated documents,\n",
|
|
" sends batches of documents to the Discovery Engine Rank API for scoring\n",
|
|
" against each query, and returns a dictionary of query-to-document-score\n",
|
|
" mappings.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" df: A Pandas DataFrame containing query and document information. It must\n",
|
|
" have columns: \"query_id\", \"query\", \"document_id\", \"title\", and\n",
|
|
" \"content\". Each row represents a document associated with a specific\n",
|
|
" query.\n",
|
|
" project_id: The ID of the Google Cloud Project.\n",
|
|
" location: The location (region) of the Discovery Engine instance.\n",
|
|
" model_name: The name of the ranking model to use in the Discovery Engine.\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" A dictionary where the outer keys are query IDs (str) and the inner values\n",
|
|
" are dictionaries mapping document IDs (str) to their ranking scores\n",
|
|
" (float),\n",
|
|
" rounded to four decimal places.\n",
|
|
" \"\"\"\n",
|
|
" results = {}\n",
|
|
"\n",
|
|
" client = discoveryengine.RankServiceClient()\n",
|
|
" ranking_config = client.ranking_config_path(\n",
|
|
" project=project_id,\n",
|
|
" location=location,\n",
|
|
" ranking_config=\"default_ranking_config\",\n",
|
|
" )\n",
|
|
" query_groups = df.groupby(\"query_id\")\n",
|
|
" dataset_name = df.iloc[0][\"name\"]\n",
|
|
" overall_passed_seconds = 0\n",
|
|
"\n",
|
|
" # Iterate through datasets\n",
|
|
" for query_id, group in tqdm(query_groups, desc=f\"Scoring {dataset_name}\"):\n",
|
|
" # Add the query id to the results\n",
|
|
" query = group.iloc[0][\"query\"]\n",
|
|
" results[query_id] = {}\n",
|
|
"\n",
|
|
" records = []\n",
|
|
" # iterate through rows of each subgroup.\n",
|
|
" for _, row in group.iterrows():\n",
|
|
" records.append(\n",
|
|
" discoveryengine.RankingRecord(\n",
|
|
" id=row[\"document_id\"],\n",
|
|
" title=row[\"title\"],\n",
|
|
" content=row[\"content\"],\n",
|
|
" )\n",
|
|
" )\n",
|
|
" # We process in batches of 200\n",
|
|
" if len(records) >= 200:\n",
|
|
" now = time.time()\n",
|
|
" request = discoveryengine.RankRequest(\n",
|
|
" ranking_config=ranking_config,\n",
|
|
" model=model_name,\n",
|
|
" query=query,\n",
|
|
" records=records,\n",
|
|
" ignore_record_details_in_response=True,\n",
|
|
" )\n",
|
|
" resp = client.rank(request=request)\n",
|
|
" overall_passed_seconds += time.time() - now\n",
|
|
" for i in resp.records:\n",
|
|
" results[query_id][i.id] = i.score\n",
|
|
" records = []\n",
|
|
"\n",
|
|
" # If any records are left process them\n",
|
|
" if len(records) > 0:\n",
|
|
" now = time.time()\n",
|
|
" request = discoveryengine.RankRequest(\n",
|
|
" ranking_config=ranking_config,\n",
|
|
" model=model_name,\n",
|
|
" query=query,\n",
|
|
" records=records,\n",
|
|
" ignore_record_details_in_response=True,\n",
|
|
" )\n",
|
|
" resp = client.rank(request=request)\n",
|
|
" overall_passed_seconds += time.time() - now\n",
|
|
"\n",
|
|
" # Round the scores and add them to the overall results.\n",
|
|
" for i in resp.records:\n",
|
|
" results[query_id][i.id] = round(i.score, 4)\n",
|
|
"\n",
|
|
" print(\n",
|
|
" f\"The time to process {len(query_groups)} queries with at total number of\"\n",
|
|
" f\" {len(df)} documents took: {round(overall_passed_seconds, 2)}s.\\n That\"\n",
|
|
" f\" is {round(1000 * overall_passed_seconds / len(query_groups), 2)}ms per\"\n",
|
|
" f\" query or {round(1000 * overall_passed_seconds / len(df), 2)}ms per\"\n",
|
|
" \" document.\"\n",
|
|
" )\n",
|
|
" return results\n",
|
|
"\n",
|
|
"\n",
|
|
"def evaluate(\n",
|
|
" query_document_ground_truth: dict[str, dict[str, int]],\n",
|
|
" query_document_prediction: dict[str, dict[str, float]],\n",
|
|
" k_values: list[int],\n",
|
|
" ignore_identical_ids: bool = True,\n",
|
|
" verbose: bool = True,\n",
|
|
") -> tuple[dict[str, float], dict[str, float], dict[str, float], dict[str, float]]:\n",
|
|
" \"\"\"Evaluates the retrieval results against the ground truth relevance judgments (query_document_ground_truth)\n",
|
|
"\n",
|
|
" using Normalized Discounted Cumulative Gain (NDCG).\n",
|
|
" This code is taken from the official BEIR evaluation code.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" query_document_ground_truth: A dictionary representing the ground truth relevance judgments. The\n",
|
|
" outer keys are query IDs (str), and the inner dictionary maps document\n",
|
|
" IDs (str) to their relevance scores (int).\n",
|
|
" query_document_prediction: A dictionary representing the retrieval results. The outer keys\n",
|
|
" are query IDs (str), and the inner dictionary maps retrieved document\n",
|
|
" IDs (str) to their retrieval scores (float).\n",
|
|
" k_values: A list of integers representing the cutoff ranks for NDCG\n",
|
|
" calculation (e.g., [1, 3, 5, 10]). NDCG will be calculated for each k in\n",
|
|
" this list.\n",
|
|
" ignore_identical_ids: If True, documents with the same ID as the query ID\n",
|
|
" will be ignored during evaluation to prevent self-ranking. Defaults to\n",
|
|
" True.\n",
|
|
" verbose: If True, detailed evaluation metrics for each k-value will be\n",
|
|
" printed. Defaults to True.\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" A tuple containing four dictionaries:\n",
|
|
" - ndcg_scores: A dictionary with the k values as keys and ndcg@k\n",
|
|
" scores as values.\n",
|
|
" \"\"\"\n",
|
|
" if ignore_identical_ids:\n",
|
|
" popped = []\n",
|
|
" for query_id, document_scores in query_document_prediction.items():\n",
|
|
" for document_id in list(document_scores):\n",
|
|
" if query_id == document_id:\n",
|
|
" document_scores.pop(document_id)\n",
|
|
" popped.append(document_id)\n",
|
|
"\n",
|
|
" ndcg = {}\n",
|
|
"\n",
|
|
" for k in k_values:\n",
|
|
" ndcg[f\"NDCG@{k}\"] = 0.0\n",
|
|
"\n",
|
|
" ndcg_string = \"ndcg_cut.\" + \",\".join([str(k) for k in k_values])\n",
|
|
" evaluator = pytrec_eval.RelevanceEvaluator(\n",
|
|
" query_document_ground_truth,\n",
|
|
" {\n",
|
|
" ndcg_string,\n",
|
|
" },\n",
|
|
" )\n",
|
|
"\n",
|
|
" scores = evaluator.evaluate(query_document_prediction)\n",
|
|
"\n",
|
|
" for query_id in scores.keys():\n",
|
|
" for k in k_values:\n",
|
|
" ndcg[f\"NDCG@{k}\"] += scores[query_id][\"ndcg_cut_\" + str(k)]\n",
|
|
"\n",
|
|
" for k in k_values:\n",
|
|
" ndcg[f\"NDCG@{k}\"] = round(ndcg[f\"NDCG@{k}\"] / len(scores), 5)\n",
|
|
"\n",
|
|
" if verbose:\n",
|
|
" for eval in [ndcg]:\n",
|
|
" for k in eval.keys():\n",
|
|
" print(f\"{k}: {eval[k]:.4f}\")\n",
|
|
"\n",
|
|
" return ndcg\n",
|
|
"\n",
|
|
"\n",
|
|
"def load_bytes_from_gcs(bucket_name: str, file_path: str):\n",
|
|
" \"\"\"Accesses a Google Cloud Storage (GCS) bucket, opens a file,\n",
|
|
"\n",
|
|
" loads its contents and returns its bytes\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" bucket_name: The name of the GCS bucket.\n",
|
|
" file_path: The path to the file within the bucket (e.g.,\n",
|
|
" 'data/my_data.json').\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" Bytes like object\n",
|
|
" \"\"\"\n",
|
|
" try:\n",
|
|
" # Initialize the GCS client\n",
|
|
" client = storage.Client()\n",
|
|
"\n",
|
|
" # Get the bucket object\n",
|
|
" bucket = client.bucket(bucket_name)\n",
|
|
"\n",
|
|
" # Get the blob (file) object\n",
|
|
" blob = bucket.blob(file_path)\n",
|
|
"\n",
|
|
" # Download the blob's content as bytes\n",
|
|
" _bytes = blob.download_as_bytes()\n",
|
|
"\n",
|
|
" return _bytes\n",
|
|
"\n",
|
|
" except Exception as e:\n",
|
|
" print(f\"An error occurred: {e}\")\n",
|
|
" return None\n",
|
|
"\n",
|
|
"\n",
|
|
"def save_dict_to_gcs_json(\n",
|
|
" data: dict, bucket_name: str, file_path: str, encoding: str = \"utf-8\"\n",
|
|
") -> None:\n",
|
|
" \"\"\"Saves a Python dictionary as a JSON file to a Google Cloud Storage (GCS) bucket.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" data: The dictionary to be saved.\n",
|
|
" bucket_name: The name of the GCS bucket.\n",
|
|
" file_path: The desired path for the JSON file within the bucket (e.g.,\n",
|
|
" 'data/my_data.json').\n",
|
|
" encoding: The encoding to use when writing the JSON file. Defaults to\n",
|
|
" 'utf-8'.\n",
|
|
" \"\"\"\n",
|
|
" try:\n",
|
|
" # Initialize the GCS client\n",
|
|
" client = storage.Client()\n",
|
|
"\n",
|
|
" # Get the bucket object\n",
|
|
" bucket = client.bucket(bucket_name)\n",
|
|
"\n",
|
|
" # Get the blob (file) object\n",
|
|
" blob = bucket.blob(file_path)\n",
|
|
"\n",
|
|
" # Serialize the dictionary to a JSON string\n",
|
|
" json_string = json.dumps(\n",
|
|
" data, indent=2, ensure_ascii=False\n",
|
|
" ) # indent for readability\n",
|
|
"\n",
|
|
" # Upload the JSON string to GCS as bytes with specified encoding\n",
|
|
" blob.upload_from_string(\n",
|
|
" json_string.encode(encoding), content_type=\"application/json\"\n",
|
|
" )\n",
|
|
"\n",
|
|
" print(f\"Dictionary successfully saved to gs://{bucket_name}/{file_path}\")\n",
|
|
"\n",
|
|
" except Exception as e:\n",
|
|
" print(f\"An error occurred while saving to GCS: {e}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "mR1elDv27JaZ"
|
|
},
|
|
"source": [
|
|
"# If you want to run the scoring for yourself run the following section. Otherwise you can skip this section"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 7,
|
|
"metadata": {
|
|
"id": "T5cRJpRI7VD6"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# @title Google Cloud Platform parameters\n",
|
|
"\n",
|
|
"# @markdown Change this to your bucket, project id and location\n",
|
|
"\n",
|
|
"\n",
|
|
"# fmt: off\n",
|
|
"PROJECT_ID = \"[your-project-id]\" # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
|
|
"if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
|
|
" PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
|
|
"\n",
|
|
"LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
|
|
"\n",
|
|
"# @markdown If you want to run the scoring for yourself you need to change this to your bucket otherwise you can leave these variables as is.\n",
|
|
"OUTPUT_SCORES_BUCKET_NAME = \"ranking_datasets\" # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
|
|
"MODEL_NAME = \"semantic-ranker-default-004\" # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
|
|
"SCORE_PATH = \"scores\" # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
|
|
"# fmt: on\n",
|
|
"\n",
|
|
"DATA_INPUT_BUCKET_NAME = \"ranking_datasets\" # leave this as is\n",
|
|
"DATASET_PATH = \"input_data/labeled_beir_with_more_than_1_label.pkl\" # leave this as is"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "YzpKt2ul-8cR"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# @title Read Dataset\n",
|
|
"df = pd.read_pickle(f\"gs://{DATA_INPUT_BUCKET_NAME}/{DATASET_PATH}\")\n",
|
|
"display(df.groupby(\"name\").agg({\"query_id\": \"nunique\", \"document_id\": \"nunique\"}))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 12,
|
|
"metadata": {
|
|
"id": "lKV9uEfOhpFJ"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# @title Test connection\n",
|
|
"docs = [\n",
|
|
" {\n",
|
|
" \"query_id\": \"1\",\n",
|
|
" \"query\": \"Why is the sky blue?\",\n",
|
|
" \"name\": \"test\",\n",
|
|
" \"document_id\": \"1\",\n",
|
|
" \"title\": \"Rayleigh Scattering and Blue Sky\",\n",
|
|
" \"content\": (\n",
|
|
" \"The sky appears blue because of a phenomenon called Rayleigh scattering. \"\n",
|
|
" \"When sunlight travels through the atmosphere, it collides with gas molecules \"\n",
|
|
" \"and scatters in all directions. Blue light is scattered more than other colors \"\n",
|
|
" \"because it travels in shorter waves.\"\n",
|
|
" ),\n",
|
|
" },\n",
|
|
" {\n",
|
|
" \"query_id\": \"1\",\n",
|
|
" \"query\": \"Why is the sky blue?\",\n",
|
|
" \"name\": \"test\",\n",
|
|
" \"document_id\": \"2\",\n",
|
|
" \"title\": \"Cooking Pasta\",\n",
|
|
" \"content\": (\n",
|
|
" \"To cook perfect pasta, bring a large pot of salted water to a rolling boil. \"\n",
|
|
" \"Add the pasta and cook until al dente, stirring occasionally. Drain and serve \"\n",
|
|
" \"with your favorite sauce.\"\n",
|
|
" ),\n",
|
|
" },\n",
|
|
"]\n",
|
|
"\n",
|
|
"res = call_model_in_batches(\n",
|
|
" df=pd.DataFrame(docs),\n",
|
|
" model_name=MODEL_NAME,\n",
|
|
" project_id=PROJECT_ID,\n",
|
|
" location=LOCATION,\n",
|
|
")\n",
|
|
"\n",
|
|
"for d in docs:\n",
|
|
" for r, v in res[\"1\"].items():\n",
|
|
" if r == d[\"document_id\"]:\n",
|
|
" print(d[\"title\"], v)\n",
|
|
" break"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "qjf6MgGe0Dta"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# @title Score Query-Document pairs\n",
|
|
"\n",
|
|
"dataset_groups = df.groupby(\"name\")\n",
|
|
"for subdataset_name, dataset in dataset_groups:\n",
|
|
" score_output_path = f\"{SCORE_PATH}/{subdataset_name.replace('-', '_')}/{MODEL_NAME.replace('-', '_')}.json\"\n",
|
|
"\n",
|
|
" print(f\"Evaluating : {subdataset_name.upper()}\")\n",
|
|
" print(f\"Shape: {dataset.shape}\")\n",
|
|
"\n",
|
|
" scores = call_model_in_batches(\n",
|
|
" df=dataset,\n",
|
|
" model_name=MODEL_NAME,\n",
|
|
" project_id=PROJECT_ID,\n",
|
|
" location=LOCATION,\n",
|
|
" )\n",
|
|
"\n",
|
|
" save_dict_to_gcs_json(\n",
|
|
" data=scores, bucket_name=OUTPUT_SCORES_BUCKET_NAME, file_path=score_output_path\n",
|
|
" )"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "53evNKYvtmRC"
|
|
},
|
|
"source": [
|
|
"# Calculate NDCG and ROC AUC Scores"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "qGbk22265OrG"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# @title Calculate...\n",
|
|
"\n",
|
|
"k_values = [1, 3, 5, 10]\n",
|
|
"\n",
|
|
"# Get datasets\n",
|
|
"client = storage.Client()\n",
|
|
"bucket = client.bucket(OUTPUT_SCORES_BUCKET_NAME)\n",
|
|
"blobs = [\n",
|
|
" \"nfcorpus\",\n",
|
|
" \"scidocs\",\n",
|
|
" \"trec-covid\",\n",
|
|
" \"webis-touche2020\",\n",
|
|
" \"dbpedia-entity\",\n",
|
|
" \"msmarco\",\n",
|
|
"]\n",
|
|
"_res = None\n",
|
|
"\n",
|
|
"for sub_name in blobs:\n",
|
|
" print(sub_name)\n",
|
|
" query_document_ground_truth_path = (\n",
|
|
" f\"{SCORE_PATH}/{sub_name.replace('-', '_')}/qrels.json\"\n",
|
|
" )\n",
|
|
" query_document_prediction_path = (\n",
|
|
" f\"{SCORE_PATH}/{sub_name.replace('-', '_')}/{MODEL_NAME.replace('-', '_')}.json\"\n",
|
|
" )\n",
|
|
"\n",
|
|
" query_document_ground_truth = json.loads(\n",
|
|
" load_bytes_from_gcs(\n",
|
|
" OUTPUT_SCORES_BUCKET_NAME, query_document_ground_truth_path\n",
|
|
" ).decode()\n",
|
|
" )\n",
|
|
" query_document_prediction = json.loads(\n",
|
|
" load_bytes_from_gcs(\n",
|
|
" OUTPUT_SCORES_BUCKET_NAME, query_document_prediction_path\n",
|
|
" ).decode()\n",
|
|
" )\n",
|
|
"\n",
|
|
" # Check if all query document pairs are present in Query Document Ground Truth and Query Document Prediction\n",
|
|
" for first, second, n1, n2 in [\n",
|
|
" (\n",
|
|
" query_document_ground_truth,\n",
|
|
" query_document_prediction,\n",
|
|
" \"Query Document Ground Truth\",\n",
|
|
" \"Query Document Prediction\",\n",
|
|
" ),\n",
|
|
" (\n",
|
|
" query_document_prediction,\n",
|
|
" query_document_ground_truth,\n",
|
|
" \"Query Document Prediction\",\n",
|
|
" \"Query Document Ground Truth\",\n",
|
|
" ),\n",
|
|
" ]:\n",
|
|
" deleted_queries = 0\n",
|
|
" deleted_docs = 0\n",
|
|
" first_query_id_keys = list(first.keys())\n",
|
|
" for query_id in first_query_id_keys:\n",
|
|
" if query_id not in second:\n",
|
|
" del first[query_id]\n",
|
|
" deleted_queries += 1\n",
|
|
" first_document_id_keys = list(first[query_id].keys())\n",
|
|
" for document_id in first_document_id_keys:\n",
|
|
" if document_id not in second[query_id]:\n",
|
|
" del first[query_id][document_id]\n",
|
|
" deleted_docs += 1\n",
|
|
" if deleted_queries > 0 or deleted_docs > 0:\n",
|
|
" print(\n",
|
|
" f\"Deleted from {n1} (because not present in {n2}):\"\n",
|
|
" f\" {deleted_queries} Queries, {deleted_docs} Documents\"\n",
|
|
" )\n",
|
|
"\n",
|
|
" # The following conversion is needed as pytrec eval only excepts positive integers as ground truth labels\n",
|
|
" y_true, y_pred = [], []\n",
|
|
" max_score = 0\n",
|
|
" for query_id, document_id_scores in query_document_ground_truth.items():\n",
|
|
" for document_id, y_true_score in document_id_scores.items():\n",
|
|
" max_score = max(max_score, y_true_score)\n",
|
|
"\n",
|
|
" for query_id, document_id_scores in query_document_ground_truth.items():\n",
|
|
" for document_id, y_pred_score in document_id_scores.items():\n",
|
|
" y_true_label = 0 if y_pred_score / max_score <= 0.5 else 1\n",
|
|
" y_true.append(y_true_label)\n",
|
|
" y_pred.append(query_document_prediction[query_id][document_id])\n",
|
|
"\n",
|
|
" ndcg = evaluate(\n",
|
|
" query_document_ground_truth, query_document_prediction, k_values, verbose=False\n",
|
|
" )\n",
|
|
" results = {k: [v] for k, v in ndcg.items()}\n",
|
|
"\n",
|
|
" results[\"ROC AUC\"] = [float(roc_auc_score(np.array(y_true), np.array(y_pred)))]\n",
|
|
"\n",
|
|
" ndcg_res = pd.DataFrame.from_dict(results, orient=\"index\", columns=[sub_name])\n",
|
|
" if _res is None:\n",
|
|
" _res = ndcg_res\n",
|
|
" else:\n",
|
|
" _res = _res.merge(ndcg_res, left_index=True, right_index=True)\n",
|
|
"\n",
|
|
"_res[\"Macro Avg.\"] = _res.mean(axis=1)\n",
|
|
"_res.index.name = MODEL_NAME"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "3vQoBzg3oklS"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"_res"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "71137aacc1d4"
|
|
},
|
|
"source": [
|
|
"## Conclusion\n",
|
|
"\n",
|
|
"In this notebook, we successfully demonstrated how to evaluate the performance of the `googles/semantic-ranker-default-004` model on several key BEIR subdatasets.\n",
|
|
"\n",
|
|
"We walked through the process of:\n",
|
|
"* Loading human-annotated ground truth data from a public Google Cloud Storage bucket.\n",
|
|
"* Utilizing the Google Ranking API (implicitly through the scoring process) to compute relevance scores for query-document pairs.\n",
|
|
"* Storing the computed scores back into GCS for persistence and accessibility.\n",
|
|
"* Loading these scores and the ground truth to calculate standard information retrieval metrics, specifically NDCG@(1, 3, 5, 10), using the `pytrec_eval` library.\n",
|
|
"\n",
|
|
"A significant takeaway from this evaluation is the strong performance of the `googles/semantic-ranker-default-004` model. At the point in time of publishing this notebook, the results indicate that this new model achieves state-of-the-art results on these specific BEIR benchmarks, showcasing its effectiveness for ranking tasks.\n",
|
|
"\n",
|
|
"### Where to Learn More\n",
|
|
"\n",
|
|
"* **Google Ranking API:** Explore the official documentation for the Ranking API to understand its capabilities, pricing, and other available models.\n",
|
|
"* **BEIR Dataset:** Visit the BEIR project website to learn more about the various datasets available, their characteristics, and how they are used for benchmarking.\n",
|
|
"* **pytrec_eval:** Refer to the documentation for the `pytrec_eval` library to understand the NDCG metric and other potential evaluation metrics in more detail."
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"colab": {
|
|
"name": "ranking_api_beir_evaluation.ipynb",
|
|
"toc_visible": true
|
|
},
|
|
"kernelspec": {
|
|
"display_name": "Python 3",
|
|
"name": "python3"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 0
|
|
}
|