from src.model.testformer import TestFormerLM
from src.model.testformer_combined import TestFormerCombinedLM
from src.model.testformer_combined_config import (
    TestFormerCombinedConfig,
    build_testformer_combined_config,
)
from src.model.motif_moe_external import (
    MotifMoEConfig,
    MotifMoETransformer,
    build_wikitext_motif_moe_config,
)
from src.model.testformer_config import (
    TESTFORMER_MOTIFS,
    TESTFORMER_SCALE_PRESETS,
    TestFormerConfig,
    TestFormerMotifConfig,
    build_testformer_config,
    build_scaled_testformer_config,
)
from src.model.qwen_motif_config import (
    DEFAULT_QWEN_ATTENTION_MOTIFS,
    DEFAULT_QWEN_MOTIFS,
    LowRankAdapterConfig,
    QwenFFNExpertLoRAConfig,
    QwenMotifAttentionPatchConfig,
    QwenMotifFullConfig,
    QwenMotifPatchConfig,
    QwenMotifRouterConfig,
    build_default_attention_patch_config,
    build_default_ffn_lora_configs,
    build_contiguous_motif_index,
    build_layer_motif_indices,
    build_layer_range,
    build_motif_index,
    build_random_motif_index,
    build_round_robin_motif_index,
)
from src.model.qwen_motif_router import (
    ContextualMotifRouter,
    StaticMotifRouter,
    build_motif_router,
)
from src.model.qwen_motif_ffn import QwenMotifSplitMLP
from src.model.qwen_motif_lora import (
    LowRankLinearAdapter,
    QwenMotifSplitLoRAMLP,
    RuntimeScaledLoRALinear,
)
from src.model.qwen_motif_attention import QwenMotifAttentionAdapter
from src.model.qwen_motif_patch import (
    apply_qwen_motif_pipeline,
    build_and_patch_qwen_attention_layers,
    build_and_patch_qwen_ffn_layers,
    build_and_patch_qwen_ffn_lora_layers,
    collect_qwen_motif_attention_adapters,
    collect_qwen_motif_mlps,
    collect_qwen_motif_trainable_names,
    freeze_model_except_qwen_motif_trainables,
    freeze_model_except_motif_routers,
    get_qwen_decoder_layers,
    partial_reinit_qwen_motif_modules,
    patch_qwen_attention_layers,
    patch_qwen_ffn_layers,
)

__all__ = [
    "TestFormerConfig",
    "TestFormerMotifConfig",
    "TestFormerCombinedConfig",
    "MotifMoEConfig",
    "TestFormerLM",
    "TestFormerCombinedLM",
    "MotifMoETransformer",
    "LowRankAdapterConfig",
    "QwenFFNExpertLoRAConfig",
    "QwenMotifAttentionPatchConfig",
    "QwenMotifFullConfig",
    "QwenMotifPatchConfig",
    "QwenMotifRouterConfig",
    "QwenMotifSplitMLP",
    "QwenMotifSplitLoRAMLP",
    "QwenMotifAttentionAdapter",
    "LowRankLinearAdapter",
    "RuntimeScaledLoRALinear",
    "StaticMotifRouter",
    "ContextualMotifRouter",
    "TESTFORMER_MOTIFS",
    "TESTFORMER_SCALE_PRESETS",
    "DEFAULT_QWEN_ATTENTION_MOTIFS",
    "DEFAULT_QWEN_MOTIFS",
    "build_scaled_testformer_config",
    "build_testformer_config",
    "build_testformer_combined_config",
    "build_wikitext_motif_moe_config",
    "build_default_ffn_lora_configs",
    "build_default_attention_patch_config",
    "build_layer_range",
    "build_motif_index",
    "build_contiguous_motif_index",
    "build_random_motif_index",
    "build_round_robin_motif_index",
    "build_layer_motif_indices",
    "build_motif_router",
    "get_qwen_decoder_layers",
    "patch_qwen_ffn_layers",
    "patch_qwen_attention_layers",
    "build_and_patch_qwen_ffn_layers",
    "build_and_patch_qwen_ffn_lora_layers",
    "build_and_patch_qwen_attention_layers",
    "apply_qwen_motif_pipeline",
    "collect_qwen_motif_mlps",
    "collect_qwen_motif_attention_adapters",
    "collect_qwen_motif_trainable_names",
    "freeze_model_except_motif_routers",
    "freeze_model_except_qwen_motif_trainables",
    "partial_reinit_qwen_motif_modules",
]