Add downstream evals to example

allenai · Oct 29, 2024 · f30358e · f30358e
1 parent 8dc0a72
commit f30358e
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/src/examples/train.py b/src/examples/train.py
@@ -31,6 +31,7 @@
     CheckpointerCallback,
     CometCallback,
     ConfigSaverCallback,
+    DownstreamEvaluatorCallbackConfig,
     GPUMemoryMonitorCallback,
     GradClipperCallback,
     LMEvaluatorCallbackConfig,
@@ -133,7 +134,7 @@ def build_config(run_name: str, overrides: List[str]) -> ExperimentConfig:
         .with_callback("config_saver", ConfigSaverCallback())
         .with_callback("profiler", ProfilerCallback(enabled=False))
         .with_callback(
-            "evaluator",
+            "lm_evaluator",
             LMEvaluatorCallbackConfig(
                 eval_dataset=NumpyDatasetConfig(
                     paths=["/net/nfs/allennlp/llm-data/c4/en/c4-validation.00000-00008.npy"],
@@ -147,6 +148,15 @@ def build_config(run_name: str, overrides: List[str]) -> ExperimentConfig:
                 eval_duration=Duration.steps(10),
             ),
         )
+        .with_callback(
+            "downstream_evaluator",
+            DownstreamEvaluatorCallbackConfig(
+                tasks=["hellaswag"],
+                tokenizer=tokenizer_config,
+                eval_rank_batch_size_instances=4,
+                eval_interval=250,
+            ),
+        )
     )
 
     return ExperimentConfig(