ViperEkura
c79d34eee1
- avoid constructing model_fn more than once when reading config - keep inference package exports focused on public entry points - rename extra strategy arguments to strategy_kwargs
96 lines
2.0 KiB
Python
96 lines
2.0 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, get_app, run_server, sample
|
|
from astrai.inference.network import ProtocolHandler
|
|
from astrai.inference.runtime.sample import SamplingPipeline
|
|
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()
|