Add trackio metrics and correlation evaluation

train_reranker.py

@@ -6,6 +6,7 @@
 #     "torch>=2.4",
 #     "transformers>=4.48",
 #     "trackio",
+#     "scipy",
 # ]
 # ///
 """
@@ -18,6 +19,8 @@ Dataset format: {"query": "...", "text": "...", "score": 0.0-1.0}
 import logging
 import os
 from collections import defaultdict
+import trackio
+import torch
 from datasets import load_dataset
 from sentence_transformers.cross_encoder import (
     CrossEncoder,
@@ -25,6 +28,7 @@ from sentence_transformers.cross_encoder import (
     CrossEncoderTrainingArguments,
 )
 from sentence_transformers.cross_encoder.evaluation import CrossEncoderNanoBEIREvaluator
+from scipy.stats import spearmanr, pearsonr
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -38,14 +42,81 @@ BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "16"))
 LEARNING_RATE = float(os.environ.get("LEARNING_RATE", "2e-5"))
 MAX_SEQ_LENGTH = int(os.environ.get("MAX_SEQ_LENGTH", "512"))
 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]
+
+    predictions = model.predict(pairs, show_progress_bar=True)
+
+    spearman = spearmanr(predictions, labels).correlation
+    pearson = pearsonr(predictions, labels).statistic
+
+    # Mean absolute error
+    mae = sum(abs(p - l) for p, l in zip(predictions, labels)) / len(labels)
+
+    return {
+        "spearman": spearman,
+        "pearson": pearson,
+        "mae": mae,
+        "pred_mean": float(predictions.mean()),
+        "pred_std": float(predictions.std()),
+        "label_mean": sum(labels) / len(labels),
+    }
+
+
+def evaluate_by_type(model, eval_dataset, type_column="type"):
+    """Evaluate correlation per content type."""
+    if type_column not in eval_dataset.column_names:
+        return {}
+
+    # Group by type
+    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)
+
+    return results
 
 
 def main():
+    # Initialize trackio with full config
+    trackio.init(
+        project="arcade-reranker",
+        name=RUN_NAME,
+        space_id=SPACE_ID,
+        config={
+            "model": BASE_MODEL,
+            "dataset": DATASET_NAME,
+            "learning_rate": LEARNING_RATE,
+            "num_epochs": NUM_EPOCHS,
+            "batch_size": BATCH_SIZE,
+            "max_seq_length": MAX_SEQ_LENGTH,
+        }
+    )
+
     logger.info(f"Configuration:")
     logger.info(f"  Dataset: {DATASET_NAME}")
     logger.info(f"  Base model: {BASE_MODEL}")
     logger.info(f"  Epochs: {NUM_EPOCHS}")
     logger.info(f"  Run name: {RUN_NAME}")
+    logger.info(f"  Trackio space: {SPACE_ID}")
 
     model = CrossEncoder(BASE_MODEL, max_length=MAX_SEQ_LENGTH)
 
@@ -53,12 +124,17 @@ def main():
     dataset = load_dataset(DATASET_NAME, split="train")
 
     # Log dataset composition
+    type_counts = defaultdict(int)
     if "type" in dataset.column_names:
-        type_counts = defaultdict(int)
         for item in dataset:
             type_counts[item["type"]] += 1
         logger.info(f"Dataset composition: {dict(type_counts)}")
 
+        # Log to trackio
+        for content_type, count in type_counts.items():
+            trackio.log({f"data/{content_type}_count": count})
+
+    trackio.log({"data/total_examples": len(dataset)})
     logger.info(f"Total examples: {len(dataset)}")
 
     # Rename columns for CrossEncoderTrainer
@@ -74,8 +150,19 @@ def main():
     train_dataset = splits["train"]
     eval_dataset = splits["test"]
 
+    trackio.log({
+        "data/train_size": len(train_dataset),
+        "data/eval_size": len(eval_dataset),
+    })
     logger.info(f"Train: {len(train_dataset)}, Eval: {len(eval_dataset)}")
 
+    # Evaluate base model before training
+    logger.info("Evaluating base model on eval set...")
+    base_metrics = evaluate_correlation(model, eval_dataset)
+    for key, value in base_metrics.items():
+        trackio.log({f"base_model/{key}": value})
+    logger.info(f"Base model metrics: {base_metrics}")
+
     # NanoBEIR for benchmark comparison
     evaluator = CrossEncoderNanoBEIREvaluator(
         dataset_names=["msmarco", "nfcorpus", "nq"],
@@ -118,8 +205,30 @@ 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)
+    for key, value in final_metrics.items():
+        trackio.log({f"final/{key}": value})
+    logger.info(f"Final metrics: {final_metrics}")
+
+    # Per-type evaluation
+    logger.info("Evaluating by content type...")
+    type_metrics = evaluate_by_type(model, eval_dataset)
+    for key, value in type_metrics.items():
+        trackio.log({f"final/by_type/{key}": value})
+    logger.info(f"Per-type metrics: {type_metrics}")
+
+    # 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
+    })
+
     logger.info(f"Pushing final model to {HUB_MODEL_ID}")
     model.push_to_hub(HUB_MODEL_ID)
+
+    trackio.finish()
     logger.info("Done!")