Upload BertForJointParsing.py

BertForJointParsing.py

@@ -8,7 +8,7 @@ from transformers import BertPreTrainedModel, BertModel, BertTokenizerFast
 from transformers.models.bert.modeling_bert import BertOnlyMLMHead
 from transformers.utils import ModelOutput
 from .BertForSyntaxParsing import BertSyntaxParsingHead, SyntaxLabels, SyntaxLogitsOutput, parse_logits as syntax_parse_logits
-from .BertForPrefixMarking import BertPrefixMarkingHead, parse_logits as prefix_parse_logits, encode_sentences_for_bert_for_prefix_marking
+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 
     
 import warnings
@@ -186,11 +186,14 @@ class BertForJointParsing(BertPreTrainedModel):
             morph_logits=morph_logits
         )
 
-    def predict(self, sentences: Union[str, List[str]], tokenizer: BertTokenizerFast, padding='longest', truncation=True, compute_syntax_mst=True, per_token_ner=False):
+    def predict(self, sentences: Union[str, List[str]], tokenizer: BertTokenizerFast, padding='longest', truncation=True, compute_syntax_mst=True, per_token_ner=False, as_iahlt_ud=False, as_htb_ud=False):
         is_single_sentence = isinstance(sentences, str)
         if is_single_sentence:
             sentences = [sentences]
-            
+        
+        if (as_htb_ud or as_iahlt_ud) and (self.prefix is None or self.morph is None or self.syntax is None or self.lex is None):
+            raise ValueError("Cannot output UD format when any of the prefix,morph,syntax,lex heads aren't loaded.")
+        
         # predict the logits for the sentence
         if self.prefix is not None:
             inputs = encode_sentences_for_bert_for_prefix_marking(tokenizer, sentences, padding)
@@ -230,10 +233,15 @@ class BertForJointParsing(BertPreTrainedModel):
                     merge_token_list(final_output[sent_idx]['tokens'], map(itemgetter(1), parsed), 'ner')
                 final_output[sent_idx]['ner_entities'] = aggregate_ner_tokens(parsed) 
         
+        if as_iahlt_ud or as_htb_ud:
+            final_output = convert_output_to_ud(final_output, htb_extras=as_htb_ud)
+        
         if is_single_sentence:
             final_output = final_output[0]
         return final_output
 
+
+
 def aggregate_ner_tokens(predictions):
     entities = []
     prev = None
@@ -302,4 +310,166 @@ def lex_parse_logits(inputs: Dict[str, torch.Tensor], sentences: List[str], toke
                 ret[-1] = (ret[-1][0] + token[2:], ret[-1][1])
                 continue
             ret.append((token, tokenizer._convert_id_to_token(predictions[batch_idx, tok_idx])))
-    return batch_ret
+    return batch_ret
+
+ud_prefixes_to_pos = {
+    'ש': ['SCONJ'],
+    'מש': ['SCONJ'],
+    'כש': ['SCONJ'],
+    'לכש': ['SCONJ'],
+    'בש': ['SCONJ'],
+    'לש': ['SCONJ'],
+    'ו': ['CCONJ'],
+    'ל': ['ADP'],
+    'ה': ['DET', 'SCONJ'],
+    'מ': ['ADP', 'SCONJ'],
+    'ב': ['ADP'],
+    'כ': ['ADP', 'ADV'],
+}
+ud_suffix_to_htb_str = {
+    'Gender=Masc|Number=Sing|Person=3': '_הוא',	
+	'Gender=Masc|Number=Plur|Person=3': '_הם',	
+	'Gender=Fem|Number=Sing|Person=3': '_היא',	
+	'Gender=Fem|Number=Plur|Person=3': '_הן',	
+	'Gender=Fem,Masc|Number=Plur|Person=1': '_אנחנו',	
+	'Gender=Fem,Masc|Number=Sing|Person=1': '_אני',	
+	'Gender=Masc|Number=Plur|Person=2': '_אתם',	
+	'Gender=Masc|Number=Sing|Person=3': '_הוא',
+	'Gender=Masc|Number=Sing|Person=2': '_אתה',	
+	'Gender=Fem|Number=Sing|Person=2': '_את',	
+	'Gender=Masc|Number=Plur|Person=3': '_הם'
+}
+def convert_output_to_ud(output_sentences, htb_extras=False):
+    final_output = []
+    for sent_idx, sentence in enumerate(output_sentences):
+        # next, go through each word and insert it in the UD format. Store in a temp format for the post process
+        intermediate_output = []
+        ranges = []
+        # store a mapping between each word index and the actual line it appears in
+        idx_to_key = {-1: 0}
+        for word_idx,word in enumerate(sentence['tokens']):
+            # handle blank lexemes
+            if word['lex'] == '[BLANK]':
+                word['lex'] = word['seg'][-1]
+            
+            start = len(intermediate_output)
+            idx_to_key[word_idx] = len(intermediate_output) + 1
+            # Add in all the prefixes
+            if len(word['seg']) > 1:
+                for pre in get_prefixes_from_str(word['seg'][0], greedy=True):
+                    # pos - just take the first valid pos that appears in the predicted prefixes list. 
+                    pos = next((pos for pos in ud_prefixes_to_pos[pre] if pos in word['morph']['prefixes']), ud_prefixes_to_pos[pre][0])
+                    dep, func = ud_get_prefix_dep(pre, word, word_idx)
+                    intermediate_output.append(dict(word=pre, lex=pre, pos=pos, dep=dep, func=func, feats='_'))
+                    
+                # if there was an implicit heh, add it in dependent on the method
+                if not 'ה' in pre and intermediate_output[-1]['pos'] == 'ADP' and 'DET' in word['morph']['prefixes']:
+                    if htb_extras:
+                        intermediate_output.append(dict(word='ה_', lex='ה', pos='DET', dep=word_idx, func='det', feats='_'))
+                    else:
+                        intermediate_output[-1]['feats'] = 'Definite=Def|PronType=Art'
+                    
+                
+            # add the main word in!
+            intermediate_output.append(dict(
+                    word=word['seg'][-1], lex=word['lex'], pos=word['morph']['pos'], 
+                    dep=word['syntax']['dep_head_idx'], func=word['syntax']['dep_func'],
+                    feats='|'.join(f'{k}={v}' for k,v in word['morph']['feats'].items())))
+            
+            # if we have suffixes, this changes things
+            if word['morph']['suffix']:
+                # first determine the dependency info:
+                # For adp, num, det - they main word points to here, and the suffix points to the dependency
+                if word['morph']['pos'] in ['ADP', 'NUM', 'DET']:
+                    intermediate_output[-1]['dep'] = len(intermediate_output)
+                    intermediate_output[-1]['absolute_dep'] = True
+                    intermediate_output[-1]['func'] = 'case'
+                    dep = word['syntax']['dep_head_idx']
+                    func = word['syntax']['dep_func']
+                else:                
+                    # if pos is verb -> obj, num -> dep, default to -> nmod:poss
+                    dep = word_idx
+                    func = {'VERB': 'obj', 'NUM': 'dep'}.get(word['morph']['pos'], 'nmod:poss')
+                
+                s_word, s_lex = word['seg'][-1], word['lex']
+                # update the word of the string and extract the string of the suffix!
+                # for IAHLT:
+                if not htb_extras:
+                    # we need to shorten the main word and extract the suffix
+                    # if it is longer than the lexeme - just take off the lexeme.
+                    if len(s_word) > len(s_lex):
+                        idx = len(s_lex)
+                    # Otherwise, try to find the last letter of the lexeme, and fail that just take the last letter
+                    else:
+                        # take either len-1, or the last occurence (which can be -1 === len-1)
+                        idx = min([len(s_word) - 1, s_word.rfind(s_lex[-1])])
+                    # extract the suffix and update the main word
+                    suf = s_word[idx:]
+                    intermediate_output[-1]['word'] = s_word[:idx]
+                # for htb:
+                else:
+                    # main word becomes the lexeme, the suffix is based on the features
+                    intermediate_output[-1]['word'] = s_lex + '_'
+                    suf_feats = word['morph']['suffix_feats']
+                    suf = ud_suffix_to_htb_str.get(f"Gender={suf_feats['Gender']}|Number={suf_feats['Number']}|Person={suf_feats['Person']}", "_הוא")
+                    # for HTB, if the function is poss, then add a shel pointing to the next word
+                    if func == 'nmod:poss':
+                        intermediate_output.append(dict(word='_של_', lex='של', pos='ADP', dep=len(intermediate_output) + 1, func='case', feats='_', absolute_dep=True))
+                # add the main suffix in
+                intermediate_output.append(dict(word=suf, lex='הוא', pos='PRON', dep=dep, func=func, feats='|'.join(f'{k}={v}' for k,v in word['morph']['suffix_feats'].items())))
+            end = len(intermediate_output)
+            ranges.append((start, end, word['token']))
+            
+        # now that we have the intermediate output, combine it to the final output
+        cur_output = []
+        final_output.append(cur_output)
+        # first, add the headers 
+        cur_output.append(f'# sent_id = {sent_idx + 1}')
+        cur_output.append(f'# text = {sentence["text"]}')
+        
+        # add in all the actual entries
+        for start,end,token in ranges:
+            if end - start > 1:
+                cur_output.append(f'{start + 1}-{end}\t{token}\t_\t_\t_\t_\t_\t_\t_\t_')
+            for idx,output in enumerate(intermediate_output[start:end], start):
+                # compute the actual dependency location
+                dep = output['dep'] if output.get('absolute_dep', False) else idx_to_key[output['dep']]
+                # and add the full ud string in
+                cur_output.append('\t'.join([
+                    str(idx),
+                    output['word'],
+                    output['lex'],
+                    output['pos'], 
+                    output['pos'],
+                    output['feats'],
+                    str(dep),
+                    output['func'],
+                    '_', '_'
+                ]))
+    return final_output
+                
+            
+def ud_get_prefix_dep(pre, word, word_idx):
+    does_follow_main = False
+    
+    # shin goes to the main word for verbs, otherwise follows the word
+    if pre.endswith('ש'): 
+        does_follow_main = word['morph']['pos'] != 'VERB'
+        func = 'mark'
+    # vuv goes to the main word if the function is in the list, otherwise follows
+    elif pre == 'ו': 
+        does_follow_main = word['syntax']['dep_func'] in ["conj", "acl:recl", "parataxis", "root", "acl", "amod", "list", "appos", "dep", "flatccomp"]
+        func = 'cc'
+    else:
+        # for adj, noun, propn, pron, verb - prefixes go to the main word
+        if word['morph']['pos'] in ["ADJ", "NOUN", "PROPN", "PRON", "VERB"]:
+            does_follow_main = True  
+        # otherwise - prefix follows the word if the function is in the list
+        else: does_follow_main = word['syntax']['dep_func'] in ["compound:affix", "det", "aux", "nummod", "advmod", "dep", "cop", "mark", "fixed"]
+        
+        func = 'case'
+        if pre == 'ה':
+            func = 'det' if 'DET' in word['morph']['prefixes'] else 'mark'
+
+    return (word['syntax']['dep_head_idx'] if does_follow_main else word_idx), func
+