diff --git a/astrai/__init__.py b/astrai/__init__.py index 8272a60..6c65541 100644 --- a/astrai/__init__.py +++ b/astrai/__init__.py @@ -1,6 +1,33 @@ __version__ = "1.3.11" __author__ = "ViperEkura" +import logging +import os + + +def setup_logging(level: str = "INFO"): + """Attach a handler to the ``astrai`` logger (only, not root). + + Call once per process, e.g. at the top of CLI scripts. + Set ``ASTR_LOG_LEVEL`` to override the default ``INFO``. + """ + _logger = logging.getLogger("astrai") + if _logger.handlers: + return + _level = getattr( + logging, os.environ.get("ASTR_LOG_LEVEL", level).upper(), logging.INFO + ) + _logger.setLevel(_level) + _handler = logging.StreamHandler() + _handler.setFormatter( + logging.Formatter( + "%(asctime)s | %(levelname)-7s | %(name)s | %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + ) + _logger.addHandler(_handler) + + from astrai.config import ( AutoRegressiveLMConfig, BaseModelConfig, @@ -94,5 +121,6 @@ __all__ = [ "only_on_rank", "run_server", "sample", + "setup_logging", "spawn_parallel_fn", ] diff --git a/scripts/tools/benchmark.py b/scripts/tools/benchmark.py index ced906a..57966b9 100644 --- a/scripts/tools/benchmark.py +++ b/scripts/tools/benchmark.py @@ -1,6 +1,7 @@ import click import torch +from astrai import setup_logging from astrai.config import AutoRegressiveLMConfig _DTYPES = ["bfloat16", "float16", "float32"] @@ -204,4 +205,5 @@ def benchmark_command( if __name__ == "__main__": + setup_logging() benchmark_command() diff --git a/scripts/tools/generate.py b/scripts/tools/generate.py index ec3131b..9848401 100644 --- a/scripts/tools/generate.py +++ b/scripts/tools/generate.py @@ -5,6 +5,7 @@ import click import torch from tqdm import tqdm +from astrai import setup_logging from astrai.inference import InferenceEngine from astrai.model import AutoModel from astrai.tokenize import AutoTokenizer @@ -162,4 +163,5 @@ def generate_command(**kwargs): if __name__ == "__main__": + setup_logging() generate_command() diff --git a/scripts/tools/preprocess.py b/scripts/tools/preprocess.py index 20bc971..3e4e4b2 100644 --- a/scripts/tools/preprocess.py +++ b/scripts/tools/preprocess.py @@ -2,6 +2,7 @@ import click +from astrai import setup_logging from astrai.config.preprocess_config import PipelineConfig from astrai.preprocessing.pipeline import Pipeline @@ -47,4 +48,5 @@ def preprocess_command(inputs, output_dir, pipeline_config, tokenizer_path, batc if __name__ == "__main__": + setup_logging() preprocess_command() diff --git a/scripts/tools/server.py b/scripts/tools/server.py index 66c62cb..31d43bf 100644 --- a/scripts/tools/server.py +++ b/scripts/tools/server.py @@ -3,6 +3,7 @@ from pathlib import Path import click import torch +from astrai import setup_logging from astrai.inference import run_server _DTYPES = ["bfloat16", "float16", "float32"] @@ -64,4 +65,5 @@ def server_command( if __name__ == "__main__": + setup_logging() server_command() diff --git a/scripts/tools/train.py b/scripts/tools/train.py index ed3be26..ff765cc 100644 --- a/scripts/tools/train.py +++ b/scripts/tools/train.py @@ -7,6 +7,7 @@ import click import torch from torch import Tensor, nn, optim +from astrai import setup_logging from astrai.config import AutoRegressiveLMConfig, TrainConfig from astrai.dataset import DatasetFactory, dpo_collate_fn, grpo_collate_fn from astrai.model import AutoRegressiveLM @@ -556,4 +557,5 @@ def train( if __name__ == "__main__": + setup_logging() train_command()