"""WhisperMLA configuration — WhisperConfig + the MLA conversion axes.

Shipped inside released model repos (trust_remote_code); loaded by AutoConfig via auto_map.
"""
from transformers import WhisperConfig


class WhisperMLAConfig(WhisperConfig):
    model_type = "whisper_mla"

    def __init__(self, mla_scheme: str = "partial_preservation", mla_scope: str = "decoder_self",
                 mla_selection: str = "uniform", mla_low_rank: int = 16, mla_rope_dim: int = 32,
                 **kwargs):
        self.mla_scheme = mla_scheme
        self.mla_scope = mla_scope
        self.mla_selection = mla_selection
        self.mla_low_rank = mla_low_rank
        self.mla_rope_dim = mla_rope_dim
        super().__init__(**kwargs)