src/diffusion_lm/__init__.py

"""Small masked discrete diffusion language model."""

from typing import TYPE_CHECKING

from diffusion_lm.config import ExperimentConfig, ModelConfig, TrainingConfig, load_config

if TYPE_CHECKING:
    from diffusion_lm.model import DiffusionTransformer

__all__ = [
    "DiffusionTransformer",
    "ExperimentConfig",
    "ModelConfig",
    "TrainingConfig",
    "load_config",
]

__version__ = "0.1.0"


def __getattr__(name: str):
    if name == "DiffusionTransformer":
        from diffusion_lm.model import DiffusionTransformer

        return DiffusionTransformer
    raise AttributeError(name)