Braindecode 1.3.2.dev170355810__tar.gz → 1.3.2.dev170805792__tar.gz

{braindecode-1.3.2.dev170355810/braindecode.egg-info → braindecode-1.3.2.dev170805792}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.3.2.dev170355810
+Version: 1.3.2.dev170805792
 Summary: Deep learning software to decode EEG, ECG or MEG signals
 Author-email: Robin Tibor Schirrmeister <robintibor@gmail.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Alexandre Gramfort <agramfort@meta.com>
 Maintainer-email: Alexandre Gramfort <agramfort@meta.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Robin Tibor Schirrmeister <robintibor@gmail.com>
@@ -22,7 +22,7 @@ Requires-Python: >=3.11
 Description-Content-Type: text/x-rst
 License-File: LICENSE.txt
 License-File: NOTICE.txt
-Requires-Dist: torch>=2.2
+Requires-Dist: torch>=2.0
 Requires-Dist: torchaudio>=2.0
 Requires-Dist: mne>=1.11.0
 Requires-Dist: pandas<3.0.0
@@ -35,7 +35,7 @@ Requires-Dist: wfdb
 Requires-Dist: linear_attention_transformer
 Requires-Dist: docstring_inheritance
 Requires-Dist: rotary_embedding_torch
-Requires-Dist: pandas<3.0.0
+Requires-Dist: pydantic>=2.0
 Provides-Extra: moabb
 Requires-Dist: moabb>=1.4.3; extra == "moabb"
 Provides-Extra: eegprep

{braindecode-1.3.2.dev170355810 → braindecode-1.3.2.dev170805792}/braindecode/augmentation/transforms.py

@@ -2,6 +2,7 @@
 #          Alexandre Gramfort <alexandre.gramfort@inria.fr>
 #          Gustavo Rodrigues <gustavenrique01@gmail.com>
 #          Bruna Lopes <brunajaflopes@gmail.com>
+#          Sarthak Tayal <sarthaktayal2@gmail.com>
 #
 # License: BSD (3-clause)
 
@@ -557,8 +558,9 @@ class BandstopFilter(Transform):
                 f" Nyquist frequency ({nyq} Hz)."
                 f" Falling back to max_freq = {nyq}."
             )
-        assert bandwidth < max_freq, (
-            f"`bandwidth` needs to be smaller than max_freq={max_freq}"
+        assert bandwidth < max_freq - 2, (
+            f"`bandwidth` needs to be smaller than max_freq - 2={max_freq - 2} "
+            f"to allow valid notch frequency sampling with 1 Hz transition bands."
         )
 
         # override bandwidth value when a magnitude is passed
@@ -600,8 +602,8 @@ class BandstopFilter(Transform):
 
         # Prevents transitions from going below 0 and above max_freq
         notched_freqs = self.rng.uniform(
-            low=1 + 2 * self.bandwidth,
-            high=self.max_freq - 1 - 2 * self.bandwidth,
+            low=1 + self.bandwidth / 2,
+            high=self.max_freq - 1 - self.bandwidth / 2,
             size=X.shape[0],
         )
         return {

{braindecode-1.3.2.dev170355810 → braindecode-1.3.2.dev170805792}/braindecode/datasets/base.py

@@ -20,18 +20,24 @@ from abc import abstractmethod
 from collections import Counter
 from collections.abc import Callable
 from glob import glob
+from pathlib import Path
 from typing import Any, Generic, Iterable, no_type_check
 
+import mne
 import mne.io
 import numpy as np
 import pandas as pd
+from mne.utils import _soft_import
 from mne.utils.docs import deprecated
 from torch.utils.data import ConcatDataset, Dataset, IterableDataset
 from typing_extensions import TypeVar
 
 from .bids.hub import HubDatasetMixin
+from .bids.hub_io import _restore_nan_from_json
 from .registry import register_dataset
 
+zarr = _soft_import("zarr", purpose="lazy loading from Zarr", strict=False)
+
 
 def _create_description(description) -> pd.Series:
     if description is not None:
@@ -176,14 +182,11 @@ def _concat_metadata_summary(datasets):
 
 
 def _channel_info(mne_obj):
-    """Extract channel summary from an mne object.
-
-    Returns (n_ch, type_str, sfreq).
-    """
-    info = mne_obj.info
+    """Extract (n_ch, type_str, sfreq) from an mne object or mne.Info."""
+    info = mne_obj if isinstance(mne_obj, mne.Info) else mne_obj.info
     n_ch = info["nchan"]
     sfreq = info["sfreq"]
-    ch_types = mne_obj.get_channel_types()
+    ch_types = info.get_channel_types()
     type_counts = Counter(ch_types)
     type_str = ", ".join(f"{cnt} {t.upper()}" for t, cnt in sorted(type_counts.items()))
     return n_ch, type_str, sfreq
@@ -294,6 +297,88 @@ def _build_windowed_repr(
     return b
 
 
+def _zarr_to_memmap(zarr_path, group_name):
+    """Decompress a zarr array to a float64 ``.npy`` memmap (one-time, atomic).
+
+    Saved as float64 so MNE objects can wrap the memmap zero-copy
+    (MNE requires float64 internally).  Copy-on-write mode (``'c'``)
+    means preprocessing writes go to RAM, the file stays untouched.
+    """
+    zarr_p = Path(zarr_path)
+    cache_dir = zarr_p.parent / f".{zarr_p.name}_memmap"
+    npy_path = cache_dir / f"{group_name}.npy"
+
+    if not npy_path.exists():
+        arr = zarr.open(zarr_path, mode="r")[group_name]["data"]
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        tmp_path = cache_dir / f"{group_name}.{os.getpid()}.tmp.npy"
+        # Write chunk-by-chunk to avoid materializing the full array in RAM
+        mm = np.lib.format.open_memmap(
+            tmp_path, mode="w+", dtype=np.float64, shape=arr.shape
+        )
+        chunk0 = arr.chunks[0] if hasattr(arr, "chunks") else arr.shape[0]
+        for start in range(0, arr.shape[0], chunk0):
+            end = min(start + chunk0, arr.shape[0])
+            mm[start:end] = np.asarray(arr[start:end], dtype=np.float64)
+        mm.flush()
+        del mm
+        try:
+            tmp_path.rename(npy_path)
+        except OSError:
+            tmp_path.unlink(missing_ok=True)
+
+    return npy_path
+
+
+class _ZarrMixin:
+    """Zarr-to-memmap lazy loading with MNE compatibility.
+
+    Decompresses zarr to float64 ``.npy`` memmaps on first access.
+    MNE objects (``raw``/``windows``) are reconstructed lazily on first
+    property access, backed by copy-on-write memmaps so preprocessing
+    writes go to RAM while the file stays untouched.
+    """
+
+    _zarr_data = None
+
+    def _open_zarr(self):
+        npy_path = _zarr_to_memmap(self._zarr_path, self._group_name)
+        # 'c' = copy-on-write: reads from file, writes to RAM
+        self._zarr_data = np.load(npy_path, mmap_mode="c")
+
+    def _make_mne_info(self):
+        """Reconstruct ``mne.Info`` from stored dict (no data I/O)."""
+        return mne.Info.from_json_dict(_restore_nan_from_json(self._info_dict))
+
+    @classmethod
+    def _init_zarr_base(cls, zarr_path, group_name, description, info_dict, transform):
+        """Shared ``_from_zarr`` preamble: create instance with zarr attrs."""
+        if zarr is False:
+            raise ImportError("zarr is required: pip install braindecode[hub]")
+        obj = object.__new__(cls)
+        RecordDataset.__init__(obj, description, transform)
+        obj._zarr_path = Path(zarr_path)
+        obj._group_name = group_name
+        obj._info_dict = info_dict
+        return obj
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state.pop("_zarr_data", None)
+        # Drop lazily-constructed mne objects; they'll be rebuilt from memmap
+        if hasattr(self, "_zarr_path"):
+            state.pop("_raw", None)
+            state.pop("_windows", None)
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        if hasattr(self, "_zarr_path"):
+            self._open_zarr()
+        else:
+            self._zarr_data = None
+
+
 class RecordDataset(Dataset[tuple[np.ndarray, int | str, tuple[int, int, int]]]):
     def __init__(
         self,
@@ -361,7 +446,7 @@ T = TypeVar("T", bound=RecordDataset)
 
 
 @register_dataset
-class RawDataset(RecordDataset):
+class RawDataset(_ZarrMixin, RecordDataset):
     """Returns samples from an mne.io.Raw object along with a target.
 
     Dataset which serves samples from an mne.io.Raw object along with a target.
@@ -395,8 +480,43 @@ class RawDataset(RecordDataset):
         self.target_name = self._target_name(target_name)
         self.raw_preproc_kwargs: list[dict[str, Any]] = []
 
+    @classmethod
+    def _from_zarr(
+        cls,
+        zarr_path,
+        group_name,
+        description,
+        info_dict,
+        target_name=None,
+        transform=None,
+    ):
+        """Create a lazy Zarr-backed RawDataset (internal API)."""
+        obj = cls._init_zarr_base(
+            zarr_path, group_name, description, info_dict, transform
+        )
+        obj._raw = None
+        obj.target_name = cls._target_name(obj, target_name)
+        obj.raw_preproc_kwargs = []
+        obj._open_zarr()
+        return obj
+
+    @property
+    def raw(self):
+        if self._raw is None and self._zarr_data is not None:
+            self._raw = mne.io.RawArray(self._zarr_data, self._make_mne_info())
+        return self._raw
+
+    @raw.setter
+    def raw(self, value):
+        self._raw = value
+
     def __getitem__(self, index):
-        X = self.raw[:, index][0]
+        if self._zarr_data is not None:
+            X = np.array(self._zarr_data[:, index], dtype="float32")
+            if X.ndim == 1:
+                X = X[:, np.newaxis]
+        else:
+            X = self._raw[:, index][0]
         y = None
         if self.target_name is not None:
             y = self.description[self.target_name]
@@ -407,13 +527,19 @@ class RawDataset(RecordDataset):
         return X, y
 
     def __len__(self):
-        return len(self.raw)
+        if self._raw is not None:
+            return len(self._raw)
+        return self._zarr_data.shape[1]
 
     def _build_repr(self):
         b = _ReprBuilder(type(self).__name__)
-        n_ch, type_str, sfreq = _channel_info(self.raw)
-        n_times = len(self.raw.times)
-        duration = n_times / sfreq
+        if self._raw is not None:
+            n_ch, type_str, sfreq = _channel_info(self._raw)
+            n_times = len(self._raw.times)
+        else:
+            n_ch, type_str, sfreq = _channel_info(self._make_mne_info())
+            n_times = self._zarr_data.shape[1]
+        duration = n_times / sfreq if sfreq > 0 else 0
         b.add_header(f"{n_ch} ch ({type_str})", "Channels", f"{n_ch} ({type_str})")
         b.add_header(f"{sfreq:.1f} Hz", "Sfreq", f"{sfreq:.1f} Hz")
         b.add_header(
@@ -424,6 +550,8 @@ class RawDataset(RecordDataset):
         if self.description is not None:
             desc_items = ", ".join(f"{k}={v}" for k, v in self.description.items())
             b.add_row("description", desc_items)
+        if self._zarr_data is not None:
+            b.add_footnote("Data loaded lazily from Zarr")
         return b
 
     def _target_name(self, target_name):
@@ -464,7 +592,7 @@ class BaseDataset(RawDataset):
 
 
 @register_dataset
-class EEGWindowsDataset(RecordDataset):
+class EEGWindowsDataset(_ZarrMixin, RecordDataset):
     """Returns windows from an mne.Raw object, its window indices, along with a target.
 
     Dataset which serves windows from an mne.Epochs object along with their
@@ -521,6 +649,47 @@ class EEGWindowsDataset(RecordDataset):
             self.y = metadata.loc[:, "target"].to_list()
         self.raw_preproc_kwargs: list[dict[str, Any]] = []
 
+    @classmethod
+    def _from_zarr(
+        cls,
+        zarr_path,
+        group_name,
+        metadata,
+        description,
+        info_dict,
+        targets_from="metadata",
+        last_target_only=True,
+        transform=None,
+    ):
+        """Create a lazy Zarr-backed EEGWindowsDataset (internal API)."""
+        obj = cls._init_zarr_base(
+            zarr_path, group_name, description, info_dict, transform
+        )
+        obj._raw = None
+        obj.metadata = metadata
+        obj.last_target_only = last_target_only
+        if targets_from not in ("metadata", "channels"):
+            raise ValueError("Wrong value for parameter `targets_from`.")
+        obj.targets_from = targets_from
+        obj.crop_inds = metadata.loc[
+            :, ["i_window_in_trial", "i_start_in_trial", "i_stop_in_trial"]
+        ].to_numpy()
+        if obj.targets_from == "metadata":
+            obj.y = metadata.loc[:, "target"].to_list()
+        obj.raw_preproc_kwargs = []
+        obj._open_zarr()
+        return obj
+
+    @property
+    def raw(self):
+        if self._raw is None and self._zarr_data is not None:
+            self._raw = mne.io.RawArray(self._zarr_data, self._make_mne_info())
+        return self._raw
+
+    @raw.setter
+    def raw(self, value):
+        self._raw = value
+
     def __getitem__(self, index: int):
         """Get a window and its target.
 
@@ -544,11 +713,13 @@ class EEGWindowsDataset(RecordDataset):
         crop_inds = self.crop_inds[index].tolist()
 
         i_window_in_trial, i_start, i_stop = crop_inds
-        X = self.raw._getitem((slice(None), slice(i_start, i_stop)), return_times=False)
-        X = X.astype("float32")
-        # ensure we don't give the user the option
-        # to accidentally modify the underlying array
-        X = X.copy()
+        if self._zarr_data is not None:
+            X = self._zarr_data[:, i_start:i_stop]
+        else:
+            X = self.raw._getitem(
+                (slice(None), slice(i_start, i_stop)), return_times=False
+            )
+        X = np.array(X, dtype="float32")
         if self.transform is not None:
             X = self.transform(X)
         if self.targets_from == "metadata":
@@ -559,10 +730,7 @@ class EEGWindowsDataset(RecordDataset):
                 y = X[misc_mask, -1]
             else:
                 y = X[misc_mask, :]
-            # ensure we don't give the user the option
-            # to accidentally modify the underlying array
             y = y.copy()
-            # remove the target channels from raw
             X = X[~misc_mask, :]
         return X, y, crop_inds
 
@@ -570,10 +738,11 @@ class EEGWindowsDataset(RecordDataset):
         return len(self.crop_inds)
 
     def _build_repr(self):
+        mne_obj = self._raw if self._raw is not None else self._make_mne_info()
         return _build_windowed_repr(
             type(self).__name__,
             len(self),
-            self.raw,
+            mne_obj,
             self.crop_inds,
             self.description,
             self.metadata,
@@ -664,7 +833,7 @@ class EEGWindowsDataset(RecordDataset):
 
 
 @register_dataset
-class WindowsDataset(RecordDataset):
+class WindowsDataset(_ZarrMixin, RecordDataset):
     """Returns windows from an mne.Epochs object along with a target.
 
     Dataset which serves windows from an mne.Epochs object along with their
@@ -717,6 +886,7 @@ class WindowsDataset(RecordDataset):
 
         metadata = self.windows.metadata
         assert metadata is not None, "WindowsDataset requires windows with metadata."
+        self.metadata = metadata
         self.crop_inds = metadata.loc[
             :, ["i_window_in_trial", "i_start_in_trial", "i_stop_in_trial"]
         ].to_numpy()
@@ -725,6 +895,67 @@ class WindowsDataset(RecordDataset):
         self.raw_preproc_kwargs: list[dict[str, Any]] = []
         self.window_preproc_kwargs: list[dict[str, Any]] = []
 
+    @classmethod
+    def _from_zarr(
+        cls,
+        zarr_path,
+        group_name,
+        metadata,
+        description,
+        info_dict,
+        target_name=None,
+        transform=None,
+        targets_from="metadata",
+        last_target_only=True,
+    ):
+        """Create a lazy Zarr-backed WindowsDataset (internal API)."""
+        obj = cls._init_zarr_base(
+            zarr_path, group_name, description, info_dict, transform
+        )
+        obj._windows = None
+        obj.target_name = target_name
+        obj._fast_disk = False
+        obj.last_target_only = last_target_only
+        if targets_from not in ("metadata", "channels"):
+            raise ValueError("Wrong value for parameter `targets_from`.")
+        obj.targets_from = targets_from
+        obj.metadata = metadata
+        obj.crop_inds = metadata.loc[
+            :, ["i_window_in_trial", "i_start_in_trial", "i_stop_in_trial"]
+        ].to_numpy()
+        if obj.targets_from == "metadata":
+            obj.y = metadata.loc[:, "target"].to_list()
+        obj.raw_preproc_kwargs = []
+        obj.window_preproc_kwargs = []
+        obj._open_zarr()
+        return obj
+
+    @property
+    def windows(self):
+        if self._windows is None and self._zarr_data is not None:
+            info = self._make_mne_info()
+            targets = self.metadata["target"].values
+            event_ids = (
+                targets
+                if np.issubdtype(targets.dtype, np.integer)
+                else np.ones(len(self.metadata), dtype=int)
+            )
+            events = np.column_stack(
+                [
+                    self.metadata["i_start_in_trial"].values.astype(int),
+                    np.zeros(len(self.metadata), dtype=int),
+                    event_ids,
+                ]
+            )
+            self._windows = mne.EpochsArray(
+                self._zarr_data, info, events=events, metadata=self.metadata
+            )
+        return self._windows
+
+    @windows.setter
+    def windows(self, value):
+        self._windows = value
+
     @staticmethod
     def _can_use_fast_get_epoch_from_raw(epochs: mne.BaseEpochs) -> bool:
         """Check if we can use the fast _get_epoch_from_raw method,
@@ -756,7 +987,9 @@ class WindowsDataset(RecordDataset):
         np.ndarray
             Crop indices.
         """
-        if self._fast_disk:
+        if self._zarr_data is not None:
+            X = self._zarr_data[index].astype("float32")
+        elif self._fast_disk:
             X = self.windows._get_epoch_from_raw(index).astype("float32")
         else:
             X = self.windows.get_data(item=index)[0].astype("float32")
@@ -770,7 +1003,6 @@ class WindowsDataset(RecordDataset):
                 y = X[misc_mask, -1]
             else:
                 y = X[misc_mask, :]
-            # remove the target channels from raw
             X = X[~misc_mask, :]
         # necessary to cast as list to get list of three tensors from batch,
         # otherwise get single 2d-tensor...
@@ -778,16 +1010,20 @@ class WindowsDataset(RecordDataset):
         return X, y, crop_inds
 
     def __len__(self) -> int:
-        return len(self.windows.events)
+        if self._windows is not None:
+            return len(self._windows.events)
+        return len(self.crop_inds)
 
     def _build_repr(self):
+        mne_obj = self._windows if self._windows is not None else self._make_mne_info()
+        md = self._windows.metadata if self._windows is not None else self.metadata
         return _build_windowed_repr(
             type(self).__name__,
             len(self),
-            self.windows,
+            mne_obj,
             self.crop_inds,
             self.description,
-            self.windows.metadata,
+            md,
         )
 
 
@@ -1014,8 +1250,8 @@ class BaseConcatDataset(ConcatDataset, HubDatasetMixin, Generic[T]):
 
         all_dfs = list()
         for ds in self.datasets:
-            if hasattr(ds, "windows"):
-                df = ds.windows.metadata
+            if hasattr(ds, "_windows") and ds._windows is not None:
+                df = ds._windows.metadata
             else:
                 df = ds.metadata
             for k, v in ds.description.items():

{braindecode-1.3.2.dev170355810 → braindecode-1.3.2.dev170805792}/braindecode/datasets/bids/hub.py

@@ -39,7 +39,6 @@ import mne
 import numpy as np
 import pandas as pd
 import scipy
-from mne._fiff.meas_info import Info
 from mne.utils import _soft_import
 
 if TYPE_CHECKING:
@@ -733,7 +732,6 @@ class HubDatasetMixin:
                     info_dict,
                     compressor,
                     target_name,
-                    chunk_size,
                 )
 
             elif dataset_type == "EEGWindowsDataset":
@@ -836,6 +834,7 @@ class HubDatasetMixin:
         datasets = []
         for i_ds in range(n_datasets):
             grp = root[f"recording_{i_ds}"]
+            group_name = f"recording_{i_ds}"
 
             if dataset_type == "WindowsDataset":
                 # Load using inlined function
@@ -843,19 +842,36 @@ class HubDatasetMixin:
                     _load_windows_from_zarr(grp, preload)
                 )
 
-                # Convert to MNE objects and create dataset
-                info = Info.from_json_dict(info_dict)
-                events = np.column_stack(
-                    [
-                        metadata["i_start_in_trial"].values,
-                        np.zeros(len(metadata), dtype=int),
-                        metadata["target"].values,
-                    ]
-                )
-                epochs = mne.EpochsArray(data, info, events=events, metadata=metadata)
-                ds = WindowsDataset(epochs, description)
-                if target_name is not None:
-                    ds.target_name = target_name
+                if preload:
+                    # Convert to MNE objects and create dataset
+                    info = mne.Info.from_json_dict(info_dict)
+                    targets = metadata["target"].values
+                    if np.issubdtype(targets.dtype, np.integer):
+                        event_ids = targets
+                    else:
+                        event_ids = np.ones(len(metadata), dtype=int)
+                    events = np.column_stack(
+                        [
+                            metadata["i_start_in_trial"].values.astype(int),
+                            np.zeros(len(metadata), dtype=int),
+                            event_ids,
+                        ]
+                    )
+                    epochs = mne.EpochsArray(
+                        data, info, events=events, metadata=metadata
+                    )
+                    ds = WindowsDataset(epochs, description)
+                    if target_name is not None:
+                        ds.target_name = target_name
+                else:
+                    ds = WindowsDataset._from_zarr(
+                        zarr_path=input_path,
+                        group_name=group_name,
+                        metadata=metadata,
+                        description=description,
+                        info_dict=info_dict,
+                        target_name=target_name,
+                    )
 
             elif dataset_type == "EEGWindowsDataset":
                 # Load using inlined function
@@ -868,17 +884,28 @@ class HubDatasetMixin:
                     last_target_only,
                 ) = _load_eegwindows_from_zarr(grp, preload)
 
-                # Convert to MNE objects and create dataset
-                # Data is already in continuous format [n_channels, n_timepoints]
-                info = Info.from_json_dict(info_dict)
-                raw = mne.io.RawArray(data, info)
-                ds = EEGWindowsDataset(
-                    raw=raw,
-                    metadata=metadata,
-                    description=description,
-                    targets_from=targets_from,
-                    last_target_only=last_target_only,
-                )
+                if preload:
+                    # Convert to MNE objects and create dataset
+                    # Data is already in continuous format [n_channels, n_timepoints]
+                    info = mne.Info.from_json_dict(info_dict)
+                    raw = mne.io.RawArray(data, info)
+                    ds = EEGWindowsDataset(
+                        raw=raw,
+                        metadata=metadata,
+                        description=description,
+                        targets_from=targets_from,
+                        last_target_only=last_target_only,
+                    )
+                else:
+                    ds = EEGWindowsDataset._from_zarr(
+                        zarr_path=input_path,
+                        group_name=group_name,
+                        metadata=metadata,
+                        description=description,
+                        info_dict=info_dict,
+                        targets_from=targets_from,
+                        last_target_only=last_target_only,
+                    )
 
             elif dataset_type == "RawDataset":
                 # Load using inlined function
@@ -886,13 +913,22 @@ class HubDatasetMixin:
                     grp, preload
                 )
 
-                # Convert to MNE objects and create dataset
-                # Data is in continuous format [n_channels, n_timepoints]
-                info = Info.from_json_dict(info_dict)
-                raw = mne.io.RawArray(data, info)
-                ds = RawDataset(raw, description)
-                if target_name is not None:
-                    ds.target_name = target_name
+                if preload:
+                    # Convert to MNE objects and create dataset
+                    # Data is in continuous format [n_channels, n_timepoints]
+                    info = mne.Info.from_json_dict(info_dict)
+                    raw = mne.io.RawArray(data, info)
+                    ds = RawDataset(raw, description)
+                    if target_name is not None:
+                        ds.target_name = target_name
+                else:
+                    ds = RawDataset._from_zarr(
+                        zarr_path=input_path,
+                        group_name=group_name,
+                        description=description,
+                        info_dict=info_dict,
+                        target_name=target_name,
+                    )
 
             else:
                 raise ValueError(f"Unsupported dataset_type: {dataset_type}")