Added major speedup

BertForJointParsing.py

@@ -208,31 +208,32 @@ class BertForJointParsing(BertPreTrainedModel):
         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)
         
-        final_output = [dict(text=sentence, tokens=combine_token_wordpieces(ids, offsets, tokenizer)) for sentence, ids, offsets in zip(sentences, inputs['input_ids'], offset_mapping)]
+        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(inputs, sentences, tokenizer, output.syntax_logits)):
+            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(inputs, sentences, tokenizer, output.prefix_logits)):
+            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(inputs, sentences, tokenizer, output.lex_logits)):
+            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(inputs, sentences, tokenizer, output.morph_logits)):
+            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(inputs, sentences, tokenizer, output.ner_logits, self.config.id2label)):
+            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) 
@@ -267,31 +268,32 @@ 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: torch.Tensor, offset_mapping: torch.Tensor, tokenizer: BertTokenizerFast):
+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 [tokenizer.cls_token, tokenizer.sep_token, tokenizer.pad_token]: continue
+        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(inputs: Dict[str, torch.Tensor], sentences: List[str], tokenizer: BertTokenizerFast, logits: torch.Tensor, id2label: Dict[int, str]):
-    input_ids = inputs['input_ids']    
-
-    predictions = torch.argmax(logits, dim=-1)
+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)
-        for tok_idx in range(input_ids.shape[1]):
-            token_id = input_ids[batch_idx, tok_idx]
-            # ignore cls, sep, pad
-            if token_id in [tokenizer.cls_token_id, tokenizer.sep_token_id, tokenizer.pad_token_id]: continue
 
-            token = tokenizer._convert_id_to_token(token_id)
+        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
@@ -299,29 +301,29 @@ def ner_parse_logits(inputs: Dict[str, torch.Tensor], sentences: List[str], toke
             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].item()]))
+            ret.append((token, id2label[predictions[batch_idx][tok_idx]]))
 
     return batch_ret
 
-def lex_parse_logits(inputs: Dict[str, torch.Tensor], sentences: List[str], tokenizer: BertTokenizerFast, logits: torch.Tensor):
-    input_ids = inputs['input_ids']
+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]
+    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 = []
-        for tok_idx in range(input_ids.shape[1]):
-            token_id = input_ids[batch_idx, tok_idx]
-            # ignore cls, sep, pad
-            if token_id in [tokenizer.cls_token_id, tokenizer.sep_token_id, tokenizer.pad_token_id]: continue
+        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
 
-            token = tokenizer._convert_id_to_token(token_id)
             # 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])))
-            
+            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)
@@ -376,10 +378,15 @@ def convert_output_to_ud(output_sentences, style: Literal['htb', 'iahlt']):
         # 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]
-            
+            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:

BertForMorphTagging.py

@@ -159,42 +159,42 @@ class BertForMorphTagging(BertPreTrainedModel):
         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, sentences, tokenizer, logits)
+        return parse_logits(inputs['input_ids'].tolist(), sentences, tokenizer, logits)
     
-def parse_logits(inputs: Dict[str, torch.Tensor], sentences: List[str], tokenizer: BertTokenizerFast, logits: MorphLogitsOutput):
+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() # Threshold at 0.5 for multi-label classification
-    pos_predictions = pos_logits.argmax(axis=-1)
-    suffix_predictions = suffix_logits.argmax(axis=-1)
-    feats_predictions = [logits.argmax(axis=-1) for logits in feats_logits]
-    suffix_feats_predictions = [logits.argmax(axis=-1) for logits in suffix_feats_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_tokens = set([tokenizer.pad_token, tokenizer.cls_token, tokenizer.sep_token])
+    special_toks = tokenizer.all_special_tokens
     ret = []
     for sent_idx,sentence in enumerate(sentences):
-        input_id_strs = tokenizer.convert_ids_to_tokens(inputs['input_ids'][sent_idx])
+        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 not token_str in special_tokens: 
-                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))
+            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
     
@@ -204,7 +204,7 @@ def get_suffix_or_false(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]]
+        val = feat_values[predictions[feat_idx][idx[0]][idx[1]]]
         if val != 'none':
             ret[feat_name] = val
     return ret

BertForPrefixMarking.py

@@ -154,15 +154,15 @@ class BertForPrefixMarking(BertPreTrainedModel):
 
         # run through bert
         logits = self.forward(**inputs, return_dict=True).logits
-        return parse_logits(inputs, sentences, tokenizer, logits)
+        return parse_logits(inputs['input_ids'].tolist(), sentences, tokenizer, logits)
 
-def parse_logits(inputs: Dict[str, torch.Tensor], sentences: List[str], tokenizer: BertTokenizerFast, logits: torch.FloatTensor):
+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)
+    logit_preds = torch.argmax(logits, axis=3).tolist()
 
     ret = []
 
-    for sent_idx,sent_ids in enumerate(inputs['input_ids']):
+    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):
@@ -176,7 +176,7 @@ def parse_logits(inputs: Dict[str, torch.Tensor], sentences: List[str], tokenize
                 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])
+            prefix_len = get_predicted_prefix_len_from_logits(token, logit_preds[sent_idx][tok_idx])
         
             if not prefix_len:
                 ret[-1].append([token])
@@ -232,7 +232,7 @@ def get_predicted_prefix_len_from_logits(token, token_logits):
         seen_prefixes.add(prefix)
 
         # check if we predicted this prefix
-        if token_logits[PREFIXES_TO_CLASS[prefix]].item():
+        if token_logits[PREFIXES_TO_CLASS[prefix]]:
             cur_len += len(prefix)
             if last_check: break
             skip_next = len(prefix) > 1

BertForSyntaxParsing.py

@@ -73,7 +73,7 @@ class BertSyntaxParsingHead(nn.Module):
             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)
+            dep_indices = compute_mst_tree(attention_scores, extended_attention_mask)
         else: 
             dep_indices = torch.argmax(attention_scores, dim=-1)
         
@@ -160,14 +160,16 @@ class BertForSyntaxParsing(BertPreTrainedModel):
         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, sentences, tokenizer, logits)
+        return parse_logits(inputs['input_ids'].tolist(), sentences, tokenizer, logits)
 
-def parse_logits(inputs: Dict[str, torch.Tensor], sentences: List[str], tokenizer: BertTokenizerFast, logits: SyntaxLogitsOutput):
+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 = tokenizer.convert_ids_to_tokens(inputs['input_ids'][i])[1:-1] # ignore cls and sep
+        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
@@ -187,6 +189,8 @@ def parse_logits(inputs: Dict[str, torch.Tensor], sentences: List[str], tokenize
                 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]]
 
@@ -200,7 +204,7 @@ def parse_logits(inputs: Dict[str, torch.Tensor], sentences: List[str], tokenize
     return outputs
 
 
-def compute_mst_tree(attention_scores: torch.Tensor):
+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]:
@@ -209,40 +213,58 @@ def compute_mst_tree(attention_scores: torch.Tensor):
     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, :] = -10000
-    attention_scores[:, -1, :] = -10000
-    attention_scores[:, :, -1] = -10000 # can never predict sep
+    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)
-    batch_indices = torch.arange(batch_size, device=root_cands.device)
-    attention_scores[batch_indices.unsqueeze(1), root_cands, 0] = -10000
-    attention_scores[batch_indices, root_cands[:, -1], 0] = 10000
-    
+    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])
+        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, attention_scores[batch_idx])
+            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])
-            
+            has_cycle, cycle_nodes = detect_cycle(indices[batch_idx], seq_lens[batch_idx])
+
     return indices
 
-def detect_cycle(indices: torch.LongTensor):
+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, len(indices) - 1): # ignore the CLS/SEP tokens
+    for node in range(1, seq_len - 1): # ignore the CLS/SEP tokens
         if node in visited:
             continue
         current_path = set()
@@ -255,31 +277,36 @@ def detect_cycle(indices: torch.LongTensor):
                 return True, current_path  # Cycle detected
     return False, None
 
-def choose_contracting_arc(indices: torch.LongTensor, sorted_indices: torch.LongTensor, cycle_nodes: set, scores: torch.FloatTensor):
+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 = float('-inf')
+    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].tolist():
+        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].item()
+            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