__init__.py

from typing import List

from mergekit.moe.arch import MoEOutputArchitecture
from mergekit.moe.deepseek import DeepseekMoE
from mergekit.moe.mixtral import MixtralMoE

ALL_OUTPUT_ARCHITECTURES: List[MoEOutputArchitecture] = [MixtralMoE(), DeepseekMoE()]

try:
    from mergekit.moe.qwen import QwenMoE
except ImportError:
    pass
else:
    ALL_OUTPUT_ARCHITECTURES.append(QwenMoE())

try:
    from mergekit.moe.qwen3 import Qwen3MoE
except ImportError:
    pass
else:
    ALL_OUTPUT_ARCHITECTURES.append(Qwen3MoE())

# --- ADD THIS SECTION START ---
try:
    from mergekit.moe.llama import LlamaMoE
except ImportError:
    # This will trigger if llama.py is missing or has a syntax error
    pass
else:
    ALL_OUTPUT_ARCHITECTURES.append(LlamaMoE())
# --- ADD THIS SECTION END ---

__all__ = [
    "ALL_OUTPUT_ARCHITECTURES",
    "MoEOutputArchitecture",
]