Sync modeling code from FinetuningBert@ebbd87d466c6

BertForJointParsing.py

@@ -10,11 +10,11 @@ from transformers.utils import ModelOutput
 # transformers' get_imports() strips the try body but keeps the except body, so the absolute
 # imports have to sit in the try or trust_remote_code demands them as pip packages
 try:
-    from BertForSyntaxParsing import BertSyntaxParsingHead, SyntaxLabels, SyntaxLogitsOutput, parse_logits as syntax_parse_logits
+    from BertForSyntaxParsing import BertSyntaxParsingHead, SyntaxLabels, SyntaxLogitsOutput, parse_logits as syntax_parse_logits, make_extended_attention_mask
     from BertForPrefixMarking import BertPrefixMarkingHead, parse_logits as prefix_parse_logits, encode_sentences_for_bert_for_prefix_marking, get_prefixes_from_str
     from BertForMorphTagging import BertMorphTaggingHead, MorphLogitsOutput, MorphLabels, parse_logits as morph_parse_logits
 except ImportError:
-    from .BertForSyntaxParsing import BertSyntaxParsingHead, SyntaxLabels, SyntaxLogitsOutput, parse_logits as syntax_parse_logits
+    from .BertForSyntaxParsing import BertSyntaxParsingHead, SyntaxLabels, SyntaxLogitsOutput, parse_logits as syntax_parse_logits, make_extended_attention_mask
     from .BertForPrefixMarking import BertPrefixMarkingHead, parse_logits as prefix_parse_logits, encode_sentences_for_bert_for_prefix_marking, get_prefixes_from_str
     from .BertForMorphTagging import BertMorphTaggingHead, MorphLogitsOutput, MorphLabels, parse_logits as morph_parse_logits
     
@@ -142,7 +142,7 @@ class BertForJointParsing(BertPreTrainedModel):
         # calculate the extended attention mask for any child that might need it
         extended_attention_mask = None
         if attention_mask is not None:
-            extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_ids.size())
+            extended_attention_mask = make_extended_attention_mask(attention_mask, bert_outputs[0].dtype)
 
         # extract the hidden states, and apply the dropout
         hidden_states = self.dropout(bert_outputs[0])

BertForPrefixMarking.py

@@ -297,3 +297,6 @@ def get_predicted_prefix_len_from_logits(token, token_logits, config: PrefixConf
             break
 
     return cur_len
+
+# name referenced by the auto_map in dicta-il/dictabert-seg config.json
+BertForPrefixMarking = BaseForPrefixMarking

BertForSyntaxParsing.py

@@ -254,6 +254,10 @@ def wrap_bert_layer_with_partial_info(layer_idx: int, bert_layer: nn.Module, par
         return hidden_states
     bert_layer.forward = new_forward
 
+# PreTrainedModel.get_extended_attention_mask was removed in transformers 5
+def make_extended_attention_mask(attention_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    return (1.0 - attention_mask[:, None, None, :].to(dtype)) * torch.finfo(dtype).min
+
 class BaseForSyntaxParsing(BertPreTrainedModel):
     base_model_prefix = ""
     def __init__(self, config, syntax_head_size=128, bert_cls=BertModel, is_partial_info_model=False, is_cls_model=False, syntax_attn_label_count=1, syntax_func_label_idx=0):
@@ -331,7 +335,7 @@ class BaseForSyntaxParsing(BertPreTrainedModel):
         
         extended_attention_mask = None
         if attention_mask is not None:
-            extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_ids.size())
+            extended_attention_mask = make_extended_attention_mask(attention_mask, bert_outputs[0].dtype)
         
         hidden_states = self.dropout(bert_outputs[0])
 
@@ -546,4 +550,6 @@ def choose_contracting_arc(indices: torch.LongTensor, sorted_indices: List[List[
     return best_base_idx, best_head_idx
 
 def replace_tensor_value(tensor, orig_value, replacement_value):
-    return torch.where(tensor == orig_value, replacement_value, tensor)
+    return torch.where(tensor == orig_value, replacement_value, tensor)
+# name referenced by the auto_map in dicta-il/dictabert-syntax config.json
+BertForSyntaxParsing = BaseForSyntaxParsing