use deduplicate_annotations() in test

ArneBinder · Nov 13, 2024 · 69e289b · 69e289b
1 parent 2a9b2d6
commit 69e289b
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 14 deletions.
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -24,7 +24,9 @@ classifiers = [
 
 [tool.poetry.dependencies]
 python = "^3.9"
-pytorch-ie = ">=0.31.2,<0.32.0"
+#pytorch-ie = ">=0.31.4,<0.32.0"
+# install from branch from https://github.com/ArneBinder/pytorch-ie/pull/436
+pytorch-ie = { git = "https://github.com/ArneBinder/pytorch-ie", branch = "document/deduplicate_annotations" }
 pytorch-lightning = "^2.1.0"
 torchmetrics = "^1"
 # >=4.35 because of BartModelWithDecoderPositionIds, <4.37 because of generation config

diff --git a/tests/taskmodules/test_pointer_network_for_end2end_re.py b/tests/taskmodules/test_pointer_network_for_end2end_re.py
@@ -568,7 +568,7 @@ def test_decode_with_add_reversed_relations():
     task_outputs = [task_encoding.targets for task_encoding in task_encodings]
     docs_with_predictions = taskmodule.decode(task_encodings, task_outputs)
     assert len(docs_with_predictions) == 1
-    doc_with_predictions: ExampleDocument = docs_with_predictions[0]
+    doc_with_predictions: ExampleDocument = docs_with_predictions[0].deduplicate_annotations()
     assert list(doc_with_predictions.entities.predictions) == list(doc_with_predictions.entities)
     assert list(doc_with_predictions.relations.predictions) == list(doc_with_predictions.relations)