Updated with fixed MST and batching from DictaBERT-Joint

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,17 @@ 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
+    special_toks.remove(tokenizer.unk_token)
     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 +190,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 +205,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 +214,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 +278,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