Replace correlation with proper ranking metrics (NDCG, MRR)

train_reranker.py

@@ -7,6 +7,7 @@
 #     "transformers>=4.48",
 #     "trackio",
 #     "scipy",
+#     "numpy",
 # ]
 # ///
 """
@@ -18,9 +19,10 @@ Dataset format: {"query": "...", "text": "...", "score": 0.0-1.0}
 
 import logging
 import os
+import math
 from collections import defaultdict
 import trackio
-import torch
+import numpy as np
 from datasets import load_dataset
 from sentence_transformers.cross_encoder import (
     CrossEncoder,
@@ -28,7 +30,7 @@ from sentence_transformers.cross_encoder import (
     CrossEncoderTrainingArguments,
 )
 from sentence_transformers.cross_encoder.evaluation import CrossEncoderNanoBEIREvaluator
-from scipy.stats import spearmanr, pearsonr
+from scipy.stats import spearmanr
 from transformers import TrainerCallback
 
 logging.basicConfig(level=logging.INFO)
@@ -46,31 +48,107 @@ RUN_NAME = os.environ.get("RUN_NAME", "reranker-03130903")
 SPACE_ID = os.environ.get("TRACKIO_SPACE_ID", "amanwithaplan/trackio")
 
 
-def evaluate_correlation(model, eval_dataset):
-    """Evaluate correlation between predicted scores and labels."""
-    pairs = [(item["sentence1"], item["sentence2"]) for item in eval_dataset]
-    labels = [item["label"] for item in eval_dataset]
+def dcg_at_k(relevances, k):
+    """Compute DCG@k."""
+    relevances = np.array(relevances)[:k]
+    if len(relevances) == 0:
+        return 0.0
+    # DCG = sum of rel_i / log2(i+2) for i in 0..k-1
+    discounts = np.log2(np.arange(len(relevances)) + 2)
+    return np.sum(relevances / discounts)
 
-    predictions = model.predict(pairs, show_progress_bar=True)
 
-    spearman = spearmanr(predictions, labels).correlation
-    pearson = pearsonr(predictions, labels).statistic
+def ndcg_at_k(predicted_order, true_relevances, k):
+    """
+    Compute NDCG@k.
 
-    # Mean absolute error
-    mae = sum(abs(p - l) for p, l in zip(predictions, labels)) / len(labels)
+    predicted_order: indices of docs sorted by model score (descending)
+    true_relevances: ground truth relevance scores for each doc
+    """
+    # Get relevances in predicted order
+    predicted_relevances = [true_relevances[i] for i in predicted_order]
+
+    # Ideal order: sort by true relevance descending
+    ideal_relevances = sorted(true_relevances, reverse=True)
+
+    dcg = dcg_at_k(predicted_relevances, k)
+    idcg = dcg_at_k(ideal_relevances, k)
+
+    if idcg == 0:
+        return 0.0
+    return dcg / idcg
+
+
+def mrr(predicted_order, true_relevances, threshold=0.5):
+    """
+    Compute MRR (Mean Reciprocal Rank).
+
+    Returns 1/rank of first relevant doc (relevance > threshold).
+    """
+    for rank, idx in enumerate(predicted_order, start=1):
+        if true_relevances[idx] > threshold:
+            return 1.0 / rank
+    return 0.0
+
+
+def evaluate_ranking(model, eval_dataset):
+    """
+    Proper ranking evaluation: group by query, compute NDCG and MRR.
+
+    This measures what we actually care about:
+    "Given a query with multiple docs, does the model rank them correctly?"
+    """
+    # Group samples by query
+    query_groups = defaultdict(list)
+    for item in eval_dataset:
+        query_groups[item["sentence1"]].append({
+            "text": item["sentence2"],
+            "label": item["label"]
+        })
+
+    # Filter to queries with multiple docs (need at least 2 to rank)
+    query_groups = {q: docs for q, docs in query_groups.items() if len(docs) >= 2}
+
+    if not query_groups:
+        return {"ndcg@3": 0.0, "ndcg@5": 0.0, "mrr": 0.0, "n_queries": 0}
+
+    ndcg_3_scores = []
+    ndcg_5_scores = []
+    mrr_scores = []
+    rank_correlations = []
+
+    for query, docs in query_groups.items():
+        # Get model predictions for this query's docs
+        pairs = [(query, d["text"]) for d in docs]
+        predictions = model.predict(pairs, show_progress_bar=False)
+
+        true_relevances = [d["label"] for d in docs]
+
+        # Get predicted order: indices sorted by prediction descending
+        predicted_order = np.argsort(predictions)[::-1].tolist()
+
+        # Compute metrics
+        ndcg_3_scores.append(ndcg_at_k(predicted_order, true_relevances, k=3))
+        ndcg_5_scores.append(ndcg_at_k(predicted_order, true_relevances, k=5))
+        mrr_scores.append(mrr(predicted_order, true_relevances, threshold=0.5))
+
+        # Rank correlation within this query
+        if len(set(true_relevances)) > 1:  # Need variance
+            corr = spearmanr(predictions, true_relevances).correlation
+            if not math.isnan(corr):
+                rank_correlations.append(corr)
 
     return {
-        "spearman": spearman,
-        "pearson": pearson,
-        "mae": mae,
-        "pred_mean": float(predictions.mean()),
-        "pred_std": float(predictions.std()),
-        "label_mean": sum(labels) / len(labels),
+        "ndcg@3": np.mean(ndcg_3_scores),
+        "ndcg@5": np.mean(ndcg_5_scores),
+        "mrr": np.mean(mrr_scores),
+        "rank_corr": np.mean(rank_correlations) if rank_correlations else 0.0,
+        "n_queries": len(query_groups),
     }
 
 
 class DomainEvalCallback(TrainerCallback):
-    """Callback to log our domain-specific correlation metrics during training."""
+    """Callback to log proper ranking metrics during training."""
 
     def __init__(self, model, eval_dataset_full):
         self.model = model
@@ -78,51 +156,53 @@ class DomainEvalCallback(TrainerCallback):
 
     def on_evaluate(self, args, state, control, **kwargs):
         """Run after each evaluation step."""
-        # Get correlation metrics
-        pairs = [(item["sentence1"], item["sentence2"]) for item in self.eval_dataset_full]
-        labels = [item["label"] for item in self.eval_dataset_full]
-
-        predictions = self.model.predict(pairs, show_progress_bar=False)
-
-        spearman = spearmanr(predictions, labels).correlation
-        pearson_val = pearsonr(predictions, labels).statistic
-        mae = sum(abs(p - l) for p, l in zip(predictions, labels)) / len(labels)
+        metrics = evaluate_ranking(self.model, self.eval_dataset_full)
 
         # Log to trackio
         trackio.log({
-            "domain/spearman": spearman,
-            "domain/pearson": pearson_val,
-            "domain/mae": float(mae),
-            "domain/pred_mean": float(predictions.mean()),
-            "domain/pred_std": float(predictions.std()),
+            "domain/ndcg@3": metrics["ndcg@3"],
+            "domain/ndcg@5": metrics["ndcg@5"],
+            "domain/mrr": metrics["mrr"],
+            "domain/rank_corr": metrics["rank_corr"],
         })
 
-        logger.info(f"Domain eval - Spearman: {spearman:.4f}, Pearson: {pearson_val:.4f}, MAE: {mae:.4f}")
+        logger.info(
+            f"Domain eval - NDCG@3: {metrics['ndcg@3']:.4f}, "
+            f"NDCG@5: {metrics['ndcg@5']:.4f}, "
+            f"MRR: {metrics['mrr']:.4f}, "
+            f"RankCorr: {metrics['rank_corr']:.4f} "
+            f"(n={metrics['n_queries']} queries)"
+        )
 
 
 def evaluate_by_type(model, eval_dataset, type_column="type"):
-    """Evaluate correlation per content type."""
+    """Evaluate ranking metrics per content type."""
     if type_column not in eval_dataset.column_names:
         return {}
 
-    # Group by type
+    # Group by type first
     by_type = defaultdict(list)
     for item in eval_dataset:
         by_type[item[type_column]].append(item)
 
     results = {}
     for content_type, items in by_type.items():
-        if len(items) < 5:
-            continue
-
-        pairs = [(item["sentence1"], item["sentence2"]) for item in items]
-        labels = [item["label"] for item in items]
-        predictions = model.predict(pairs)
-
-        if len(set(labels)) > 1:  # Need variance for correlation
-            results[f"{content_type}_spearman"] = spearmanr(predictions, labels).correlation
-            results[f"{content_type}_mae"] = sum(abs(p - l) for p, l in zip(predictions, labels)) / len(labels)
-            results[f"{content_type}_n"] = len(items)
+        # Create a mini dataset for this type
+        class TypeDataset:
+            def __init__(self, items):
+                self.items = items
+            def __iter__(self):
+                return iter(self.items)
+            @property
+            def column_names(self):
+                return ["sentence1", "sentence2", "label"]
+
+        type_metrics = evaluate_ranking(model, TypeDataset(items))
+
+        if type_metrics["n_queries"] >= 2:
+            results[f"{content_type}_ndcg@5"] = type_metrics["ndcg@5"]
+            results[f"{content_type}_mrr"] = type_metrics["mrr"]
+            results[f"{content_type}_n_queries"] = type_metrics["n_queries"]
 
     return results
 
@@ -193,9 +273,9 @@ def main():
     })
     logger.info(f"Train: {len(train_dataset)}, Eval: {len(eval_dataset)}")
 
-    # Evaluate base model before training
+    # Evaluate base model before training with proper ranking metrics
     logger.info("Evaluating base model on eval set...")
-    base_metrics = evaluate_correlation(model, eval_dataset)
+    base_metrics = evaluate_ranking(model, eval_dataset_full)
     for key, value in base_metrics.items():
         trackio.log({f"base_model/{key}": value})
     logger.info(f"Base model metrics: {base_metrics}")
@@ -231,7 +311,7 @@ def main():
         run_name=RUN_NAME,
     )
 
-    # Custom callback to log domain-specific metrics during training
+    # Custom callback to log domain-specific ranking metrics during training
     domain_callback = DomainEvalCallback(model, eval_dataset_full)
 
     trainer = CrossEncoderTrainer(
@@ -246,14 +326,14 @@ def main():
     logger.info("Starting training...")
     trainer.train()
 
-    # Final evaluation on our eval set
-    logger.info("Running final correlation evaluation...")
-    final_metrics = evaluate_correlation(model, eval_dataset)
+    # Final evaluation with proper ranking metrics
+    logger.info("Running final ranking evaluation...")
+    final_metrics = evaluate_ranking(model, eval_dataset_full)
     for key, value in final_metrics.items():
         trackio.log({f"final/{key}": value})
     logger.info(f"Final metrics: {final_metrics}")
 
-    # Per-type evaluation (use full eval dataset with type column)
+    # Per-type evaluation
     logger.info("Evaluating by content type...")
     type_metrics = evaluate_by_type(model, eval_dataset_full)
     for key, value in type_metrics.items():
@@ -262,8 +342,8 @@ def main():
 
     # Log improvement
     trackio.log({
-        "improvement/spearman_delta": final_metrics["spearman"] - base_metrics["spearman"],
-        "improvement/mae_delta": base_metrics["mae"] - final_metrics["mae"],  # Lower is better
+        "improvement/ndcg5_delta": final_metrics["ndcg@5"] - base_metrics["ndcg@5"],
+        "improvement/mrr_delta": final_metrics["mrr"] - base_metrics["mrr"],
     })
 
     logger.info(f"Pushing final model to {HUB_MODEL_ID}")