Add domain metrics callback + fix push_to_hub exist_ok

train_reranker.py

@@ -29,6 +29,7 @@ from sentence_transformers.cross_encoder import (
 )
 from sentence_transformers.cross_encoder.evaluation import CrossEncoderNanoBEIREvaluator
 from scipy.stats import spearmanr, pearsonr
+from transformers import TrainerCallback
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -68,6 +69,37 @@ def evaluate_correlation(model, eval_dataset):
     }
 
 
+class DomainEvalCallback(TrainerCallback):
+    """Callback to log our domain-specific correlation metrics during training."""
+
+    def __init__(self, model, eval_dataset_full):
+        self.model = model
+        self.eval_dataset_full = eval_dataset_full
+
+    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)
+
+        # 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()),
+        })
+
+        logger.info(f"Domain eval - Spearman: {spearman:.4f}, Pearson: {pearson_val:.4f}, MAE: {mae:.4f}")
+
+
 def evaluate_by_type(model, eval_dataset, type_column="type"):
     """Evaluate correlation per content type."""
     if type_column not in eval_dataset.column_names:
@@ -199,12 +231,16 @@ def main():
         run_name=RUN_NAME,
     )
 
+    # Custom callback to log domain-specific metrics during training
+    domain_callback = DomainEvalCallback(model, eval_dataset_full)
+
     trainer = CrossEncoderTrainer(
         model=model,
         args=args,
         train_dataset=train_dataset,
         eval_dataset=eval_dataset,
         evaluator=evaluator,
+        callbacks=[domain_callback],
     )
 
     logger.info("Starting training...")
@@ -231,7 +267,7 @@ def main():
     })
 
     logger.info(f"Pushing final model to {HUB_MODEL_ID}")
-    model.push_to_hub(HUB_MODEL_ID)
+    model.push_to_hub(HUB_MODEL_ID, exist_ok=True)
 
     trackio.finish()
     logger.info("Done!")