color to logger

PhilChina · Jul 25, 2021 · 64ad8df · 64ad8df
1 parent 263cdb1
commit 64ad8df
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 3 deletions.
diff --git a/scripts/consolidate.py b/scripts/consolidate.py
@@ -138,7 +138,12 @@ def main():
     target_dir = model_dir / "consolidated"
 
     logger.remove()
-    logger.add(sys.stdout, format="{level} {message}", level="INFO")
+    logger.add(
+        sys.stdout,
+        format="<level>{level} {message}</level>",
+        level="INFO",
+        colorize=True,
+        )
     logger.add(Path(target_dir) / "consolidate.log", level="DEBUG")
 
     logger.info(f"looking for models in {model_dir}")

diff --git a/scripts/predict.py b/scripts/predict.py
@@ -63,7 +63,12 @@ def run(cfg: dict,
     prediction_dir = training_dir / "test_predictions"
 
     logger.remove()
-    logger.add(sys.stdout, format="{level} {message}", level="INFO")
+    logger.add(
+        sys.stdout,
+        format="<level>{level} {message}</level>",
+        level="INFO",
+        colorize=True,
+        )
     logger.add(Path(training_dir) / "inference.log", level="INFO")
 
     if process:

diff --git a/scripts/train.py b/scripts/train.py
@@ -201,7 +201,12 @@ def _train(
         {"trainer": OmegaConf.to_container(cfg["trainer_cfg"], resolve=True)}))
 
     logger.remove()
-    logger.add(sys.stdout, format="{level} {message}", level="INFO")
+    logger.add(
+        sys.stdout,
+        format="<level>{level} {message}</level>",
+        level="INFO",
+        colorize=True,
+        )
     log_file = Path(os.getcwd()) / "train.log"
     logger.add(log_file, level="INFO")
     logger.info(f"Log file at {log_file}")