From 1d73f8a42a93fc0ff6f8a2ca192f99eb09c41bea Mon Sep 17 00:00:00 2001
From: Arne Binder <arne.binder@dfki.de>
Date: Sun, 15 Sep 2024 22:04:03 +0200
Subject: [PATCH] add_negative_coref_relations does not return docs without
 relations

---
 .../document/processing/text_pair.py          |  7 +++----
 tests/document/processing/test_text_pair.py   | 20 +++----------------
 2 files changed, 6 insertions(+), 21 deletions(-)

diff --git a/src/pie_modules/document/processing/text_pair.py b/src/pie_modules/document/processing/text_pair.py
index 6b85e51a8..e052342cc 100644
--- a/src/pie_modules/document/processing/text_pair.py
+++ b/src/pie_modules/document/processing/text_pair.py
@@ -225,10 +225,9 @@ def add_negative_coref_relations(
     for rel in positive_rels:
         new_rels2new_docs[rel].binary_coref_relations.append(rel)
 
-    # TODO: implement down sampling
+    # TODO: implement downsampling
     for rel in negative_rels:
         new_rels2new_docs[rel].binary_coref_relations.append(rel)
 
-    # docs_with_rels = [doc for doc in new_docs if len(doc.binary_coref_relations) > 0]
-    # return docs_with_rels
-    return new_docs
+    docs_with_rels = [doc for doc in new_docs if len(doc.binary_coref_relations) > 0]
+    return docs_with_rels
diff --git a/tests/document/processing/test_text_pair.py b/tests/document/processing/test_text_pair.py
index cd016a117..e452219de 100644
--- a/tests/document/processing/test_text_pair.py
+++ b/tests/document/processing/test_text_pair.py
@@ -204,7 +204,7 @@ def positive_and_negative_documents(positive_documents):
 
 
 def test_construct_negative_documents(positive_and_negative_documents):
-    assert len(positive_and_negative_documents) == 16
+    assert len(positive_and_negative_documents) == 8
     TEXTS = [
         "Entity A works at B.",
         "And she founded C.",
@@ -229,7 +229,6 @@ def test_construct_negative_documents(positive_and_negative_documents):
     ]
 
     assert all_rels_and_scores == [
-        (("And she founded C.", "And she founded C."), []),
         (
             ("And she founded C.", "Bob loves his cat."),
             [(0.0, ("coref", (("PERSON", "she"), ("PERSON", "Bob"))))],
@@ -241,12 +240,10 @@ def test_construct_negative_documents(positive_and_negative_documents):
                 (0.0, ("coref", (("COMPANY", "C"), ("COMPANY", "B")))),
             ],
         ),
-        (("And she founded C.", "She sleeps a lot."), []),
         (
             ("Bob loves his cat.", "And she founded C."),
             [(0.0, ("coref", (("PERSON", "Bob"), ("PERSON", "she"))))],
         ),
-        (("Bob loves his cat.", "Bob loves his cat."), []),
         (
             ("Bob loves his cat.", "Entity A works at B."),
             [(0.0, ("coref", (("PERSON", "Bob"), ("PERSON", "Entity A"))))],
@@ -266,15 +263,10 @@ def test_construct_negative_documents(positive_and_negative_documents):
             ("Entity A works at B.", "Bob loves his cat."),
             [(0.0, ("coref", (("PERSON", "Entity A"), ("PERSON", "Bob"))))],
         ),
-        (("Entity A works at B.", "Entity A works at B."), []),
-        (("Entity A works at B.", "She sleeps a lot."), []),
-        (("She sleeps a lot.", "And she founded C."), []),
         (
             ("She sleeps a lot.", "Bob loves his cat."),
             [(1.0, ("coref", (("ANIMAL", "She"), ("ANIMAL", "his cat"))))],
         ),
-        (("She sleeps a lot.", "Entity A works at B."), []),
-        (("She sleeps a lot.", "She sleeps a lot."), []),
     ]
 
 
@@ -289,14 +281,8 @@ def test_construct_text_document_from_text_pair_coref_document(positive_and_nega
         )
         for doc in positive_and_negative_documents
     ]
-    assert len(docs) == 16
+    assert len(docs) == 8
     doc = docs[0]
-    assert doc.text == "And she founded C."
-    assert doc.labeled_spans.resolve() == [("PERSON", "she"), ("COMPANY", "C")]
-    assert doc.binary_relations.resolve() == []
-    assert [rel.score for rel in doc.binary_relations] == []
-
-    doc = docs[1]
     assert doc.text == "And she founded C.<s><s>Bob loves his cat."
     assert doc.labeled_spans.resolve() == [
         ("PERSON", "she"),
@@ -309,7 +295,7 @@ def test_construct_text_document_from_text_pair_coref_document(positive_and_nega
     ]
     assert [rel.score for rel in doc.binary_relations] == [1.0]
 
-    doc = docs[7]
+    doc = docs[4]
     assert doc.text == "Bob loves his cat.<s><s>She sleeps a lot."
     assert doc.labeled_spans.resolve() == [
         ("PERSON", "Bob"),