From b712d9868ff178239d0aa90f55d8ca45da924f2b Mon Sep 17 00:00:00 2001 From: epwalsh Date: Mon, 25 Nov 2024 16:15:59 -0800 Subject: [PATCH] add more tasks --- src/olmo_core/model_ladder.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/olmo_core/model_ladder.py b/src/olmo_core/model_ladder.py index 2ae443dd..b332bbfa 100644 --- a/src/olmo_core/model_ladder.py +++ b/src/olmo_core/model_ladder.py @@ -263,6 +263,8 @@ def get_trainer_config( :param size: The target model size. :param gpu_type: The type of GPU as given by ``torch.cuda.get_device_name()``. """ + from olmo_eval import list_tasks + rank_mbz = self.get_rank_microbatch_size(size=size, gpu_type=gpu_type) if rank_mbz % self.sequence_length != 0: raise OLMoConfigurationError( @@ -303,7 +305,9 @@ def get_trainer_config( .with_callback( "downstream_evaluator", DownstreamEvaluatorCallbackConfig( - tasks=["hellaswag"], # TODO: which other tasks? + tasks=[ + task for task in list_tasks() if "_mc" not in task and "_var" not in task + ], tokenizer=self.tokenizer, eval_interval=250, ),