AstrAI/astrai/__init__.py

36 lines
768 B
Python

__version__ = "1.3.2"
__author__ = "ViperEkura"
from astrai.config import (
ModelConfig,
TrainConfig,
)
from astrai.model.transformer import Transformer
from astrai.data import DatasetLoader, BpeTokenizer
from astrai.inference.generator import (
GenerationRequest,
LoopGenerator,
StreamGenerator,
BatchGenerator,
EmbeddingEncoder,
GeneratorFactory,
)
from astrai.trainer import Trainer, StrategyFactory, SchedulerFactory
__all__ = [
"Transformer",
"ModelConfig",
"TrainConfig",
"DatasetLoader",
"BpeTokenizer",
"GenerationRequest",
"LoopGenerator",
"StreamGenerator",
"BatchGenerator",
"EmbeddingEncoder",
"GeneratorFactory",
"Trainer",
"StrategyFactory",
"SchedulerFactory",
]