cleanup

lm_eval/tasks/__init__.py

@@ -19,5 +19,12 @@ for file in os.listdir(tasks_dir):
 ALL_TASKS = sorted(list(TASK_REGISTRY.registry))
 
 
-def get_task(model_name):
-    return TASK_REGISTRY.registry[model_name]
+def get_task(task_name):
+    return TASK_REGISTRY.registry[task_name]
+
+
+def get_task_dict(task_name_list):
+    return {
+        task_name: get_task(task_name)()
+        for task_name in task_name_list
+    }

main.py

@@ -28,10 +28,7 @@ def main():
         task_names = tasks.ALL_TASKS
     else:
         task_names = args.tasks.split(",")
-    task_dict = {
-        task_name: tasks.get_task(task_name)()
-        for task_name in task_names
-    }
+    task_dict = tasks.get_task_dict(task_names)
     results = {}
     for task_name, task in task_dict.items():
         if not task.has_validation_docs():
@@ -43,6 +40,7 @@ def main():
             num_fewshot=args.num_fewshot,
         )
         results[task_name] = result
+                
     dumped = json.dumps(results, indent=2)
     print(dumped)
     if args.output_path:

write_out.py

@@ -28,10 +28,7 @@ def main():
         task_names = tasks.ALL_TASKS
     else:
         task_names = args.tasks.split(",")
-    task_dict = {
-        task_name: tasks.get_task(task_name)()
-        for task_name in task_names
-    }
+    task_dict = tasks.get_task_dict(task_names)
     os.makedirs(args.output_base_path, exist_ok=True)
     for task_name, task in task_dict.items():
         if not task.has_validation_docs():