Enhanced dependency parsing

combo/utils/graph.py

@@ -18,10 +18,12 @@ def sdp_to_dag_deps(arc_scores, rel_scores, tree_tokens: List, root_idx=0, vocab
         index = heads.index(head)
         deprel = tree_tokens[i - 1]["deprel"]
         deprel = deprel.split('>')[-1]
-        # TODO is this necessary?
-        if len(heads) >= 2:
-            heads.pop(index)
-            rels.pop(index)
+        # TODO - Consider if there should be a condition,
+        # It doesn't seem to make any sense as DEPS should contain DEPREL
+        # (although sometimes with different/more detailed label)
+        # if len(heads) >= 2:
+        #     heads.pop(index)
+        #     rels.pop(index)
         deps = '|'.join(f'{h}:{r}' for h, r in zip(heads, rels))
         tree_tokens[i - 1]["deps"] = deps
         tree_tokens[i - 1]["deprel"] = deprel
@@ -32,6 +34,8 @@ def adjust_root_score_then_add_secondary_arcs(arc_scores, rel_scores, tree_heads
     if len(arc_scores) != tree_heads:
         arc_scores = arc_scores[:len(tree_heads)][:len(tree_heads)]
         rel_scores = rel_scores[:len(tree_heads)][:len(tree_heads)]
+    # Self-loops aren't allowed, mask with 0. This is an in-place operation.
+    np.fill_diagonal(arc_scores, 0)
     parse_preds = np.array(arc_scores) > 0
     parse_preds[:, 0] = False  # set heads to False
     rel_scores[:, :, root_idx] = -float('inf')