fix slow test: test_train_predict() (#157)

ArneBinder · Mar 6, 2024 · 1fd445c · 1fd445c
1 parent 961ee64
commit 1fd445c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_predict.py b/tests/test_predict.py
@@ -75,7 +75,7 @@ def test_train_predict(tmp_path, cfg_train, cfg_predict):
     HydraConfig().set_config(cfg_predict)
     _, object_dict = predict(cfg_predict)
 
-    predicted_entities = [list(doc.entities.predictions) for doc in object_dict["documents"]]
+    predicted_entities = [list(doc.labeled_spans.predictions) for doc in object_dict["documents"]]
     num_predicted_entities = sum([len(preds) for preds in predicted_entities])
     assert num_predicted_entities > 0