"""Configuration for GLEAN-pruned OLMoE: variable-width, variable-count experts.
"""

from transformers.models.olmoe.configuration_olmoe import OlmoeConfig


class PrunedOlmoeConfig(OlmoeConfig):
    """OlmoeConfig plus a per-(layer, expert) width table.

    ``expert_widths[l]`` lists the SwiGLU intermediate width of each surviving
    expert in decoder layer ``l``, in expert order. Lists are ragged: layers
    may keep different numbers of experts (deleted experts simply don't
    appear — the router in layer ``l`` has ``len(expert_widths[l])`` rows),
    and each width may differ (multiples of the GEMM block size, 128, for
    variable-MegaBlocks execution). ``None`` means an unpruned model
    (uniform ``num_experts`` × ``intermediate_size``).

    The inherited ``num_experts`` / ``intermediate_size`` keep their ORIGINAL
    (pre-pruning) values for provenance; the width table is authoritative for
    the built architecture.
    """

    model_type = "pruned_olmoe"

    def __init__(self, expert_widths: list[list[int]] | None = None, **kwargs):
        super().__init__(**kwargs)
        self.expert_widths = expert_widths