color to logger

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}")

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:

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}")