# modeling_solon_moe.py — 자립형 Solon-MoE 모델
# =================================================
# 네이티브 Gemma4ForCausalLM을 상속하되, config.moe_layers에 지정된 층에만
# MoE 블록을 생성하도록 DecoderLayer 초기화를 재정의한다.
# config.json의 auto_map에 등록하면 trust_remote_code=True 로 어떤 도구(lm-eval,
# vLLM, llm-compressor 등)에서도 별도 plugin 없이 15층 MoE가 정확히 로드된다.
#
# 원리: 네이티브 enable_moe_block은 전역 flag라 48층 전부 MoE를 만든다.
#       여기서는 layer_idx가 config.moe_layers에 있을 때만 MoE를 만든다.

from transformers.models.gemma4.modeling_gemma4 import (
    Gemma4TextDecoderLayer,
    Gemma4ForCausalLM,
    Gemma4TextModel,
)
from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig


class SolonMoEConfig(Gemma4TextConfig):
    model_type = "solon_moe"
    # moe_layers는 config.json에서 로드됨 (커스텀 필드)


class SolonMoEDecoderLayer(Gemma4TextDecoderLayer):
    def __init__(self, config, layer_idx):
        moe_layers = getattr(config, "moe_layers", None)
        want = (moe_layers is None and config.enable_moe_block) or \
               (moe_layers is not None and layer_idx in set(moe_layers))
        saved = config.enable_moe_block
        config.enable_moe_block = bool(want)
        try:
            super().__init__(config, layer_idx)
        finally:
            config.enable_moe_block = saved
        self.enable_moe_block = bool(want)


class SolonMoETextModel(Gemma4TextModel):
    config_class = SolonMoEConfig

    def __init__(self, config):
        super().__init__(config)
        # 네이티브가 만든 층을 층별 MoE 버전으로 교체
        import torch.nn as nn
        self.layers = nn.ModuleList(
            [SolonMoEDecoderLayer(config, i) for i in range(config.num_hidden_layers)]
        )
        self.post_init()


class SolonMoEForCausalLM(Gemma4ForCausalLM):
    config_class = SolonMoEConfig

    def __init__(self, config):
        super().__init__(config)
        self.model = SolonMoETextModel(config)
        self.post_init()