Fix: Reduce CLI loading time from 2.2s to 0.05s (#3099)

* Lazy-load submodules to reduce import time

* pacify pre-commit

---------

Co-authored-by: Baber <[email protected]>

Author: stakodiak

lm_eval/__init__.py

@@ -1,7 +1,21 @@
 import logging
 import os
 
-from .evaluator import evaluate, simple_evaluate
-
 
 __version__ = "0.4.9"
+
+
+# Lazy-load .evaluator module to improve CLI startup
+def __getattr__(name):
+    if name == "evaluate":
+        from .evaluator import evaluate
+
+        return evaluate
+    elif name == "simple_evaluate":
+        from .evaluator import simple_evaluate
+
+        return simple_evaluate
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = ["evaluate", "simple_evaluate", "__version__"]

lm_eval/__main__.py

@@ -7,16 +7,6 @@
 from pathlib import Path
 from typing import Union
 
-from lm_eval import evaluator, utils
-from lm_eval.evaluator import request_caching_arg_to_dict
-from lm_eval.loggers import EvaluationTracker, WandbLogger
-from lm_eval.tasks import TaskManager
-from lm_eval.utils import (
-    handle_non_serializable,
-    make_table,
-    simple_parse_args_string,
-)
-
 
 def try_parse_json(value: str) -> Union[str, dict, None]:
     if value is None:
@@ -314,6 +304,17 @@ def cli_evaluate(args: Union[argparse.Namespace, None] = None) -> None:
         parser = setup_parser()
         args = parse_eval_args(parser)
 
+    # defer loading `lm_eval` submodules for faster CLI load
+    from lm_eval import evaluator, utils
+    from lm_eval.evaluator import request_caching_arg_to_dict
+    from lm_eval.loggers import EvaluationTracker, WandbLogger
+    from lm_eval.tasks import TaskManager
+    from lm_eval.utils import (
+        handle_non_serializable,
+        make_table,
+        simple_parse_args_string,
+    )
+
     if args.wandb_args:
         wandb_args_dict = simple_parse_args_string(args.wandb_args)
         wandb_config_args_dict = simple_parse_args_string(args.wandb_config_args)