Linked python files to main dictabert-joint

BertForJointParsing.py

@@ -1,523 +0,0 @@
-from dataclasses import dataclass
-import re
-from operator import itemgetter
-import torch
-from torch import nn
-from typing import Any, Dict, List, Literal, Optional, Tuple, Union
-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, get_prefixes_from_str
-from .BertForMorphTagging import BertMorphTaggingHead, MorphLogitsOutput, MorphLabels, parse_logits as morph_parse_logits 
-    
-import warnings
-
-@dataclass
-class JointParsingOutput(ModelOutput):
-    loss: Optional[torch.FloatTensor] = None
-    # logits will contain the optional predictions for the given labels
-    logits: Optional[Union[SyntaxLogitsOutput, None]] = None
-    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
-    attentions: Optional[Tuple[torch.FloatTensor]] = None
-    # if no labels are given, we will always include the syntax logits separately
-    syntax_logits: Optional[SyntaxLogitsOutput] = None
-    ner_logits: Optional[torch.FloatTensor] = None
-    prefix_logits: Optional[torch.FloatTensor] = None
-    lex_logits: Optional[torch.FloatTensor] = None
-    morph_logits: Optional[MorphLogitsOutput] = None
-
-# wrapper class to wrap a torch.nn.Module so that you can store a module in multiple linked
-# properties without registering the parameter multiple times
-class ModuleRef:
-    def __init__(self, module: torch.nn.Module):
-        self.module = module
-
-    def forward(self, *args, **kwargs):
-        return self.module.forward(*args, **kwargs)
-
-    def __call__(self, *args, **kwargs):
-        return self.module(*args, **kwargs)
-
-class BertForJointParsing(BertPreTrainedModel):
-    _tied_weights_keys = ["predictions.decoder.bias", "cls.predictions.decoder.weight"]
-
-    def __init__(self, config, do_syntax=None, do_ner=None, do_prefix=None, do_lex=None, do_morph=None, syntax_head_size=64):
-        super().__init__(config)
-
-        self.bert = BertModel(config, add_pooling_layer=False)
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
-        # create all the heads as None, and then populate them as defined
-        self.syntax, self.ner, self.prefix, self.lex, self.morph = (None,)*5
-
-        if do_syntax is not None: 
-            config.do_syntax = do_syntax
-            config.syntax_head_size = syntax_head_size
-        if do_ner is not None: config.do_ner = do_ner
-        if do_prefix is not None: config.do_prefix = do_prefix
-        if do_lex is not None: config.do_lex = do_lex
-        if do_morph is not None: config.do_morph = do_morph
-        
-        # add all the individual heads
-        if config.do_syntax:
-            self.syntax = BertSyntaxParsingHead(config)
-        if config.do_ner:
-            self.num_labels = config.num_labels
-            self.classifier = nn.Linear(config.hidden_size, config.num_labels) # name it same as in BertForTokenClassification 
-            self.ner = ModuleRef(self.classifier)
-        if config.do_prefix:
-            self.prefix = BertPrefixMarkingHead(config)
-        if config.do_lex:
-            self.cls = BertOnlyMLMHead(config) # name it the same as in BertForMaskedLM
-            self.lex = ModuleRef(self.cls)
-        if config.do_morph:
-            self.morph = BertMorphTaggingHead(config)
-        
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def get_output_embeddings(self):
-        return self.cls.predictions.decoder if self.lex is not None else None
-
-    def set_output_embeddings(self, new_embeddings):
-        if self.lex is not None:
-
-            self.cls.predictions.decoder = new_embeddings
-
-    def forward(
-        self,
-        input_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        token_type_ids: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-        prefix_class_id_options: Optional[torch.Tensor] = None,
-        labels: Optional[Union[SyntaxLabels, MorphLabels, torch.Tensor]] = None,
-        labels_type: Optional[Literal['syntax', 'ner', 'prefix', 'lex', 'morph']] = None,
-        head_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        compute_syntax_mst: Optional[bool] = None
-    ):
-        if return_dict is False:
-            warnings.warn("Specified `return_dict=False` but the flag is ignored and treated as always True in this model.")
-        
-        if labels is not None and labels_type is None:
-            raise ValueError("Cannot specify labels without labels_type")
-        
-        if labels_type == 'seg' and prefix_class_id_options is None:
-            raise ValueError('Cannot calculate prefix logits without prefix_class_id_options')
-        
-        if compute_syntax_mst is not None and self.syntax is None:
-            raise ValueError("Cannot compute syntax MST when the syntax head isn't loaded")
-
-
-        bert_outputs = self.bert(
-            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_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=True,
-        )
-        
-        # 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())
-
-        # extract the hidden states, and apply the dropout
-        hidden_states = self.dropout(bert_outputs[0])
-
-        logits = None    
-        syntax_logits = None
-        ner_logits = None
-        prefix_logits = None
-        lex_logits = None
-        morph_logits = None
-
-        # Calculate the syntax
-        if self.syntax is not None and (labels is None or labels_type == 'syntax'):
-            # apply the syntax head
-            loss, syntax_logits = self.syntax(hidden_states, extended_attention_mask, labels, compute_syntax_mst)
-            logits = syntax_logits
-
-        # Calculate the NER
-        if self.ner is not None and (labels is None or labels_type == 'ner'):
-            ner_logits = self.ner(hidden_states)
-            logits = ner_logits
-            if labels is not None:
-                loss_fct = nn.CrossEntropyLoss()
-                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-
-        # Calculate the segmentation
-        if self.prefix is not None and (labels is None or labels_type == 'prefix'):
-            loss, prefix_logits = self.prefix(hidden_states, prefix_class_id_options, labels)
-            logits = prefix_logits
-        
-        # Calculate the lexeme
-        if self.lex is not None and (labels is None or labels_type == 'lex'):
-            lex_logits = self.lex(hidden_states)
-            logits = lex_logits
-            if labels is not None:
-                loss_fct = nn.CrossEntropyLoss()  # -100 index = padding token
-                loss = loss_fct(lex_logits.view(-1, self.config.vocab_size), labels.view(-1))
-
-        if self.morph is not None and (labels is None or labels_type == 'morph'):
-            loss, morph_logits = self.morph(hidden_states, labels)
-            logits = morph_logits
-
-        # no labels => logits = None
-        if labels is None: logits = None
-
-        return JointParsingOutput(
-            loss,
-            logits,
-            hidden_states=bert_outputs.hidden_states,
-            attentions=bert_outputs.attentions,
-            # all the predicted logits section
-            syntax_logits=syntax_logits,
-            ner_logits=ner_logits,
-            prefix_logits=prefix_logits,
-            lex_logits=lex_logits,
-            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, output_style: Literal['json', 'ud', 'iahlt_ud'] = 'json'):
-        is_single_sentence = isinstance(sentences, str)
-        if is_single_sentence:
-            sentences = [sentences]
-        
-        if output_style not in ['json', 'ud', 'iahlt_ud']:
-            raise ValueError('output_style must be in json/ud/iahlt_ud')
-        if output_style in ['ud', '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, and 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)
-        else:
-            inputs = tokenizer(sentences, padding=padding, truncation=truncation, return_offsets_mapping=True, return_tensors='pt')
-
-        offset_mapping = inputs.pop('offset_mapping')
-        # Copy the tensors to the right device, and parse!
-        inputs = {k:v.to(self.device) for k,v in inputs.items()}
-        output = self.forward(**inputs, return_dict=True, compute_syntax_mst=compute_syntax_mst)
-        
-        input_ids = inputs['input_ids'].tolist() # convert once
-        final_output = [dict(text=sentence, tokens=combine_token_wordpieces(ids, offsets, tokenizer)) for sentence, ids, offsets in zip(sentences, input_ids, offset_mapping)]
-        # Syntax logits: each sentence gets a dict(tree: List[dict(word,dep_head,dep_head_idx,dep_func)], root_idx: int)
-        if output.syntax_logits is not None:
-            for sent_idx,parsed in enumerate(syntax_parse_logits(input_ids, sentences, tokenizer, output.syntax_logits)):
-                merge_token_list(final_output[sent_idx]['tokens'], parsed['tree'], 'syntax')
-                final_output[sent_idx]['root_idx'] = parsed['root_idx']
-                
-        # Prefix logits: each sentence gets a list([prefix_segment, word_without_prefix]) - **WITH CLS & SEP**
-        if output.prefix_logits is not None:
-            for sent_idx,parsed in enumerate(prefix_parse_logits(input_ids, sentences, tokenizer, output.prefix_logits)):
-                merge_token_list(final_output[sent_idx]['tokens'], map(tuple, parsed[1:-1]), 'seg')
-            
-        # Lex logits each sentence gets a list(tuple(word, lexeme))
-        if output.lex_logits is not None:
-            for sent_idx, parsed in enumerate(lex_parse_logits(input_ids, sentences, tokenizer, output.lex_logits)):
-                merge_token_list(final_output[sent_idx]['tokens'], map(itemgetter(1), parsed), 'lex')
-                
-        # morph logits each sentences get a dict(text=str, tokens=list(dict(token, pos, feats, prefixes, suffix, suffix_feats?)))
-        if output.morph_logits is not None:
-            for sent_idx,parsed in enumerate(morph_parse_logits(input_ids, sentences, tokenizer, output.morph_logits)):
-                merge_token_list(final_output[sent_idx]['tokens'], parsed['tokens'], 'morph')
-            
-        # NER logits each sentence gets a list(tuple(word, ner))
-        if output.ner_logits is not None:
-            for sent_idx,parsed in enumerate(ner_parse_logits(input_ids, sentences, tokenizer, output.ner_logits, self.config.id2label)):
-                if per_token_ner:
-                    merge_token_list(final_output[sent_idx]['tokens'], map(itemgetter(1), parsed), 'ner')
-                final_output[sent_idx]['ner_entities'] = aggregate_ner_tokens(final_output[sent_idx], parsed) 
-        
-        if output_style in ['ud', 'iahlt_ud']:
-            final_output = convert_output_to_ud(final_output, style='htb' if output_style == 'ud' else 'iahlt')
-        
-        if is_single_sentence:
-            final_output = final_output[0]
-        return final_output
-
-
-
-def aggregate_ner_tokens(final_output, parsed):
-    entities = []
-    prev = None
-    for token_idx, (d, (word, pred)) in enumerate(zip(final_output['tokens'], parsed)):
-        # O does nothing
-        if pred == 'O': prev = None
-        # B- || I-entity != prev (different entity or none)
-        elif pred.startswith('B-') or pred[2:] != prev:
-            prev = pred[2:]
-            entities.append([[word], dict(label=prev, start=d['offsets']['start'], end=d['offsets']['end'], token_start=token_idx, token_end=token_idx)])
-        else: 
-            entities[-1][0].append(word)
-            entities[-1][1]['end'] = d['offsets']['end']
-            entities[-1][1]['token_end'] = token_idx
-
-    return [dict(phrase=' '.join(words), **d) for words, d in entities]
-
-def merge_token_list(src, update, key):
-    for token_src, token_update in zip(src, update):
-        token_src[key] = token_update
-        
-def combine_token_wordpieces(input_ids: List[int], offset_mapping: torch.Tensor, tokenizer: BertTokenizerFast):
-    offset_mapping = offset_mapping.tolist()
-    ret = []
-    special_toks = tokenizer.all_special_tokens
-    for token, offsets in zip(tokenizer.convert_ids_to_tokens(input_ids), offset_mapping):
-        if token in special_toks: continue
-        if token.startswith('##'):
-            ret[-1]['token'] += token[2:]
-            ret[-1]['offsets']['end'] = offsets[1]
-        else: ret.append(dict(token=token, offsets=dict(start=offsets[0], end=offsets[1])))
-    return ret
-
-def ner_parse_logits(input_ids: List[List[int]], sentences: List[str], tokenizer: BertTokenizerFast, logits: torch.Tensor, id2label: Dict[int, str]):
-    predictions = torch.argmax(logits, dim=-1).tolist()
-    batch_ret = []
-
-    special_toks = tokenizer.all_special_tokens
-    for batch_idx in range(len(sentences)):
-
-        ret = []
-        batch_ret.append(ret)
-
-        tokens = tokenizer.convert_ids_to_tokens(input_ids[batch_idx])
-        for tok_idx in range(len(tokens)):
-            token = tokens[tok_idx]
-            if token in special_toks: continue
-
-            # wordpieces should just be appended to the previous word
-            # we modify the last token in ret
-            # by discarding the original end position and replacing it with the new token's end position
-            if token.startswith('##'):
-                continue
-            # for each token, we append a tuple containing: token, label, start position, end position
-            ret.append((token, id2label[predictions[batch_idx][tok_idx]]))
-
-    return batch_ret
-
-def lex_parse_logits(input_ids: List[List[int]], sentences: List[str], tokenizer: BertTokenizerFast, logits: torch.Tensor):
-    
-    predictions = torch.argsort(logits, dim=-1, descending=True)[..., :3].tolist()
-    batch_ret = []
-
-    special_toks = tokenizer.all_special_tokens
-    for batch_idx in range(len(sentences)):
-        intermediate_ret = []
-        tokens = tokenizer.convert_ids_to_tokens(input_ids[batch_idx])
-        for tok_idx in range(len(tokens)):
-            token = tokens[tok_idx]
-            if token in special_toks: continue
-
-            # wordpieces should just be appended to the previous word
-            if token.startswith('##'):
-                intermediate_ret[-1] = (intermediate_ret[-1][0] + token[2:], intermediate_ret[-1][1])
-                continue
-            intermediate_ret.append((token, tokenizer.convert_ids_to_tokens(predictions[batch_idx][tok_idx])))
-        
-        # build the final output taking into account valid letters
-        ret = []
-        batch_ret.append(ret)
-        for (token, lexemes) in intermediate_ret:
-            # must overlap on at least 2 non אהוי letters
-            possible_lets = set(c for c in token if c not in 'אהוי')
-            final_lex = '[BLANK]'
-            for lex in lexemes:
-                if sum(c in possible_lets for c in lex) >= min([2, len(possible_lets), len([c for c in lex if c not in 'אהוי'])]): 
-                    final_lex = lex
-                    break
-            ret.append((token, final_lex))
-        
-    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, style: Literal['htb', 'iahlt']):
-    if style not in ['htb', 'iahlt']:
-        raise ValueError('style must be htb/iahlt')
-        
-    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']):
-            try:
-                # handle blank lexemes
-                if word['lex'] == '[BLANK]':
-                    word['lex'] = word['seg'][-1]
-            except KeyError:
-                import json
-                print(json.dumps(sentence, ensure_ascii=False, indent=2))
-                exit(0)
-
-            start = len(intermediate_output)
-            # 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 style == 'htb':
-                        intermediate_output.append(dict(word='ה_', lex='ה', pos='DET', dep=word_idx, func='det', feats='_'))
-                    elif style == 'iahlt':
-                        intermediate_output[-1]['feats'] = 'Definite=Def|PronType=Art'
-                    
-                
-            idx_to_key[word_idx] = len(intermediate_output) + 1
-            # 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
-                entry_to_assign_suf_dep = None
-                if word['morph']['pos'] in ['ADP', 'NUM', 'DET']:
-                    entry_to_assign_suf_dep = intermediate_output[-1]
-                    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 style == 'iahlt':
-                    # 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:
-                elif style == 'htb':
-                    # main word becomes the lexeme, the suffix is based on the features
-                    intermediate_output[-1]['word'] = (s_lex if s_lex != s_word else s_word[:-1]) + '_'
-                    suf_feats = word['morph']['suffix_feats']
-                    suf = ud_suffix_to_htb_str.get(f"Gender={suf_feats.get('Gender', 'Fem,Masc')}|Number={suf_feats.get('Number', 'Sing')}|Person={suf_feats.get('Person', '3')}", "_הוא")
-                    # for HTB, if the function is poss, then add a shel pointing to the next word
-                    if func == 'nmod:poss' and s_lex != 'של':
-                        intermediate_output.append(dict(word='_של_', lex='של', pos='ADP', dep=len(intermediate_output) + 2, 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())))
-                if entry_to_assign_suf_dep:
-                    entry_to_assign_suf_dep['dep'] = len(intermediate_output)
-                    entry_to_assign_suf_dep['absolute_dep'] = True
-                    
-            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 + 1):
-                # compute the actual dependency location
-                dep = output['dep'] if output.get('absolute_dep', False) else idx_to_key[output['dep']]
-                func = normalize_dep_rel(output['func'], style)
-                # 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),
-                    func,
-                    '_', '_'
-                ]))
-    return final_output
-                
-def normalize_dep_rel(dep, style: Literal['htb', 'iahlt']):
-    if style == 'iahlt':
-        if dep == 'compound:smixut': return 'compound'
-        if dep == 'nsubj:cop': return 'nsubj'
-        if dep == 'mark:q': return 'mark'
-        if dep == 'case:gen' or dep == 'case:acc': return 'case'
-    return dep
-    
-    
-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'] not 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 = False
-        # 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
-    

BertForMorphTagging.py

@@ -1,212 +0,0 @@
-from collections import OrderedDict
-from operator import itemgetter
-from transformers.utils import ModelOutput
-import torch
-from torch import nn
-from typing import Dict, List, Tuple, Optional
-from dataclasses import dataclass
-from transformers import BertPreTrainedModel, BertModel, BertTokenizerFast
-
-ALL_POS = ['DET', 'NOUN', 'VERB', 'CCONJ', 'ADP', 'PRON', 'PUNCT', 'ADJ', 'ADV', 'SCONJ', 'NUM', 'PROPN', 'AUX', 'X', 'INTJ', 'SYM']
-ALL_PREFIX_POS = ['SCONJ', 'DET', 'ADV', 'CCONJ', 'ADP', 'NUM']
-ALL_SUFFIX_POS = ['none', 'ADP_PRON', 'PRON']
-ALL_FEATURES = [
-    ('Gender', ['none', 'Masc', 'Fem', 'Fem,Masc']),
-    ('Number', ['none', 'Sing', 'Plur', 'Plur,Sing', 'Dual', 'Dual,Plur']),
-    ('Person', ['none', '1', '2', '3', '1,2,3']),
-    ('Tense', ['none', 'Past', 'Fut', 'Pres', 'Imp'])
-]
-
-@dataclass
-class MorphLogitsOutput(ModelOutput):
-    prefix_logits: torch.FloatTensor = None
-    pos_logits: torch.FloatTensor = None
-    features_logits: List[torch.FloatTensor] = None
-    suffix_logits: torch.FloatTensor = None
-    suffix_features_logits: List[torch.FloatTensor] = None
-
-    def detach(self):
-        return MorphLogitsOutput(self.prefix_logits.detach(), self.pos_logits.detach(), [logits.deatch() for logits in self.features_logits], self.suffix_logits.detach(), [logits.deatch() for logits in self.suffix_features_logits])
-
-
-@dataclass
-class MorphTaggingOutput(ModelOutput):
-    loss: Optional[torch.FloatTensor] = None
-    logits: Optional[MorphLogitsOutput] = None
-    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
-    attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-@dataclass
-class MorphLabels(ModelOutput):
-    prefix_labels: Optional[torch.FloatTensor] = None
-    pos_labels: Optional[torch.FloatTensor] = None
-    features_labels: Optional[List[torch.FloatTensor]] = None
-    suffix_labels: Optional[torch.FloatTensor] = None
-    suffix_features_labels: Optional[List[torch.FloatTensor]] = None
-
-    def detach(self):
-        return MorphLabels(self.prefix_labels.detach(), self.pos_labels.detach(), [labels.detach() for labels in self.features_labels], self.suffix_labels.detach(), [labels.detach() for labels in self.suffix_features_labels])
-    
-    def to(self, device):
-        return MorphLabels(self.prefix_labels.to(device), self.pos_labels.to(device), [feat.to(device) for feat in self.features_labels], self.suffix_labels.to(device), [feat.to(device) for feat in self.suffix_features_labels])
-
-class BertMorphTaggingHead(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-
-        self.num_prefix_classes = len(ALL_PREFIX_POS)
-        self.num_pos_classes = len(ALL_POS)
-        self.num_suffix_classes = len(ALL_SUFFIX_POS)
-        self.num_features_classes = list(map(len, map(itemgetter(1), ALL_FEATURES)))
-        # we need a classifier for prefix cls and POS cls
-        # the prefix will use BCEWithLogits for multiple labels cls
-        self.prefix_cls = nn.Linear(config.hidden_size, self.num_prefix_classes)
-        # and pos + feats will use good old cross entropy for single label
-        self.pos_cls = nn.Linear(config.hidden_size, self.num_pos_classes)
-        self.features_cls = nn.ModuleList([nn.Linear(config.hidden_size, len(features)) for _, features in ALL_FEATURES])
-        # and suffix + feats will also be cross entropy
-        self.suffix_cls = nn.Linear(config.hidden_size, self.num_suffix_classes)
-        self.suffix_features_cls = nn.ModuleList([nn.Linear(config.hidden_size, len(features)) for _, features in ALL_FEATURES])
-
-    def forward(
-            self, 
-            hidden_states: torch.Tensor,
-            labels: Optional[MorphLabels] = None):
-        # run each of the classifiers on the transformed output
-        prefix_logits = self.prefix_cls(hidden_states)
-        pos_logits = self.pos_cls(hidden_states)
-        suffix_logits = self.suffix_cls(hidden_states)
-        features_logits = [cls(hidden_states) for cls in self.features_cls]
-        suffix_features_logits = [cls(hidden_states) for cls in self.suffix_features_cls]
-
-        loss = None
-        if labels is not None:
-            # step 1: prefix labels loss
-            loss_fct = nn.BCEWithLogitsLoss(weight=(labels.prefix_labels != -100).float())
-            loss = loss_fct(prefix_logits, labels.prefix_labels)
-            # step 2: pos labels loss
-            loss_fct = nn.CrossEntropyLoss()
-            loss += loss_fct(pos_logits.view(-1, self.num_pos_classes), labels.pos_labels.view(-1))
-            # step 2b: features
-            for feat_logits,feat_labels,num_features in zip(features_logits, labels.features_labels, self.num_features_classes):
-                loss += loss_fct(feat_logits.view(-1, num_features), feat_labels.view(-1))
-            # step 3: suffix logits loss
-            loss += loss_fct(suffix_logits.view(-1, self.num_suffix_classes), labels.suffix_labels.view(-1))
-            # step 3b: suffix features
-            for feat_logits,feat_labels,num_features in zip(suffix_features_logits, labels.suffix_features_labels, self.num_features_classes):
-                loss += loss_fct(feat_logits.view(-1, num_features), feat_labels.view(-1))
-
-        return loss, MorphLogitsOutput(prefix_logits, pos_logits, features_logits, suffix_logits, suffix_features_logits)
-
-class BertForMorphTagging(BertPreTrainedModel):
-
-    def __init__(self, config):
-        super().__init__(config)
-
-        self.bert = BertModel(config, add_pooling_layer=False)
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
-        self.morph = BertMorphTaggingHead(config)
-        
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def forward(
-        self,
-        input_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        token_type_ids: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-        labels: Optional[MorphLabels] = None,
-        head_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ):
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        bert_outputs = self.bert(
-            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_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-        hidden_states = bert_outputs[0]
-        hidden_states = self.dropout(hidden_states)
-
-        loss, logits = self.morph(hidden_states, labels)
-        
-        if not return_dict:
-            return (loss,logits) + bert_outputs[2:]
-        
-        return MorphTaggingOutput(
-            loss=loss,
-            logits=logits,
-            hidden_states=bert_outputs.hidden_states,
-            attentions=bert_outputs.attentions,
-        )
-
-    def predict(self, sentences: List[str], tokenizer: BertTokenizerFast, padding='longest'):
-        # tokenize the inputs and convert them to relevant device
-        inputs = tokenizer(sentences, padding=padding, truncation=True, return_tensors='pt')
-        inputs = {k:v.to(self.device) for k,v in inputs.items()}
-        # calculate the logits
-        logits = self.forward(**inputs, return_dict=True).logits
-        return parse_logits(inputs['input_ids'].tolist(), sentences, tokenizer, logits)
-    
-def parse_logits(input_ids: List[List[int]], sentences: List[str], tokenizer: BertTokenizerFast, logits: MorphLogitsOutput):
-    prefix_logits, pos_logits, feats_logits, suffix_logits, suffix_feats_logits = \
-                logits.prefix_logits, logits.pos_logits, logits.features_logits, logits.suffix_logits, logits.suffix_features_logits
-
-    prefix_predictions = (prefix_logits > 0.5).int().tolist() # Threshold at 0.5 for multi-label classification
-    pos_predictions = pos_logits.argmax(axis=-1).tolist()
-    suffix_predictions = suffix_logits.argmax(axis=-1).tolist()
-    feats_predictions = [logits.argmax(axis=-1).tolist() for logits in feats_logits]
-    suffix_feats_predictions = [logits.argmax(axis=-1).tolist() for logits in suffix_feats_logits]
-
-    # create the return dictionary 
-    # for each sentence, return a dict object with the following files { text, tokens }
-    # Where tokens is a list of dicts, where each dict is: 
-    #       { pos: str, feats: dict, prefixes: List[str], suffix: str | bool, suffix_feats: dict | None}
-    special_toks = tokenizer.all_special_tokens
-    ret = []
-    for sent_idx,sentence in enumerate(sentences):
-        input_id_strs = tokenizer.convert_ids_to_tokens(input_ids[sent_idx])
-        # iterate through each token in the sentence, ignoring special tokens
-        tokens = []
-        for token_idx,token_str in enumerate(input_id_strs):
-            if token_str in special_toks: continue
-            if token_str.startswith('##'):
-                tokens[-1]['token'] += token_str[2:]
-                continue
-            tokens.append(dict(
-                token=token_str,
-                pos=ALL_POS[pos_predictions[sent_idx][token_idx]],
-                feats=get_features_dict_from_predictions(feats_predictions, (sent_idx, token_idx)),
-                prefixes=[ALL_PREFIX_POS[idx] for idx,i in enumerate(prefix_predictions[sent_idx][token_idx]) if i > 0],
-                suffix=get_suffix_or_false(ALL_SUFFIX_POS[suffix_predictions[sent_idx][token_idx]]),
-            ))
-            if tokens[-1]['suffix']:
-                tokens[-1]['suffix_feats'] = get_features_dict_from_predictions(suffix_feats_predictions, (sent_idx, token_idx))
-        ret.append(dict(text=sentence, tokens=tokens))
-    return ret
-    
-def get_suffix_or_false(suffix):
-    return False if suffix == 'none' else suffix
-
-def get_features_dict_from_predictions(predictions, idx):
-    ret = {}
-    for (feat_idx, (feat_name, feat_values)) in enumerate(ALL_FEATURES):
-        val = feat_values[predictions[feat_idx][idx[0]][idx[1]]]
-        if val != 'none':
-            ret[feat_name] = val
-    return ret
-
-

BertForPrefixMarking.py

@@ -1,248 +0,0 @@
-from transformers.utils import ModelOutput
-import torch
-from torch import nn
-from typing import Dict, List, Tuple, Optional
-from dataclasses import dataclass
-from transformers import BertPreTrainedModel, BertModel, BertTokenizerFast
-
-# define the classes, and the possible prefixes for each class
-POSSIBLE_PREFIX_CLASSES =  [ ['לכש', 'כש', 'מש', 'בש', 'לש'], ['מ'], ['ש'], ['ה'], ['ו'], ['כ'], ['ל'], ['ב'] ]
-# map each individual prefix to it's class number
-PREFIXES_TO_CLASS = {w:i for i,l in enumerate(POSSIBLE_PREFIX_CLASSES) for w in l}
-# keep a list of all the prefixes, sorted by length, so that we can decompose
-# a given prefixes and figure out the classes
-ALL_PREFIX_ITEMS = list(sorted(PREFIXES_TO_CLASS.keys(), key=len, reverse=True))
-TOTAL_POSSIBLE_PREFIX_CLASSES = len(POSSIBLE_PREFIX_CLASSES)    
-
-def get_prefixes_from_str(s, greedy=False):
-    # keep trimming prefixes from the string
-    while len(s) > 0 and s[0] in PREFIXES_TO_CLASS:
-        # find the longest string to trim
-        next_pre = next((pre for pre in ALL_PREFIX_ITEMS if s.startswith(pre)), None)
-        if next_pre is None:
-            return
-        yield next_pre
-        # if the chosen prefix is more than one letter, there is always an option that the 
-        # prefix is actually just the first letter of the prefix - so offer that up as a valid prefix
-        # as well. We will still jump to the length of the longer one, since if the next two/three
-        # letters are a prefix, they have to be the longest one
-        if not greedy and len(next_pre) > 1:
-            yield next_pre[0]
-        s = s[len(next_pre):]
-
-def get_prefix_classes_from_str(s, greedy=False):
-    for pre in get_prefixes_from_str(s, greedy):
-        yield PREFIXES_TO_CLASS[pre]
-
-@dataclass
-class PrefixesClassifiersOutput(ModelOutput):
-    loss: Optional[torch.FloatTensor] = None
-    logits: Optional[torch.FloatTensor] = None
-    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
-    attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-class BertPrefixMarkingHead(nn.Module):
-    def __init__(self, config) -> None:
-        super().__init__()
-        self.config = config
-
-        # an embedding table containing an embedding for each prefix class + 1 for NONE
-        # we will concatenate either the embedding/NONE for each class - and we want the concatenate
-        # size to be the hidden_size
-        prefix_class_embed = config.hidden_size // TOTAL_POSSIBLE_PREFIX_CLASSES
-        self.prefix_class_embeddings = nn.Embedding(TOTAL_POSSIBLE_PREFIX_CLASSES + 1, prefix_class_embed)
-        
-        # one layer for transformation, apply an activation, then another N classifiers for each prefix class
-        self.transform = nn.Linear(config.hidden_size + prefix_class_embed * TOTAL_POSSIBLE_PREFIX_CLASSES, config.hidden_size)
-        self.activation = nn.Tanh()
-        self.classifiers = nn.ModuleList([nn.Linear(config.hidden_size, 2) for _ in range(TOTAL_POSSIBLE_PREFIX_CLASSES)])
-
-    def forward(
-            self,
-            hidden_states: torch.Tensor,
-            prefix_class_id_options: torch.Tensor,
-            labels: Optional[torch.Tensor] = None) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
-        
-        # encode the prefix_class_id_options
-        # If input_ids is batch x seq_len
-        # Then sequence_output is batch x seq_len x hidden_dim
-        # So prefix_class_id_options is batch x seq_len x TOTAL_POSSIBLE_PREFIX_CLASSES
-        # Looking up the embeddings should give us batch x seq_len x TOTAL_POSSIBLE_PREFIX_CLASSES x hidden_dim / N
-        possible_class_embed = self.prefix_class_embeddings(prefix_class_id_options)
-        # then flatten the final dimension - now we have batch x seq_len x hidden_dim_2
-        possible_class_embed = possible_class_embed.reshape(possible_class_embed.shape[:-2] + (-1,))
-
-        # concatenate the new class embed into the sequence output before the transform
-        pre_transform_output = torch.cat((hidden_states, possible_class_embed), dim=-1) # batch x seq_len x (hidden_dim + hidden_dim_2)
-        pre_logits_output = self.activation(self.transform(pre_transform_output))# batch x seq_len x hidden_dim
-
-        # run each of the classifiers on the transformed output
-        logits = torch.cat([cls(pre_logits_output).unsqueeze(-2) for cls in self.classifiers], dim=-2)
-
-        loss = None
-        if labels is not None:
-            loss_fct = nn.CrossEntropyLoss()
-            loss = loss_fct(logits.view(-1, 2), labels.view(-1))
-        
-        return (loss, logits)
-        
-
-
-class BertForPrefixMarking(BertPreTrainedModel):
-
-    def __init__(self, config):
-        super().__init__(config)
-
-        self.bert = BertModel(config, add_pooling_layer=False)
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
-        self.prefix = BertPrefixMarkingHead(config)
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def forward(
-        self,
-        input_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        token_type_ids: Optional[torch.Tensor] = None,
-        prefix_class_id_options: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-        labels: Optional[torch.Tensor] = None,
-        head_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ):
-        r"""
-        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
-            Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
-        """
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        bert_outputs = self.bert(
-            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_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-        hidden_states = bert_outputs[0]
-        hidden_states = self.dropout(hidden_states)
-
-        loss, logits = self.prefix.forward(hidden_states, prefix_class_id_options, labels)
-        if not return_dict:
-            return (loss,logits,) + bert_outputs[2:]
-
-        return PrefixesClassifiersOutput(
-            loss=loss,
-            logits=logits,
-            hidden_states=bert_outputs.hidden_states,
-            attentions=bert_outputs.attentions,
-        )
-    
-    def predict(self, sentences: List[str], tokenizer: BertTokenizerFast, padding='longest'):
-        # step 1: encode the sentences through using the tokenizer, and get the input tensors + prefix id tensors
-        inputs = encode_sentences_for_bert_for_prefix_marking(tokenizer, sentences, padding)
-        inputs.pop('offset_mapping')
-        inputs = {k:v.to(self.device) for k,v in inputs.items()}
-
-        # run through bert
-        logits = self.forward(**inputs, return_dict=True).logits
-        return parse_logits(inputs['input_ids'].tolist(), sentences, tokenizer, logits)
-
-def parse_logits(input_ids: List[List[int]], sentences: List[str], tokenizer: BertTokenizerFast, logits: torch.FloatTensor):
-    # extract the predictions by argmaxing the final dimension (batch x sequence x prefixes x prediction)
-    logit_preds = torch.argmax(logits, axis=3).tolist()
-
-    ret = []
-
-    for sent_idx,sent_ids in enumerate(input_ids):
-        tokens = tokenizer.convert_ids_to_tokens(sent_ids)
-        ret.append([])
-        for tok_idx,token in enumerate(tokens):
-            # If we've reached the pad token, then we are at the end
-            if token == tokenizer.pad_token: continue
-            if token.startswith('##'): continue
-
-            # combine the next tokens in? only if it's a breakup
-            next_tok_idx = tok_idx + 1
-            while next_tok_idx < len(tokens) and tokens[next_tok_idx].startswith('##'):
-                token += tokens[next_tok_idx][2:]
-                next_tok_idx += 1
-
-            prefix_len = get_predicted_prefix_len_from_logits(token, logit_preds[sent_idx][tok_idx])
-        
-            if not prefix_len:
-                ret[-1].append([token])
-            else:
-                ret[-1].append([token[:prefix_len], token[prefix_len:]])
-    return ret
-
-def encode_sentences_for_bert_for_prefix_marking(tokenizer: BertTokenizerFast, sentences: List[str], padding='longest', truncation=True):
-    inputs = tokenizer(sentences, padding=padding, truncation=truncation, return_offsets_mapping=True, return_tensors='pt')
-    # create our prefix_id_options array which will be like the input ids shape but with an addtional
-    # dimension containing for each prefix whether it can be for that word
-    prefix_id_options = torch.full(inputs['input_ids'].shape + (TOTAL_POSSIBLE_PREFIX_CLASSES,), TOTAL_POSSIBLE_PREFIX_CLASSES, dtype=torch.long)
-
-    # go through each token, and fill in the vector accordingly
-    for sent_idx, sent_ids in enumerate(inputs['input_ids']):
-        tokens = tokenizer.convert_ids_to_tokens(sent_ids)
-        for tok_idx, token in enumerate(tokens):
-            # if the first letter isn't a valid prefix letter, nothing to talk about
-            if len(token) < 2 or not token[0] in PREFIXES_TO_CLASS: continue
-
-            # combine the next tokens in? only if it's a breakup
-            next_tok_idx = tok_idx + 1
-            while next_tok_idx < len(tokens) and tokens[next_tok_idx].startswith('##'):
-                token += tokens[next_tok_idx][2:]
-                next_tok_idx += 1
-
-            # find all the possible prefixes - and mark them as 0 (and in the possible mark it as it's value for embed lookup)
-            for pre_class in get_prefix_classes_from_str(token):
-                prefix_id_options[sent_idx, tok_idx, pre_class] = pre_class
-        
-    inputs['prefix_class_id_options'] = prefix_id_options
-    return inputs
-
-def get_predicted_prefix_len_from_logits(token, token_logits):
-    # Go through each possible prefix, and check if the prefix is yes - and if
-    # so increase the counter of the matched length, otherwise break out. That will solve cases
-    # of predicting prefix combinations that don't exist on the word.
-    # For example, if we have the word ושכשהלכתי and the model predict ו & כש, then we will only
-    # take the vuv because in order to get the כש we need the ש as well.
-    # Two extra items:
-    # 1] Don't allow the same prefix multiple times
-    # 2] Always check that the word starts with that prefix - otherwise it's bad 
-    #    (except for the case of multi-letter prefix, where we force the next to be last)
-    cur_len, skip_next, last_check, seen_prefixes = 0, False, False, set()
-    for prefix in get_prefixes_from_str(token):
-        # Are we skipping this prefix? This will be the case where we matched כש, don't allow ש
-        if skip_next:
-            skip_next = False
-            continue
-        # check for duplicate prefixes, we don't allow two of the same prefix
-        # if it predicted two of the same, then we will break out
-        if prefix in seen_prefixes: break
-        seen_prefixes.add(prefix)
-
-        # check if we predicted this prefix
-        if token_logits[PREFIXES_TO_CLASS[prefix]]:
-            cur_len += len(prefix)
-            if last_check: break
-            skip_next = len(prefix) > 1
-        # Otherwise, we predicted no. If we didn't, then this is the end of the prefix
-        # and time to break out. *Except* if it's a multi letter prefix, then we allow
-        # just the next letter - e.g., if כש doesn't match, then we allow כ, but then we know
-        # the word continues with a ש, and if it's not כש, then it's not כ-ש- (invalid)
-        elif len(prefix) > 1:
-            last_check = True
-        else:
-            break
-
-    return cur_len

BertForSyntaxParsing.py

@@ -1,312 +0,0 @@
-import math
-from transformers.utils import ModelOutput
-import torch
-from torch import nn
-from typing import Dict, List, Tuple, Optional, Union
-from dataclasses import dataclass
-from transformers import BertPreTrainedModel, BertModel, BertTokenizerFast
-
-ALL_FUNCTION_LABELS = ["nsubj", "nsubj:cop", "punct", "mark", "mark:q", "case", "case:gen", "case:acc", "fixed", "obl", "det", "amod", "acl:relcl", "nmod", "cc", "conj", "root", "compound:smixut", "cop", "compound:affix", "advmod", "nummod", "appos", "nsubj:pass", "nmod:poss", "xcomp", "obj", "aux", "parataxis", "advcl", "ccomp", "csubj", "acl", "obl:tmod", "csubj:pass", "dep", "dislocated", "nmod:tmod", "nmod:npmod", "flat", "obl:npmod", "goeswith", "reparandum", "orphan", "list", "discourse", "iobj", "vocative", "expl", "flat:name"]
-
-@dataclass
-class SyntaxLogitsOutput(ModelOutput):
-    dependency_logits: torch.FloatTensor = None
-    function_logits: torch.FloatTensor = None
-    dependency_head_indices: torch.LongTensor = None
-
-    def detach(self):
-        return SyntaxTaggingOutput(self.dependency_logits.detach(), self.function_logits.detach(), self.dependency_head_indices.detach())
-
-@dataclass
-class SyntaxTaggingOutput(ModelOutput):
-    loss: Optional[torch.FloatTensor] = None
-    logits: Optional[SyntaxLogitsOutput] = None
-    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
-    attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-@dataclass
-class SyntaxLabels(ModelOutput):
-    dependency_labels: Optional[torch.LongTensor] = None
-    function_labels: Optional[torch.LongTensor] = None
-
-    def detach(self):
-        return SyntaxLabels(self.dependency_labels.detach(), self.function_labels.detach())
-    
-    def to(self, device):
-        return SyntaxLabels(self.dependency_labels.to(device), self.function_labels.to(device))
-
-class BertSyntaxParsingHead(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-
-        # the attention query & key values
-        self.head_size = config.syntax_head_size# int(config.hidden_size / config.num_attention_heads * 2)
-        self.query = nn.Linear(config.hidden_size, self.head_size)
-        self.key = nn.Linear(config.hidden_size, self.head_size)
-        # the function classifier gets two encoding values and predicts the labels
-        self.num_function_classes = len(ALL_FUNCTION_LABELS)
-        self.cls = nn.Linear(config.hidden_size * 2, self.num_function_classes)
-
-    def forward(
-            self, 
-            hidden_states: torch.Tensor, 
-            extended_attention_mask: Optional[torch.Tensor],
-            labels: Optional[SyntaxLabels] = None,
-            compute_mst: bool = False) -> Tuple[torch.Tensor, SyntaxLogitsOutput]:
-        
-        # Take the dot product between "query" and "key" to get the raw attention scores.
-        query_layer = self.query(hidden_states)
-        key_layer = self.key(hidden_states)
-        attention_scores = torch.bmm(query_layer, key_layer.transpose(-1, -2)) / math.sqrt(self.head_size)
-
-        # add in the attention mask
-        if extended_attention_mask is not None:
-            if extended_attention_mask.ndim == 4:
-                extended_attention_mask = extended_attention_mask.squeeze(1)
-            attention_scores += extended_attention_mask# batch x seq x seq
-
-        # At this point take the hidden_state of the word and of the dependency word, and predict the function
-        # If labels are provided, use the labels.
-        if self.training and labels is not None:
-            # Note that the labels can have -100, so just set those to zero with a max
-            dep_indices = labels.dependency_labels.clamp_min(0)
-        # Otherwise - check if he wants the MST or just the argmax
-        elif compute_mst:
-            dep_indices = compute_mst_tree(attention_scores, extended_attention_mask)
-        else: 
-            dep_indices = torch.argmax(attention_scores, dim=-1)
-        
-        # After we retrieved the dependency indicies, create a tensor of teh batch indices, and and retrieve the vectors of the heads to calculate the function
-        batch_indices = torch.arange(dep_indices.size(0)).view(-1, 1).expand(-1, dep_indices.size(1)).to(dep_indices.device)
-        dep_vectors = hidden_states[batch_indices, dep_indices, :] # batch x seq x dim
-
-        # concatenate that with the last hidden states, and send to the classifier output
-        cls_inputs = torch.cat((hidden_states, dep_vectors), dim=-1)
-        function_logits = self.cls(cls_inputs)
-
-        loss = None
-        if labels is not None:
-            loss_fct = nn.CrossEntropyLoss()
-            # step 1: dependency scores loss - this is applied to the attention scores
-            loss = loss_fct(attention_scores.view(-1, hidden_states.size(-2)), labels.dependency_labels.view(-1))
-            # step 2: function loss
-            loss += loss_fct(function_logits.view(-1, self.num_function_classes), labels.function_labels.view(-1))
-        
-        return (loss, SyntaxLogitsOutput(attention_scores, function_logits, dep_indices))
-
-
-class BertForSyntaxParsing(BertPreTrainedModel):
-
-    def __init__(self, config):
-        super().__init__(config)
-
-        self.bert = BertModel(config, add_pooling_layer=False)
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
-        self.syntax = BertSyntaxParsingHead(config)
-        
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def forward(
-        self,
-        input_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        token_type_ids: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-        labels: Optional[SyntaxLabels] = None,
-        head_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        compute_syntax_mst: Optional[bool] = None,
-    ):
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        bert_outputs = self.bert(
-            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_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-        
-        extended_attention_mask = None
-        if attention_mask is not None:
-            extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_ids.size())
-        # apply the syntax head
-        loss, logits = self.syntax(self.dropout(bert_outputs[0]), extended_attention_mask, labels, compute_syntax_mst)
-        
-        if not return_dict:
-            return (loss,(logits.dependency_logits, logits.function_logits)) + bert_outputs[2:]
-        
-        return SyntaxTaggingOutput(
-            loss=loss,
-            logits=logits,
-            hidden_states=bert_outputs.hidden_states,
-            attentions=bert_outputs.attentions,
-        )
-
-    def predict(self, sentences: Union[str, List[str]], tokenizer: BertTokenizerFast, compute_mst=True):
-        if isinstance(sentences, str):
-            sentences = [sentences]
-            
-        # predict the logits for the sentence
-        inputs = tokenizer(sentences, padding='longest', truncation=True, return_tensors='pt')
-        inputs = {k:v.to(self.device) for k,v in inputs.items()}
-        logits = self.forward(**inputs, return_dict=True, compute_syntax_mst=compute_mst).logits
-        return parse_logits(inputs['input_ids'].tolist(), sentences, tokenizer, logits)
-
-def parse_logits(input_ids: List[List[int]], sentences: List[str], tokenizer: BertTokenizerFast, logits: SyntaxLogitsOutput):
-    outputs = []        
-
-    special_toks = tokenizer.all_special_tokens
-    for i in range(len(sentences)):
-        deps = logits.dependency_head_indices[i].tolist()
-        funcs = logits.function_logits.argmax(-1)[i].tolist()
-        toks = [tok for tok in tokenizer.convert_ids_to_tokens(input_ids[i]) if tok not in special_toks]
-        
-        # first, go through the tokens and create a mapping between each dependency index and the index without wordpieces
-        # wordpieces. At the same time, append the wordpieces in
-        idx_mapping = {-1:-1} # default root 
-        real_idx = -1
-        for i in range(len(toks)):
-            if not toks[i].startswith('##'):
-                real_idx += 1
-            idx_mapping[i] = real_idx
-            
-        # build our tree, keeping tracking of the root idx
-        tree = []
-        root_idx = 0
-        for i in range(len(toks)):
-            if toks[i].startswith('##'):
-                tree[-1]['word'] += toks[i][2:]
-                continue 
-            
-            dep_idx = deps[i + 1] - 1 # increase 1 for cls, decrease 1 for cls
-            if dep_idx == len(toks): dep_idx = i - 1 # if he predicts sep, then just point to the previous word
-            
-            dep_head = 'root' if dep_idx == -1 else toks[dep_idx]
-            dep_func = ALL_FUNCTION_LABELS[funcs[i + 1]]
-
-            if dep_head == 'root': root_idx = len(tree)
-            tree.append(dict(word=toks[i], dep_head_idx=idx_mapping[dep_idx], dep_func=dep_func))
-        # append the head word
-        for d in tree:
-            d['dep_head'] = tree[d['dep_head_idx']]['word']
-        
-        outputs.append(dict(tree=tree, root_idx=root_idx))
-    return outputs
-
-
-def compute_mst_tree(attention_scores: torch.Tensor, extended_attention_mask: torch.LongTensor):
-    # attention scores should be 3 dimensions - batch x seq x seq (if it is 2 - just unsqueeze)
-    if attention_scores.ndim == 2: attention_scores = attention_scores.unsqueeze(0)
-    if attention_scores.ndim != 3 or attention_scores.shape[1] != attention_scores.shape[2]:
-        raise ValueError(f'Expected attention scores to be of shape batch x seq x seq, instead got {attention_scores.shape}')
-    
-    batch_size, seq_len, _ = attention_scores.shape
-    # start by softmaxing so the scores are comparable
-    attention_scores = attention_scores.softmax(dim=-1)
-    
-    batch_indices = torch.arange(batch_size, device=attention_scores.device)
-    seq_indices = torch.arange(seq_len, device=attention_scores.device)
-
-    seq_lens = torch.full((batch_size,), seq_len)
-    
-    if extended_attention_mask is not None:
-        seq_lens = torch.argmax((extended_attention_mask != 0).int(), dim=2).squeeze(1)     
-        # zero out any padding
-        attention_scores[extended_attention_mask.squeeze(1) != 0] = 0
-
-    # set the values for the CLS and sep to all by very low, so they never get chosen as a replacement arc
-    attention_scores[:, 0, :] = 0
-    attention_scores[batch_indices, seq_lens - 1, :] = 0
-    attention_scores[batch_indices, :, seq_lens - 1] = 0 # can never predict sep
-    # set the values for each token pointing to itself be 0
-    attention_scores[:, seq_indices, seq_indices] = 0
-    
-    # find the root, and make him super high so we never have a conflict
-    root_cands = torch.argsort(attention_scores[:, :, 0], dim=-1)
-    attention_scores[batch_indices.unsqueeze(1), root_cands, 0] = 0
-    attention_scores[batch_indices, root_cands[:, -1], 0] = 1.0
-
-    # we start by getting the argmax for each score, and then computing the cycles and contracting them
-    sorted_indices = torch.argsort(attention_scores, dim=-1, descending=True)
-    indices = sorted_indices[:, :, 0].clone() # take the argmax
-    
-    attention_scores = attention_scores.tolist()
-    seq_lens = seq_lens.tolist()
-    sorted_indices = [[sub_l[:slen] for sub_l in l[:slen]] for l,slen in zip(sorted_indices.tolist(), seq_lens)]
-
-
-    # go through each batch item and make sure our tree works
-    for batch_idx in range(batch_size):
-        # We have one root - detect the cycles and contract them. A cycle can never contain the root so really
-        # for every cycle, we look at all the nodes, and find the highest arc out of the cycle for any values. Replace that and tada
-        has_cycle, cycle_nodes = detect_cycle(indices[batch_idx], seq_lens[batch_idx])
-        contracted_arcs = set()
-        while has_cycle:
-            base_idx, head_idx = choose_contracting_arc(indices[batch_idx], sorted_indices[batch_idx], cycle_nodes, contracted_arcs, seq_lens[batch_idx], attention_scores[batch_idx])
-            indices[batch_idx, base_idx] = head_idx
-            contracted_arcs.add(base_idx)
-            # find the next cycle
-            has_cycle, cycle_nodes = detect_cycle(indices[batch_idx], seq_lens[batch_idx])
-
-    return indices
-
-def detect_cycle(indices: torch.LongTensor, seq_len: int):
-    # Simple cycle detection algorithm
-    # Returns a boolean indicating if a cycle is detected and the nodes involved in the cycle
-    visited = set()
-    for node in range(1, seq_len - 1): # ignore the CLS/SEP tokens
-        if node in visited:
-            continue
-        current_path = set()
-        while node not in visited:
-            visited.add(node)
-            current_path.add(node)
-            node = indices[node].item()
-            if node == 0: break # roots never point to anything
-            if node in current_path:
-                return True, current_path  # Cycle detected
-    return False, None
-
-def choose_contracting_arc(indices: torch.LongTensor, sorted_indices: List[List[int]], cycle_nodes: set, contracted_arcs: set, seq_len: int, scores: List[List[float]]):
-    # Chooses the highest-scoring, non-cycling arc from a graph. Iterates through 'cycle_nodes' to find
-    # the best arc based on 'scores', avoiding cycles and zero node connections.
-    # For each node, we only look at the next highest scoring non-cycling arc 
-    best_base_idx, best_head_idx = -1, -1
-    score = 0
-    
-    # convert the indices to a list once, to avoid multiple conversions (saves a few seconds)
-    currents = indices.tolist()
-    for base_node in cycle_nodes:
-        if base_node in contracted_arcs: continue
-        # we don't want to take anything that has a higher score than the current value - we can end up in an endless loop
-        # Since the indices are sorted, as soon as we find our current item, we can move on to the next. 
-        current = currents[base_node]
-        found_current = False
-        
-        for head_node in sorted_indices[base_node]:
-            if head_node == current:
-                found_current = True
-                continue
-            if head_node in contracted_arcs: continue
-            if not found_current or head_node in cycle_nodes or head_node == 0: 
-                continue
-            
-            current_score = scores[base_node][head_node]
-            if current_score > score:
-                best_base_idx, best_head_idx, score = base_node, head_node, current_score
-            break
-    
-    if best_base_idx == -1:
-        raise ValueError('Stuck in endless loop trying to compute syntax mst. Please try again setting compute_syntax_mst=False')
-
-    return best_base_idx, best_head_idx

config.json

@@ -3,7 +3,7 @@
     "BertForJointParsing"
   ],
   "auto_map": {
-    "AutoModel": "BertForJointParsing.BertForJointParsing"
+    "AutoModel": "dicta-il/dictabert-joint--BertForJointParsing.BertForJointParsing"
   },
   "attention_probs_dropout_prob": 0.1,
   "classifier_dropout": null,