import datasets only when needed

huggingface · Nov 15, 2024 · ee90d06 · ee90d06
1 parent d9e9e4a
commit ee90d06
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/tests/utils/test_task_processors.py b/tests/utils/test_task_processors.py
@@ -20,9 +20,9 @@
 from unittest import TestCase
 
 import pytest
-from datasets import DatasetDict
 from transformers import AutoConfig, AutoFeatureExtractor, AutoTokenizer
 
+from optimum.utils.import_utils import is_datasets_available
 from optimum.utils.preprocessing import TaskProcessorsManager
 from optimum.utils.testing_utils import require_datasets
 
@@ -31,6 +31,9 @@
     from transformers import PretrainedConfig, PreTrainedTokenizerBase
     from transformers.image_processing_utils import BaseImageProcessor
 
+if is_datasets_available():
+    from datasets import DatasetDict
+
 
 TEXT_MODEL_NAME = "bert-base-uncased"
 CONFIG = AutoConfig.from_pretrained(TEXT_MODEL_NAME)