diff --git a/dataset_builders/pie/argmicro/requirements.txt b/dataset_builders/pie/argmicro/requirements.txt
index 2c289fa6..6e799260 100644
--- a/dataset_builders/pie/argmicro/requirements.txt
+++ b/dataset_builders/pie/argmicro/requirements.txt
@@ -1 +1 @@
-pie-datasets>=0.3.3,<0.5.0
+pie-datasets>=0.3.3,<0.6.0
diff --git a/tests/dataset_builders/pie/test_argmicro.py b/tests/dataset_builders/pie/test_argmicro.py
index e81fece3..3a13c5f5 100644
--- a/tests/dataset_builders/pie/test_argmicro.py
+++ b/tests/dataset_builders/pie/test_argmicro.py
@@ -4,6 +4,7 @@
 
 import pytest
 from datasets import disable_caching, load_dataset
+from pie_models.document.processing import tokenize_document
 from pytorch_ie.core import Document
 from pytorch_ie.documents import TextDocumentWithLabeledSpansAndBinaryRelations
 from transformers import AutoTokenizer, PreTrainedTokenizer
@@ -16,7 +17,6 @@
     example_to_document,
 )
 from pie_datasets import DatasetDict
-from pie_datasets.document.processing import tokenize_document
 from pie_datasets.document.types import TokenDocumentWithLabeledSpansAndBinaryRelations
 from tests import FIXTURES_ROOT
 from tests.dataset_builders.common import HF_DS_FIXTURE_DATA_PATH, PIE_BASE_PATH