660 lines
24 KiB
Plaintext
660 lines
24 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "2zyrrTp0QwkE"
|
|
},
|
|
"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": "4yz5pBb8RHi2"
|
|
},
|
|
"source": [
|
|
"# Clearbox for Ranking Tuning"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "iGLQTYerXKa0"
|
|
},
|
|
"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/custom-ranking/clearbox.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%2Fcustom-ranking%2Fclearbox.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/custom-ranking/clearbox.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/custom-ranking/clearbox.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/custom-ranking/clearbox.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/custom-ranking/clearbox.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/custom-ranking/clearbox.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/custom-ranking/clearbox.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/custom-ranking/clearbox.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": "jZy3Q-RagNKm"
|
|
},
|
|
"source": [
|
|
"| Author |\n",
|
|
"| --- |\n",
|
|
"| [Andrei Papou](https://github.com/andrei-papou) |"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "8tGlF6DhStHp"
|
|
},
|
|
"source": [
|
|
"## Overview"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "eMZE2a9vSwsR"
|
|
},
|
|
"source": [
|
|
"In this notebook we will use the [ClearBox](https://github.com/GoogleCloudPlatform/clearbox) library to improve the default [Vertex AI Search](https://cloud.google.com/generative-ai-app-builder/docs/enterprise-search-introduction) ranking for a given query set."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "i-iwIpJqRcke"
|
|
},
|
|
"source": [
|
|
"## Get Started"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "nxI7hhthRf6L"
|
|
},
|
|
"source": [
|
|
"### Install ClearBox from GitHub"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "kzPb2Ju8RmhD"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"%pip install \"git+https://github.com/GoogleCloudPlatform/clearbox\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "OLreVAILTBAL"
|
|
},
|
|
"source": [
|
|
"## Ranking Tuning"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "dFy2MsD-TFNl"
|
|
},
|
|
"source": [
|
|
"### Imports"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "ODC6BF12n6VA"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"import math\n",
|
|
"import random\n",
|
|
"import warnings\n",
|
|
"\n",
|
|
"from clearbox import features as F\n",
|
|
"from clearbox.features import signals as S\n",
|
|
"from clearbox.metrics.recall import RecallAtK\n",
|
|
"from clearbox.models.probabilistic import (\n",
|
|
" BayesianRidgeModel,\n",
|
|
" BayesOptLinearModel,\n",
|
|
")\n",
|
|
"from clearbox.models.probabilistic import ExpectedImprovement as EI\n",
|
|
"from clearbox.models.regression import LinRegModel\n",
|
|
"from clearbox.training import Trainer\n",
|
|
"from clearbox.visualization import Visualizer\n",
|
|
"import matplotlib.pyplot as plt\n",
|
|
"import numpy as np\n",
|
|
"import pandas as pd"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "7fGvXIH4DbQj"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"warnings.filterwarnings(\"ignore\", category=DeprecationWarning)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "_F0-noZcTHmz"
|
|
},
|
|
"source": [
|
|
"### Pin the random seed"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "0Zv36-pQsobR"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"def seed_everything(seed: int):\n",
|
|
" random.seed(seed)\n",
|
|
" np.random.seed(seed)\n",
|
|
"\n",
|
|
"\n",
|
|
"seed_everything(42)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "vJMQd5XwTMfx"
|
|
},
|
|
"source": [
|
|
"### Load and explore the dataset"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "M2eLSvNPmjYv"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"qs_df = pd.read_csv(\n",
|
|
" \"https://storage.googleapis.com/github-repo/search/clearbox/beir_fiqa_50.csv\"\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "YGJumCMEmqfi"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"qs_df.head(3)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "a9MZiVShTQpu"
|
|
},
|
|
"source": [
|
|
"### Preprocess the dataset and generate features"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "AKVClzm49g9v"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"COL_QUERY_CODE = \"query_code\"\n",
|
|
"COL_TARGET = \"is_match\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "mDfVntkE9qc6"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"qs_df[COL_QUERY_CODE] = F.encode(qs_df[\"query\"])"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "G9y3f7oT_uuy"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"def _generate_is_match_col(df: pd.DataFrame) -> pd.DataFrame:\n",
|
|
" query_to_matched_titles = {}\n",
|
|
" for _, row in df.iterrows():\n",
|
|
" query_code = row[COL_QUERY_CODE]\n",
|
|
" if query_code in query_to_matched_titles:\n",
|
|
" continue\n",
|
|
" query_to_matched_titles[query_code] = {\n",
|
|
" targets_str.removeprefix(\"gs://beir_fiqa/\").removesuffix(\".txt\")\n",
|
|
" for targets_str in row[\"targets\"].split(\" \")\n",
|
|
" }\n",
|
|
" df[COL_TARGET] = np.zeros((len(df),), dtype=np.uint8)\n",
|
|
" for query_code, matched_titles in query_to_matched_titles.items():\n",
|
|
" df.loc[\n",
|
|
" (df[COL_QUERY_CODE] == query_code)\n",
|
|
" & (df[\"title\"].astype(str).isin(matched_titles)),\n",
|
|
" COL_TARGET,\n",
|
|
" ] = 1\n",
|
|
"\n",
|
|
"\n",
|
|
"_generate_is_match_col(qs_df)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "PnH-R_JjB0XL"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"qs_df.head(3)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "GNnDCy_NCqA7"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"len(qs_df[qs_df[\"is_match\"] == 1])"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "2tsCIFS0CDaH"
|
|
},
|
|
"source": [
|
|
" ### Explore the dataset"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "wORGCOf4CHej"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"def _plot_signal_distribution(\n",
|
|
" df: pd.DataFrame, signal_list: list[str], n_per_row: int = 3\n",
|
|
"):\n",
|
|
" n = len(signal_list)\n",
|
|
" n_rows = math.ceil(n / n_per_row)\n",
|
|
" f, axes = plt.subplots(n_rows, n_per_row, squeeze=False)\n",
|
|
" f.set_figheight(5 * n_rows)\n",
|
|
" f.set_figwidth(7 * n_per_row)\n",
|
|
" for i, signal in enumerate(signal_list):\n",
|
|
" ax = axes[i // n_per_row][i % n_per_row]\n",
|
|
" ax.set_title(signal)\n",
|
|
" ax.hist(df[signal], bins=50)\n",
|
|
"\n",
|
|
"\n",
|
|
"_plot_signal_distribution(\n",
|
|
" qs_df,\n",
|
|
" [\n",
|
|
" \"gecko_score\",\n",
|
|
" \"bm25_score\",\n",
|
|
" \"jetstream_score\",\n",
|
|
" \"freshness_rank\",\n",
|
|
" \"base_rank\",\n",
|
|
" ],\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "F1xu_uVtDBiM"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"for signal_name in [\n",
|
|
" \"gecko_score\",\n",
|
|
" \"bm25_score\",\n",
|
|
" \"jetstream_score\",\n",
|
|
" \"freshness_rank\",\n",
|
|
" \"base_rank\",\n",
|
|
"]:\n",
|
|
" if qs_df[signal_name].isna().any():\n",
|
|
" print(\n",
|
|
" f\"{signal_name}: {len(qs_df[qs_df[signal_name].isna()]) / len(qs_df) * 100:.2f}% are NaNs\"\n",
|
|
" )"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "BtpiU9gkbERh"
|
|
},
|
|
"source": [
|
|
"Looks like we have some NaN values in `bm25_score` column. We'll handle that later."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "7U1o1AMDC_wt"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"plt.hist(qs_df[\"is_match\"], bins=50)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "d-s2nczJTetv"
|
|
},
|
|
"source": [
|
|
"### Train the model"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "jJ8XcHqxToSi"
|
|
},
|
|
"source": [
|
|
"First we will create instances of `Trainer` and `Visualizer` classes. We will use `trainer` to train and validate the model and `visualizer` will help us explore the results.\n",
|
|
"\n",
|
|
"The `Trainer` has a number of parameters to customize the training procedure, let's look at some of them.\n",
|
|
"\n",
|
|
"First of all, under the hood `Trainer` does a number\n",
|
|
"of random cross-validation splits and trains and validates a new model on each of them. `seeds` argument specifies how many CV splits we should do and which random seed we should use for each. `n_folds` specifies the number of folds we generate for each split.\n",
|
|
"\n",
|
|
"Once a new model is trained for a given seed and fold combination, the `Trainer` will calculate metric value on both train and validation parts of the fold for each metric object from `metrics` list. All those metric values are then aggregated and added to the return value of `train` method, you can access them using `.metrics` attribute.\n",
|
|
"\n",
|
|
"`target_col` specifies the column of the `qs_df` data frame to use as a target. `query_col` is used to identify the query column, please note that we use the integer encoded version of the query we've generated above."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "GJScuHenA7w9"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"trainer = Trainer(\n",
|
|
" df=qs_df,\n",
|
|
" seeds=[7, 15, 21, 42, 81],\n",
|
|
" n_folds=3,\n",
|
|
" metrics=[RecallAtK(k) for k in [1, 3, 5]],\n",
|
|
" target_col=COL_TARGET,\n",
|
|
" query_col=COL_QUERY_CODE,\n",
|
|
")\n",
|
|
"visualizer = Visualizer(metrics=trainer.metrics)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "LlrDrP7nUDJo"
|
|
},
|
|
"source": [
|
|
"It's a good practice in ranking to compare the model against the individual signal, just to make sure our model is an improvement compared to the features we train on. Also, as the ranking signals Vertex AI exposes take different aspects of query-document relevance into account (e.g. topicality, relevance, semantic similarity, etc.), looking at the baselines is a good way to understand which if the signals are more important for your query set."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "L2V8grQ0BV6K"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"baseline_list = [\n",
|
|
" (\"base_rank\", trainer.get_feature_baseline(feature=-S.base_rank)),\n",
|
|
" (\"gecko_score\", trainer.get_feature_baseline(feature=S.gecko_score)),\n",
|
|
" (\n",
|
|
" \"bm25_score\",\n",
|
|
" trainer.get_feature_baseline(feature=F.FillNaN(S.bm25_score, F.Constant(0.0))),\n",
|
|
" ),\n",
|
|
" (\"jetstream_score\", trainer.get_feature_baseline(feature=S.jetstream_score)),\n",
|
|
" (\"freshness_rank\", trainer.get_feature_baseline(feature=-S.freshness_rank)),\n",
|
|
"]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "12qHHlKLVIWD"
|
|
},
|
|
"source": [
|
|
"Now we will train a linear regression model on reciprocal ranks of input signals we've explored above.\n",
|
|
"\n",
|
|
"`features` argument of the `train` method accepts a list of input features to train the model on. Notice how the features are created:\n",
|
|
"\n",
|
|
"- We use `clearbox.features` module (conveniently imported as `F` alias) to do all the feature engineering we need. Doing it this way guarantees we'll be able to seamlessly deploy the formula to production.\n",
|
|
"\n",
|
|
"- We use `S` utility object to create a signal node. `S.gecko_score` tells the trainer that we should read the `gecko_score` column of the input data frame. Under the hood, it's just a syntactic sugar over `F.Signal(\"gecko_score\")`.\n",
|
|
"\n",
|
|
"- We use `F.RR` class to generate reciprocal ranks. Although for some of the signals (like `gecko_score` or `bm25_score`) we have raw values, not just ranks, it's always recommended to train the model on reciprocal ranks. Doing it this way guarantees better stability of the formula across signal updates (e.g. a new version of Jetstream model is released and the distribution of the score changes slightly).\n",
|
|
"\n",
|
|
"- When computing reciprocal ranks, we make sure the input signal is monotonically increasing (e.g. better match corresponds to higher signal value). As the ranks are usually monotonically decreasing, we preprocess them using negation operation: `-S.base_rank`.\n",
|
|
"\n",
|
|
"- We use `F.FillNaN` class to replace `NaN` values with `0`. It's important to do that here using `F` module utilities as this way we can guarantee same logic will be used for any NaNs we encounter while serving the formula in production."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "OkcfwMF4FO5d"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"reg_training_results = trainer.train(\n",
|
|
" LinRegModel(),\n",
|
|
" features=[\n",
|
|
" F.RR(-S.base_rank, 40.0, group_by=S.query_code),\n",
|
|
" F.RR(S.gecko_score, 40.0, group_by=S.query_code),\n",
|
|
" F.RR(F.FillNaN(S.bm25_score, F.Constant(0.0)), 40.0, group_by=S.query_code),\n",
|
|
" F.RR(S.jetstream_score, 40.0, group_by=S.query_code),\n",
|
|
" F.RR(-S.freshness_rank, 40.0, group_by=S.query_code),\n",
|
|
" ],\n",
|
|
" num_parallel_workers=4,\n",
|
|
" print_progress=True,\n",
|
|
")\n",
|
|
"visualizer.visualize_training_results(reg_training_results, baseline_list);"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "-cV8yG98cBSM"
|
|
},
|
|
"source": [
|
|
"Looks like we've improved over the baselines quite a bit!\n",
|
|
"\n",
|
|
"`training_results` object which is returned by the `train` method provides a number of useful properties:\n",
|
|
"\n",
|
|
"- `metrics` data frame contains all the metric values computed on train and validation parts for each seed & fold pair.\n",
|
|
"- `ranking_formula` represents the final ranking model we've trained.\n",
|
|
"\n",
|
|
"The ranking formula object has a number of useful methods, but the most important one is `serialize_to_ranking_expression`. It serializes the formula into the deployment format. To deploy the formula, put the string returned by the method into the `ranking_expression` field of the request and make sure `ranking_expression_backend` field is set to `CLEARBOX`. And that's it, the formula is now used for reranking the results of this particular request.\n",
|
|
"\n",
|
|
"Now let's look at the formula we've just tuned."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "n4o3RKTefINo"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"print(reg_training_results.ranking_formula.serialize_to_ranking_expression())"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "-xSEAMJXfghj"
|
|
},
|
|
"source": [
|
|
"Now let's try a different kind of model and see if we can beat the current result. Intuitively it does seem like optimizing against a 0/1 target using a regression model isn't a great idea.\n",
|
|
"\n",
|
|
"What we'll do instead is optimize the recall directly. Of course, the metric is non-differentiable, but we still can try different combinations of weights and see which of those give us better results. But even if we are going to try just 10 values from 0.0 to 1.0 for each of the 5 signal, we'll have to wait a lot as we'll need to compute recall 5 seed x 3 folds x 10 ** 5 models = 1500000 times. Instead, we gonna use Bayesian Optimization technique to pick the most promising combinations of weights. `BayesOptLinearModel` does just that."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "u-WehEiOk2ub"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"bayes_training_results = trainer.train(\n",
|
|
" BayesOptLinearModel(\n",
|
|
" surrogate_model=BayesianRidgeModel(),\n",
|
|
" acquisition_function=EI(xi=1e-6),\n",
|
|
" metric=RecallAtK(1),\n",
|
|
" grid_size=21,\n",
|
|
" seed_batch_size=512,\n",
|
|
" batch_size=32,\n",
|
|
" n_opt_steps=8,\n",
|
|
" ),\n",
|
|
" features=[\n",
|
|
" F.RR(-S.base_rank, 40.0, group_by=S.query_code),\n",
|
|
" F.RR(S.gecko_score, 40.0, group_by=S.query_code),\n",
|
|
" F.RR(F.FillNaN(S.bm25_score, F.Constant(0.0)), 40.0, group_by=S.query_code),\n",
|
|
" F.RR(S.jetstream_score, 40.0, group_by=S.query_code),\n",
|
|
" F.RR(-S.freshness_rank, 40.0, group_by=S.query_code),\n",
|
|
" ],\n",
|
|
" num_parallel_workers=4,\n",
|
|
" print_progress=True,\n",
|
|
")\n",
|
|
"visualizer.visualize_training_results(\n",
|
|
" bayes_training_results,\n",
|
|
" [\n",
|
|
" *baseline_list,\n",
|
|
" (\"regression_model\", reg_training_results.metrics),\n",
|
|
" ],\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "03zXmeIOtkwO"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"print(bayes_training_results.ranking_formula.serialize_to_ranking_expression())"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "F3gyRLipxiv_"
|
|
},
|
|
"source": [
|
|
"The intuition worked and we were able to improve all 3 recall metrics, `~0.04` improvement in `recall@1` looks decent. Notice how the formula is very close to the previous one, the only substantial difference is that it uses `base_rank` instead of `freshness_rank`."
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"colab": {
|
|
"name": "clearbox.ipynb",
|
|
"toc_visible": true
|
|
},
|
|
"kernelspec": {
|
|
"display_name": "Python 3",
|
|
"name": "python3"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 0
|
|
}
|