diff --git a/requirements.txt b/requirements.txt index a04901e5b..ed7e0fb7f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -12,6 +12,7 @@ einops xformers optimum hf_transfer +colorama numba numpy==1.24.4 # qlora things diff --git a/src/axolotl/logging_config.py b/src/axolotl/logging_config.py index 546e64a77..bec20cbad 100644 --- a/src/axolotl/logging_config.py +++ b/src/axolotl/logging_config.py @@ -1,16 +1,42 @@ -"""Logging configuration settings""" +""" +Common logging module for axolotl +""" import os import sys +from logging import Formatter from logging.config import dictConfig from typing import Any, Dict +from colorama import Fore, Style, init + + +class ColorfulFormatter(Formatter): + """ + Formatter to add coloring to log messages by log type + """ + + COLORS = { + "WARNING": Fore.YELLOW, + "ERROR": Fore.RED, + "CRITICAL": Fore.RED + Style.BRIGHT, + } + + def format(self, record): + log_message = super().format(record) + return self.COLORS.get(record.levelname, "") + log_message + Fore.RESET + + DEFAULT_LOGGING_CONFIG: Dict[str, Any] = { "version": 1, "formatters": { "simple": { "format": "[%(asctime)s] [%(levelname)s] [%(name)s.%(funcName)s:%(lineno)d] [PID:%(process)d] %(message)s", }, + "colorful": { + "()": ColorfulFormatter, + "format": "[%(asctime)s] [%(levelname)s] [%(name)s.%(funcName)s:%(lineno)d] [PID:%(process)d] %(message)s", + }, }, "filters": {}, "handlers": { @@ -20,14 +46,25 @@ DEFAULT_LOGGING_CONFIG: Dict[str, Any] = { "filters": [], "stream": sys.stdout, }, + "color_console": { + "class": "logging.StreamHandler", + "formatter": "colorful", + "filters": [], + "stream": sys.stdout, + }, }, "root": {"handlers": ["console"], "level": os.getenv("LOG_LEVEL", "INFO")}, "loggers": { - "axolotl": {"handlers": ["console"], "level": "DEBUG", "propagate": False}, + "axolotl": { + "handlers": ["color_console"], + "level": "DEBUG", + "propagate": False, + }, }, } def configure_logging(): """Configure with default logging""" + init() # Initialize colorama dictConfig(DEFAULT_LOGGING_CONFIG)