Generative AI Google Brain

search-ranking-api-ranking-api-beir-evaluation

3 min read

Setup

Source notebook

Repo path: search/ranking-api/ranking_api_beir_evaluation.ipynb · Open on GitHub · advanced

Evaluates Discovery Engine semantic ranking on BEIR datasets with NDCG and ROC AUC metrics.

Summary

This notebook evaluates the Ranking API model googles/semantic-ranker-default-004 on selected labeled BEIR subdatasets. It loads query, document, and relevance data from Google Cloud Storage, scores query-document pairs with the Discovery Engine Rank API, saves scores to GCS, then computes NDCG@1/3/5/10 and ROC AUC with pytrec_eval and scikit-learn.

Key code patterns

Rank API client setup

client = discoveryengine.RankServiceClient()
ranking_config = client.ranking_config_path(
    project=project_id,
    location=location,
    ranking_config="default_ranking_config",
)

Builds the Discovery Engine ranking config path used by RankRequest.

Batch ranking records

request = discoveryengine.RankRequest(
    ranking_config=ranking_config,
    model=model_name,
    query=query,
    records=records,
    ignore_record_details_in_response=True,
)
resp = client.rank(request=request)

Scores batches of up to 200 query-document records with the ranking model.

Persist scores to GCS

score_output_path = f"{SCORE_PATH}/{subdataset_name.replace('-', '_')}/{MODEL_NAME.replace('-', '_')}.json"
save_dict_to_gcs_json(
    data=scores,
    bucket_name=OUTPUT_SCORES_BUCKET_NAME,
    file_path=score_output_path,
)

Stores computed ranking scores by BEIR subdataset for later evaluation.

Evaluate ranking metrics

k_values = [1, 3, 5, 10]
ndcg = evaluate(
    query_document_ground_truth,
    query_document_prediction,
    k_values,
    verbose=False,
)
results["ROC AUC"] = [float(roc_auc_score(np.array(y_true), np.array(y_pred)))]

Computes NDCG cutoffs and ROC AUC from ground truth and prediction scores.

Models & APIs used

  • Models: googles/semantic-ranker-default-004, semantic-ranker-default-004
  • APIs / services: Discovery Engine, Google Cloud Storage
  • SDKs / libraries: google-cloud-discoveryengine, google-cloud-storage, pytrec_eval, numpy, pandas, scikit-learn, tqdm, gcsfs

When to use this

Use this pattern to benchmark Discovery Engine ranking quality against labeled query-document relevance datasets.

Gotchas & caveats

  • Discovery Engine API must be enabled before scoring.
  • Colab requires explicit user authentication with google.colab.auth.authenticate_user().
  • PROJECT_ID defaults to GOOGLE_CLOUD_PROJECT if the placeholder is unchanged.
  • LOCATION defaults to us-central1 from GOOGLE_CLOUD_REGION when unset.
  • To rerun scoring, OUTPUT_SCORES_BUCKET_NAME should point to a bucket in the user’s project.
  • pytrec_eval expects positive integer ground truth labels, so labels are converted before ROC AUC calculation.
  • The notebook restricts evaluation to labeled datapoints and BEIR datasets with more than one unique label.

Best practices

  • Processes records in batches of 200 before calling the Rank API.
  • Uses ignore_record_details_in_response=True when only scores are needed.
  • Persists computed scores to Cloud Storage for reuse and inspection.
  • Filters evaluation to labeled datapoints to avoid underestimating model performance from unlabeled relevant examples.
  • Removes query-document pairs missing from either ground truth or predictions before metric calculation.
  • Reports NDCG@1, NDCG@3, NDCG@5, NDCG@10, ROC AUC, and macro averages across datasets.

Graph View

Backlinks