src/model/__init__.py

"""Atlas model module."""

from .topomlp_adapter import TopoMLPToAtlasMapTokens
from .streampetr_adapter import extract_streampetr_topk_tokens

# Atlas (LLM) side depends on `transformers`. StreamPETR/TopoMLP pretraining does not.
try:
    from .configuration_atlas import AtlasConfig
    from .modeling_atlas import AtlasProjector, AtlasForCausalLM
    _ATLAS_AVAILABLE = True
except Exception:
    AtlasConfig = None  # type: ignore[assignment]
    AtlasProjector = None  # type: ignore[assignment]
    AtlasForCausalLM = None  # type: ignore[assignment]
    _ATLAS_AVAILABLE = False

__all__ = [
    "TopoMLPToAtlasMapTokens",
    "extract_streampetr_topk_tokens",
]

if _ATLAS_AVAILABLE:
    __all__ += [
        "AtlasConfig",
        "AtlasProjector",
        "AtlasForCausalLM",
    ]