Deferred imports for datasets libraries (torchvision, tensorflow-datasets, and HF datasets)

oodeel/datasets/tf_data_handler.py

@@ -23,8 +23,6 @@
 from typing import get_args
 
 import tensorflow as tf
-import tensorflow_datasets as tfds
-from datasets import load_dataset as hf_load_dataset
 
 from ..types import Callable
 from ..types import ItemType
@@ -240,6 +238,14 @@ def load_from_huggingface(
         Returns:
             tf.data.Dataset: dataset
         """
+        try:
+            from datasets import load_dataset as hf_load_dataset
+        except ImportError:
+            raise ImportError(
+                "datasets is not installed. Please install it to use "
+                'hub="huggingface" in load_dataset.'
+            )
+
         dataset = hf_load_dataset(dataset_id, **load_kwargs)
         dataset = dataset.to_tf_dataset()
         return dataset
@@ -259,6 +265,14 @@ def load_from_tensorflow_datasets(
         Returns:
             tf.data.Dataset
         """
+        try:
+            import tensorflow_datasets as tfds
+        except ImportError:
+            raise ImportError(
+                "tensorflow_datasets is not installed. Please install it to use "
+                'hub="tensorflow-datasets" in load_dataset.'
+            )
+
         assert (
             dataset_id in tfds.list_builders()
         ), "Dataset not available on tensorflow datasets catalog"

oodeel/datasets/torch_data_handler.py

@@ -26,8 +26,6 @@
 import numpy as np
 import torch
 import torch.nn.functional as F
-import torchvision
-from datasets import load_dataset as hf_load_dataset
 from torch.utils.data import DataLoader
 from torch.utils.data import Dataset
 from torch.utils.data import Subset
@@ -377,6 +375,14 @@ def load_from_huggingface(
         Returns:
             DictDataset: dataset
         """
+        try:
+            from datasets import load_dataset as hf_load_dataset
+        except ImportError:
+            raise ImportError(
+                "datasets is not installed. Please install it to use "
+                'hub="huggingface" in load_dataset.'
+            )
+
         if "transform" in load_kwargs.keys():
             transform = load_kwargs["transform"]
             load_kwargs.pop("transform")
@@ -421,6 +427,13 @@ def load_from_torchvision(
         Returns:
             DictDataset: dataset
         """
+        try:
+            import torchvision
+        except ImportError:
+            raise ImportError(
+                "torchvision is not installed. Please install it to use "
+                'hub="torchvision" in load_dataset.'
+            )
         assert (
             dataset_id in torchvision.datasets.__all__
         ), "Dataset not available on torchvision datasets catalog"

oodeel/types/__init__.py

@@ -23,6 +23,7 @@
 """
 Typing module
 """
+
 from typing import Any
 from typing import Callable
 from typing import Dict

setup.cfg

@@ -51,6 +51,7 @@ deps =
     pandas
     seaborn
     requests
+    datasets
     plotly == 5.15.0
     tf211: protobuf>=3.9.2,<3.20
     tf213: protobuf>=3.20.3,<4.21.0
@@ -103,15 +104,19 @@ deps =
     torch200: torch == 2.0.0+cpu
     torch200: torchvision == 0.15.0+cpu
     torch200: transformers == 4.49.0
+    torch200: datasets==2.19.1
     torch22: torch == 2.2.0+cpu
     torch22: torchvision == 0.17.0+cpu
     torch22: transformers == 4.49.0
+    torch24: datasets==2.19.1
     torch24: torch == 2.4.0+cpu
     torch24: torchvision == 0.19.0+cpu
     torch24: transformers == 4.49.0
+    torch24: datasets==2.19.1
     torch26: torch == 2.6.0+cpu
     torch26: torchvision == 0.21.0+cpu
     torch26: transformers == 4.49.0
+    torch26: datasets
 install_command = uv pip install --extra-index-url https://download.pytorch.org/whl/cpu {opts} {packages}
 setenv =
     DL_LIB = torch
@@ -154,9 +159,11 @@ deps =
     torch22: torch == 2.2.0+cpu
     torch22: torchvision == 0.17.0+cpu
     torch22: transformers == 4.49.0
+    torch22: datasets==2.19.1
     tf215: tensorflow ~= 2.15.0
     tf215: tensorflow_datasets
     tf215: tensorflow_probability ~= 0.23.0
+    tf215: datasets
 install_command = uv pip install --extra-index-url https://download.pytorch.org/whl/cpu {opts} {packages}
 setenv =
     DL_LIB = both

setup.py

@@ -40,16 +40,16 @@
     "seaborn",
     "plotly",
     "tqdm",
-    "datasets",
 ]
 
 tensorflow_requirements = [
     "tensorflow==2.11.0",
     "tensorflow_datasets",
     "tensorflow_probability==0.19.0",
+    "datasets",
 ]
 
-torch_requirements = ["timm", "torch==1.13.1", "torchvision==0.14.1"]
+torch_requirements = ["timm", "torch==1.13.1", "torchvision==0.14.1", "datasets"]
 
 dev_requirements = [
     "mypy",