Add transformers v5 support on main; persist full sliding_window; tidy auto_map

README.md

@@ -119,11 +119,21 @@ mlm = AutoModelForMaskedLM.from_pretrained("BidirLM/BidirLM-270M-Base", trust_re
 ## Requirements
 
 ```
-transformers>=4.57.6,<5.0.0
+transformers>=5.0
 ```
 
 This model requires `trust_remote_code=True`.
 
+> **Note:** This model was trained with `transformers==4.57.6` (transformers 4.x). The version on `main` was patched to work with `transformers>=5.0`. For the original (pre-patch) version, which is compatible with `transformers>=4.57.6,<5.0.0`, use the `transformers-v4` branch:
+> ```python
+> from transformers import AutoModel
+> model = AutoModel.from_pretrained(
+>     "BidirLM/BidirLM-270M-Base",
+>     trust_remote_code=True,
+>     revision="transformers-v4",
+> )
+> ```
+
 ## Citation
 
 ```bibtex

config.json

@@ -7,10 +7,9 @@
   "attention_dropout": 0.0,
   "attn_logit_softcapping": null,
   "auto_map": {
-    "AutoConfig": "configuration_bidirlm.BidirLMTextConfig",
-    "AutoModel": "modeling_bidirlm.BidirLMTextModel",
+    "AutoConfig": "configuration_bidirlm.BidirLMConfig",
+    "AutoModel": "modeling_bidirlm.BidirLMModel",
     "AutoModelForMaskedLM": "modeling_bidirlm.BidirLMForMaskedLM",
-    "AutoModelForPreTraining": "modeling_bidirlm.BidirLMPreTrainedModel",
     "AutoModelForSequenceClassification": "modeling_bidirlm.BidirLMForSequenceClassification",
     "AutoModelForTokenClassification": "modeling_bidirlm.BidirLMForTokenClassification"
   },
@@ -56,8 +55,8 @@
   "rope_scaling": null,
   "rope_theta": 1000000.0,
   "sliding_window": 512,
-  "transformers_version": "4.57.3",
+  "transformers_version": "5.9.0",
   "use_bidirectional_attention": true,
   "use_cache": true,
   "vocab_size": 262144
-}
+}

configuration_bidirlm.py

@@ -23,13 +23,14 @@ from typing import Any, Optional, Union
 
 import transformers
 _v = transformers.__version__
-if _v < "4.57.6" or _v >= "5.0.0":
+if _v < "5.0.0":
     raise ImportError(
-        f"BidirLM requires transformers>=4.57.6,<5.0.0 (found {_v}). "
-        f"Install a compatible version: pip install 'transformers>=4.57.6,<5.0.0'"
+        f"BidirLM requires transformers>=5.0.0 on this branch (found {_v}). "
+        f"Install a compatible version: pip install 'transformers>=5.0.0'. "
+        f"For transformers 4.x, use the `transformers-v4` branch instead."
     )
 
-from transformers.configuration_utils import PretrainedConfig, layer_type_validation
+from transformers.configuration_utils import PretrainedConfig
 from transformers.modeling_rope_utils import rope_config_validation
 from transformers.utils import logging
 from transformers.models.siglip import SiglipVisionConfig
@@ -174,6 +175,20 @@ class BidirLMConfig(PretrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
 
+    @property
+    def effective_sliding_window(self):
+        # Per-side sliding window used at attention time. For bidirectional
+        # models the configured (total) window is split symmetrically, so
+        # each side is half. Derived at runtime so the full `sliding_window`
+        # is what gets persisted (no halving on every save/load).
+        sw = getattr(self, "sliding_window", None)
+        if sw is None:
+            return None
+        if getattr(self, "use_bidirectional_attention", False):
+            return sw // 2
+        return sw
+
+
     def __init__(
         self,
         vocab_size=262_208,
@@ -235,8 +250,6 @@ class BidirLMConfig(PretrainedConfig):
         self.layer_types = layer_types
         self.use_bidirectional_attention = use_bidirectional_attention
         self.classifier_pooling = classifier_pooling
-        if use_bidirectional_attention:
-            self.sliding_window = self.sliding_window // 2
 
         self.rope_local_base_freq = rope_local_base_freq
         self.rope_scaling = rope_scaling
@@ -250,7 +263,7 @@ class BidirLMConfig(PretrainedConfig):
                 "sliding_attention" if bool((i + 1) % self._sliding_window_pattern) else "full_attention"
                 for i in range(self.num_hidden_layers)
             ]
-        layer_type_validation(self.layer_types, self.num_hidden_layers)
+        self.validate_layer_type()
 
 
 class Gemma3Config(PretrainedConfig):

modeling_bidirlm.py

@@ -3,10 +3,11 @@ from typing import Optional
 
 import transformers
 _v = transformers.__version__
-if _v < "4.57.6" or _v >= "5.0.0":
+if _v < "5.0.0":
     raise ImportError(
-        f"BidirLM requires transformers>=4.57.6,<5.0.0 (found {_v}). "
-        f"Install a compatible version: pip install 'transformers>=4.57.6,<5.0.0'"
+        f"BidirLM requires transformers>=5.0.0 on this branch (found {_v}). "
+        f"Install a compatible version: pip install 'transformers>=5.0.0'. "
+        f"For transformers 4.x, use the `transformers-v4` branch instead."
     )
 
 import torch
@@ -126,7 +127,7 @@ class Gemma3Attention(nn.Module):
             bias=config.attention_bias,
         )
         self.attn_logit_softcapping = self.config.attn_logit_softcapping
-        self.sliding_window = config.sliding_window if self.is_sliding else None
+        self.sliding_window = config.effective_sliding_window if self.is_sliding else None
 
         self.q_norm = Gemma3RMSNorm(dim=config.head_dim, eps=config.rms_norm_eps)
         self.k_norm = Gemma3RMSNorm(dim=config.head_dim, eps=config.rms_norm_eps)
@@ -334,13 +335,12 @@ class BidirLMPreTrainedModel(PreTrainedModel):
 
     def _init_weights(self, module):
         super()._init_weights(module)
-        # if isinstance(module, Gemma3MultiModalProjector):
-        #     module.mm_input_projection_weight.data.zero_()
-        # # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight)
-        # elif "RMSNorm" in module.__class__.__name__:
-        #     module.weight.data.zero_()
         if "RMSNorm" in module.__class__.__name__:
             module.weight.data.zero_()
+        elif isinstance(module, Gemma3TextScaledWordEmbedding):
+            # transformers 5.x resets non-persistent buffers in this hook;
+            # restore embed_scale to the value chosen at construction.
+            torch.nn.init.constant_(module.embed_scale, module.scalar_embed_scale)
 
 
 class Gemma3TextScaledWordEmbedding(nn.Embedding):
@@ -356,6 +356,10 @@ class Gemma3TextScaledWordEmbedding(nn.Embedding):
         embed_scale: float = 1.0,
     ):
         super().__init__(num_embeddings, embedding_dim, padding_idx)
+        # transformers 5.x calls _init_weights on every module post-load and
+        # resets non-persistent buffers to zero; keep the scalar so the
+        # PreTrainedModel._init_weights below can re-initialize embed_scale.
+        self.scalar_embed_scale = embed_scale
         self.register_buffer("embed_scale", torch.tensor(embed_scale), persistent=False)
 
     def forward(self, input_ids: torch.Tensor):
@@ -411,12 +415,21 @@ class Gemma3RotaryEmbedding(nn.Module):
         self.original_max_seq_len = config.max_position_embeddings
 
         self.config = config
-        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        # transformers 5.x removed 'default' from ROPE_INIT_FUNCTIONS
+        rope_init_fn = self.compute_default_rope_parameters if self.rope_type == "default" else ROPE_INIT_FUNCTIONS[self.rope_type]
 
-        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
         self.register_buffer("inv_freq", inv_freq, persistent=False)
         self.original_inv_freq = self.inv_freq
 
+    @staticmethod
+    def compute_default_rope_parameters(config, device=None, **kwargs):
+        base = config.rope_theta
+        head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        dim = int(head_dim * getattr(config, "partial_rotary_factor", 1.0))
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim))
+        return inv_freq, 1.0
+
     @torch.no_grad()
     @dynamic_rope_update
     def forward(self, x, position_ids):
@@ -541,7 +554,7 @@ class BidirLMModel(BidirLMPreTrainedModel):
             else self.config.output_hidden_states
         )
         return_dict = (
-            return_dict if return_dict is not None else self.config.use_return_dict
+            return_dict if return_dict is not None else True
         )
         all_hidden_states = () if output_hidden_states else None
         all_self_attns = () if output_attentions else None
@@ -573,10 +586,11 @@ class BidirLMModel(BidirLMPreTrainedModel):
         position_embeddings_global = self.rotary_emb(hidden_states, position_ids)
         position_embeddings_local = self.rotary_emb_local(hidden_states, position_ids)
 
+        _eff_sw = self.config.effective_sliding_window
         window_size = (
             (
-                self.config.sliding_window,
-                self.config.sliding_window if self.config.use_bidirectional_attention else 0
+                _eff_sw,
+                _eff_sw if self.config.use_bidirectional_attention else 0
             )
             if self.config.sliding_window is not None
             else None
@@ -645,7 +659,7 @@ class BidirLMModel(BidirLMPreTrainedModel):
 
 
 class BidirLMForMaskedLM(BidirLMPreTrainedModel):
-    _tied_weights_keys = ["lm_head.weight"]
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
     config: BidirLMConfig
 
     def __init__(self, config):
@@ -670,7 +684,7 @@ class BidirLMForMaskedLM(BidirLMPreTrainedModel):
         **kwargs,
     ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
         return_dict = (
-            return_dict if return_dict is not None else self.config.use_return_dict
+            return_dict if return_dict is not None else True
         )
         encoder_output = self.model(
             input_ids=input_ids,
@@ -726,7 +740,7 @@ class BidirLMForSequenceClassification(BidirLMPreTrainedModel):
         **kwargs,
     ) -> tuple[torch.Tensor] | SequenceClassifierOutput:
         return_dict = (
-            return_dict if return_dict is not None else self.config.use_return_dict
+            return_dict if return_dict is not None else True
         )
 
         encoder_output = self.model(
@@ -823,7 +837,7 @@ class BidirLMForTokenClassification(BidirLMPreTrainedModel):
         return_dict: Optional[bool] = None,
     ) -> tuple[torch.Tensor] | TokenClassifierOutput:
         return_dict = (
-            return_dict if return_dict is not None else self.config.use_return_dict
+            return_dict if return_dict is not None else True
         )
 
         outputs = self.model(
@@ -976,7 +990,7 @@ class BidirLMForTokenClassification(BidirLMPreTrainedModel):
 #         output_hidden_states = (
 #             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
 #         )
-#         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+#         return_dict = return_dict if return_dict is not None else True
 
 #         # Replace image id with PAD if the image token if OOV, to avoid index-errors
 #         if input_ids is not None and self.config.image_token_id >= self.vocab_size: