Update configuration_neollm.py

configuration_neollm.py

@@ -1,3 +1,4 @@
+# mypy: ignore-errors
 # ==================== configuration_neollm.py ====================
 from transformers.configuration_utils import PretrainedConfig
 from transformers.modeling_rope_utils import rope_config_validation
@@ -636,9 +637,9 @@ class NeoLLMConfig(PretrainedConfig):
         attn_res_num_blocks=4,
         # ── StackMemory / STACKTRANS (Zhang et al., NeurIPS 2025) ───────
         use_stack_memory=False,
-        stack_d_model=32,          # H * d_s = 4 * 16
-        num_mem_heads=4,           # H = 4
-        stack_slots=16,            # S = 24
+        stack_d_model=32,  # H * d_s = 4 * 16
+        num_mem_heads=4,  # H = 4
+        stack_slots=16,  # S = 24
         stack_memory_cache_size=2048,
         # ── ResFormer cross-layer FAN residual (He et al., 2023) ─────────
         use_fan_residual=False,
@@ -693,9 +694,9 @@ class NeoLLMConfig(PretrainedConfig):
         laurel_lr_rank=32,
         # ── Interleaved Head Attention (Duvvuri et al., 2026) ─────────────
         use_iha=True,
-        iha_num_pseudo_heads=2,          # P=2 → 2×2=4 patrones por head
-        iha_local_global_pattern="LLLLG", # 4 locales + 1 global (paper §5.1)
-        iha_sliding_window=None,          # auto = N // (2*P^2) usando la longitud real del batch
+        iha_num_pseudo_heads=2,  # P=2 → 2×2=4 patrones por head
+        iha_local_global_pattern="LLLLG",  # 4 locales + 1 global (paper §5.1)
+        iha_sliding_window=None,  # auto = N // (2*P^2) usando la longitud real del batch
         iha_global_layers_use_iha=False,  # False replica el paper: la capa G es global estándar, sin IHA
         **kwargs,
     ):
@@ -758,7 +759,9 @@ class NeoLLMConfig(PretrainedConfig):
             repo_start_layer = num_hidden_layers // 3
         if repo_d_p is None:
             repo_d_p = hidden_size // 8
-        if (use_repo or use_repo_grape or use_repo_goat_prior) and not (0 <= repo_start_layer < num_hidden_layers):
+        if (use_repo or use_repo_grape or use_repo_goat_prior) and not (
+            0 <= repo_start_layer < num_hidden_layers
+        ):
             raise ValueError(
                 f"`repo_start_layer` must satisfy "
                 f"0 <= repo_start_layer < num_hidden_layers, "
@@ -775,8 +778,7 @@ class NeoLLMConfig(PretrainedConfig):
                 )
             if repo_goat_sink_decay <= 0.0:
                 raise ValueError(
-                    f"`repo_goat_sink_decay` must be > 0, "
-                    f"got {repo_goat_sink_decay}."
+                    f"`repo_goat_sink_decay` must be > 0, got {repo_goat_sink_decay}."
                 )
             _repo_goat_prior_dim = 2 * int(repo_goat_num_frequencies) + 2
             if (head_dim + _repo_goat_prior_dim) % 8 != 0:
@@ -785,7 +787,9 @@ class NeoLLMConfig(PretrainedConfig):
                     "head_dim=%d, giving internal attention dim=%d. Some "
                     "FlashAttention builds prefer dimensions divisible by 8; "
                     "the default repo_goat_num_frequencies=3 gives +8.",
-                    _repo_goat_prior_dim, head_dim, head_dim + _repo_goat_prior_dim,
+                    _repo_goat_prior_dim,
+                    head_dim,
+                    head_dim + _repo_goat_prior_dim,
                 )
         # ── SiameseNorm: RMS-only topology by construction ────────────────
 
@@ -803,7 +807,9 @@ class NeoLLMConfig(PretrainedConfig):
         if use_iha:
             _iha_pattern = str(iha_local_global_pattern).upper().strip()
             if not _iha_pattern:
-                raise ValueError("`iha_local_global_pattern` must not be empty when `use_iha=True`.")
+                raise ValueError(
+                    "`iha_local_global_pattern` must not be empty when `use_iha=True`."
+                )
             _bad_iha_pattern_chars = sorted(set(_iha_pattern) - {"L", "G"})
             if _bad_iha_pattern_chars:
                 raise ValueError(
@@ -821,7 +827,9 @@ class NeoLLMConfig(PretrainedConfig):
             "use_stack_memory": bool(use_stack_memory),
             "use_siamesenorm": bool(use_siamesenorm),
         }
-        _active_residual_flow = [name for name, enabled in _residual_flow_flags.items() if enabled]
+        _active_residual_flow = [
+            name for name, enabled in _residual_flow_flags.items() if enabled
+        ]
         if len(_active_residual_flow) > 1:
             raise ValueError(
                 "Enable at most one residual-flow mechanism among "
@@ -854,45 +862,45 @@ class NeoLLMConfig(PretrainedConfig):
         super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
 
         # ── Core Transformer ──────────────────────────────────────────────
-        self.vocab_size              = vocab_size
+        self.vocab_size = vocab_size
         self.max_position_embeddings = max_position_embeddings
-        self.hidden_size             = hidden_size
-        self.intermediate_size       = intermediate_size
-        self.num_hidden_layers       = num_hidden_layers
-        self.num_attention_heads     = num_attention_heads
-        self.num_key_value_heads     = num_key_value_heads
-        self.hidden_act              = hidden_act
-        self.initializer_range       = initializer_range
-        self.rms_norm_eps            = rms_norm_eps
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
 
         # ── Positional encoding ───────────────────────────────────────────
-        self.rope_theta              = rope_theta
-        self.rope_scaling            = rope_scaling
-        self.partial_rotary_factor   = partial_rotary_factor
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.partial_rotary_factor = partial_rotary_factor
 
         # ── Attention ─────────────────────────────────────────────────────
-        self.attention_bias                = attention_bias
-        self.attention_dropout             = attention_dropout
-        self.head_dim                      = head_dim
-        self.use_momentum_attention        = use_momentum_attention
-        self.momentum_gamma                = momentum_gamma
-        self.use_mea_attention             = use_mea_attention
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.head_dim = head_dim
+        self.use_momentum_attention = use_momentum_attention
+        self.momentum_gamma = momentum_gamma
+        self.use_mea_attention = use_mea_attention
         self.mea_component_key_value_heads = (
             num_key_value_heads
             if mea_component_key_value_heads is None
             else int(mea_component_key_value_heads)
         )
-        self.mea_groupnorm_eps             = mea_groupnorm_eps
-        self.use_lucid_attention           = use_lucid_attention
-        self.lucid_attention_eps           = lucid_attention_eps
-        self.use_affine_scaled_attention   = use_affine_scaled_attention
-        self.affine_momentum               = affine_momentum
-        self.use_xsa                       = use_xsa
-        self.xsa_eps                       = xsa_eps
-# ── Directional Routing ───────────────────────────────────────────
-        self.use_directional_routing       = use_directional_routing
-        self.directional_routing_k         = directional_routing_k
-        self.directional_routing_temp      = directional_routing_temp
+        self.mea_groupnorm_eps = mea_groupnorm_eps
+        self.use_lucid_attention = use_lucid_attention
+        self.lucid_attention_eps = lucid_attention_eps
+        self.use_affine_scaled_attention = use_affine_scaled_attention
+        self.affine_momentum = affine_momentum
+        self.use_xsa = use_xsa
+        self.xsa_eps = xsa_eps
+        # ── Directional Routing ───────────────────────────────────────────
+        self.use_directional_routing = use_directional_routing
+        self.directional_routing_k = directional_routing_k
+        self.directional_routing_temp = directional_routing_temp
 
         # ── Attention Residuals ───────────────────────────────────────────
         # use_attn_res=True:  replace fixed residual accumulation with learned
@@ -903,36 +911,36 @@ class NeoLLMConfig(PretrainedConfig):
         #   block_size = num_hidden_layers // 4 = 3 layers per block.
         #   Memory cost: O(num_blocks × batch × seq × hidden) instead of
         #   O(num_layers × batch × seq × hidden).
-        self.use_attn_res                  = use_attn_res
-        self.attn_res_num_blocks           = attn_res_num_blocks
+        self.use_attn_res = use_attn_res
+        self.attn_res_num_blocks = attn_res_num_blocks
 
         # ── StackMemory / STACKTRANS ──────────────────────────────────────
-        self.use_stack_memory              = use_stack_memory
-        self.stack_d_model                 = stack_d_model
-        self.num_mem_heads                 = num_mem_heads
-        self.stack_slots                   = stack_slots
-        self.stack_memory_cache_size       = stack_memory_cache_size
+        self.use_stack_memory = use_stack_memory
+        self.stack_d_model = stack_d_model
+        self.num_mem_heads = num_mem_heads
+        self.stack_slots = stack_slots
+        self.stack_memory_cache_size = stack_memory_cache_size
         rope_config_validation(self)
 
         # ── FANformer periodicity ─────────────────────────────────────────
-        self.use_fan_residual              = use_fan_residual
-        self.fan_ratio                     = fan_ratio
-        self.fan_ratio_ffn                 = fan_ratio_ffn
+        self.use_fan_residual = use_fan_residual
+        self.fan_ratio = fan_ratio
+        self.fan_ratio_ffn = fan_ratio_ffn
 
         # ── Regularization ────────────────────────────────────────────────
-        self.dropout_rate                  = dropout_rate
+        self.dropout_rate = dropout_rate
 
         # ── Learnable Multipliers ─────────────────────────────────────────
-        self.use_learnable_multipliers     = use_learnable_multipliers
-        self.use_embedding_multipliers     = use_embedding_multipliers
+        self.use_learnable_multipliers = use_learnable_multipliers
+        self.use_embedding_multipliers = use_embedding_multipliers
 
         # ── Normalization controls ────────────────────────────────────────
-        self.use_lns                       = bool(use_lns)
-        self.use_gpas                      = bool(use_gpas)
-        self.use_siamesenorm               = bool(use_siamesenorm)
-        self.siamese_normalized_input      = bool(siamese_normalized_input)
-        self.siamese_depth_scaling         = bool(siamese_depth_scaling)
-        self.siamese_attn_x_scale_init     = float(siamese_attn_x_scale_init)
+        self.use_lns = bool(use_lns)
+        self.use_gpas = bool(use_gpas)
+        self.use_siamesenorm = bool(use_siamesenorm)
+        self.siamese_normalized_input = bool(siamese_normalized_input)
+        self.siamese_depth_scaling = bool(siamese_depth_scaling)
+        self.siamese_attn_x_scale_init = float(siamese_attn_x_scale_init)
 
         # ── Embedding input normalization ─────────────────────────────────
         # True preserves the previous layer-0 pre-norm path, using the active
@@ -941,48 +949,48 @@ class NeoLLMConfig(PretrainedConfig):
         # False removes that first input-norm module entirely, so raw embeddings
         # enter the first attention block and the layer-0 input-norm parameters
         # are not counted; layers >= 1 keep their pre-norm unchanged.
-        self.use_embedding_input_norm      = use_embedding_input_norm
+        self.use_embedding_input_norm = use_embedding_input_norm
 
         # ── Leviathan generator ───────────────────────────────────────────
-        self.use_token_generator           = use_token_generator
-        self.generator_d_seed              = generator_d_seed
-        self.generator_num_modes           = generator_num_modes
-        self.generator_num_knots           = generator_num_knots
-        self.generator_spline_degree       = generator_spline_degree
-        self.generator_k                   = generator_k
-        self.generator_krank               = generator_krank
+        self.use_token_generator = use_token_generator
+        self.generator_d_seed = generator_d_seed
+        self.generator_num_modes = generator_num_modes
+        self.generator_num_knots = generator_num_knots
+        self.generator_spline_degree = generator_spline_degree
+        self.generator_k = generator_k
+        self.generator_krank = generator_krank
 
         # ── Leviathan-JTok-M ─────────────────────────────────────────────
-        self.use_jtokm                     = use_jtokm
-        self.jtokm_num_experts             = jtokm_num_experts
-        self.jtokm_top_k                   = jtokm_top_k
-        self.jtokm_num_modes               = jtokm_num_modes
-        self.jtokm_aux_loss_weight         = jtokm_aux_loss_weight
-        self.jtokm_norm_eps                = jtokm_norm_eps
+        self.use_jtokm = use_jtokm
+        self.jtokm_num_experts = jtokm_num_experts
+        self.jtokm_top_k = jtokm_top_k
+        self.jtokm_num_modes = jtokm_num_modes
+        self.jtokm_aux_loss_weight = jtokm_aux_loss_weight
+        self.jtokm_norm_eps = jtokm_norm_eps
 
         # ── Hadamard output projection (Aggarwal & Kumar, 2026) ───────────
-        self.use_hadamard_o_proj           = use_hadamard_o_proj
+        self.use_hadamard_o_proj = use_hadamard_o_proj
 
         # ── PolyNorm exclusivity ────────────���─────────────────────────────
-        self.polynorm_exclusive            = polynorm_exclusive
+        self.polynorm_exclusive = polynorm_exclusive
 
         # ── Spelling Bee Embeddings (Rabe et al., 2026) ───────────────────
-        self.use_spelling_bee_embeddings   = use_spelling_bee_embeddings
+        self.use_spelling_bee_embeddings = use_spelling_bee_embeddings
 
         # ── Context Re-Positioning (Li et al., 2026) ──────────────────────
-        self.use_repo                      = use_repo
-        self.repo_start_layer              = repo_start_layer
-        self.repo_d_p                      = repo_d_p
-        self.use_repo_grape                = use_repo_grape
-        self.use_repo_goat_prior           = use_repo_goat_prior
-        self.repo_goat_num_frequencies     = int(repo_goat_num_frequencies)
-        self.repo_goat_sink_decay          = float(repo_goat_sink_decay)
+        self.use_repo = use_repo
+        self.repo_start_layer = repo_start_layer
+        self.repo_d_p = repo_d_p
+        self.use_repo_grape = use_repo_grape
+        self.use_repo_goat_prior = use_repo_goat_prior
+        self.repo_goat_num_frequencies = int(repo_goat_num_frequencies)
+        self.repo_goat_sink_decay = float(repo_goat_sink_decay)
 
         # ── LAuReL: Learned Augmented Residual Layer (Menghani et al., 2025) ─
-        self.use_laurel                    = use_laurel
-        self.use_laurel_rw                 = use_laurel_rw
-        self.use_laurel_lr                 = use_laurel_lr
-        self.laurel_lr_rank                = laurel_lr_rank
+        self.use_laurel = use_laurel
+        self.use_laurel_rw = use_laurel_rw
+        self.use_laurel_lr = use_laurel_lr
+        self.laurel_lr_rank = laurel_lr_rank
 
         # ── Interleaved Head Attention (Duvvuri et al., 2026) ─────────────
         # use_iha=True: enables learned cross-head mixing of Q, K, V.
@@ -1011,17 +1019,17 @@ class NeoLLMConfig(PretrainedConfig):
         #   True  → G layers instantiate and run global IHA over the expanded
         #           sequence; useful only as an explicit high-cost ablation.
         # Init: identity (IHA ≡ MHA at step 0, Theorem 2 inclusion proof).
-        self.use_iha                       = use_iha
-        self.iha_num_pseudo_heads          = iha_num_pseudo_heads
-        self.iha_local_global_pattern      = iha_local_global_pattern
-        self.iha_sliding_window            = iha_sliding_window
-        self.iha_global_layers_use_iha     = bool(iha_global_layers_use_iha)
+        self.use_iha = use_iha
+        self.iha_num_pseudo_heads = iha_num_pseudo_heads
+        self.iha_local_global_pattern = iha_local_global_pattern
+        self.iha_sliding_window = iha_sliding_window
+        self.iha_global_layers_use_iha = bool(iha_global_layers_use_iha)
 
         self.auto_map = {
-            "AutoConfig":           "configuration_neollm.NeoLLMConfig",
-            "AutoModel":            "modeling_neollm.NeoLLMModel",
+            "AutoConfig": "configuration_neollm.NeoLLMConfig",
+            "AutoModel": "modeling_neollm.NeoLLMModel",
             "AutoModelForCausalLM": "modeling_neollm.NeoLLMForCausalLM",
         }
 
 
-__all__ = ["NeoLLMConfig"]
+__all__ = ["NeoLLMConfig"]