make the metrics work for ft as well

Author: yawenzzzz

scripts/get_max_eval_metrics_from_wandb.py

@@ -16,7 +16,6 @@
 from olmoearth_pretrain.train.callbacks.evaluator_callback import EvalMode
 
 WANDB_ENTITY = "eai-ai2"
-METRICS = list(EVAL_TASKS.keys())
 
 # Dataset partitions to consider (excluding default)
 PARTITIONS = [
@@ -375,21 +374,18 @@ def save_metrics_to_csv(metrics_dict: dict[str, dict[str, float]], filename: str
         help="Aggregate metrics per dataset partition instead of grouping by '_step'",
     )
     parser.add_argument(
-        "--get_test_metrics",
+        "--finetune",
         action="store_true",
-        help="Report test metrics based on the configuration of the validation results witht the highest score",
+        help="Use finetune evaluation tasks when determining metrics",
     )
     parser.add_argument(
-        "--finetune",
+        "--get_test_metrics",
         action="store_true",
-        help="Use finetune evaluation tasks when determining metrics",
+        help="Report test metrics based on the configuration of the validation results witht the highest score",
     )
 
     args = parser.parse_args()
-
-    global METRICS
-    selected_tasks = FT_EVAL_TASKS if args.finetune else EVAL_TASKS
-    METRICS = list(selected_tasks.keys())
+    metrics = list(FT_EVAL_TASKS.keys()) if args.finetune else list(EVAL_TASKS.keys())
 
     if args.per_partition:
         if not args.run_prefix:
@@ -404,7 +400,7 @@ def save_metrics_to_csv(metrics_dict: dict[str, dict[str, float]], filename: str
         for partition in PARTITIONS:
             if partition in partition_metrics:
                 print(f"\n{partition}:")
-                for metric in METRICS:
+                for metric in metrics:
                     # Try original name
                     key = f"eval/{metric}"
                     val = partition_metrics[partition].get(key)
@@ -445,7 +441,7 @@ def save_metrics_to_csv(metrics_dict: dict[str, dict[str, float]], filename: str
         print("\nFinal Results:")
         for group_name, metrics in group_metrics.items():
             print(f"\n{group_name}:")
-            for metric in METRICS:
+            for metric in metrics:
                 try:
                     k = f"eval/{metric}"
                     print(f"  {metric}: {metrics[k]}")
@@ -460,7 +456,7 @@ def save_metrics_to_csv(metrics_dict: dict[str, dict[str, float]], filename: str
             print("\nFinal Test Results:")
             for group_name, metrics in group_test_metrics.items():
                 print(f"\n{group_name}:")
-                for metric in METRICS:
+                for metric in metrics:
                     try:
                         k = f"eval/test/{metric}"
                         print(f"  {metric}: {metrics[k]}")