Fix Transformers 5 compatibility and make xFormers optional

README.md

@@ -55,7 +55,7 @@ For detailed benchmarking, see https://acr.ps/neoarabert.
 ### How to Use
 Install these libraries:
 ```
-pip install fast-disambig hf_transfer==0.1.9 torch==2.5.1 transformers==4.49.0 xformers==0.0.28.post3 torchvision torchaudio
+pip install fast-disambig hf_transfer==0.1.9 torch==2.14.1 transformers==5.20.2 xformers==0.0.35 torchvision torchaudio
 ```
 Load the model and use it to generate embeddings:
 ```python

model.py

@@ -1,15 +1,3 @@
-import warnings
-
-# xformers 0.0.27 still calls `torch.library.impl_abstract`, which torch
-# deprecated in favor of `register_fake`. Suppress the FutureWarning before
-# importing xformers so anyone doing `AutoModel.from_pretrained(...)` gets
-# clean output. xformers 0.0.28+ fixed it; this filter is a no-op there.
-warnings.filterwarnings(
-    "ignore",
-    message=r"`torch\.library\.impl_abstract` was renamed.*",
-    category=FutureWarning,
-)
-
 import torch
 from torch import nn
 
@@ -19,8 +7,6 @@ from torch.nn.functional import scaled_dot_product_attention
 from typing import Optional, Tuple
 import numpy as np
 
-from xformers.ops import SwiGLU
-
 try:
     from flash_attn.flash_attn_interface import flash_attn_varlen_func
 
@@ -38,6 +24,7 @@ from transformers.modeling_outputs import (
     MaskedLMOutput,
     SequenceClassifierOutput,
 )
+from transformers import initialization as init
 
 try:
     import logging as _std_logging
@@ -52,6 +39,7 @@ except Exception:
     pass
 
 from .rotary import precompute_freqs_cis, apply_rotary_emb
+from .swiglu import SwiGLU
 
 
 class DataCollatorWithPacking(DataCollatorForLanguageModeling):
@@ -226,9 +214,9 @@ class NeoBERTPreTrainedModel(PreTrainedModel):
 
     def _init_weights(self, module):
         if isinstance(module, nn.Linear):
-            module.weight.data.uniform_(-self.config.decoder_init_range, self.config.decoder_init_range)
+            init.uniform_(module.weight, -self.config.decoder_init_range, self.config.decoder_init_range)
         elif isinstance(module, nn.Embedding):
-            module.weight.data.uniform_(-self.config.embedding_init_range, self.config.embedding_init_range)
+            init.uniform_(module.weight, -self.config.embedding_init_range, self.config.embedding_init_range)
 
 
 class NeoBERT(NeoBERTPreTrainedModel):
@@ -318,6 +306,13 @@ class NeoBERT(NeoBERTPreTrainedModel):
             attentions=attentions if output_attentions else None,
         )
 
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, NeoBERT):
+            config = module.config
+            freqs_cis = precompute_freqs_cis(config.hidden_size // config.num_attention_heads, config.max_length)
+            init.copy_(module.freqs_cis, freqs_cis)
+
 
 class NeoBERTLMHead(NeoBERTPreTrainedModel):
     config_class = NeoBERTConfig
@@ -384,9 +379,9 @@ class NeoBERTForSequenceClassification(NeoBERTPreTrainedModel):
 
     def _init_weights(self, module):
         if isinstance(module, nn.Linear):
-            module.weight.data.normal_(mean=0.0, std=self.classifier_init_range)
+            init.normal_(module.weight, mean=0.0, std=self.classifier_init_range)
             if module.bias is not None:
-                module.bias.data.zero_()
+                init.zeros_(module.bias)
 
     def forward(
         self,

swiglu.py

@@ -0,0 +1,49 @@
+import warnings
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class _SwiGLU(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: int,
+        out_features: int | None = None,
+        bias: bool = True,
+    ):
+        super().__init__()
+
+        out_features = out_features or in_features
+
+        self.in_features = in_features
+        self.hidden_features = hidden_features
+        self.out_features = out_features or in_features
+
+        self.w12 = nn.Linear(
+            in_features, 2 * hidden_features, bias=bias
+        )
+        self.w3 = nn.Linear(
+            hidden_features, out_features, bias=bias
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x1, x2 = self.w12(x).chunk(2, dim=-1)
+        return self.w3(F.silu(x1) * x2)
+
+
+# xformers 0.0.27 still calls `torch.library.impl_abstract`, which torch
+# deprecated in favor of `register_fake`. Suppress the FutureWarning before
+# importing xformers so anyone doing `AutoModel.from_pretrained(...)` gets
+# clean output. xformers 0.0.28+ fixed it; this filter is a no-op there.
+warnings.filterwarnings(
+    "ignore",
+    message=r"`torch\.library\.impl_abstract` was renamed.*",
+    category=FutureWarning,
+)
+
+try:
+    from xformers.ops import SwiGLU
+except ImportError:
+    SwiGLU = _SwiGLU