File size: 793 Bytes
9fe982a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | """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",
]
|