Refactoring the add_metrics_to_eval_loaders function to accept list of metric names instead of a dictionary of metrics.

llmfoundry/utils/builders.py

@@ -28,7 +28,6 @@
 from omegaconf import DictConfig, ListConfig
 from omegaconf import OmegaConf as om
 from torch.optim.optimizer import Optimizer
-from torchmetrics import Metric
 from transformers import AutoTokenizer, PreTrainedTokenizerBase
 
 from llmfoundry.callbacks import (AsyncEval, EvalGauntlet, FDiffMetrics,
@@ -108,9 +107,8 @@ def build_eval_loaders(
 
 def add_metrics_to_eval_loaders(
     evaluators: List[Evaluator],
-    metrics: Dict[str, Metric],
+    metric_names: List[str],
 ) -> List[Evaluator]:
-    metric_names = list(metrics.keys())
     eval_loaders, other_evaluators = [], []
     for evaluator in evaluators:
         if evaluator.metric_names == []:

scripts/eval/eval.py

@@ -184,7 +184,8 @@ def evaluate_model(
     # Now add the eval metrics
     if eval_loader_config is not None:
         train_metrics = composer_model.get_metrics(is_train=True)
-        evaluators = add_metrics_to_eval_loaders(evaluators, train_metrics)
+        evaluators = add_metrics_to_eval_loaders(evaluators,
+                                                 list(train_metrics.keys()))
 
     if eval_gauntlet_df is None and eval_gauntlet_callback is not None:
         eval_gauntlet_df = pd.DataFrame(

scripts/train/train.py

@@ -544,7 +544,8 @@ def main(cfg: DictConfig) -> Trainer:
     # Now add the eval metrics
     if eval_loader_config is not None and not use_async_eval:
         train_metrics = model.get_metrics(is_train=True)
-        evaluators = add_metrics_to_eval_loaders(evaluators, train_metrics)
+        evaluators = add_metrics_to_eval_loaders(evaluators,
+                                                 list(train_metrics.keys()))
 
     # Build the Trainer
     log.info('Building trainer...')

tests/utils/test_builders.py

@@ -335,13 +335,7 @@ def test_add_metrics_to_eval_loaders():
         )
     ]
 
-    new_evaluators = add_metrics_to_eval_loaders(
-        evaluators,
-        {
-            'new1': 'foo',
-            'new2': 'bar'
-        },  # type: ignore
-    )
+    new_evaluators = add_metrics_to_eval_loaders(evaluators, ['new1', 'new2'])
     assert len(new_evaluators) == 3
 
     assert new_evaluators[0].label == 'second'