ViperEkura
cbc584470d
- 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
101 lines
1.9 KiB
Python
101 lines
1.9 KiB
Python
__version__ = "1.3.13"
|
|
__author__ = "ViperEkura"
|
|
|
|
from astrai.config import (
|
|
AutoRegressiveLMConfig,
|
|
BaseModelConfig,
|
|
ConfigFactory,
|
|
EncoderConfig,
|
|
PipelineConfig,
|
|
TrainConfig,
|
|
)
|
|
from astrai.dataset import (
|
|
BaseDataset,
|
|
DatasetFactory,
|
|
RDSampler,
|
|
Store,
|
|
StoreFactory,
|
|
)
|
|
from astrai.factory import BaseFactory
|
|
from astrai.inference import (
|
|
InferenceEngine,
|
|
ProtocolHandler,
|
|
SamplingPipeline,
|
|
get_app,
|
|
run_server,
|
|
sample,
|
|
)
|
|
from astrai.logging import setup_logging
|
|
from astrai.model import (
|
|
AutoModel,
|
|
AutoRegressiveLM,
|
|
EmbeddingEncoder,
|
|
LoRAConfig,
|
|
inject_lora,
|
|
)
|
|
from astrai.parallel import (
|
|
ExecutorFactory,
|
|
get_rank,
|
|
get_world_size,
|
|
only_on_rank,
|
|
spawn_parallel_fn,
|
|
)
|
|
from astrai.preprocessing import Pipeline, filter_by_length
|
|
from astrai.serialization import Checkpoint
|
|
from astrai.tokenize import AutoTokenizer, ChatTemplate
|
|
from astrai.trainer import (
|
|
BaseScheduler,
|
|
BaseStrategy,
|
|
CallbackFactory,
|
|
SchedulerFactory,
|
|
StrategyFactory,
|
|
TrainCallback,
|
|
Trainer,
|
|
)
|
|
|
|
__all__ = [
|
|
"AutoRegressiveLM",
|
|
"AutoRegressiveLMConfig",
|
|
"AutoModel",
|
|
"AutoTokenizer",
|
|
"BaseDataset",
|
|
"BaseFactory",
|
|
"BaseModelConfig",
|
|
"BaseScheduler",
|
|
"BaseStrategy",
|
|
"CallbackFactory",
|
|
"ChatTemplate",
|
|
"Checkpoint",
|
|
"ConfigFactory",
|
|
"DatasetFactory",
|
|
"EmbeddingEncoder",
|
|
"EncoderConfig",
|
|
"ExecutorFactory",
|
|
"InferenceEngine",
|
|
"LoRAConfig",
|
|
"Pipeline",
|
|
"PipelineConfig",
|
|
"ProtocolHandler",
|
|
"RDSampler",
|
|
"SamplingPipeline",
|
|
"SchedulerFactory",
|
|
"Store",
|
|
"StoreFactory",
|
|
"StrategyFactory",
|
|
"TrainCallback",
|
|
"TrainConfig",
|
|
"Trainer",
|
|
"filter_by_length",
|
|
"get_app",
|
|
"get_rank",
|
|
"get_world_size",
|
|
"inject_lora",
|
|
"only_on_rank",
|
|
"run_server",
|
|
"sample",
|
|
"setup_logging",
|
|
"spawn_parallel_fn",
|
|
]
|
|
|
|
setup_logging()
|