corefud.GuessSpan: add empty nodes that are causing gaps

martinpopel · martinpopel · commit 301b80808225 · 2023-02-04T01:36:25.000+01:00
diff --git a/udapi/block/corefud/guessspan.py b/udapi/block/corefud/guessspan.py
@@ -4,6 +4,30 @@ class GuessSpan(Block):
     """Block corefud.GuessSpan heuristically fills mention spans, while keeping mention.head"""
 
     def process_coref_mention(self, mention):
-        mention.words = mention.head.descendants(add_self=True)
-        # TODO add empty nodes that are causing gaps
+        mwords = mention.head.descendants(add_self=True)
         # TODO add heuristics from corefud.PrintMentions almost_forest=1
+
+        # Add empty nodes that are causing gaps.
+        # A node "within the span" whose enhanced parent is in the mentions
+        # must be added to the mention as well.
+        # "within the span" includes also empty nodes "on the boundary".
+        # However, don't add empty nodes which are in a gap cause by non-empty nodes.
+        to_add = []
+        min_ord = int(mwords[0].ord) if mwords[0].is_empty() else mwords[0].ord - 1
+        max_ord = int(mwords[-1].ord) + 1
+        root = mention.head.root
+        for empty in root.empty_nodes:
+            if empty in mwords:
+                continue
+            if empty.ord > max_ord:
+                break
+            if empty.ord > min_ord:
+                if any(enh['parent'] in mwords for enh in empty.deps):
+                    to_add.append(empty)
+                elif empty.ord > min_ord + 1 and empty.ord < max_ord - 1:
+                    prev_nonempty = root.descendants[int(empty.ord) - 1]
+                    next_nonempty = root.descendants[int(empty.ord)]
+                    if prev_nonempty in mwords and next_nonempty in mwords:
+                        to_add.append(empty)
+                    #else: empty.misc['Mark'] = f'not_in_treelet_of_{mention.entity.eid}'
+        mention.words = sorted(mwords + to_add)
