Repository: ktonal/mimikit (language: python)

Issue:
Refactor/data.factory
- simplify factory.py
- refactor Metadata into Regions


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 012823b..5b2f5cf 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,4 +3,4 @@
 *.ipynb_checkpoints*
 dist/*
 build/*
-*egg-info*
\ No newline at end of file
+*egg-info*
diff --git a/README.md b/README.md
index 7affe19..325b086 100644
--- a/README.md
+++ b/README.md
@@ -1,54 +1,16 @@
 # mimikit
 
-Do deep-learning on your own audio files like a pro with just a google account.
+`mimikit` is an open-source python package for audio and music data
+modeling with neural networks.
 
-`mimikit` is a music modelling kit that lets you mimic / transform your own audio files with generative neural-networks.
-
-It contains a collection of models in pytorch and pytorch-lightning as well as powerful ways to :
-- prepare & store your data for these models
-- train the models online by free gpu providers
-- store and track every experiment you make & every sound bits you generate on neptune.ai - also for free
+It contains models out-of-the-box and simple utilities for dealing with audio data.
+The aim of mimikit is double : it will lower the barrier to deep-learning for non-developers while providing tools to developers that enable fast experimentation.
 
 -----
 
-**Table of Contents**
-
-* [Installation](#installation)
-* [Quickstart](#quickstart)
-* [Usage](#usage)
-* [Documentation](#documentation)
-* [Contribute](#contribute)
-* [License](#license)
-
-## Installation
-
-mimikit is available as a `pip` package. Open a terminal and type :
-
-```bash
-$ pip install mimikit
-```
-
-## Quickstart
-
-If you never did deep-learning before, we recommend you start with the [quickest intro to practical deep-learning ever]()
-
-- Have a [google account](https://accounts.google.com/signup/v2/webcreateaccount?flowName=GlifWebSignIn&flowEntry=SignUp) and register with it to [neptune.ai](https://neptune.ai/)
-- Put some audio files in your google drive or [make a database]() on your computer
-- Open the [FreqNet starter notebook](https://colab.research.google.com/github/k-tonal/mmk-notebooks/blob/main/FreqNet.ipynb) in colab and follow the instructions
-
-For more, check out the [mimikit-notebooks](https://github.com/k-tonal/mimikit-notebooks), the [mmikit docs]() or the documentation for the [freqnet package]() 
-
 ## Usage 
 
-Check out the [mimikit-notebooks](https://github.com/k-tonal/mimikit-notebooks) for client code examples 
-
-## Documentation
-
-TODO !
-
-## Contribute
-
-`mimikit` welcomes all kinds of contributions! From bug-fixes to new cool experimental models or improving coverage of tests and docs : get in touch and/or make a pull request. 
+Check out the [mimikit-notebooks](https://github.com/ktonal/mimikit-notebooks) for examples of client code.
 
 
 ## License
diff --git a/mimikit/__init__.py b/mimikit/__init__.py
index d9f27bf..696d661 100644
--- a/mimikit/__init__.py
+++ b/mimikit/__init__.py
@@ -1,4 +1,4 @@
-__version__ = '0.1.9'
+__version__ = '0.1.10'
 
 from .kit import get_trainer
 from .connectors.neptune import NeptuneConnector
diff --git a/mimikit/data/__init__.py b/mimikit/data/__init__.py
index 0e52d95..2e4e4ef 100644
--- a/mimikit/data/__init__.py
+++ b/mimikit/data/__init__.py
@@ -1,6 +1,6 @@
 from .api import Database, FeatureProxy
 from .data_object import DataObject
 from .factory import make_root_db, AudioFileWalker
-from .metadata import Metadata
+from .regions import Regions
 from .transforms import file_to_fft, default_extract_func, N_FFT, HOP_LENGTH, SR
 from .freqnet_db import freqnet_db
diff --git a/mimikit/data/api.py b/mimikit/data/api.py
index 2682f30..9037ef9 100644
--- a/mimikit/data/api.py
+++ b/mimikit/data/api.py
@@ -6,20 +6,23 @@
 
 from torch.utils.data.dataset import Subset
 
-from .metadata import Metadata
+from .regions import Regions
 from ..data.data_object import DataObject
 
 
 class FeatureProxy(object):
     """
-    Interface to the numerical data (array) stored in a .h5 file created with mimikit.
+    Interface to the numerical data (array) stored in a .h5 file created with ``mimikit.data.make_root_db``
+
+    ``FeatureProxy`` objects are like ``numpy`` arrays (on disk) and can be used as map-style ``Dataset``
+    because they implement ``__len__`` and ``__getitem__``.
 
     Parameters
     ----------
     h5_file : str
         name of the file
     ds_name : str
-        name of the ``h5py.Dataset`` containing the data
+        name of the ``h5py.Dataset`` containing the data to be served
 
     Attributes
     ----------
@@ -52,13 +55,13 @@ def __getitem__(self, item):
             rv = f[self.name][item]
         return rv
 
-    def get(self, metadata):
+    def get(self, regions):
         """
-        get the data (numpy array) corresponding to the rows of `metadata`
+        get the data (numpy array) corresponding to the rows of `regions`
 
         Parameters
         ----------
-        metadata : Metadata
+        regions : Regions
             the files you want to get as array
 
         Returns
@@ -71,36 +74,73 @@ def get(self, metadata):
         >>>from mimikit import Database
         >>>db = Database("my-db.h5")
         # only get files 0, 3 & 7
-        >>>db.fft.get(db.metadata.iloc[[0, 3, 7]])
+        >>>db.fft.get(db.regions.iloc[[0, 3, 7]])
 
         """
         t_axis = self.attrs.get("time_axis", 0)
-        slices = metadata.slices(t_axis)
+        slices = regions.slices(t_axis)
         return np.concatenate(tuple(self[slice_i] for slice_i in slices), axis=t_axis)
 
     def add(self, array, filename=None):
-        new = add_data(self.h5_file, self.name, array, filename)
+        """
+        EXPERIMENTAL! append `array` to the feature and fill the `"name"` column of `db.regions` with 'filename'
+        Parameters
+        ----------
+        array : np.ndarray
+            the array to append at the end of the feature
+        filename : str, optional
+            a name for the array being added
+
+        Returns
+        -------
+        new : FeatureProxy
+            the updated object
+        """
+        new = _add_data(self.h5_file, self.name, array, filename)
         return new
 
     def subset(self, indices):
         """
-        transform self into a torch `Subset` (`Dataset`) containing only `indices` from the original data
+        transform self into a torch ``Subset`` (``Dataset``) containing only ``indices`` from the original data
 
         Parameters
         ----------
-        indices : Metadata or any object that `Subset` accepts as indices
-            if `indices` is of type `Metadata`, the returned `Subset` will only contain the files (rows) present
-            in `indices`.
+        indices : ``Regions`` or any object that ``Subset`` accepts as indices
+            if ``indices`` is of type ``Regions``, the returned ``Subset`` will only contain the files (rows) present
+            in ``indices``.
 
         Returns
         -------
         subset : torch.utils.data.Subset
             the obtained subset
         """
-        if isinstance(indices, Metadata):
+        if isinstance(indices, Regions):
             indices = indices.all_indices
         return Subset(DataObject(self), indices)
 
+    def regions(self, indices):
+        """
+        transform self into a torch ``Subset`` containing the files whose indices have been passed
+        as ``ints`` in ``indices``
+
+        Parameters
+        ----------
+        indices : list of ints
+            correspond to the rows of db.regions you want to keep
+
+        Returns
+        -------
+        subset : torch.utils.data.Subset
+            the obtained subset
+
+        Examples
+        --------
+        >>>>db = Database("test.h5")
+        >>>>sub = db.fft.regions([1, 2, 3, 9])
+        """
+        regions = Database(self.h5_file).regions.iloc[indices]
+        return self.subset(regions)
+
     def __repr__(self):
         return "<FeatureProxy: '%s/%s'>" % (self.h5_file, self.name)
 
@@ -116,9 +156,9 @@ class Database(object):
 
     Attributes
     ----------
-    metadata : Metadata
+    regions : Regions
         pandas DataFrame where each row contains information about one stored file.
-        see ``Metadata`` for more information
+        see ``Regions`` for more information
 
     <feature_proxy> : FeatureProxy
         each feature created by the extracting function passed to ``make_root_db``
@@ -128,31 +168,26 @@ class Database(object):
     """
     def __init__(self, h5_file: str):
         self.h5_file = h5_file
-        self.info = self._get_dataframe("/info")
-        self.metadata = Metadata(self._get_dataframe("/metadata"))
+        self.regions = Regions(self._get_dataframe("/regions"))
         with h5py.File(h5_file, "r") as f:
+            self.attrs = {k: v for k, v in f.attrs.items()}
+            self.features = f.attrs["features"]
             # add found features as self.feature_name = FeatureProxy(self.h5_file, feature_name)
             self._register_features(self.features)
-            self._register_dataframes(self.dataframes)
-
-    @property
-    def features(self):
-        names = self.info.iloc[:, 2:].T.index.get_level_values(0)
-        return list(set(names))
-
-    @property
-    def dataframes(self):
-        keys = set()
 
-        def func(k, v):
-            if "pandas_type" in v.attrs.keys() and k.split("/")[0] not in ("layouts", "info", "metadata"):
-                keys.add(k)
-            return None
+    def visit(self, func=print):
+        """
+        wrapper for ``h5py.File.visititems()``
 
-        self.visit(func)
-        return list(keys)
+        Parameters
+        ----------
+        func : function
+            a function to be applied recursively
 
-    def visit(self, func=print):
+        Returns
+        -------
+        None
+        """
         with h5py.File(self.h5_file, "r") as f:
             f.visititems(func)
 
@@ -163,18 +198,27 @@ def _get_dataframe(self, key):
             return pd.DataFrame()
 
     def save_dataframe(self, key, df):
+        """
+        stores a ``pd.DataFrame`` object under ``key``
+
+        Parameters
+        ----------
+        key : str
+            the key under which ``df`` will be stored
+        df : pd.DataFrame
+            the ``DataFrame`` to be stored
+
+        Returns
+        -------
+        df : pd.DataFrame
+            the ``DataFrame`` as it has been stored
+        """
         with h5py.File(self.h5_file, "r+") as f:
             if key in f:
                 f.pop(key)
         df.to_hdf(self.h5_file, key=key, mode="r+")
         return self._get_dataframe(key)
 
-    def layout_for(self, feature):
-        with h5py.File(self.h5_file, "r") as f:
-            if "layouts" not in f.keys():
-                return pd.DataFrame()
-        return self._get_dataframe("layouts/" + feature)
-
     def _register_features(self, names):
         for name in names:
             setattr(self, name, FeatureProxy(self.h5_file, name))
@@ -194,28 +238,39 @@ def add_feature(h5_file, feature_name, array):
     pass
 
 
-def add_metadata(h5_file, start, stop, duration, ds_name, filename=None):
-    meta = pd.read_hdf(h5_file, "metadata")
-    layout = pd.read_hdf(h5_file, "layouts/" + ds_name)
-    info = pd.read_hdf(h5_file, "info")
-    new = Metadata.from_start_stop([start], [stop], [duration])
+def _add_regions(h5_file, start, stop, duration, filename=None):
+    """
+    adds a row of regions to ``db.regions``
+    """
+    meta = pd.read_hdf(h5_file, "regions")
+    new = Regions.from_start_stop([start], [stop], [duration])
     filename = datetime.now() if filename is None else filename
     new["name"] = filename
     new_meta = pd.concat((meta, new), axis=0, ignore_index=True)
-    new_layout = pd.concat((layout, new), axis=0, ignore_index=True)
-    new_info = info.iloc[info.index.max()]
-    new_info.loc[:] = ("added", filename,
-                       new_info[(ds_name, "dtype")],
-                       (new.duration.item(), *new_info[(ds_name, "shape")][1:]),
-                       "xMb")
-    new_info = info.append(new_info, ignore_index=True)
-    new_meta.to_hdf(h5_file, "metadata")
-    new_layout.to_hdf(h5_file, "layouts/" + ds_name)
-    new_info.to_hdf(h5_file, "info")
-    return new_info
-
-
-def add_data(h5_file, ds_name, array, filename):
+    new_meta.to_hdf(h5_file, "regions")
+    return new_meta
+
+
+def _add_data(h5_file, ds_name, array, filename):
+    """
+    EXPERIMENTAL!! adds data at the end of a ``Dataset`` in a .h5 file.
+
+    Parameters
+    ----------
+    h5_file : str
+        path to the file
+    ds_name : str
+        name of the ``h5py.Dataset`` to which ``array`` will be added
+    array : np.ndarray
+        the array to add
+    filename : str
+        a name for the array in ``db.regions``
+
+    Returns
+    -------
+    arr : np.ndarray
+        the array as it has been stored
+    """
     N = array.shape[0]
     with h5py.File(h5_file, "r+") as f:
         if f[ds_name].shape[1:] != array.shape[1:]:
@@ -226,5 +281,5 @@ def add_data(h5_file, ds_name, array, filename):
         f[ds_name].resize(M + N, axis=0)
         f[ds_name][-N:] = array
         rv = f[ds_name][-N:]
-    add_metadata(h5_file, M, N + M, N, ds_name, filename)
+    _add_regions(h5_file, M, N + M, N, filename)
     return rv
\ No newline at end of file
diff --git a/mimikit/data/data_object.py b/mimikit/data/data_object.py
index 5020794..ae94dad 100644
--- a/mimikit/data/data_object.py
+++ b/mimikit/data/data_object.py
@@ -6,6 +6,19 @@
 
 
 def map_if_multi(attr):
+    """
+    decorator for applying a method recursively to ``getattr(self, attr)`` if ``getattr(self, attr)`` is a tuple.
+
+    Parameters
+    ----------
+    attr : str
+        name of the ``attribute`` to which the method will apply and that can possibly be a ``tuple``.
+
+    Returns
+    -------
+    decorator : function
+        the wrapping function for decorating the method
+    """
     def decorator(fn):
         def wrapper(self, *arg, **kwargs):
             if isinstance(getattr(self, attr), tuple):
@@ -23,6 +36,30 @@ def wrapper(self, *arg, **kwargs):
 
 
 class DataObject(TorchDataset):
+    """
+    the goal of the ``DataObject`` class is to provide a single universal object for making smart ``Datasets``.
+
+    In its constructor, ``DataObject`` accepts ``numpy`` arrays, ``python iterables``, ``torch.Tensors`` and ``Datasets``,
+    ``FeatureProxies`` from ``mimikit``'s ``Databases`` and generally, any object that implements ``__len__`` and ``__getitem__``
+    or ``__iter__`` - i.e. objects that behaves like 'map-style' and 'iter-style' ``Datasets``.
+
+    Furthermore, ``DataObject`` handles nested ``tuples`` of such "data objects" as if they were single containers with
+    multiple fields, thus allowing you to easily define the shape and content of your batches at the very top-level.
+    For example, if you have raw images and labels, you just need
+    ```
+    dataloader = Dataloader(DataObject((images, labels)), ....)
+    ```
+    in order to later do
+    ```
+    images, labels = batch
+    ```
+    in your ``training_step()``.
+
+    ``DataObject`` instances also implement some of the most common ``numpy`` and ``torch`` properties like ``shape``
+    or ``dtype`` and some handy methods like ``to_tensor()``, ``to(device)`` or ``split([...])`` thus removing a lot of
+    boilerplate imports and code-lines that deal with low-level data management while maintaining access to the parts
+    of the low-level interfaces that are always needed when defining models and trainings.
+    """
 
     @property
     def data(self):
diff --git a/mimikit/data/factory.py b/mimikit/data/factory.py
index 5f1db11..856a167 100644
--- a/mimikit/data/factory.py
+++ b/mimikit/data/factory.py
@@ -7,12 +7,10 @@
 import warnings
 
 from .api import Database
-from .metadata import Metadata
+from .regions import Regions
 from .transforms import default_extract_func
 
 warnings.filterwarnings("ignore", message="PySoundFile failed.")
-warnings.filterwarnings("ignore", message="PerformanceWarning")
-warnings.filterwarnings("ignore", message="Creating an ndarray from ragged nested sequences")
 
 
 class AudioFileWalker:
@@ -78,13 +76,13 @@ def walk_root(root):
 
     @staticmethod
     def is_audio_file(filename):
-        # filter out hidden files (isn't cross-platform, but, it's a start!...)
-        if filename.startswith("."):
+        # filter out hidden files
+        if os.path.split(filename.strip("/"))[-1].startswith("."):
             return False
         return os.path.splitext(filename)[-1].strip(".") in AudioFileWalker.AUDIO_EXTENSIONS
 
 
-def sizeof_fmt(num, suffix='b'):
+def _sizeof_fmt(num, suffix='b'):
     for unit in ['', 'k', 'M', 'G', 'T', 'P', 'E', 'Z']:
         if abs(num) < 1024.0:
             return "%3.1f%s%s" % (num, unit, suffix)
@@ -92,25 +90,11 @@ def sizeof_fmt(num, suffix='b'):
     return "%.1f%s%s" % (num, 'Yi', suffix)
 
 
-def _empty_info(features_names):
-    tuples = [("directory", ""), ("name", ""),
-              *[t for feat in features_names for t in [(feat, "dtype"), (feat, "shape"), (feat, "size")]
-                if feat != "metadata"]
-              ]
-    idx = pd.MultiIndex.from_tuples(tuples)
-    return pd.DataFrame([], columns=idx)
-
-
-def split_path(path):
-    prefix, file_name = os.path.split(path)
-    return prefix, file_name
-
-
 # Core function
 
-def file_to_db(abs_path, extract_func=default_extract_func, mode="w"):
+def file_to_db(abs_path, extract_func=default_extract_func, output_path=None, mode="w"):
     """
-    apply `extract_func` to `abs_path` and write the result in a .h5 file
+    apply ``extract_func`` to ``abs_path`` and write the result in a .h5 file.
 
     Parameters
     ----------
@@ -118,155 +102,170 @@ def file_to_db(abs_path, extract_func=default_extract_func, mode="w"):
         path to the file to be extracted
     extract_func : function
         the function to use for the extraction - should take exactly one argument
+    output_path : str or None
+        name of the created .h5 file if not None, else the name of the created .h5 file
+        will be of the form : "<abs_path_without_extension>.h5"
     mode : str
         the mode to use when opening the .h5 file. default is "w".
 
     Returns
     -------
-    created : str
-        the name of the created .h5 file
+    created_file, infos : str, dict
+        the name of the created .h5 file and a ``dict`` with the keys ``"dtype"`` and ``"shape"``
     """
-    print("making db for %s" % abs_path)
-    tmp_db = os.path.splitext(abs_path)[0] + ".h5"
+    print("making .h5 for %s" % abs_path)
+    if output_path is None:
+        output_path = os.path.splitext(abs_path)[0] + ".h5"
+    else:
+        if output_path[-3:] != ".h5":
+            output_path += ".h5"
     rv = extract_func(abs_path)
-    if "metadata" not in rv:
-        raise ValueError("Expected `extract_func` to return a ('metadata', Metadata) item. Found none")
-    info = _empty_info(rv.keys())
-    info.loc[0, [("directory", ""), ("name", "")]] = split_path(abs_path)
-    f = h5py.File(tmp_db, mode)
+    if "regions" not in rv:
+        raise ValueError("Expected `extract_func` to return a ('regions', Regions) item. Found none")
+    info = {}
+    f = h5py.File(output_path, mode)
     for name, (attrs, data) in rv.items():
         if issubclass(type(data), np.ndarray):
             ds = f.create_dataset(name=name, shape=data.shape, data=data)
             ds.attrs.update(attrs)
-            info.at[0, [(name, "dtype"), (name, "shape"), (name, "size")]] = tuple([ds.dtype, ds.shape, sizeof_fmt(data.nbytes)])
+            info[name] = {"dtype": ds.dtype, "shape": ds.shape}
         elif issubclass(type(data), pd.DataFrame):
             f.close()
-            pd.DataFrame(data).to_hdf(tmp_db, name, "r+")
-            f = h5py.File(tmp_db, "r+")
+            pd.DataFrame(data).to_hdf(output_path, name, "r+")
+            f = h5py.File(output_path, "r+")
     f.flush()
     f.close()
-    if "info" in f.keys():
-        prior = pd.read_hdf(tmp_db, "info", "r")
-        info = pd.concat((prior, info.iloc[:, 2:]), axis=1)
-    info.to_hdf(tmp_db, "info", "r+")
-    return tmp_db
+    return output_path, info
 
 
 # Multiprocessing routine
 
-def make_db_for_each_file(file_walker,
-                          extract_func=default_extract_func,
-                          n_cores=cpu_count()):
-    args = [(file, extract_func) for file in file_walker]
-    with Pool(n_cores) as p:
-        tmp_dbs = p.starmap(file_to_db, args)
-    return tmp_dbs
-
-
-# Aggregating sub-tasks
-
-def collect_infos(tmp_dbs):
-    infos = []
-    for db in tmp_dbs:
-        infos += [Database(db).info]
-    return pd.concat(infos, ignore_index=True)
-
+def _make_db_for_each_file(file_walker,
+                           extract_func=default_extract_func,
+                           n_cores=cpu_count()):
+    """
+    apply ``extract_func`` to the files found by ``file_walker``
 
-def collect_metadatas(tmp_dbs):
-    metadatas = []
-    offset = 0
-    for db in tmp_dbs:
-        meta = Database(db).metadata
-        meta.loc[:, ("start", "stop")] = meta.loc[:, ("start", "stop")].values + offset
-        meta.loc[:, "name"] = os.path.splitext(db)[0]
-        metadatas += [meta]
-        offset = meta.last_stop
-    return pd.DataFrame(pd.concat(metadatas, ignore_index=True))
+    If file_walker found more than ``n_cores`` files, a multiprocessing ``Pool`` is used to speed up the process.
 
+    Parameters
+    ----------
+    file_walker : iterable of str
+        collection of files to be processed
+    extract_func : function
+        the function to apply to each file. Must take only one argument (the path to the file)
+    n_cores : int, optional
+        the number of cores used in the ``Pool``, default is the number of available cores on the system.
 
-def zip_prev_next(iterable):
-    return zip(iterable[:-1], iterable[1:])
+    Returns
+    -------
+    temp_dbs : list of tuples
+        each tuple in the list is of the form ``("<created_file>.h5", dict(feature_name=dict(dtype=..., shape=...), ...))``
+    """
+    # add ".tmp_" prefix to the output_paths
+    args = [(file, extract_func, os.path.join(os.path.split(file)[0], ".tmp_" + os.path.split(file)[1]))
+            for file in file_walker]
+    if len(args) > n_cores:
+        with Pool(n_cores) as p:
+            tmp_dbs_infos = p.starmap(file_to_db, args)
+    else:
+        tmp_dbs_infos = [file_to_db(*arg) for arg in args]
+    return tmp_dbs_infos
+
+
+def _aggregate_db_infos(infos):
+    """
+    helper to reduce shapes and dtypes of several files and to concatenate their regions
 
+    Parameters
+    ----------
+    infos : list of tuples
+        see the returned value of ``_make_db_for_each_file`` for the expected type
 
-def ds_definitions_from_infos(infos):
-    tb = infos.iloc[:, 2:].T
-    paths = [os.path.join(*parts) for parts in infos.iloc[:, :2].values]
-    # change the paths' extensions
-    paths = [os.path.splitext(path)[0] + ".h5" for path in paths]
-    features = set(tb.index.get_level_values(0))
+    Returns
+    -------
+    ds_definitions : dict
+        the keys are the name of the H5Datasets (features)
+        the values are dictionaries with keys ("shape", "dtype", "regions") and corresponding values
+    """
+    paths = [x[0] for x in infos]
+    features = set([feature for db in infos for feature in db[1].keys()])
     ds_definitions = {}
     for f in features:
-        dtype = tb.loc[(f, "dtype"), :].unique().item()
-        shapes = tb.loc[(f, "shape"), :].values
+        dtype = set([db[1][f]["dtype"] for db in infos])
+        assert len(dtype) == 1, "aggregated features must be of a unique dtype"
+        dtype = dtype.pop()
+        shapes = [db[1][f]["shape"] for db in infos]
         dims = shapes[0][1:]
         assert all(shp[1:] == dims for shp in
                    shapes[1:]), "all features should have the same dimensions but for the first axis"
-        layout = Metadata.from_duration([s[0] for s in shapes])
-        ds_shape = (layout.last_stop, *dims)
-        layout.index = paths
-        ds_definitions[f] = {"shape": ds_shape, "dtype": dtype, "layout": layout}
+        # collect the regions for the files
+        regions = Regions.from_duration([s[0] for s in shapes])
+        ds_shape = (regions.last_stop, *dims)
+        regions.index = paths
+        ds_definitions[f] = {"shape": ds_shape, "dtype": dtype, "meta_regions": regions}
     return ds_definitions
 
 
-def create_datasets_from_defs(target, defs, mode="w"):
-    f = h5py.File(target, mode)
-    for name, params in defs.items():
-        f.create_dataset(name, shape=params["shape"], dtype=params["dtype"],
-                         chunks=True, maxshape=(None, *params["shape"][1:]))
-        layout = params["layout"]
-        layout.reset_index(drop=False, inplace=True)
-        layout = layout.rename(columns={"index": "name"})
-        f.flush()
-        f.close()
-        pd.DataFrame(layout).to_hdf(target, "layouts/" + name, "r+", format="table")
-        f = h5py.File(target, "r+")
-    f.flush()
-    f.close()
-    return
+# Aggregating function
+
+def _aggregate_dbs(target, tmp_dbs_infos, mode="w"):
+    """
+    copy a list of .h5 files as defined in ``tmp_dbs_infos`` into ``target``
+
+    Parameters
+    ----------
+    target : str
+        name for the target file
+    tmp_dbs_infos : list
+        see the returned value of ``_make_db_for_each_file`` for the expected type
+    mode : str
+        the mode with which ``target`` is opened the first time. Must be 'w' if ``target`` doesn't exist.
+    """
+    features_infos = _aggregate_db_infos(tmp_dbs_infos)
 
+    # open the file & create the master datasets to host the features
+    with h5py.File(target, mode) as f:
+        # add the list of features
+        f.attrs["features"] = list(features_infos.keys())
+        for name, params in features_infos.items():
+            f.create_dataset(name, shape=params["shape"], dtype=params["dtype"],
+                             chunks=True, maxshape=(None, *params["shape"][1:]))
 
-def make_integration_args(target):
+    # prepare args for copying the tmp_files into the appropriate regions
     args = []
-    with h5py.File(target, "r") as f:
-        for feature in f["layouts"].keys():
-            df = Metadata(pd.read_hdf(target, "layouts/" + feature))
-            args += [(target, source, feature, indices) for source, indices in
-                     zip(df.name, df.slices(time_axis=0))]
-    return args
-
-
-def integrate(target, source, key, indices):
-    with h5py.File(source, "r") as src:
-        data = src[key][()]
-        attrs = {k: v for k, v in src[key].attrs.items()}
-    with h5py.File(target, "r+") as trgt:
-        trgt[key][indices] = data
-        trgt[key].attrs.update(attrs)
-    return
-
-
-# Aggregating function and main client
-
-def aggregate_dbs(target, dbs, mode="w", remove_sources=False):
-    infos = collect_infos(dbs)
-    metadata = collect_metadatas(dbs)
-    definitions = ds_definitions_from_infos(infos)
-    create_datasets_from_defs(target, definitions, mode)
-    args = make_integration_args(target)
-    for arg in args: integrate(*arg)
-    if remove_sources:
-        for src in dbs:
-            if src != target:
-                os.remove(src)
-    infos = infos.astype(object)
-    infos.to_hdf(target, "info", "r+")
-    metadata.to_hdf(target, "metadata", "r+")
+    for feature, info in features_infos.items():
+        # collect the arguments for copying each file into its regions in the master datasets/features
+        args += [(source, feature, indices) for source, indices in
+                 zip(info["meta_regions"].index, info["meta_regions"].slices(time_axis=0))]
+
+    # copy the data
+    intra_regions = None
+    for source, key, indices in args:
+        with h5py.File(source, "r") as src:
+            data = src[key][()]
+            attrs = {k: v for k, v in src[key].attrs.items()}
+        # concat the regions :
+        regions = pd.read_hdf(source, key="regions", mode="r")
+        regions.loc[:, ("start", "stop")] += indices[0].start
+        # remove ".tmp_" from the source's name
+        regions.loc[:, "name"] = "".join(source.split(".tmp_"))
+        intra_regions = pd.concat(([intra_regions] if intra_regions is not None else []) + [regions])
+        with h5py.File(target, "r+") as trgt:
+            trgt[key][indices] = data
+            trgt[key].attrs.update(attrs)
+    pd.DataFrame(intra_regions).to_hdf(target, "regions", mode="r+")
+
+    # remove the temp_dbs
+    for src, _ in tmp_dbs_infos:
+        if src != target:
+            os.remove(src)
 
 
 def make_root_db(db_name, roots='./', files=None, extract_func=default_extract_func,
                  n_cores=cpu_count()):
     """
-    extract and aggregate several files into a .h5 Database
+    extract and aggregate several files into a single .h5 Database
 
     Parameters
     ----------
@@ -280,7 +279,7 @@ def make_root_db(db_name, roots='./', files=None, extract_func=default_extract_f
         default is `None`
     extract_func : function, optional
         the function to use for the extraction - should take exactly one argument.
-        the default transforms the file to a stft with n_fft=2048 and hop_length=512.
+        the default transforms the file to the stft with n_fft=2048 and hop_length=512.
     n_cores : int, optional
         the number of cores to use to parallelize the extraction process.
         default is the number of available cores on the system.
@@ -291,6 +290,6 @@ def make_root_db(db_name, roots='./', files=None, extract_func=default_extract_f
         the created db
     """
     walker = AudioFileWalker(roots, files)
-    dbs = make_db_for_each_file(walker, extract_func, n_cores)
-    aggregate_dbs(db_name, dbs, "w", True)
+    tmp_dbs_infos = _make_db_for_each_file(walker, extract_func, n_cores)
+    _aggregate_dbs(db_name, tmp_dbs_infos, "w")
     return Database(db_name)
diff --git a/mimikit/data/metadata.py b/mimikit/data/regions.py
similarity index 89%
rename from mimikit/data/metadata.py
rename to mimikit/data/regions.py
index f8a8ad0..7149a2e 100644
--- a/mimikit/data/metadata.py
+++ b/mimikit/data/regions.py
@@ -8,11 +8,11 @@ def ssts(item, axis=0):
     return slices
 
 
-class Metadata(pd.DataFrame):
+class Regions(pd.DataFrame):
     """
     subclass of ``pandas.DataFrame`` to store a slice structure. This is handy when concatenating arrays
 
-    Instances of ``Metadata`` created by the constructors mentioned below will automatically have
+    Instances of ``Regions`` created by the constructors mentioned below will automatically have
     three columns : "start", "stop", "duration".
 
     :no-inherited-members:
@@ -20,7 +20,7 @@ class Metadata(pd.DataFrame):
 
     @property
     def _constructor(self):
-        return Metadata
+        return Regions
 
     @staticmethod
     def _validate(obj):
@@ -85,7 +85,7 @@ def is_contiguous(self):
 
     @staticmethod
     def from_start_stop(starts, stops, durations):
-        return Metadata((starts, stops, durations), index=["start", "stop", "duration"]).T
+        return Regions((starts, stops, durations), index=["start", "stop", "duration"]).T
 
     @staticmethod
     def from_stop(stop):
@@ -100,25 +100,25 @@ def from_stop(stop):
         else:
             starts = np.r_[0, stop[:-1]]
         durations = stop - starts
-        return Metadata.from_start_stop(starts, stop, durations)
+        return Regions.from_start_stop(starts, stop, durations)
 
     @staticmethod
     def from_duration(duration):
         stops = np.cumsum(duration)
         starts = np.r_[0, stops[:-1]]
-        return Metadata.from_start_stop(starts, stops, duration)
+        return Regions.from_start_stop(starts, stops, duration)
 
     @staticmethod
     def from_data(sequence, time_axis=1):
         duration = np.array([x.shape[time_axis] for x in sequence])
-        return Metadata.from_duration(duration)
+        return Regions.from_duration(duration)
 
     @staticmethod
     def from_frame_definition(total_duration, frame_length, stride=1, butlasts=0):
         starts = np.arange(total_duration - frame_length - butlasts + 1, step=stride)
         durations = frame_length + np.zeros_like(starts, dtype=np.int)
         stops = starts + durations
-        return Metadata.from_start_stop(starts, stops, durations)
+        return Regions.from_start_stop(starts, stops, durations)
 
 
 @pd.api.extensions.register_dataframe_accessor("soft_q")
diff --git a/mimikit/data/transforms.py b/mimikit/data/transforms.py
index 1f15db9..5f66ee0 100644
--- a/mimikit/data/transforms.py
+++ b/mimikit/data/transforms.py
@@ -1,5 +1,5 @@
 import librosa
-from .metadata import Metadata
+from .regions import Regions
 
 N_FFT = 2048
 HOP_LENGTH = 512
@@ -17,9 +17,9 @@ def stft(file, n_fft=N_FFT, hop_length=HOP_LENGTH, sr=SR):
 def file_to_fft(abs_path, n_fft=N_FFT, hop_length=HOP_LENGTH, sr=SR):
     fft, params = stft(abs_path, n_fft, hop_length, sr)
     fft = abs(fft)
-    metadata = Metadata.from_duration([fft.shape[1]])
+    regions = Regions.from_duration([fft.shape[1]])
     params.update(dict(time_axis=0))
-    return dict(fft=(params, fft.T), metadata=({}, metadata))
+    return dict(fft=(params, fft.T), regions=({}, regions))
 
 
 def mu_law_compress(file, mu=MU, sr=SR):
@@ -32,8 +32,8 @@ def mu_law_compress(file, mu=MU, sr=SR):
 
 def file_to_qx(abs_path, mu=MU, sr=SR):
     qx, params = mu_law_compress(abs_path, mu, sr)
-    metadata = Metadata.from_duration([qx.shape[0]])
-    return dict(qx=(params, qx.reshape(-1, 1)), metadata=({}, metadata))
+    regions = Regions.from_duration([qx.shape[0]])
+    return dict(qx=(params, qx.reshape(-1, 1)), regions=({}, regions))
 
 
 default_extract_func = file_to_fft
diff --git a/mimikit/extract/segment.py b/mimikit/extract/segment.py
index 170ef05..4890b8b 100644
--- a/mimikit/extract/segment.py
+++ b/mimikit/extract/segment.py
@@ -2,7 +2,7 @@
 from librosa.segment import recurrence_matrix
 from librosa.util import localmax
 from scipy.ndimage import convolve
-from ..data.metadata import Metadata
+from ..data.regions import Regions
 
 
 def from_recurrence_matrix(X,
@@ -32,4 +32,4 @@ def checker(N):
     mx = mx * (np.diff(mx, append=R.shape[0]) >= min_dur)
     mx = mx[mx > 0]
     stops = np.r_[mx, R.shape[0]]
-    return Metadata.from_stop(stops)
\ No newline at end of file
+    return Regions.from_stop(stops)
\ No newline at end of file
diff --git a/mimikit/freqnet/base.py b/mimikit/freqnet/base.py
index dea3d38..aef1155 100644
--- a/mimikit/freqnet/base.py
+++ b/mimikit/freqnet/base.py
@@ -18,7 +18,13 @@ def step(self, epoch=None):
 
 
 class FreqOptim:
+    """
+    simple class to modularize the optimization setup used in a ``FreqNetModel``.
 
+    Here the method ``configure_optimizers()`` returns an ``Adam`` optimizer and a slightly modified
+    ``OneCycleLR`` scheduler (the only modification being that it won't raise a ``ValueError`` if you use for more
+    steps than what it expects and will instead restarts its cycle)
+    """
     def __init__(self,
                  model,
                  max_lr=1e-3,
@@ -52,12 +58,19 @@ def configure_optimizers(self):
 
 
 class FreqData(LightningDataModule):
+    """
+    boilerplate subclass of ``pytorch_lightning.LightningDataModule`` to handle the data of a ``FreqNetModel``.
+
+    the data is passed through ``data_object``, wrapped in a ``ShiftedSeqsPair`` in ``prepare_data()`` and
+    in ``setup("fit")`` it is moved to the gpu if ``in_mem_data`` is ``True`` and split into train, val and test sets
+    according to ``splits``
+    """
     def __init__(self,
                  model,
                  data_object=None,
                  input_seq_length=64,
                  batch_size=64,
-                 to_gpu=True,
+                 in_mem_data=True,
                  splits=None,
                  **loader_kwargs,
                  ):
@@ -66,7 +79,7 @@ def __init__(self,
         self.ds = DataObject(data_object)
         self.input_seq_length = input_seq_length
         self.batch_size = batch_size
-        self.to_gpu = to_gpu
+        self.in_mem_data = in_mem_data
         self.splits = splits
         loader_kwargs.setdefault("drop_last", False)
         loader_kwargs.setdefault("shuffle", True)
@@ -83,7 +96,7 @@ def prepare_data(self, *args, **kwargs):
 
     def setup(self, stage=None):
         if stage == "fit":
-            if self.to_gpu and torch.cuda.is_available():
+            if self.in_mem_data and torch.cuda.is_available():
                 self.ds.to_tensor()
                 self.ds.to("cuda")
             if self.splits is None:
@@ -128,7 +141,9 @@ class FreqNetModel(MMKHooks,
                    LoggingHooks,
                    LightningModule,
                    ABC):
-
+    """
+    base class for all ``FreqNets`` that handles optim, data and a few handy methods like ``generate()``
+    """
     @property
     def data(self):
         """short-hand to quickly access the data object passed to the constructor"""
@@ -138,7 +153,7 @@ def __init__(self,
                  data_object=None,
                  input_seq_length=64,
                  batch_size=64,
-                 to_gpu=True,
+                 in_mem_data=True,
                  splits=[.8, .2],
                  max_lr=1e-3,
                  betas=(.9, .9),
@@ -152,7 +167,7 @@ def __init__(self,
         if data_object is not None and not isinstance(data_object, str):
             self.input_dim = data_object.shape[-1]
             self.datamodule = FreqData(self, data_object, input_seq_length, batch_size,
-                                       to_gpu, splits, **loaders_kwargs)
+                                       in_mem_data, splits, **loaders_kwargs)
         else:
             raise ValueError("Please pass a valid data_object to instantiate a FreqNetModel."
                              " If you are loading from a checkpoint, you can do so by modifying your method call to :\n"
diff --git a/mimikit/freqnet/freaksnet.py b/mimikit/freqnet/freaksnet.py
index fe73e2c..08adfa5 100644
--- a/mimikit/freqnet/freaksnet.py
+++ b/mimikit/freqnet/freaksnet.py
@@ -17,6 +17,7 @@ def __init__(self,
                  model_dim=512,
                  groups=((1, 1, 1),),
                  n_layers=(3,),
+                 kernel_size=((2, 2, 2),),
                  strict=((False, True, False),),
                  accum_outputs=((-1, 1, -1),),
                  concat_outputs=((1, -1, 1),),
@@ -32,6 +33,7 @@ def __init__(self,
         self.groups = groups
         self.n_layers = n_layers
         self.strict = strict
+        self.kernel_size = kernel_size
         self.accum_outputs = accum_outputs
         self.concat_outputs = concat_outputs
         self.pad_input = pad_input
diff --git a/mimikit/freqnet/freq_layer.py b/mimikit/freqnet/freq_layer.py
index baa8676..34309f9 100644
--- a/mimikit/freqnet/freq_layer.py
+++ b/mimikit/freqnet/freq_layer.py
@@ -33,7 +33,6 @@ def concat(x, y, shift=1):
 
 
 class FreqLayer(nn.Module):
-    kernel_size = 2
     stride = 1
     bias = True
     sides = {1: 1, -1: -1, 0: 0, "left": 1, "right": -1}
@@ -44,6 +43,7 @@ def __init__(self,
                  layer_dim=512,
                  groups=1,
                  strict=False,
+                 kernel_size=2,
                  accum_outputs=0,
                  concat_outputs=0,
                  pad_input=0,
@@ -57,6 +57,7 @@ def __init__(self,
         self.layer_dim = layer_dim
         self.groups = groups
         self.strict = strict
+        self.kernel_size = kernel_size
         if pad_input != 0 and concat_outputs != 0:
             raise ValueError("pad_input and concat_outputs can not be both non zero")
         # accum and concat have opposite signs :
@@ -114,7 +115,7 @@ def forward(self, x, skip=None):
 
     @property
     def dilation(self):
-        return 2 ** self.layer_index
+        return self.kernel_size ** self.layer_index
 
     @property
     def padding(self):
@@ -128,7 +129,7 @@ def receptive_field(self):
         """
         amount of inputs necessary for 1 output (this is independent of the shift!)
         """
-        return 2 ** (self.layer_index + 1)
+        return self.kernel_size ** (self.layer_index + 1)
 
     def shift(self):
         """total shift at this layer wrt. to the beginning of its block"""
@@ -136,7 +137,7 @@ def shift(self):
 
     def rel_shift(self):
         """relative shift at this layer wrt. to the previous layer"""
-        return (int(self.strict) + self.receptive_field() // 2) if self.pad_input == 0 else int(self.strict)
+        return (int(self.strict) + (self.kernel_size - 1) * self.dilation) if self.pad_input == 0 else int(self.strict)
 
     def output_length(self, input_length):
         if abs(self.concat_outputs):
diff --git a/mimikit/freqnet/freqnet.py b/mimikit/freqnet/freqnet.py
index 9e3001f..78f39d4 100644
--- a/mimikit/freqnet/freqnet.py
+++ b/mimikit/freqnet/freqnet.py
@@ -7,7 +7,7 @@
 
 
 class FreqNet(FreqNetModel):
-    LAYER_KWARGS = ["groups", "strict", "accum_outputs", "concat_outputs",
+    LAYER_KWARGS = ["groups", "strict", "accum_outputs", "concat_outputs", "kernel_size",
                     "pad_input", "learn_padding", "with_skip_conv", "with_residual_conv"]
 
     def __init__(self,
@@ -15,6 +15,7 @@ def __init__(self,
                  model_dim=512,
                  groups=1,
                  n_layers=(2,),
+                 kernel_size=2,
                  strict=False,
                  accum_outputs=0,
                  concat_outputs=0,
@@ -28,6 +29,7 @@ def __init__(self,
         self.model_dim = model_dim
         self.groups = groups
         self.n_layers = n_layers
+        self.kernel_size = kernel_size
         self.strict = strict
         self.accum_outputs = accum_outputs
         self.concat_outputs = concat_outputs
diff --git a/mimikit/freqnet/hierarchical_kernel.py b/mimikit/freqnet/hierarchical_kernel.py
index c887559..5c24f6d 100644
--- a/mimikit/freqnet/hierarchical_kernel.py
+++ b/mimikit/freqnet/hierarchical_kernel.py
@@ -19,6 +19,7 @@ def __init__(self,
                  groups=1,
                  n_layers=2,
                  strict=False,
+                 kernel_size=2,
                  accum_outputs=0,
                  concat_outputs=0,
                  pad_input=0,
@@ -26,7 +27,7 @@ def __init__(self,
                  with_skip_conv=True,
                  with_residual_conv=True,
                  **data_optim_kwargs):
-        super(HKFreqNet, self).__init__(loss_fn, model_dim, groups, (n_layers,), strict,
+        super(HKFreqNet, self).__init__(loss_fn, model_dim, groups, (n_layers,), kernel_size, strict,
                                         accum_outputs, concat_outputs, pad_input, learn_padding,
                                         with_skip_conv, with_residual_conv,
                                         **data_optim_kwargs)
diff --git a/mimikit/freqnet/layer_wise_loss.py b/mimikit/freqnet/layer_wise_loss.py
index 84c48d9..3d6d785 100644
--- a/mimikit/freqnet/layer_wise_loss.py
+++ b/mimikit/freqnet/layer_wise_loss.py
@@ -21,6 +21,10 @@ def forward(self, x):
             return tuple(outputs)
         # else redirect to a special method for inference
         return self.infer_layer_wise(outputs)
+    
+    def training_step(self, batch, batch_idx):
+        batch = batch[0], batch[1:]
+        return super(LayerWiseLossFreqNet, self).training_step(batch, batch_idx)
 
     def infer_layer_wise(self, outputs):
         """
diff --git a/mimikit/freqnet/wavenet.py b/mimikit/freqnet/wavenet.py
index d8947eb..9ede450 100644
--- a/mimikit/freqnet/wavenet.py
+++ b/mimikit/freqnet/wavenet.py
@@ -2,12 +2,11 @@
 import torch.nn.functional as F
 from functools import partial
 from multiprocessing import cpu_count
-import os
 
 from . import FreqNet
 from .modules.io import PermuteTF
 from ..data import Database, make_root_db
-from ..kit import NeptuneConnector
+from ..connectors.neptune import NeptuneConnector
 from ..data.transforms import file_to_qx
 
 
diff --git a/mimikit/kit/ds_wrappers.py b/mimikit/kit/ds_wrappers.py
index a0516b9..1fa2f73 100644
--- a/mimikit/kit/ds_wrappers.py
+++ b/mimikit/kit/ds_wrappers.py
@@ -40,9 +40,7 @@ def __init__(self, input_length, targets, stride=1):
         @param targets: [(shift, length)... ]
         @param stride:
         """
-        self.input_length = input_length
-        self.shifts = list(zip(*targets))[0]
-        self.lengths = list(zip(*targets))[1]
+        self.shifts, self.lengths = list(zip(*[(0, input_length), *targets]))
         self.stride = stride
         self.N = None
 
@@ -52,16 +50,13 @@ def __call__(self, dataset: DataObject):
         return super(ShiftedSeqsPair, self).__call__(dataset)
 
     def __len__(self):
-        ln = (self.N - max(self.lengths) - max(self.shifts) + 1) // self.stride
+        ln = (self.N - max(shift + ln for shift, ln in zip(self.shifts, self.lengths)) + 1) // self.stride
         return ln
 
     def __getitem__(self, item):
         i = item * self.stride
-        input_slice = slice(i, i + self.input_length)
-        target_slices = [slice(i + shift, i + shift + length)
+        items_slices = [slice(i + shift, i + shift + length)
                          for shift, length in zip(self.shifts, self.lengths)]
-        inputs = self.data[input_slice]
-        targets = tuple(self.data[idx] for idx in target_slices)
-        return inputs, targets[0] if len(targets) == 1 else targets
+        return tuple(self.data[idx] for idx in items_slices)
 
 


Test patch:
diff --git a/tests/test_data.py b/tests/test_data.py
index 74e5a22..43950f2 100644
--- a/tests/test_data.py
+++ b/tests/test_data.py
@@ -38,9 +38,9 @@ def test_audio_file_walker(audio_tree):
 
 def test_default_extract_func(audio_tree):
     dic = default_extract_func(audio_tree + "/dir2/test1.wav")
-    # we must have metadata
-    assert "metadata" in dic and isinstance(dic["metadata"][1], pd.DataFrame), dic["metadata"]
-    # metadata + feature
+    # we must have regions
+    assert "regions" in dic and isinstance(dic["regions"][1], pd.DataFrame), dic["regions"]
+    # regions + feature
     assert len(dic) == 2
 
 
@@ -50,27 +50,25 @@ def test_file_to_db(audio_tree):
 
     def faulty_extract_func(path):
         dic = default_extract_func(path)
-        dic.pop("metadata")
+        dic.pop("regions")
         return dic
 
-    with pytest.raises(ValueError, match=r".*metadata.*"):
+    with pytest.raises(ValueError, match=r".*regions.*"):
         file_to_db(file_path, faulty_extract_func)
 
     file_to_db(file_path)
     with h5py.File(file_path.split(".")[0] + ".h5", "r") as f:
-        assert f["metadata"] is not None
-        assert f["info"] is not None
+        assert f["regions"] is not None
 
 
 def test_make_root_db_and_Database(audio_tree):
     make_root_db(audio_tree + "/test_db.h5", roots=audio_tree, extract_func=default_extract_func, n_cores=1)
 
     db = Database(audio_tree + "/test_db.h5")
-    assert db.metadata is not pd.DataFrame()
-    assert db.layout_for("fft") is not pd.DataFrame()
-    assert len(db.metadata) == 4, len(db.metadata)
+    assert db.regions is not pd.DataFrame()
+    assert len(db.regions) == 4, len(db.regions)
     assert np.any(db.fft[:4] != 0), db.fft[:4]
-    assert isinstance(db.fft.get(db.metadata.iloc[:1]), np.ndarray), db.fft.get(db.metadata.iloc[:1])
+    assert isinstance(db.fft.get(db.regions.iloc[:1]), np.ndarray), db.fft.get(db.regions.iloc[:1])
 
     # test Dataset Integration
     ds = DataObject(db.fft)
diff --git a/tests/test_freqnets.py b/tests/test_freqnets.py
index c40e616..f72b9be 100644
--- a/tests/test_freqnets.py
+++ b/tests/test_freqnets.py
@@ -54,7 +54,7 @@ def targets_shifts_and_lengths(self, input_length):
     base = TestModel(data_object=np.random.randn(32, 1025),
                      input_seq_length=16,
                      batch_size=10,
-                     to_gpu=False,
+                     in_mem_data=False,
                      splits=[1.])
 
     assert isinstance(base.datamodule, FreqData)
diff --git a/tests/test_kit/test_ds_wrappers.py b/tests/test_kit/test_ds_wrappers.py
index 49c2697..893ab90 100644
--- a/tests/test_kit/test_ds_wrappers.py
+++ b/tests/test_kit/test_ds_wrappers.py
@@ -48,11 +48,11 @@ def test_overrides_getitem(self):
         # returns input and target
         assert len(returned) == 2, returned
         # returns correct input shape
-        assert returned[0].shape == (self.wrapper.input_length, self.data.shape[-1]), returned[0].shape
+        assert returned[0].shape == (self.wrapper.lengths[0], self.data.shape[-1]), returned[0].shape
         # returns correct target shape
-        assert returned[1].shape == (self.wrapper.lengths[0], self.data.shape[-1]), returned[1].shape
+        assert returned[1].shape == (self.wrapper.lengths[1], self.data.shape[-1]), returned[1].shape
         # shifts sequences correctly
-        assert np.all(returned[0][self.wrapper.shifts] == returned[1][0]), (returned[0][self.wrapper.shifts],
+        assert np.all(returned[0][self.wrapper.shifts[1]] == returned[1][0]), (returned[0][self.wrapper.shifts[1]],
                                                                             returned[1][0])
 
     def test_overrides_len(self):