feat : add setup_logging with hierarchical astrai logger

- setup_logging(): attach handler only to astrai logger, not root
- all astrai.* sub-module loggers inherit automatically
- controlled by ASTR_LOG_LEVEL env var (default INFO)
- called in if __name__ == '__main__' of each CLI script

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()

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()

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()

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()

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()