Fix model correctness and Hugging Face compatibility

README.md

@@ -152,8 +152,7 @@ model = AutoModel.from_pretrained("Taykhoom/SpliceBERT-1024nt",
 
 Original model and code by Chen et al. Source:
 [GitHub](https://github.com/biomed-AI/SpliceBERT).
-The HF conversion code was authored primarily by [Claude Code](https://claude.ai/code)
-and reviewed manually by Taykhoom Dalal.
+Hugging Face port maintained by Taykhoom Dalal.
 
 ## License
 

config.json

@@ -5,9 +5,9 @@
   ],
   "model_type": "bert_updated",
   "auto_map": {
-    "AutoConfig": "Taykhoom/BERT-updated--configuration_bert_updated.BertUpdatedConfig",
-    "AutoModel": "Taykhoom/BERT-updated--modeling_bert.BertModel",
-    "AutoModelForMaskedLM": "Taykhoom/BERT-updated--modeling_bert.BertForMaskedLM"
+    "AutoConfig": "configuration_bert_updated.BertUpdatedConfig",
+    "AutoModel": "modeling_bert.BertModel",
+    "AutoModelForMaskedLM": "modeling_bert.BertForMaskedLM"
   },
   "vocab_size": 10,
   "hidden_size": 512,
@@ -24,4 +24,4 @@
   "pad_token_id": 0,
   "model_max_length": 1024,
   "transformers_version": "4.57.6"
-}
+}

configuration_bert_updated.py

@@ -18,6 +18,7 @@ class BertUpdatedConfig(PretrainedConfig):
         num_attention_heads=12,
         intermediate_size=3072,
         hidden_act="gelu",
+        mlm_hidden_act=None,
         hidden_dropout_prob=0.1,
         attention_probs_dropout_prob=0.1,
         max_position_embeddings=512,
@@ -33,6 +34,7 @@ class BertUpdatedConfig(PretrainedConfig):
         self.num_attention_heads = num_attention_heads
         self.intermediate_size = intermediate_size
         self.hidden_act = hidden_act
+        self.mlm_hidden_act = mlm_hidden_act
         self.hidden_dropout_prob = hidden_dropout_prob
         self.attention_probs_dropout_prob = attention_probs_dropout_prob
         self.max_position_embeddings = max_position_embeddings

modeling_bert.py

@@ -4,6 +4,8 @@ from typing import Optional, Tuple, Union
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import torch.utils.checkpoint
+from transformers.activations import ACT2FN
 from transformers import PreTrainedModel, PretrainedConfig
 from transformers.modeling_outputs import BaseModelOutputWithPooling, MaskedLMOutput
 
@@ -32,6 +34,7 @@ class BertSelfAttention(nn.Module):
         hidden_states: torch.Tensor,
         key_padding_mask: Optional[torch.Tensor] = None,
         output_attentions: bool = False,
+        head_mask: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
         q = self._split_heads(self.query(hidden_states))
         k = self._split_heads(self.key(hidden_states))
@@ -41,9 +44,14 @@ class BertSelfAttention(nn.Module):
         scores = torch.matmul(q, k.transpose(-1, -2)) / scale
         if key_padding_mask is not None:
             scores = scores.masked_fill(key_padding_mask[:, None, None, :], float("-inf"))
-        probs = F.softmax(scores, dim=-1)
-        probs = self.dropout(probs)
-        context = torch.matmul(probs, v)
+        probs = torch.nan_to_num(
+            F.softmax(scores.float(), dim=-1),
+            nan=0.0,
+        )
+        context_probs = self.dropout(probs)
+        if head_mask is not None:
+            context_probs = context_probs * head_mask
+        context = torch.matmul(context_probs.to(v.dtype), v)
 
         B, _, T, _ = context.shape
         context = context.permute(0, 2, 1, 3).contiguous().view(B, T, self.all_head_size)
@@ -60,9 +68,15 @@ class BertSdpaSelfAttention(BertSelfAttention):
         hidden_states: torch.Tensor,
         key_padding_mask: Optional[torch.Tensor] = None,
         output_attentions: bool = False,
+        head_mask: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
-        if output_attentions:
-            return super().forward(hidden_states, key_padding_mask, output_attentions=True)
+        if output_attentions or head_mask is not None:
+            return super().forward(
+                hidden_states,
+                key_padding_mask,
+                output_attentions=output_attentions,
+                head_mask=head_mask,
+            )
 
         B, T, _ = hidden_states.shape
         q = self._split_heads(self.query(hidden_states))
@@ -74,7 +88,13 @@ class BertSdpaSelfAttention(BertSelfAttention):
             attn_mask = torch.zeros(B, 1, 1, T, dtype=q.dtype, device=q.device)
             attn_mask = attn_mask.masked_fill(key_padding_mask[:, None, None, :], float("-inf"))
 
-        context = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+        context = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=attn_mask,
+            dropout_p=self.dropout.p if self.training else 0.0,
+        )
         context = context.permute(0, 2, 1, 3).contiguous().view(B, T, self.all_head_size)
         return context, None
 
@@ -86,9 +106,22 @@ class BertFlashSelfAttention(BertSelfAttention):
         hidden_states: torch.Tensor,
         key_padding_mask: Optional[torch.Tensor] = None,
         output_attentions: bool = False,
+        head_mask: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
-        if output_attentions:
-            return super().forward(hidden_states, key_padding_mask, output_attentions=True)
+        if (
+            output_attentions
+            or head_mask is not None
+            or (
+                key_padding_mask is not None
+                and (~key_padding_mask).sum(dim=-1).eq(0).any()
+            )
+        ):
+            return super().forward(
+                hidden_states,
+                key_padding_mask,
+                output_attentions=output_attentions,
+                head_mask=head_mask,
+            )
 
         try:
             from flash_attn import flash_attn_func, flash_attn_varlen_func
@@ -104,9 +137,11 @@ class BertFlashSelfAttention(BertSelfAttention):
         k = self._split_heads(self.key(hidden_states)).permute(0, 2, 1, 3)
         v = self._split_heads(self.value(hidden_states)).permute(0, 2, 1, 3)
 
-        orig_dtype = q.dtype
-        if orig_dtype not in (torch.float16, torch.bfloat16):
-            q, k, v = q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16)
+        if q.dtype not in (torch.float16, torch.bfloat16):
+            raise ValueError(
+                "flash_attention_2 requires float16 or bfloat16 model weights. "
+                f"Received {q.dtype}."
+            )
 
         if key_padding_mask is not None and key_padding_mask.any():
             attend = ~key_padding_mask
@@ -117,13 +152,20 @@ class BertFlashSelfAttention(BertSelfAttention):
                 q_u, k_u, v_u,
                 cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens,
                 max_seqlen_q=max_seqlen, max_seqlen_k=max_seqlen,
+                dropout_p=self.dropout.p if self.training else 0.0,
                 causal=False,
             )
             out = pad_input(out_u, indices, B, T)
         else:
-            out = flash_attn_func(q, k, v, causal=False)
+            out = flash_attn_func(
+                q,
+                k,
+                v,
+                dropout_p=self.dropout.p if self.training else 0.0,
+                causal=False,
+            )
 
-        out = out.to(orig_dtype).reshape(B, T, self.all_head_size)
+        out = out.reshape(B, T, self.all_head_size)
         return out, None
 
 
@@ -158,8 +200,14 @@ class BertAttention(nn.Module):
         hidden_states: torch.Tensor,
         key_padding_mask: Optional[torch.Tensor],
         output_attentions: bool = False,
+        head_mask: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
-        self_out, attn_weights = self.self(hidden_states, key_padding_mask, output_attentions)
+        self_out, attn_weights = self.self(
+            hidden_states,
+            key_padding_mask,
+            output_attentions,
+            head_mask,
+        )
         return self.output(self_out, hidden_states), attn_weights
 
 
@@ -167,9 +215,15 @@ class BertIntermediate(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        if config.hidden_act in (None, "identity", "linear"):
+            self.intermediate_act_fn = lambda value: value
+        elif isinstance(config.hidden_act, str):
+            self.intermediate_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.intermediate_act_fn = config.hidden_act
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        return F.gelu(self.dense(hidden_states))
+        return self.intermediate_act_fn(self.dense(hidden_states))
 
 
 class BertOutput(nn.Module):
@@ -196,8 +250,14 @@ class BertLayer(nn.Module):
         hidden_states: torch.Tensor,
         key_padding_mask: Optional[torch.Tensor],
         output_attentions: bool = False,
+        head_mask: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
-        attn_out, attn_weights = self.attention(hidden_states, key_padding_mask, output_attentions)
+        attn_out, attn_weights = self.attention(
+            hidden_states,
+            key_padding_mask,
+            output_attentions,
+            head_mask,
+        )
         return self.output(self.intermediate(attn_out), attn_out), attn_weights
 
 
@@ -205,6 +265,7 @@ class BertEncoder(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = False
 
     def forward(
         self,
@@ -212,12 +273,31 @@ class BertEncoder(nn.Module):
         key_padding_mask: Optional[torch.Tensor],
         output_hidden_states: bool = False,
         output_attentions: bool = False,
+        head_mask: Optional[torch.Tensor] = None,
     ) -> Tuple:
         all_hidden_states = (hidden_states,) if output_hidden_states else None
         all_attentions = () if output_attentions else None
 
-        for layer in self.layer:
-            hidden_states, attn_weights = layer(hidden_states, key_padding_mask, output_attentions)
+        for layer_index, layer in enumerate(self.layer):
+            layer_head_mask = (
+                head_mask[layer_index] if head_mask is not None else None
+            )
+            if self.gradient_checkpointing and self.training:
+                hidden_states, attn_weights = torch.utils.checkpoint.checkpoint(
+                    layer.__call__,
+                    hidden_states,
+                    key_padding_mask,
+                    output_attentions,
+                    layer_head_mask,
+                    use_reentrant=False,
+                )
+            else:
+                hidden_states, attn_weights = layer(
+                    hidden_states,
+                    key_padding_mask,
+                    output_attentions,
+                    layer_head_mask,
+                )
             if output_hidden_states:
                 all_hidden_states = all_hidden_states + (hidden_states,)
             if output_attentions:
@@ -236,13 +316,39 @@ class BertEmbeddings(nn.Module):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False)
 
-    def forward(self, input_ids: torch.LongTensor, token_type_ids: Optional[torch.LongTensor] = None) -> torch.Tensor:
-        B, T = input_ids.shape
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+    ) -> torch.Tensor:
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Specify exactly one of input_ids or inputs_embeds.")
+        B, T = (
+            input_ids.shape
+            if input_ids is not None
+            else inputs_embeds.shape[:2]
+        )
         if token_type_ids is None:
-            token_type_ids = torch.zeros_like(input_ids)
-        x = self.word_embeddings(input_ids)
-        x = x + self.position_embeddings(self.position_ids[:, :T])
+            token_type_ids = torch.zeros(
+                (B, T),
+                dtype=torch.long,
+                device=(
+                    input_ids.device
+                    if input_ids is not None
+                    else inputs_embeds.device
+                ),
+            )
+        if position_ids is None:
+            position_ids = self.position_ids[:, :T]
+        x = (
+            self.word_embeddings(input_ids)
+            if inputs_embeds is None
+            else inputs_embeds
+        )
         x = x + self.token_type_embeddings(token_type_ids)
+        x = x + self.position_embeddings(position_ids)
         return self.dropout(self.LayerNorm(x))
 
 
@@ -261,9 +367,20 @@ class BertPredictionHeadTransform(nn.Module):
         super().__init__()
         self.dense = nn.Linear(config.hidden_size, config.hidden_size)
         self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        hidden_act = (
+            config.mlm_hidden_act
+            if getattr(config, "mlm_hidden_act", None) is not None
+            else config.hidden_act
+        )
+        if hidden_act in (None, "identity", "linear"):
+            self.transform_act_fn = lambda value: value
+        elif isinstance(hidden_act, str):
+            self.transform_act_fn = ACT2FN[hidden_act]
+        else:
+            self.transform_act_fn = hidden_act
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        return self.LayerNorm(F.gelu(self.dense(hidden_states)))
+        return self.LayerNorm(self.transform_act_fn(self.dense(hidden_states)))
 
 
 class BertModel(PreTrainedModel):
@@ -271,12 +388,14 @@ class BertModel(PreTrainedModel):
     base_model_prefix = "bert"
     _supports_sdpa = True
     _supports_flash_attn_2 = True
+    supports_gradient_checkpointing = True
+    _keys_to_ignore_on_load_missing = [r"pooler\."]
 
-    def __init__(self, config):
+    def __init__(self, config, add_pooling_layer=True):
         super().__init__(config)
         self.embeddings = BertEmbeddings(config)
         self.encoder = BertEncoder(config)
-        self.pooler = BertPooler(config)
+        self.pooler = BertPooler(config) if add_pooling_layer else None
         self.post_init()
 
     def get_input_embeddings(self):
@@ -287,9 +406,12 @@ class BertModel(PreTrainedModel):
 
     def forward(
         self,
-        input_ids: torch.LongTensor,
+        input_ids: Optional[torch.LongTensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
         token_type_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
         output_hidden_states: Optional[bool] = None,
         output_attentions: Optional[bool] = None,
         return_dict: Optional[bool] = None,
@@ -298,22 +420,51 @@ class BertModel(PreTrainedModel):
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Specify exactly one of input_ids or inputs_embeds.")
+        input_shape = (
+            input_ids.shape
+            if input_ids is not None
+            else inputs_embeds.shape[:2]
+        )
+        input_device = (
+            input_ids.device
+            if input_ids is not None
+            else inputs_embeds.device
+        )
         if attention_mask is None:
-            attention_mask = torch.ones_like(input_ids)
+            attention_mask = torch.ones(input_shape, device=input_device)
         key_padding_mask = attention_mask.eq(0)
         if not key_padding_mask.any():
             key_padding_mask = None
-
-        x = self.embeddings(input_ids, token_type_ids)
+        head_mask = self.get_head_mask(
+            head_mask,
+            self.config.num_hidden_layers,
+        )
+        if all(mask is None for mask in head_mask):
+            head_mask = None
+
+        x = self.embeddings(
+            input_ids=input_ids,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+        )
         last_hidden_state, all_hidden_states, all_attentions = self.encoder(
             x, key_padding_mask,
             output_hidden_states=output_hidden_states,
             output_attentions=output_attentions,
+            head_mask=head_mask,
         )
-        pooled = self.pooler(last_hidden_state)
+        pooled = self.pooler(last_hidden_state) if self.pooler is not None else None
 
         if not return_dict:
-            return tuple(v for v in [last_hidden_state, pooled, all_hidden_states, all_attentions] if v is not None)
+            output = (last_hidden_state, pooled)
+            if output_hidden_states:
+                output += (all_hidden_states,)
+            if output_attentions:
+                output += (all_attentions,)
+            return output
 
         return BaseModelOutputWithPooling(
             last_hidden_state=last_hidden_state,
@@ -328,10 +479,12 @@ class BertForMaskedLM(PreTrainedModel):
     base_model_prefix = "bert"
     _supports_sdpa = True
     _supports_flash_attn_2 = True
+    supports_gradient_checkpointing = True
+    _keys_to_ignore_on_load_unexpected = [r"bert\.pooler\."]
 
     def __init__(self, config):
         super().__init__(config)
-        self.bert = BertModel(config)
+        self.bert = BertModel(config, add_pooling_layer=False)
         self.transform = BertPredictionHeadTransform(config)
         self.cls = nn.Linear(config.hidden_size, config.vocab_size)
         self.post_init()
@@ -339,11 +492,23 @@ class BertForMaskedLM(PreTrainedModel):
     def get_input_embeddings(self):
         return self.bert.embeddings.word_embeddings
 
+    def set_input_embeddings(self, value):
+        self.bert.embeddings.word_embeddings = value
+
+    def get_output_embeddings(self):
+        return self.cls
+
+    def set_output_embeddings(self, value):
+        self.cls = value
+
     def forward(
         self,
-        input_ids: torch.LongTensor,
+        input_ids: Optional[torch.LongTensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
         token_type_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
         labels: Optional[torch.LongTensor] = None,
         output_hidden_states: Optional[bool] = None,
         output_attentions: Optional[bool] = None,
@@ -352,8 +517,14 @@ class BertForMaskedLM(PreTrainedModel):
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         outputs = self.bert(
-            input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids,
-            output_hidden_states=output_hidden_states, output_attentions=output_attentions,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
             return_dict=True,
         )
         logits = self.cls(self.transform(outputs.last_hidden_state))
@@ -363,7 +534,11 @@ class BertForMaskedLM(PreTrainedModel):
             loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), labels.view(-1), ignore_index=-100)
 
         if not return_dict:
-            output = (logits,) + outputs[2:]
+            output = (logits,)
+            if output_hidden_states:
+                output += (outputs.hidden_states,)
+            if output_attentions:
+                output += (outputs.attentions,)
             return (loss,) + output if loss is not None else output
 
         return MaskedLMOutput(

tokenization_splicebert.py

@@ -79,7 +79,7 @@ class SpliceBERTTokenizer(PreTrainedTokenizer):
         sep = [self.sep_token_id]
         if token_ids_1 is None:
             return cls + token_ids_0 + sep
-        return cls + token_ids_0 + sep + cls + token_ids_1 + sep
+        return cls + token_ids_0 + sep + token_ids_1 + sep
 
     def get_special_tokens_mask(self, token_ids_0, token_ids_1=None,
                                 already_has_special_tokens=False):
@@ -89,10 +89,13 @@ class SpliceBERTTokenizer(PreTrainedTokenizer):
             )
         mask = [1] + [0] * len(token_ids_0) + [1]
         if token_ids_1 is not None:
-            mask += [1] + [0] * len(token_ids_1) + [1]
+            mask += [0] * len(token_ids_1) + [1]
         return mask
 
     def create_token_type_ids_from_sequences(self, token_ids_0, token_ids_1=None):
         if token_ids_1 is None:
-            return [0] + token_ids_0 + [0]
-        return [0] + token_ids_0 + [0, 0] + token_ids_1 + [0]
+            return [0] * (len(token_ids_0) + 2)
+        return (
+            [0] * (len(token_ids_0) + 2)
+            + [1] * (len(token_ids_1) + 1)
+        )