fix for transformers >= 4.35.2

README.md

@@ -9,7 +9,7 @@ pipeline_tag: fill-mask
 ---
 
 # LSG model 
-**Transformers >= 4.23.1**\
+**Transformers >= 4.35.2**\
 **This model relies on a custom modeling file, you need to add trust_remote_code=True**\
 **See [\#13467](https://github.com/huggingface/transformers/pull/13467)**
 

modeling_lsg_pegasus.py

@@ -1,7 +1,7 @@
 from logging import warn
 import torch
 from transformers.models.pegasus.modeling_pegasus import *
-from transformers.models.pegasus.modeling_pegasus import _expand_mask
+from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_causal_attention_mask
 import torch.nn as nn
 import sys
 
@@ -895,7 +895,7 @@ class LSGPegasusEncoder(LSGPegasusPreTrainedModel, PegasusEncoder):
         # expand attention_mask
         if attention_mask is not None:
             # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
-            attention_mask = _expand_mask(attention_mask, inputs_embeds.dtype)
+            attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype)
 
         encoder_states = () if output_hidden_states else None
         all_attentions = () if output_attentions else None
@@ -1122,4 +1122,4 @@ try:
         str_to_class(value.split(".")[-1]).register_for_auto_class(key)
 except:
     warn("AutoRegister isn't available, you'll have to manually copy modeling.py after .save_pretrained(...).")
-    warn("Update to transformers >= 4.23.1 to fix.")
+    warn("Update to transformers >= 4.35.2 to fix.")