Upgrade transformers version

README.md

@@ -67,10 +67,10 @@ We trained the model on public datasets described in the Dataset and Training se
 
 ### **Installation**
 
-The model requires transformers version 4.47.1.
+The model requires transformers version >=4.47.1.
 
 ```bash
-pip install transformers==4.47.1
+pip install transformers>=4.47.1
 ```
 
 ### **Usage**

llama_bidirectional_model.py

@@ -5,7 +5,7 @@ import torch.nn.functional as F
 from torch import Tensor, nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from transformers.cache_utils import Cache, HybridCache
-from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
+from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_attention_mask_for_sdpa
 from transformers.modeling_outputs import (
     BaseModelOutputWithPast,
     SequenceClassifierOutputWithPast,
@@ -21,6 +21,24 @@ from transformers.utils import logging
 logger = logging.get_logger(__name__)
 
 
+def create_bidirectional_attention_mask(
+    attn_implementation: str,
+    attention_mask: torch.Tensor,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+
+    if attn_implementation == "flash_attention_2":
+        if attention_mask is not None and (attention_mask == 0.0).any():
+            return attention_mask
+        return None
+    elif attn_implementation == "eager":
+        return _prepare_4d_attention_mask(attention_mask, dtype=dtype)
+    elif attn_implementation == "sdpa":
+        return _prepare_4d_attention_mask_for_sdpa(attention_mask, dtype=dtype)
+    else:
+        raise ValueError(f"Unsupported attention implementation: {attn_implementation}, only support flash_attention_2, eager or sdpa")
+
+
 def pool(last_hidden_states: Tensor, attention_mask: Tensor, pool_type: str) -> Tensor:
     last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
 
@@ -64,31 +82,17 @@ class LlamaBidirectionalModel(LlamaModel):
         super().__init__(config)
         for layer in self.layers:
             layer.self_attn.is_causal = False
-        self.config._attn_implementation = "eager"
-
-    def _update_causal_mask(
-        self,
-        attention_mask: torch.Tensor,
-        input_tensor: torch.Tensor,
-        cache_position: torch.Tensor,
-        past_key_values: Cache,
-        output_attentions: bool,
-    ):
-        # Generates bi-directional attention.
-        causal_mask = _prepare_4d_attention_mask(attention_mask, input_tensor.dtype)
-        return causal_mask
 
 
 class LlamaBidirectionalForSequenceClassification(LlamaForSequenceClassification):
     config_class = LlamaBidirectionalConfig
 
     def __init__(self, config):
-        super().__init__(config)
-        # Releasing the parameters of LlamaModel
-        # created by parent LlamaForSequenceClassification
-        del self.model
+        LlamaPreTrainedModel.__init__(self, config)
 
+        self.num_labels = config.num_labels
         self.model = LlamaBidirectionalModel(config)
+        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
 
         # Initialize weights and apply final processing
         self.post_init()
@@ -105,6 +109,7 @@ class LlamaBidirectionalForSequenceClassification(LlamaForSequenceClassification
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        **kwargs,
     ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
         r"""
         labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
@@ -116,6 +121,16 @@ class LlamaBidirectionalForSequenceClassification(LlamaForSequenceClassification
             return_dict if return_dict is not None else self.config.use_return_dict
         )
 
+        # Keep original 2D mask for pooling
+        attention_mask_2d = attention_mask
+
+        # Create 4D bidirectional attention mask
+        attention_mask = create_bidirectional_attention_mask(
+            attn_implementation=self.config._attn_implementation,
+            attention_mask=attention_mask,
+            dtype=self.config.torch_dtype,
+        )
+
         transformer_outputs = self.model(
             input_ids,
             attention_mask=attention_mask,
@@ -126,12 +141,13 @@ class LlamaBidirectionalForSequenceClassification(LlamaForSequenceClassification
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
+            **kwargs,
         )
         hidden_states = transformer_outputs[0]
 
         pooled_hidden_states = pool(
             last_hidden_states=hidden_states,
-            attention_mask=attention_mask,
+            attention_mask=attention_mask_2d,
             pool_type=self.config.pooling,
         )
 
@@ -140,7 +156,7 @@ class LlamaBidirectionalForSequenceClassification(LlamaForSequenceClassification
 
         loss = None
         if labels is not None:
-            labels = labels.to(logits.device)
+            labels = labels.to(pooled_logits.device)
             if self.config.problem_type is None:
                 if self.num_labels == 1:
                     self.config.problem_type = "regression"