Clearbox for Ranking Tuning

Source notebook

Repo path: search/custom-ranking/clearbox.ipynb · Open on GitHub · intermediate

Tunes Vertex AI Search ranking with ClearBox using BEIR FIQA signals and recall-based validation.

Summary

This notebook teaches how to use the ClearBox library to improve default Vertex AI Search ranking for a query set. It loads a BEIR FIQA sample from Cloud Storage, creates match labels, explores ranking signals, trains ClearBox ranking formulas with cross-validation, compares baselines, and serializes formulas for deployment with the CLEARBOX ranking backend.

Key code patterns

Install ClearBox

%pip install "git+https://github.com/GoogleCloudPlatform/clearbox"

Installs the ClearBox ranking tuning library directly from GitHub.

Encode queries and label matches

qs_df[COL_QUERY_CODE] = F.encode(qs_df["query"])
_generate_is_match_col(qs_df)

Creates an integer query key and binary target column for supervised ranking evaluation.

Trainer setup

trainer = Trainer(
    df=qs_df,
    seeds=[7, 15, 21, 42, 81],
    n_folds=3,
    metrics=[RecallAtK(k) for k in [1, 3, 5]],
    target_col=COL_TARGET,
    query_col=COL_QUERY_CODE,
)

Runs repeated cross-validation and evaluates recall at multiple cutoffs.

Reciprocal-rank features

features=[
    F.RR(-S.base_rank, 40.0, group_by=S.query_code),
    F.RR(S.gecko_score, 40.0, group_by=S.query_code),
    F.RR(F.FillNaN(S.bm25_score, F.Constant(0.0)), 40.0, group_by=S.query_code),
]

Uses ClearBox feature nodes so the same feature logic can be serialized for serving.

Serialize ranking formula

print(reg_training_results.ranking_formula.serialize_to_ranking_expression())

Produces the ranking expression to send with ranking_expression_backend set to CLEARBOX.

Models & APIs used

  • APIs / services: Vertex AI Search, Cloud Storage
  • SDKs / libraries: clearbox, matplotlib, numpy, pandas

When to use this

Use this pattern when you have query-result judgments and want to tune Vertex AI Search ranking signals into a deployable ClearBox expression.

Gotchas & caveats

  • Install requires pulling ClearBox from GitHub.
  • The dataset is loaded from a public Cloud Storage URL.
  • bm25_score contains NaN values and is filled with 0 using ClearBox feature utilities.
  • Ranking expression deployment requires ranking_expression_backend set to CLEARBOX.
  • Training uses multiple seeds, folds, and parallel workers, so runtime depends on model and optimization settings.

Best practices

  • Pin random seeds for reproducible training.
  • Compare trained models against individual signal baselines.
  • Train on reciprocal ranks for better stability across signal distribution changes.
  • Use ClearBox feature utilities such as FillNaN so production serving uses the same logic as training.
  • Make rank-like signals monotonically increasing before reciprocal-rank computation.