track metadata in predictions

geneformer/classifier.py

@@ -801,7 +801,7 @@ class Classifier:
                     # 5-fold cross-validate
                     num_cells = len(data)
                     fifth_cells = int(np.floor(num_cells * 0.2))
-                    num_eval = min((self.eval_size * num_cells), fifth_cells)
+                    num_eval = int(min((self.eval_size * num_cells), fifth_cells))
                     start = i * fifth_cells
                     end = start + num_eval
                     eval_indices = [j for j in range(start, end)]
@@ -1313,6 +1313,7 @@ class Classifier:
         predict=False,
         output_directory=None,
         output_prefix=None,
+        predict_metadata=None,
     ):
         """
         Evaluate the fine-tuned model.
@@ -1338,9 +1339,11 @@ class Classifier:
 
         ##### Evaluate the model #####
         labels = id_class_dict.keys()
-        y_pred, y_true, logits_list = eu.classifier_predict(
-            model, self.classifier, eval_data, self.forward_batch_size, self.gene_token_dict
+
+        y_pred, y_true, logits_list, predict_metadata_all = eu.classifier_predict(
+            model, self.classifier, eval_data, self.forward_batch_size, self.gene_token_dict, predict_metadata
         )
+
         conf_mat, macro_f1, acc, roc_metrics = eu.get_metrics(
             y_pred, y_true, logits_list, num_classes, labels
         )
@@ -1350,6 +1353,9 @@ class Classifier:
                 "label_ids": y_true,
                 "predictions": logits_list,
             }
+            if predict_metadata is not None:
+                pred_dict["prediction_metadata"] = predict_metadata_all
+                
             pred_dict_output_path = (
                 Path(output_directory) / f"{output_prefix}_pred_dict"
             ).with_suffix(".pkl")
@@ -1370,6 +1376,7 @@ class Classifier:
         output_directory,
         output_prefix,
         predict=True,
+        predict_metadata=None,
     ):
         """
         Evaluate the fine-tuned model.
@@ -1389,6 +1396,8 @@ class Classifier:
             | Prefix for output files
         predict : bool
             | Whether or not to save eval predictions
+        predict_metadata : None | list
+            | Metadata labels to output with predictions (columns in test_data_file)
         """
 
         # load numerical id to class dictionary (id:class)
@@ -1401,6 +1410,15 @@ class Classifier:
         # load previously filtered and prepared data
         test_data = pu.load_and_filter(None, self.nproc, test_data_file)
 
+        if predict_metadata is not None:
+            absent_metadata = []
+            for predict_metadata_x in predict_metadata:
+                if predict_metadata_x not in test_data.features.keys():
+                    absent_metadata += [predict_metadata_x]
+            if len(absent_metadata)>0:
+                logger.error(f"Following predict_metadata was not found as column in test_data_file: {absent_metadata}")
+                raise
+                    
         # load previously fine-tuned model
         model = pu.load_model(
             self.model_type,
@@ -1419,6 +1437,7 @@ class Classifier:
             predict=predict,
             output_directory=output_directory,
             output_prefix=output_prefix,
+            predict_metadata=predict_metadata,
         )
 
         all_conf_mat_df = pd.DataFrame(

geneformer/classifier_utils.py

@@ -570,6 +570,27 @@ def compute_metrics(pred):
     return {"accuracy": acc, "macro_f1": macro_f1}
 
 
+def robust_compute_objective(metrics: dict):
+    # tries both prefixed ("eval_") and raw metric names to support different transformers versions
+    metric_name = "macro_f1"
+    
+    # check for the prefixed version
+    prefixed_metric_name = f"eval_{metric_name}"
+    if prefixed_metric_name in metrics:
+        return metrics[prefixed_metric_name]
+    
+    # fall back to the raw name
+    elif metric_name in metrics:
+        return metrics[metric_name]
+    
+    # if neither is found, raise a clear error to help with debugging
+    raise KeyError(
+        f"Could not find '{prefixed_metric_name}' or '{metric_name}' in the reported metrics. "
+        f"Please check your `compute_metrics` function and `TrainingArguments`. "
+        f"Available metrics: {list(metrics.keys())}"
+    )
+
+
 def get_default_train_args(model, classifier, data, output_dir):
     num_layers = pu.quant_layers(model)
     freeze_layers = 0

geneformer/evaluation_utils.py

@@ -77,7 +77,7 @@ def py_softmax(vector):
     return e / e.sum()
 
 
-def classifier_predict(model, classifier_type, evalset, forward_batch_size, gene_token_dict):
+def classifier_predict(model, classifier_type, evalset, forward_batch_size, gene_token_dict, predict_metadata=None):
     if classifier_type == "gene":
         label_name = "labels"
     elif classifier_type == "cell":
@@ -85,6 +85,14 @@ def classifier_predict(model, classifier_type, evalset, forward_batch_size, gene
 
     predict_logits = []
     predict_labels = []
+
+    predict_metadata_all = None
+    
+    if predict_metadata is not None:
+        predict_metadata_all = dict()
+        for metadata_name in predict_metadata:
+            predict_metadata_all[metadata_name] = []
+    
     model.eval()
 
     # ensure there is at least 2 examples in each batch to avoid incorrect tensor dims
@@ -99,9 +107,15 @@ def classifier_predict(model, classifier_type, evalset, forward_batch_size, gene
     for i in trange(0, evalset_len, forward_batch_size):
         max_range = min(i + forward_batch_size, evalset_len)
         batch_evalset = evalset.select([i for i in range(i, max_range)])
+
+        if predict_metadata is not None:
+            for metadata_name in predict_metadata:
+                predict_metadata_all[metadata_name] += batch_evalset[metadata_name]
+        
         padded_batch = preprocess_classifier_batch(
             batch_evalset, max_evalset_len, label_name, gene_token_dict
         )
+            
         padded_batch.set_format(type="torch")
 
         # For datasets>=4.0.0, convert to dict to avoid format issues
@@ -134,7 +148,8 @@ def classifier_predict(model, classifier_type, evalset, forward_batch_size, gene
     y_pred = [vote(item[0]) for item in logit_label_paired]
     y_true = [item[1] for item in logit_label_paired]
     logits_list = [item[0] for item in logit_label_paired]
-    return y_pred, y_true, logits_list
+
+    return y_pred, y_true, logits_list, predict_metadata_all
 
 
 def get_metrics(y_pred, y_true, logits_list, num_classes, labels):