"""Configuration for MinistralDualRope — a Ministral with a per-layer sliding-window / full
attention pattern and a second, UNSCALED rotary for the sliding layers.

The base model is byte-for-byte a Ministral (== Llama math: SiLU-gated MLP, RMSNorm
without unit offset, GQA, rotary embeddings, no biases, no QK-norm, no sandwich norms,
no logit softcapping). Ministral already provides the per-layer sliding/full attention
pattern via ``layer_types`` (mask routing identical to how this model was trained). The
*only* addition over stock Ministral is dual RoPE:

  * full-attention layers use the standard ``rope_theta`` (+ ``rope_scaling``, e.g. the
    llama3-scaled cache of a context-extended checkpoint);
  * sliding (local) layers use an UNSCALED rotary at ``rope_local_base_freq``.

This mirrors our pretraining stack (``pretraining/flexattn_patch.py``) and is the exact
Gemma-3 dual-RoPE design. For a factor-1 checkpoint with equal bases and no scaling the
two rotaries collapse to one, so a single ``rope_theta`` is exact; for a context-extended
checkpoint (``rope_scaling`` set) the global and local caches diverge and both are needed.
"""

from __future__ import annotations

from transformers.models.ministral.configuration_ministral import MinistralConfig


class MinistralDualRopeConfig(MinistralConfig):
    r"""MinistralConfig + an unscaled local rotary base for the sliding layers.

    Extra arg beyond :class:`~transformers.MinistralConfig`:
        rope_local_base_freq (`float`, *optional*):
            RoPE base for the sliding (local) layers, applied with NO rope scaling.
            Defaults to ``rope_theta`` (i.e. single-RoPE, byte-identical to Ministral).

    (``layer_types`` and ``sliding_window`` are inherited from Ministral and drive the
    per-layer full/sliding attention pattern.)
    """

    model_type = "ministral_dual_rope"

    def __init__(self, rope_local_base_freq=None, **kwargs):
        super().__init__(**kwargs)
        if rope_local_base_freq is None:
            rope_local_base_freq = self._global_rope_theta()
        self.rope_local_base_freq = float(rope_local_base_freq)

    def _global_rope_theta(self) -> float:
        """The GLOBAL rotary base, robust across transformers versions (5.x stores it in
        ``rope_parameters['rope_theta']``; 4.x in ``rope_theta``)."""
        rp = getattr(self, "rope_parameters", None)
        if isinstance(rp, dict) and "rope_theta" in rp:
            return float(rp["rope_theta"])
        return float(getattr(self, "rope_theta", 10000.0))


__all__ = ["MinistralDualRopeConfig"]