refactor: centralize logging in astrai.logging, replace ASTRAI_TIMED with log level

- move setup_logging to astrai/logging.py
- timed() now uses logger.isEnabledFor(DEBUG) instead of separate env var
- enable ASTR_LOG_LEVEL=DEBUG to see per-step timing logs
- call setup_logging() in stream_chat.py

scripts/tools/benchmark.py

@@ -5,7 +5,6 @@ from typing import Optional, Union
 import click
 import torch
 
-from astrai import setup_logging
 from astrai.config import BaseModelConfig, ConfigFactory
 from astrai.extension import ATTN_BACKEND, AttentionBackendFactory, attn_backend
 from astrai.inference.core.cache import PagePool
@@ -478,5 +477,4 @@ def benchmark_command(
 
 
 if __name__ == "__main__":
-    setup_logging()
     benchmark_command()

scripts/tools/generate.py

@@ -6,7 +6,6 @@ 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
@@ -157,5 +156,4 @@ def generate_command(**kwargs):
 
 
 if __name__ == "__main__":
-    setup_logging()
     generate_command()

scripts/tools/preprocess.py

@@ -2,7 +2,6 @@
 
 import click
 
-from astrai import setup_logging
 from astrai.config.preprocess_config import PipelineConfig
 from astrai.preprocessing.pipeline import Pipeline
 
@@ -48,5 +47,4 @@ def preprocess_command(inputs, output_dir, pipeline_config, tokenizer_path, batc
 
 
 if __name__ == "__main__":
-    setup_logging()
     preprocess_command()

scripts/tools/server.py

@@ -3,7 +3,6 @@ from pathlib import Path
 import click
 import torch
 
-from astrai import setup_logging
 from astrai.inference import run_server
 
 _DTYPES = ["bfloat16", "float16", "float32"]
@@ -65,5 +64,4 @@ def server_command(
 
 
 if __name__ == "__main__":
-    setup_logging()
     server_command()

scripts/tools/train.py

@@ -8,7 +8,6 @@ import torch
 from click.core import ParameterSource
 from torch import 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
@@ -828,5 +827,4 @@ def train(
 
 
 if __name__ == "__main__":
-    setup_logging()
     train_command()