From 40f2d19fc1c2b2c39313ad62697009272a758f3d Mon Sep 17 00:00:00 2001 From: ayulockin Date: Thu, 22 Feb 2024 17:33:14 +0000 Subject: [PATCH] track system metrics + total eval time + stdout --- lm_eval/__main__.py | 11 ++++++++--- lm_eval/logging_utils.py | 11 ++++++----- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/lm_eval/__main__.py b/lm_eval/__main__.py index eb6fbeaf0f..8b02446d3b 100644 --- a/lm_eval/__main__.py +++ b/lm_eval/__main__.py @@ -201,6 +201,9 @@ def cli_evaluate(args: Union[argparse.Namespace, None] = None) -> None: # we allow for args to be passed externally, else we parse them ourselves args = parse_eval_args() + if args.wandb_args: + wandb_logger = WandbLogger(args) + eval_logger = utils.eval_logger eval_logger.setLevel(getattr(logging, f"{args.verbosity}")) eval_logger.info(f"Verbosity set to {args.verbosity}") @@ -318,12 +321,10 @@ def cli_evaluate(args: Union[argparse.Namespace, None] = None) -> None: # Add W&B logging if args.wandb_args: try: - wandb_logger = WandbLogger(results, args) + wandb_logger.post_init(results) wandb_logger.log_eval_result() if args.log_samples: wandb_logger.log_eval_samples(samples) - # Tear down wandb run once all the logging is done. - wandb_logger.run.finish() except Exception as e: eval_logger.info(f"Logging to Weights and Biases failed due to {e}") @@ -352,6 +353,10 @@ def cli_evaluate(args: Union[argparse.Namespace, None] = None) -> None: if "groups" in results: print(make_table(results, "groups")) + if args.wandb_args: + # Tear down wandb run once all the logging is done. + wandb_logger.run.finish() + if __name__ == "__main__": cli_evaluate() diff --git a/lm_eval/logging_utils.py b/lm_eval/logging_utils.py index f7ae0efeb8..464ba9f732 100644 --- a/lm_eval/logging_utils.py +++ b/lm_eval/logging_utils.py @@ -78,21 +78,17 @@ def get_wandb_printer() -> Literal["Printer"]: class WandbLogger: - def __init__(self, results: Dict[str, Any], args: Any) -> None: + def __init__(self, args: Any) -> None: """Initialize the WandbLogger. Args: results (Dict[str, Any]): The results dictionary. args (Any): Arguments for configuration. """ - self.results: Dict[str, Any] = copy.deepcopy(results) self.wandb_args: Dict[str, Any] = utils.simple_parse_args_string( args.wandb_args ) - self.task_names: List[str] = list(results.get("results", {}).keys()) - self.group_names: List[str] = list(results.get("groups", {}).keys()) - # initialize a W&B run if wandb.run is None: self.run = wandb.init(**self.wandb_args) @@ -101,6 +97,11 @@ def __init__(self, results: Dict[str, Any], args: Any) -> None: self.printer = get_wandb_printer() + def post_init(self, results: Dict[str, Any]) -> None: + self.results: Dict[str, Any] = copy.deepcopy(results) + self.task_names: List[str] = list(results.get("results", {}).keys()) + self.group_names: List[str] = list(results.get("groups", {}).keys()) + def _get_config(self) -> Dict[str, Any]: """Get configuration parameters.""" self.task_configs = self.results.get("configs", {})