# Agnes 3.0 Flash configuration for sglang.
#
# The HF checkpoint says model_type "agnes", names its layer types
# agnes_delta_attention / agnes_global_attention and carries a parallel FFN
# branch per layer.  The server runs it on its built-in hybrid
# (delta-rule + global attention) implementation, so this class maps those
# onto the fields that implementation reads; the checkpoint's tensor names
# are translated while loading (see the model file patched by apply_patch.py).
from sglang.srt.configs.qwen3_5 import Qwen3_5Config

AGNES_DELTA = "agnes_delta_attention"
AGNES_GLOBAL = "agnes_global_attention"


class AgnesConfig(Qwen3_5Config):
    model_type = "agnes"

    def __init__(self, text_config=None, vision_config=None, **kwargs):
        kwargs.pop("auto_map", None)  # the transformers remote code is not used in the server
        if isinstance(text_config, dict):
            text_config = dict(text_config)
            text_config["model_type"] = "qwen3_5_text"
            width = int(text_config.pop("parallel_ffn_intermediate_size", 0) or 0)
            plan = text_config.pop("layer_types", None)
            interval = text_config.pop("global_attention_interval", None)
            if interval is None:
                interval = text_config.pop("full_attention_interval", None)
            if interval is None and plan:
                interval = next(i + 1 for i, t in enumerate(plan) if t == AGNES_GLOBAL)
            text_config["full_attention_interval"] = int(interval or 4)
            main = int(text_config["intermediate_size"])
            # the parallel branch is folded into the main MLP at load time
            text_config["intermediate_size"] = main + width
            text_config["agnes_main_intermediate_size"] = main
            text_config["agnes_parallel_ffn_intermediate_size"] = width
        if isinstance(vision_config, dict):
            vision_config = dict(vision_config)
            vision_config["model_type"] = "qwen3_5"
        kwargs["architectures"] = ["Qwen3_5ForConditionalGeneration"]
        super().__init__(text_config=text_config, vision_config=vision_config, **kwargs)
        # every downstream check sees the built-in hybrid architecture
        self.model_type = "qwen3_5"

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
        # The weight loader needs the checkpoint directory to pick up the parallel
        # branch tensors.  The path is written into the config *dict* before the
        # object is built: the config reaches the worker processes through a
        # to_dict round trip, which keeps fields that came in through __init__
        # and drops attributes set afterwards (from_pretrained's own kwargs only
        # override known fields, so they cannot carry it either).
        path = str(pretrained_model_name_or_path)
        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
        config_dict["agnes_model_path"] = path
        if isinstance(config_dict.get("text_config"), dict):
            config_dict["text_config"]["agnes_model_path"] = path
        return cls.from_dict(config_dict, **kwargs)