Fix: select only required columns for training

train_reranker.py

@@ -144,11 +144,16 @@ def main():
         "score": "label"
     })
 
-    # Split for evaluation
+    # Split for evaluation (before removing extra columns so we keep type for eval)
     eval_size = min(400, int(len(dataset) * 0.15))
     splits = dataset.train_test_split(test_size=eval_size, seed=42)
-    train_dataset = splits["train"]
-    eval_dataset = splits["test"]
+
+    # Keep full eval dataset with type column for per-type evaluation
+    eval_dataset_full = splits["test"]
+
+    # Remove extra columns for training (CrossEncoderTrainer only wants sentence1, sentence2, label)
+    train_dataset = splits["train"].select_columns(["sentence1", "sentence2", "label"])
+    eval_dataset = splits["test"].select_columns(["sentence1", "sentence2", "label"])
 
     trackio.log({
         "data/train_size": len(train_dataset),
@@ -212,9 +217,9 @@ def main():
         trackio.log({f"final/{key}": value})
     logger.info(f"Final metrics: {final_metrics}")
 
-    # Per-type evaluation
+    # Per-type evaluation (use full eval dataset with type column)
     logger.info("Evaluating by content type...")
-    type_metrics = evaluate_by_type(model, eval_dataset)
+    type_metrics = evaluate_by_type(model, eval_dataset_full)
     for key, value in type_metrics.items():
         trackio.log({f"final/by_type/{key}": value})
     logger.info(f"Per-type metrics: {type_metrics}")