fix verb assigment for srl_frames in document_to_example()

ArneBinder · Mar 18, 2024 · c2d7188 · c2d7188
1 parent 1e54498
commit c2d7188
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 8 deletions.
diff --git a/dataset_builders/pie/conll2012_ontonotesv5/conll2012_ontonotesv5.py b/dataset_builders/pie/conll2012_ontonotesv5/conll2012_ontonotesv5.py
@@ -289,15 +289,19 @@ def document_to_example(
             span_start = min([span.start for span in srl_rel.arguments])
             span_end = max([span.end for span in srl_rel.arguments])
             if sent_start <= span_start and span_end <= sent_end:
+                verb = None
                 frames = ["O"] * sent_len
                 for arg, role in zip(srl_rel.arguments, srl_rel.roles):
                     frames[arg.start - sent_start] = "B-" + role
                     if arg.end - arg.start > 1:
                         frames[arg.start - sent_start + 1 : arg.end - sent_start] = [
                             "I-" + role
                         ] * (arg.end - arg.start - 1)
-                    if role == "V":
+                    # english_v4 and arabic_v4 contain some weird role names (in addition to "V") for the verb
+                    if role in ["V", "ARG1(V", "C-ARG0(V", "C-ARG1(V", "C-ARG2(V", "R-ARG1(V"]:
                         verb = document.tokens[arg.start]
+                if verb is None:
+                    raise ValueError(f"Verb not found for SRL relation: {srl_rel}")
                 srl_frames.append({"verb": verb, "frames": frames})
 
         coref_spans = []

diff --git a/tests/dataset_builders/pie/test_conll2012_ontonotesv5.py b/tests/dataset_builders/pie/test_conll2012_ontonotesv5.py
@@ -231,15 +231,9 @@ def test_generate_example_and_back_all(
         assert hf_ex_back["document_id"] == hf_ex["document_id"]
         for ex, ex_back in zip(hf_ex["sentences"], hf_ex_back["sentences"]):
             for key in ex.keys():
-                # 'coref_spans' must be sorted before compare
+                # 'coref_spans' must be sorted before comparison
                 if key == "coref_spans":
                     assert sorted(ex[key]) == sorted(ex_back[key])
-                # TODO: srl_frames are not equal, but they should be
-                elif key == "srl_frames":
-                    if not ex[key] == ex_back[key]:
-                        print(
-                            f"Error for example {hf_ex['document_id']}: {key}. Expected: {ex[key]}, Got: {ex_back[key]}"
-                        )
                 else:
                     assert ex[key] == ex_back[key]
         i += 1