coref_labels now correctly ignore mention longer than max_span_size

Author: Aethor

tests/test_bertcoref.py

@@ -13,7 +13,7 @@ def bert_tokenizer() -> BertTokenizerFast:
 # we suppress the `function_scoped_fixture` health check since we want
 # to execute the `bert_tokenizer` fixture only once.
 @settings(deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture])
-@given(doc=coref_docs(min_size=5, max_size=10))
+@given(doc=coref_docs(min_size=5, max_size=10, max_span_size=4))
 def test_doc_is_reconstructed(
     doc: CoreferenceDocument, bert_tokenizer: BertTokenizerFast
 ):

tibert/bertcoref.py

@@ -94,21 +94,23 @@ def coref_labels(self, max_span_size: int) -> List[List[int]]:
 
         # spans in a coref chain : mark all antecedents
         for chain in self.coref_chains:
+            # mentions can be longer than max_span_size. We filter
+            # these mentions so that they do not appear in the labels.
+            chain = [m for m in chain if len(m.tokens) <= max_span_size]
+            if len(chain) == 0:
+                continue
             for mention in chain:
-                try:
-                    mention_idx = spans_idx[(mention.start_idx, mention.end_idx)]
-                    for other_mention in chain:
-                        if other_mention == mention:
-                            continue
-                        other_mention_idx = spans_idx[
-                            (other_mention.start_idx, other_mention.end_idx)
-                        ]
-                        labels[mention_idx][other_mention_idx] = 1
-                # ValueError happens if the mention does not exist in
-                # spans_idx. This is possible since the mention can be
-                # larger than max_span_size
-                except ValueError:
-                    continue
+                mention_idx = spans_idx[(mention.start_idx, mention.end_idx)]
+                for other_mention in chain:
+                    if other_mention == mention:
+                        continue
+                    key = (other_mention.start_idx, other_mention.end_idx)
+                    if not key in spans_idx:
+                        continue
+                    other_mention_idx = spans_idx[
+                        (other_mention.start_idx, other_mention.end_idx)
+                    ]
+                    labels[mention_idx][other_mention_idx] = 1
 
         # spans without preceding mentions : mark preceding mention to
         # be the null span