Repository: Strong-AI-Lab/emotion (language: python)

Issue:
Add documentation
Need API and CLI documentation, with examples.


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index e9acf6c..c954a21 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,6 +6,7 @@
 !/datasets/CMU-MOSEI/splits.py
 
 # Misc. project related
+.env
 
 # Python related
 .venv*
@@ -14,11 +15,13 @@
 /build
 /dist
 /docs/_build
+/docs/api/generated
 __pycache__
 .mypy_cache
 .coverage*
 *.lcov
 .virtual_documents/
+.hypothesis/
 
 # NumPy
 *.npy
@@ -31,6 +34,7 @@ __pycache__
 # PyTorch
 *.pt
 whisper_checkpoints/
+pretrained_models/
 
 
 # Misc general
diff --git a/.readthedocs.yaml b/.readthedocs.yaml
new file mode 100644
index 0000000..4cc4aa0
--- /dev/null
+++ b/.readthedocs.yaml
@@ -0,0 +1,16 @@
+# Read the Docs configuration file
+version: 2
+build:
+  os: ubuntu-22.04
+  tools:
+    python: "3.10"
+sphinx:
+  builder: html
+  configuration: docs/conf.py
+python:
+  install:
+  - requirements: requirements.txt
+  - requirements: docs/requirements.txt
+submodules:
+  include: []
+  recursive: false
diff --git a/datasets/README.md b/datasets/README.md
index 445f129..d0e377f 100644
--- a/datasets/README.md
+++ b/datasets/README.md
@@ -53,12 +53,10 @@ For most datasets this involves resampling to 16 kHz 16-bit WAV, and
 outputting annotations. The scripts for some datasets (e.g. IEMOCAP,
 MSP-IMPROV) print additional statistics about perceptual ratings.
 
-All scripts will output any annotation files in CSV format, as well as
-print the emotion label distribution and any inter-rater agreement
-metrics that can be calculated (e.g. Krippendorf's alpha, Fleiss'
-kappa).
+All scripts will output annotations in CSV format, and any inter-rater
+agreement metrics that can be calculated (e.g. Krippendorf's alpha,
+Fleiss' kappa) from ratings data.
 
-Once these steps are done, you can use the feature extraction scripts.
 
 ## File lists and subsets
 A list of file paths to the resampled audio clips will be created for
diff --git a/docs/Makefile b/docs/Makefile
new file mode 100644
index 0000000..68b6cad
--- /dev/null
+++ b/docs/Makefile
@@ -0,0 +1,22 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = .
+BUILDDIR      = _build
+export CUDA_VISIBLE_DEVICES=
+export TF_CPP_MIN_LOG_LEVEL=1
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
diff --git a/docs/_templates/autosummary/class.rst b/docs/_templates/autosummary/class.rst
new file mode 100644
index 0000000..03c0355
--- /dev/null
+++ b/docs/_templates/autosummary/class.rst
@@ -0,0 +1,39 @@
+{{ fullname | escape | underline}}
+
+.. currentmodule:: {{ module }}
+
+.. autoclass:: {{ objname }}
+    :members:
+    :show-inheritance:
+
+    {% block methods %}
+    .. automethod:: __init__
+
+    {% if methods %}
+    .. rubric:: {{ _('Methods') }}
+
+    .. autosummary::
+    {% for item in methods if item not in inherited_members %}
+        ~{{ name }}.{{ item }}
+    {%- endfor %}
+
+    {% if inherited_members %}
+    .. rubric:: {{ _('Inherited Methods') }}
+    .. autosummary::
+    {% for item in methods if item in inherited_members %}
+        ~{{ name }}.{{ item }}
+    {%- endfor %}
+    {% endif %}
+    {% endif %}
+    {% endblock %}
+
+    {% block attributes %}
+    {% if attributes %}
+    .. rubric:: {{ _('Attributes') }}
+
+    .. autosummary::
+    {% for item in attributes %}
+        ~{{ name }}.{{ item }}
+    {%- endfor %}
+    {% endif %}
+    {% endblock %}
diff --git a/docs/_templates/autosummary/module.rst b/docs/_templates/autosummary/module.rst
new file mode 100644
index 0000000..df86398
--- /dev/null
+++ b/docs/_templates/autosummary/module.rst
@@ -0,0 +1,3 @@
+{{ fullname | escape | underline}}
+
+.. automodule:: {{ fullname }}
diff --git a/docs/api/cli.rst b/docs/api/cli.rst
new file mode 100644
index 0000000..1e9ffbd
--- /dev/null
+++ b/docs/api/cli.rst
@@ -0,0 +1,1 @@
+.. automodule:: ertk.cli
diff --git a/docs/api/config.rst b/docs/api/config.rst
new file mode 100644
index 0000000..da65b9c
--- /dev/null
+++ b/docs/api/config.rst
@@ -0,0 +1,1 @@
+.. automodule:: ertk.config
diff --git a/docs/api/dataset.rst b/docs/api/dataset.rst
new file mode 100644
index 0000000..1bc293e
--- /dev/null
+++ b/docs/api/dataset.rst
@@ -0,0 +1,1 @@
+.. automodule:: ertk.dataset
diff --git a/docs/api/preprocessing.rst b/docs/api/preprocessing.rst
new file mode 100644
index 0000000..7f6e885
--- /dev/null
+++ b/docs/api/preprocessing.rst
@@ -0,0 +1,1 @@
+.. automodule:: ertk.preprocessing
diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst
new file mode 100644
index 0000000..a96b840
--- /dev/null
+++ b/docs/api/pytorch.rst
@@ -0,0 +1,1 @@
+.. automodule:: ertk.pytorch
diff --git a/docs/api/sklearn.rst b/docs/api/sklearn.rst
new file mode 100644
index 0000000..18ad782
--- /dev/null
+++ b/docs/api/sklearn.rst
@@ -0,0 +1,1 @@
+.. automodule:: ertk.sklearn
diff --git a/docs/api/tensorflow.rst b/docs/api/tensorflow.rst
new file mode 100644
index 0000000..fc2aa2a
--- /dev/null
+++ b/docs/api/tensorflow.rst
@@ -0,0 +1,1 @@
+.. automodule:: ertk.tensorflow
diff --git a/docs/api/train.rst b/docs/api/train.rst
new file mode 100644
index 0000000..fb14d19
--- /dev/null
+++ b/docs/api/train.rst
@@ -0,0 +1,26 @@
+Training
+========
+
+
+Backends
+--------
+.. toctree::
+   :maxdepth: 1
+
+   sklearn
+   pytorch
+   tensorflow
+
+Classes and utilities
+---------------------
+.. automodule:: ertk.train
+
+
+Classification
+--------------
+.. automodule:: ertk.classification
+
+
+Transforms
+----------
+.. automodule:: ertk.transform
diff --git a/docs/api/utils.rst b/docs/api/utils.rst
new file mode 100644
index 0000000..c5fb7b6
--- /dev/null
+++ b/docs/api/utils.rst
@@ -0,0 +1,1 @@
+.. automodule:: ertk.utils
diff --git a/docs/cli_ref.rst b/docs/cli_ref.rst
new file mode 100644
index 0000000..ff0ae05
--- /dev/null
+++ b/docs/cli_ref.rst
@@ -0,0 +1,92 @@
+Command-line programs reference
+===============================
+
+A number of command-line tools are installed along with the ERTK Python
+library. Three executables are installed:
+
+* `ertk-cli`_
+* `ertk-dataset`_
+* `ertk-util`_
+
+These each have a number of subcommands available, shown with the
+``--help`` option.
+
+ertk-cli
+--------
+
+Experiments
+^^^^^^^^^^^
+.. click:: ertk.cli.cli.exp:main
+    :prog: ertk-cli exp
+
+.. click:: ertk.cli.cli.exp2:main
+    :prog: ertk-cli exp2
+
+Model training and inference
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+.. click:: ertk.cli.cli.classify:main
+    :prog: ertk-cli classify
+
+.. click:: ertk.cli.cli.train:main
+    :prog: ertk-cli train
+
+
+ertk-dataset
+------------
+
+Info
+^^^^
+.. click:: ertk.cli.dataset.annotation_stats:main
+    :prog: ertk-dataset annotation
+
+.. click:: ertk.cli.dataset.info:main
+    :prog: ertk-dataset info
+
+
+Running processors and feature extractors
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+.. click:: ertk.cli.dataset.process:main
+    :prog: ertk-dataset process
+
+
+Features
+^^^^^^^^
+.. click:: ertk.cli.dataset.combine:main
+    :prog: ertk-dataset combine
+
+.. click:: ertk.cli.dataset.convert:main
+    :prog: ertk-dataset convert
+
+.. click:: ertk.cli.dataset.remove_instances:main
+    :prog: ertk-dataset remove_instances
+
+.. click:: ertk.cli.dataset.vis:main
+    :prog: ertk-dataset vis
+
+
+ertk-util
+---------
+
+Run parallel CPU or GPU jobs with a simple command:
+
+.. click:: ertk.cli.util.parallel_jobs:main
+    :prog: ertk-util parallel_jobs
+
+Other utilities
+^^^^^^^^^^^^^^^
+.. click:: ertk.cli.util.create_cv_dirs:main
+    :prog: ertk-util create_cv_dirs
+
+.. click:: ertk.cli.util.grid_to_conf:main
+    :prog: ertk-util grid_to_conf
+
+.. click:: ertk.cli.util.names_to_filenames:main
+    :prog: ertk-util names_to_filenames
+
+CHAT and ELAN file formats
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+.. click:: ertk.cli.util.split_chat:main
+    :prog: ertk-util split_chat
+
+.. click:: ertk.cli.util.split_elan:main
+    :prog: ertk-util split_elan
diff --git a/docs/conf.py b/docs/conf.py
new file mode 100644
index 0000000..c54308f
--- /dev/null
+++ b/docs/conf.py
@@ -0,0 +1,74 @@
+# Configuration file for the Sphinx documentation builder.
+#
+# This file only contains a selection of the most common options. For a full
+# list see the documentation:
+# https://www.sphinx-doc.org/en/master/usage/configuration.html
+
+# -- Path setup --------------------------------------------------------------
+
+import sys
+from pathlib import Path
+
+import sphinx
+
+import ertk
+
+_docs_dir = Path(__file__).parent
+_src_dir = _docs_dir.parent / "src"
+sys.path.insert(0, str(_src_dir.resolve()))
+
+
+# -- Project information -----------------------------------------------------
+
+project = "ERTK"
+copyright = "2019, Aaron Keesing"
+author = "Aaron Keesing"
+version = ertk.__version__
+
+
+# -- General configuration ---------------------------------------------------
+extensions = [
+    "sphinx.ext.autodoc",
+    "sphinx.ext.autosummary",
+    "sphinx.ext.intersphinx",
+    "numpydoc",
+    "sphinx_rtd_theme",
+    "sphinx_click",
+]
+
+templates_path = ["_templates"]
+exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
+
+
+# -- Options for HTML output -------------------------------------------------
+html_theme = "sphinx_rtd_theme"
+html_theme_options = {
+    "collapse_navigation": False,
+    "display_version": True,
+    "navigation_depth": 4,
+    "sticky_navigation": True,
+}
+html_static_path = ["_static"]
+
+
+primary_domain = "py"
+default_role = "py:obj"
+
+add_module_names = True
+autodoc_default_options = {
+    "member-order": "alphabetical",
+}
+
+autosummary_generate = True
+
+intersphinx_mapping = {
+    "python": ("https://docs.python.org/3", None),
+    "numpy": ("https://numpy.org/doc/stable/", None),
+    "np": ("https://numpy.org/doc/stable/", None),
+    "sklearn": ("https://scikit-learn.org/stable/", None),
+    "torch": ("https://pytorch.org/docs/stable/", None),
+}
+
+numpydoc_show_class_members = False
+
+add_module_names = True
diff --git a/docs/datasets.rst b/docs/datasets.rst
new file mode 100644
index 0000000..66a080c
--- /dev/null
+++ b/docs/datasets.rst
@@ -0,0 +1,70 @@
+Emotion Datasets
+================
+
+A large number of emotional speech datasets have processing scripts and
+metadata available in the ERTK repository. Currently the scripts only
+deal with audio data, but in future we intend to support video data.
+
+
+Currently supported datasets:
+-----------------------------
+* `AESDD
+  <http://m3c.web.auth.gr/research/aesdd-speech-emotion-recognition/>`_
+* `ASED <https://github.com/wgwangang/ASED_V1>`_
+* `BAVED
+  <https://www.kaggle.com/a13x10/basic-arabic-vocal-emotions-dataset>`_
+* `CaFE <https://zenodo.org/record/1478765>`_
+* `CMU-MOSEI
+  <http://multicomp.cs.cmu.edu/resources/cmu-mosei-dataset/>`_
+* `CREMA-D <https://github.com/CheyneyComputerScience/CREMA-D>`_
+* `DEMoS <https://zenodo.org/record/2544829>`_
+* `EESC
+  <https://metashare.ut.ee/repository/browse/estonian-emotional-speech-corpus/4d42d7a8463411e2a6e4005056b40024a19021a316b54b7fb707757d43d1a889/>`_
+* `EMO-DB <http://emodb.bilderbar.info/>`_
+* `EmoFilm <https://zenodo.org/record/1326428>`_
+* `EmoryNLP <https://github.com/declare-lab/MELD/>`_
+* `EmoV-DB <https://github.com/numediart/EmoV-DB/>`_
+* `EMOVO <http://voice.fub.it/activities/corpora/emovo/index.html>`_
+* `ESD <https://github.com/HLTSingapore/Emotional-Speech-Data/>`_
+* `eNTERFACE <http://www.enterface.net/results/>`_
+* `IEMOCAP <https://sail.usc.edu/iemocap/>`_
+* `JL-corpus <https://www.kaggle.com/tli725/jl-corpus>`_
+* `MELD <https://github.com/declare-lab/MELD/>`_
+* `MESD <https://data.mendeley.com/datasets/cy34mh68j9/3>`_
+* `MESS <https://zenodo.org/record/3813437>`_
+* `MLendSND
+  <https://www.kaggle.com/datasets/jesusrequena/mlend-spoken-numerals>`_
+* `MSP-IMPROV
+  <https://ecs.utdallas.edu/research/researchlabs/msp-lab/MSP-Improv.html>`_
+* `MSP-Podcast
+  <https://ecs.utdallas.edu/research/researchlabs/msp-lab/MSP-Podcast.html>`_
+* `Oréau <https://zenodo.org/record/4405783>`_
+* `Portuguese <https://link.springer.com/article/10.3758/BRM.42.1.74>`_
+* `RAVDESS <https://zenodo.org/record/1188976>`_
+* `SAVEE <http://kahlan.eps.surrey.ac.uk/savee/>`_
+* `SEMAINE <https://semaine-db.eu/>`_
+* `ShEMO <https://github.com/mansourehk/ShEMO>`_
+* `SmartKom
+  <https://clarin.phonetik.uni-muenchen.de/BASRepository/index.php>`_
+* `SUBESCO <https://zenodo.org/record/4526477>`_
+* `TESS <https://tspace.library.utoronto.ca/handle/1807/24487/>`_
+* `URDU <https://github.com/siddiquelatif/URDU-Dataset/>`_
+* `VENEC (Public subset)
+  <https://www.nature.com/articles/s41562-019-0533-6>`_. The full
+  dataset is available on request.
+* `VIVAE <https://zenodo.org/record/4066235>`_
+
+
+Standardising datasets
+----------------------
+Each dataset has a subdirectory in the ``datasets`` directory. This
+subdirectory contains a ``process.py`` script, which takes a path to the
+original dataset, and converts it to a simple standard format, consisting
+of annotation CSVs and audio files. The script also generates a
+``corpus.yaml`` file, which contains metadata about the dataset.
+
+You can run the ``process.py`` script for a dataset, such as EMO-DB, as
+follows::
+
+    cd datasets/EMO-DB
+    python process.py /path/to/EMO-DB
diff --git a/docs/experiments.rst b/docs/experiments.rst
new file mode 100644
index 0000000..6a566c5
--- /dev/null
+++ b/docs/experiments.rst
@@ -0,0 +1,75 @@
+Running experiments
+===================
+
+Experiments can be specified with a YAML configuration file. The
+configuration file specifies the dataset(s) to load, the model to use,
+the training parameters, and the evaluation parameters.
+
+The `ertk-cli exp2 <api/cli>`_ command is designed to run experiments
+from a configuration file and takes optional parameter overrides on the
+command line::
+
+    ertk-cli exp2 <config_file> <overrides>
+
+Example with EMO-DB
+-------------------
+Suppose we want to run a quick experiment with `EMO-DB <datasets>`_.
+After downloading the dataset, standardising using the build-in script,
+and extracting features, we can run an experiment.
+
+The following is an example of such a configuration file:
+
+.. code-block:: yaml
+    :caption: exp_config.yaml
+
+    # Configuration file for an experiment
+    # ------------------------------------
+    # This file specifies the dataset(s) to load, the model to use,
+    # the training parameters, and the evaluation parameters.
+    name: EMO-DB
+    data:
+        datasets:
+            EMO-DB:
+                path: datasets/EMO-DB/corpus.yaml
+        features: eGeMAPS
+
+    # Model parameters
+    model:
+        # Multinomial logistic regression
+        type: sk/lr
+        config:
+            penalty: l2
+            class_weight: null
+            solver: lbfgs
+            max_iter: 100
+            C: 1
+            multi_class: multinomial
+
+    # Evaluation parameters
+    eval:
+        # Speaker-independent cross-validtion
+        cv:
+            part: speaker
+            kfold: -1
+        # Two-fold speaker-independent inner cross-validation for
+        # hyperparameter tuning
+        inner_kfold: 2
+        inner_part: speaker
+
+    results: results.csv
+
+    # Training parameters
+    train:
+        # Run cross-validation folds in parallel
+        n_jobs: -1
+
+If we use this configuration file, we can run the experiment with the
+following command::
+
+    ertk-cli exp2 exp_config.yaml
+
+Overrides can be specified on the command line. For example, to run
+the experiment with a different cost parameter C, we can use the
+following command::
+
+    ertk-cli exp2 exp_config.yaml train.config.C=0.1
diff --git a/docs/genindex.rst b/docs/genindex.rst
new file mode 100644
index 0000000..44f5c34
--- /dev/null
+++ b/docs/genindex.rst
@@ -0,0 +1,3 @@
+Index
+=====
+:ref:`genindex`
diff --git a/docs/index.rst b/docs/index.rst
new file mode 100644
index 0000000..bec6b6f
--- /dev/null
+++ b/docs/index.rst
@@ -0,0 +1,65 @@
+ERTK documentation
+==================
+
+ERTK: Emotion Recognition ToolKit is a python library for processing
+emotional speech datasets, extracting common features, and
+training/testing models. It is built upon scikit-learn, with
+infrastructure for PyTorch and TensorFlow models.
+
+Citing
+------
+If you use ERTK in your research, please cite our paper [1]_:
+
+.. code-block:: bibtex
+
+    @inproceedings{keesingAcousticFeaturesNeural2021,
+        title = {Acoustic Features and Neural Representations for Categorical Rmotion Recognition from Speech},
+        booktitle = {Interspeech 2021},
+        author = {Keesing, Aaron and Koh, Yun Sing and Witbrock, Michael},
+        year = {2021},
+        month = aug,
+        pages = {3415--3419},
+        publisher = {{ISCA}},
+        doi = {10.21437/Interspeech.2021-2217}
+    }
+
+.. [1] A. Keesing, Y. S. Koh, and M. Witbrock, "Acoustic Features and Neural
+       Representations for Categorical Emotion Recognition from Speech," in
+       Interspeech 2021, ISCA, Aug. 2021, pp. 3415-3419. doi:
+       10.21437/Interspeech.2021-2217
+
+
+.. toctree::
+    :maxdepth: 1
+    :caption: Contents
+
+    installation
+    datasets
+    processors
+    experiments
+
+
+.. toctree::
+    :maxdepth: 1
+    :caption: CLI reference
+
+    cli_ref
+
+
+.. toctree::
+    :maxdepth: 1
+    :caption: API reference
+
+    api/dataset
+    api/preprocessing
+    api/train
+    api/config
+    api/cli
+    api/utils
+
+
+.. toctree::
+    :caption: Index
+    :maxdepth: 1
+
+    genindex
diff --git a/docs/installation.rst b/docs/installation.rst
new file mode 100644
index 0000000..95f8ac5
--- /dev/null
+++ b/docs/installation.rst
@@ -0,0 +1,51 @@
+Installation
+============
+
+ERTK
+----
+
+ERTK can be installed from `PyPI <https://pypi.org/project/ertk/>`_
+using pip::
+
+    pip install ertk
+
+Alternatively, you can clone the repository and install it manually::
+
+    git clone https://github.com/Strong-AI-Lab/emotion.git
+    cd emotion
+    pip install .
+
+To install in editable mode, use the ``-e`` flag::
+
+    pip install -e .
+
+You can also install directly from GitHub with pip::
+
+    pip install git+https://github.com/Strong-AI-Lab/emotion.git
+
+
+Optional third-party libraries
+------------------------------
+
+ERTK can use the following third-party libraries if they are installed:
+
+* `HuggingFace Transformers
+  <https://github.com/huggingface/transformers>`_ for using
+  transformer-based models for feature extraction, ASR, etc.
+* `fairseq <https://github.com/facebookresearch/fairseq>`_ for
+  using fairseq models for feature extraction, etc.
+* `phonemizer <https://github.com/bootphon/phonemizer>`_ for
+  phonemizing text.
+* `openSMILE <https://github.com/audeering/opensmile-python>`_ for
+  feature extraction.
+* `resampy <https://github.com/bmcfee/resampy>`_ for audio resampling.
+* `SpeechBrain <https://github.com/speechbrain/speechbrain>`_ for
+  SpeechBrain models.
+* `Encodec <https://github.com/facebookresearch/encodec>`_ for feature
+  extraction.
+* `TensorFlow-Slim <https://github.com/google-research/tf-slim>`_ to run
+  the audioset models VGGish and YAMNet.
+
+To install ERTK with these dependencies, use the ``[all]`` extra::
+
+    pip install ertk[all]
diff --git a/docs/make.bat b/docs/make.bat
new file mode 100644
index 0000000..153be5e
--- /dev/null
+++ b/docs/make.bat
@@ -0,0 +1,35 @@
+@ECHO OFF
+
+pushd %~dp0
+
+REM Command file for Sphinx documentation
+
+if "%SPHINXBUILD%" == "" (
+	set SPHINXBUILD=sphinx-build
+)
+set SOURCEDIR=.
+set BUILDDIR=_build
+
+if "%1" == "" goto help
+
+%SPHINXBUILD% >NUL 2>NUL
+if errorlevel 9009 (
+	echo.
+	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
+	echo.installed, then set the SPHINXBUILD environment variable to point
+	echo.to the full path of the 'sphinx-build' executable. Alternatively you
+	echo.may add the Sphinx directory to PATH.
+	echo.
+	echo.If you don't have Sphinx installed, grab it from
+	echo.https://www.sphinx-doc.org/
+	exit /b 1
+)
+
+%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+goto end
+
+:help
+%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+
+:end
+popd
diff --git a/docs/processors.rst b/docs/processors.rst
new file mode 100644
index 0000000..68f9015
--- /dev/null
+++ b/docs/processors.rst
@@ -0,0 +1,31 @@
+Instance processing
+===================
+
+A number of built-in processors are available to process the data,
+depending on the installed packages. Additionally, custom processors
+can be created through a plugin interface.
+
+The full list of available processors can be found in the `preprocessor
+documentation <api/preprocessing>`_.
+
+
+Examples
+--------
+Suppose we want to extract eGeMAPS features for a dataset such as
+EMO-DB. We can use the `ertk-cli process` command to process the
+instances in EMO-DB and extract eGeMAPS features using OpenSMILE::
+
+    ertk-cli process \
+        --processor opensmile \
+        datasets/EMO-DB/files_all.txt \
+        datasets/EMO-DB/features/eGeMAPS.nc \
+        opensmile_config=eGeMAPS
+
+We can run an automatic speech recogniser from HuggingFace using the
+``hugginface`` processor::
+
+    ertk-cli process \
+        --processor huggingface \
+        datasets/EMO-DB/files_all.txt \
+        datasets/EMO-DB/transcripts.txt \
+        model=facebook/wav2vec2-base-960h
diff --git a/docs/requirements.txt b/docs/requirements.txt
new file mode 100644
index 0000000..a22e5f6
--- /dev/null
+++ b/docs/requirements.txt
@@ -0,0 +1,4 @@
+numpydoc==1.5.0
+sphinx==6.1.3
+sphinx-click==4.4.0
+sphinx-rtd-theme==1.2.1
diff --git a/src/ertk/classification.py b/src/ertk/classification.py
index cbe47c5..e9b73c5 100644
--- a/src/ertk/classification.py
+++ b/src/ertk/classification.py
@@ -1,3 +1,21 @@
+"""
+Classification functions.
+
+This module contains functions for performing classification tasks.
+
+.. autosummary::
+    :toctree: generated/
+
+    binary_accuracy_score
+    standard_class_scoring
+    cross_validate
+    dataset_cross_validation
+    dataset_train_val_test
+    get_balanced_class_weights
+    get_balanced_sample_weights
+    class_ratings_to_probs
+"""
+
 import logging
 import time
 from functools import partial
@@ -18,6 +36,18 @@
 
 from ertk.dataset import Dataset
 from ertk.train import ExperimentResult, get_cv_splitter, scores_to_df
+from ertk.utils import ScoreFunction
+
+__all__ = [
+    "binary_accuracy_score",
+    "standard_class_scoring",
+    "cross_validate",
+    "dataset_cross_validation",
+    "dataset_train_val_test",
+    "get_balanced_class_weights",
+    "get_balanced_sample_weights",
+    "class_ratings_to_probs",
+]
 
 
 def binary_accuracy_score(
@@ -121,10 +151,90 @@ def standard_class_scoring(classes: Sequence[str]):
     return scoring
 
 
+def cross_validate(
+    clf_lib: str,
+    clf,
+    x: np.ndarray,
+    y: Optional[np.ndarray] = None,
+    *,
+    groups: Optional[np.ndarray] = None,
+    cv: BaseCrossValidator = None,
+    scoring: Union[
+        str, List[str], Dict[str, ScoreFunction], Callable[..., float]
+    ] = "accuracy",
+    verbose: int = 0,
+    n_jobs: int = 1,
+    fit_params: Dict[str, Any] = {},
+):
+    """Cross validate a classifier.
+
+    Parameters
+    ----------
+    clf_lib: str
+        Classifier library to use. Must be one of "sk", "tf", or "pt".
+    clf:
+        Classifier to cross validate.
+    x: np.ndarray
+        Input data.
+    y: np.ndarray, optional
+        Target labels.
+    groups: np.ndarray, optional
+        Groups to split data into for cross validation.
+    cv: BaseCrossValidator, optional
+        Cross validation strategy. If None, use StratifiedKFold.
+    scoring: str, list, dict, or callable, optional
+        Scoring metric(s) to use. If a string, must be a valid
+        scikit-learn scorer.
+    verbose: int, optional
+        Verbosity level.
+    n_jobs: int, optional
+        Number of jobs to run in parallel. Only used for scikit-learn.
+    fit_params: dict, optional
+        Parameters to pass to the classifier's fit method.
+
+    Returns
+    -------
+    result: ExperimentResult
+        Cross validation results.
+    """
+    cross_validate_fn: Callable[..., ExperimentResult]
+    if clf_lib == "sk":
+        from ertk.sklearn.classification import sk_cross_validate
+
+        # We have to set n_jobs here because using a
+        # `with joblib.Parallel...` clause doesn't work properly
+        cross_validate_fn = partial(sk_cross_validate, n_jobs=n_jobs)
+    elif clf_lib == "tf":
+        from ertk.tensorflow.classification import tf_cross_validate
+
+        n_jobs = 1
+        cross_validate_fn = tf_cross_validate
+    elif clf_lib == "pt":
+        from ertk.pytorch.classification import pt_cross_validate
+
+        n_jobs = 1
+        cross_validate_fn = pt_cross_validate
+    else:
+        raise ValueError(f"Unknown classifier type: {clf_lib}")
+
+    result = cross_validate_fn(
+        clf,
+        x,
+        y,
+        groups=groups,
+        cv=cv,
+        scoring=scoring,
+        verbose=verbose,
+        n_jobs=n_jobs,
+        fit_params=fit_params,
+    )
+    return result
+
+
 def dataset_cross_validation(
     clf,
     dataset: Dataset,
-    clf_lib: Optional[str] = None,
+    clf_lib: str,
     partition: Optional[str] = None,
     label: str = "label",
     cv: Union[BaseCrossValidator, int] = 10,
@@ -176,34 +286,18 @@ def dataset_cross_validation(
     if scoring is None:
         scoring = standard_class_scoring(dataset.get_group_names(label))
 
-    cross_validate_fn: Callable[..., ExperimentResult]
-    if clf_lib == "sk":
-        from ertk.sklearn.classification import sk_cross_validate
-
-        # We have to set n_jobs here because using a
-        # `with joblib.Parallel...` clause doesn't work properly
-        cross_validate_fn = partial(sk_cross_validate, n_jobs=n_jobs)
-    elif clf_lib == "tf":
-        from ertk.tensorflow.classification import tf_cross_validate
-
-        n_jobs = 1
-        cross_validate_fn = tf_cross_validate
-    elif clf_lib == "pt":
-        from ertk.pytorch.classification import pt_cross_validate
-
-        n_jobs = 1
-        cross_validate_fn = pt_cross_validate
-
-    start_time = time.perf_counter()
     logging.info(f"Starting cross-validation with n_jobs={n_jobs}")
-    result = cross_validate_fn(
+    start_time = time.perf_counter()
+    result = cross_validate(
+        clf_lib,
         clf,
         dataset.x,
         dataset.get_group_indices(label),
+        groups=groups,
         cv=cv,
         scoring=scoring,
-        groups=groups,
         verbose=verbose,
+        n_jobs=n_jobs,
         fit_params=fit_params,
     )
     total_time = time.perf_counter() - start_time
@@ -224,7 +318,7 @@ def dataset_train_val_test(
     test_idx: Union[Sequence[int], np.ndarray, None] = None,
     label: str = "label",
     clf_lib: Optional[str] = None,
-    sample_weight: np.ndarray = None,
+    sample_weight: Optional[np.ndarray] = None,
     verbose: int = 0,
     scoring=None,
     fit_params: Dict[str, Any] = {},
diff --git a/src/ertk/cli/__init__.py b/src/ertk/cli/__init__.py
index e69de29..c0b6614 100644
--- a/src/ertk/cli/__init__.py
+++ b/src/ertk/cli/__init__.py
@@ -0,0 +1,13 @@
+"""
+Command-line programs
+=====================
+
+This module contains the command-line programs for ERTK.
+
+.. autosummary::
+    :toctree: generated/
+
+    cli
+    dataset
+    util
+"""
diff --git a/src/ertk/cli/cli/__init__.py b/src/ertk/cli/cli/__init__.py
index c06da36..2bbf05c 100644
--- a/src/ertk/cli/cli/__init__.py
+++ b/src/ertk/cli/cli/__init__.py
@@ -1,3 +1,25 @@
+"""
+``ertk-cli`` scripts
+====================
+
+This module contains the ``ertk-cli`` scripts.
+
+.. autosummary::
+    :toctree:
+
+    cli
+
+
+Commands
+--------
+.. autosummary::
+    :toctree:
+
+    classify_main
+    exp_main
+    exp2_main
+"""
+
 import click
 
 from .classify import main as classify_main
diff --git a/src/ertk/cli/dataset/__init__.py b/src/ertk/cli/dataset/__init__.py
index dab1d07..2711200 100644
--- a/src/ertk/cli/dataset/__init__.py
+++ b/src/ertk/cli/dataset/__init__.py
@@ -1,3 +1,30 @@
+"""
+``ertk-dataset`` scripts
+========================
+
+This module contains the ``ertk-dataset`` scripts.
+
+.. autosummary::
+    :toctree:
+
+    cli
+
+
+Commands
+--------
+.. autosummary::
+    :toctree:
+
+    annotation_stats_main
+    combine_main
+    convert_main
+    filter_clips_main
+    info_main
+    process_main
+    remove_instances_main
+    vis_main
+"""
+
 import click
 
 from .annotation_stats import main as annotation_stats_main
diff --git a/src/ertk/cli/util/__init__.py b/src/ertk/cli/util/__init__.py
index a11c94c..e2f84d5 100644
--- a/src/ertk/cli/util/__init__.py
+++ b/src/ertk/cli/util/__init__.py
@@ -1,3 +1,29 @@
+"""
+``ertk-util`` scripts
+=====================
+
+This module contains the ``ertk-util`` scripts.
+
+.. autosummary::
+    :toctree:
+
+    cli
+
+
+Commands
+--------
+.. autosummary::
+    :toctree:
+
+    convert_audeep_main
+    create_cv_dirs_main
+    grid_to_conf_main
+    names_to_filenames_main
+    parallel_jobs_main
+    split_chat_main
+    split_elan_main
+"""
+
 import click
 
 from .convert_audeep import main as convert_audeep_main
diff --git a/src/ertk/config.py b/src/ertk/config.py
index 35161b6..483d8c7 100644
--- a/src/ertk/config.py
+++ b/src/ertk/config.py
@@ -1,3 +1,27 @@
+"""
+Configuration utilities
+=======================
+
+This module provides utilities for working with configuration
+dataclasses and YAML files.
+
+
+Base classes
+------------
+.. autosummary::
+    :toctree: generated/
+
+    ERTKConfig
+
+
+Functions
+---------
+.. autosummary::
+    :toctree: generated/
+
+    get_arg_mapping
+"""
+
 from abc import ABC
 from collections import defaultdict
 from dataclasses import dataclass
@@ -9,6 +33,9 @@
 
 from ertk.utils import PathOrStr
 
+__all__ = ["ERTKConfig", "get_arg_mapping"]
+
+
 T = TypeVar("T", bound="ERTKConfig")
 
 
@@ -29,11 +56,30 @@ class ERTKConfig(ABC):
     """Base class for ERTK configuration dataclasses."""
 
     def to_dictconfig(self) -> DictConfig:
+        """Convert config to DictConfig.
+
+        Returns
+        -------
+        omegaconf.DictConfig
+            The structured config which behaves like the corresponding
+            dataclass.
+        """
         return DictConfig(self)
 
     @classmethod
     def from_config(cls: Type[T], config: Any) -> T:
-        """Create config from any omegaconf config."""
+        """Create config object from any compatible config.
+
+        Parameters
+        ----------
+        config: omegaconf.DictConfig or dict
+            The object to create the config from.
+
+        Returns
+        -------
+        ERKConfig
+            The resulting config.
+        """
         schema = OmegaConf.structured(cls)
         return cast(T, OmegaConf.merge(schema, config))
 
@@ -53,23 +99,48 @@ def from_file(
 
         Returns
         -------
-        omegaconf.DictConfig
-            The structured config which behaves like the corresponding
-            config dataclass.
+        ERKConfig
+            The resulting config.
         """
         config = OmegaConf.load(Path(path))
         if override is not None:
             config = OmegaConf.merge(config, OmegaConf.from_dotlist(override))
         return cls.from_config(config)
 
-    def to_string(self):
+    def to_string(self) -> str:
+        """Generate YAML string representation of config.
+
+        Returns
+        -------
+        str
+            YAML string representation of config.
+        """
         return OmegaConf.to_yaml(self)
 
-    def to_file(self, path: PathOrStr):
+    def to_file(self, path: PathOrStr) -> None:
+        """Write config to YAML file.
+
+        Parameters
+        ----------
+        path: os.Pathlike or str
+            The path to YAML file to write config to.
+        """
         with open(path, "w") as fid:
             fid.write(OmegaConf.to_yaml(self))
 
     def merge_with_args(self: T, args: Optional[List[str]] = None) -> T:
+        """Merge config with command-line arguments.
+
+        Parameters
+        ----------
+        args: list of str, optional
+            Argument overrides in the form of key=value pairs.
+
+        Returns
+        -------
+        ERKConfig
+            The resulting config.
+        """
         return cast(T, OmegaConf.merge(self, OmegaConf.from_cli(args)))
 
 
diff --git a/src/ertk/dataset/__init__.py b/src/ertk/dataset/__init__.py
index 401049e..047608a 100644
--- a/src/ertk/dataset/__init__.py
+++ b/src/ertk/dataset/__init__.py
@@ -1,3 +1,59 @@
+"""
+Dataset classes and functions
+=============================
+
+This module contains classes and functions for working with datasets.
+
+
+Features
+--------
+.. autosummary::
+    :toctree: generated/
+
+    read_features
+    read_features_iterable
+    register_format
+    write_features
+
+
+Annotations
+-----------
+.. autosummary::
+    :toctree: generated/
+
+    read_annotations
+    write_annotations
+
+
+Dataset loading and manipulation
+--------------------------------
+.. autosummary::
+    :toctree: generated/
+
+    CombinedDataset
+    CorpusInfo
+    DataLoadConfig
+    DataSelector
+    Dataset
+    DatasetConfig
+    MapGroups
+    RemoveGroups
+    SubsetInfo
+    load_datasets_config
+    load_multiple
+
+
+Utilities
+---------
+.. autosummary::
+    :toctree: generated/
+
+    get_audio_paths
+    resample_audio
+    resample_rename_clips
+    write_filelist
+"""
+
 from .annotation import read_annotations, write_annotations
 from .dataset import (
     CombinedDataset,
diff --git a/src/ertk/dataset/annotation.py b/src/ertk/dataset/annotation.py
index 1329892..b6f78a7 100644
--- a/src/ertk/dataset/annotation.py
+++ b/src/ertk/dataset/annotation.py
@@ -1,3 +1,5 @@
+"""Functions for reading and writing annotations."""
+
 from typing import Dict, Mapping, Optional, Type, Union
 
 import pandas as pd
@@ -5,6 +7,8 @@
 
 from ertk.utils import PathOrStr
 
+__all__ = ["read_annotations", "write_annotations"]
+
 
 def read_annotations(
     filename: PathOrStr, dtype: Optional[Union[Type, Literal["category"]]] = None
diff --git a/src/ertk/dataset/dataset.py b/src/ertk/dataset/dataset.py
index 1338d27..51ae992 100644
--- a/src/ertk/dataset/dataset.py
+++ b/src/ertk/dataset/dataset.py
@@ -1,3 +1,5 @@
+"""Dataset classes and related functions."""
+
 import copy
 import logging
 import os
@@ -37,6 +39,20 @@
 from ertk.transform import group_transform
 from ertk.utils import PathOrStr, clip_arrays, frame_arrays, pad_arrays, transpose_time
 
+__all__ = [
+    "Dataset",
+    "CombinedDataset",
+    "DatasetConfig",
+    "DataLoadConfig",
+    "DataSelector",
+    "SubsetInfo",
+    "MapGroups",
+    "RemoveGroups",
+    "CorpusInfo",
+    "load_multiple",
+    "load_datasets_config",
+]
+
 logger = logging.getLogger(__name__)
 
 
diff --git a/src/ertk/dataset/features.py b/src/ertk/dataset/features.py
index 4d7fd4a..8502224 100644
--- a/src/ertk/dataset/features.py
+++ b/src/ertk/dataset/features.py
@@ -1,3 +1,5 @@
+"""Dataset features backend."""
+
 import csv
 import itertools
 import logging
@@ -36,68 +38,22 @@
     make_array_array,
 )
 
-logger = logging.getLogger(__name__)
-
-
-def read_arff(path: PathOrStr, label: bool = False):
-    path = Path(path)
-    with open(path) as fid:
-        data = arff.load(fid)
-
-    attr_names = [x[0] for x in data["attributes"]]
-    counts = Counter([x[0] for x in data["data"]])
-    idx = slice(1, -1) if label else slice(1, None)
-
-    return FeaturesData(
-        corpus=data["relation"],
-        names=list(counts.keys()),
-        features=np.array([x[idx] for x in data["data"]]),
-        slices=np.array(list(counts.values())),
-        feature_names=attr_names[idx],
-    )
-
+__all__ = [
+    "read_features",
+    "read_features_iterable",
+    "write_features",
+    "register_format",
+    "FeaturesData",
+    "SequentialReader",
+    "RandomAccessReader",
+    "SequentialWriter",
+    "NetCDFReader",
+    "RawAudioReader",
+    "ARFFSequentialReader",
+    "CSVSequentialReader",
+]
 
-def read_csv(path: PathOrStr, header: bool = True, label: bool = False):
-    df = pd.read_csv(path, header=0 if header else None, converters={0: str})
-    counts: typing.Counter[str] = Counter(df.iloc[:, 0])
-    idx = slice(1, -1) if label else slice(1, None)
-
-    return FeaturesData(
-        names=list(counts.keys()),
-        features=np.array(df.iloc[:, idx]),
-        slices=np.array(list(counts.values())),
-        feature_names=list(df.columns[idx]),
-    )
-
-
-def read_netcdf(path: PathOrStr):
-    reader = NetCDFReader(path)
-    features = reader[:]
-    reader.close()
-    return FeaturesData(features, reader.names, reader.corpus, reader.feature_names)
-
-
-def _read_audio_file(path, sample_rate: int = 16000):
-    warnings.filterwarnings("ignore", message="PySoundFile", category=UserWarning)
-    audio, _ = librosa.load(path, sr=sample_rate, mono=True, res_type="kaiser_fast")
-    return np.expand_dims(audio, -1)
-
-
-def read_raw(path: PathOrStr, sample_rate: int = 16000):
-    logger.info(f"Reading raw audio with sample_rate: {sample_rate}")
-    path = Path(path)
-    filepaths = get_audio_paths(path)
-    _audio = TqdmParallel(len(filepaths), desc="Reading audio", leave=False, n_jobs=-1)(
-        delayed(_read_audio_file)(filepath, sample_rate) for filepath in filepaths
-    )
-
-    return FeaturesData(
-        corpus=path.parent.stem,
-        names=[x.stem for x in filepaths],
-        features=np.concatenate(_audio),
-        slices=[len(x) for x in _audio],
-        feature_names=["pcm"],
-    )
+logger = logging.getLogger(__name__)
 
 
 class FeaturesData:
@@ -218,7 +174,72 @@ def __getitem__(self, index) -> np.ndarray:
         return self.features[index]
 
 
+def read_arff_mem(path: PathOrStr, label: bool = False) -> FeaturesData:
+    path = Path(path)
+    with open(path) as fid:
+        data = arff.load(fid)
+
+    attr_names = [x[0] for x in data["attributes"]]
+    counts = Counter([x[0] for x in data["data"]])
+    idx = slice(1, -1) if label else slice(1, None)
+
+    return FeaturesData(
+        corpus=data["relation"],
+        names=list(counts.keys()),
+        features=np.array([x[idx] for x in data["data"]]),
+        slices=np.array(list(counts.values())),
+        feature_names=attr_names[idx],
+    )
+
+
+def read_csv_mem(
+    path: PathOrStr, header: bool = True, label: bool = False
+) -> FeaturesData:
+    df = pd.read_csv(path, header=0 if header else None, converters={0: str})
+    counts: typing.Counter[str] = Counter(df.iloc[:, 0])
+    idx = slice(1, -1) if label else slice(1, None)
+
+    return FeaturesData(
+        names=list(counts.keys()),
+        features=np.array(df.iloc[:, idx]),
+        slices=np.array(list(counts.values())),
+        feature_names=list(df.columns[idx]),
+    )
+
+
+def read_netcdf_mem(path: PathOrStr) -> FeaturesData:
+    reader = NetCDFReader(path)
+    features = reader[:]
+    reader.close()
+    return FeaturesData(features, reader.names, reader.corpus, reader.feature_names)
+
+
+def _read_audio_file(path, sample_rate: int = 16000):
+    warnings.filterwarnings("ignore", message="PySoundFile", category=UserWarning)
+    audio, _ = librosa.load(path, sr=sample_rate, mono=True, res_type="kaiser_fast")
+    return np.expand_dims(audio, -1)
+
+
+def read_raw_mem(path: PathOrStr, sample_rate: int = 16000) -> FeaturesData:
+    logger.info(f"Reading raw audio with sample_rate: {sample_rate}")
+    path = Path(path)
+    filepaths = get_audio_paths(path)
+    _audio = TqdmParallel(len(filepaths), desc="Reading audio", leave=False, n_jobs=-1)(
+        delayed(_read_audio_file)(filepath, sample_rate) for filepath in filepaths
+    )
+
+    return FeaturesData(
+        corpus=path.parent.stem,
+        names=[x.stem for x in filepaths],
+        features=np.concatenate(_audio),
+        slices=[len(x) for x in _audio],
+        feature_names=["pcm"],
+    )
+
+
 class SequentialReader(Iterable[np.ndarray], Sized, ABC):
+    """Base class for readers that read features sequentially."""
+
     feature_names: List[str]
     names: List[str]
     corpus: str = ""
@@ -243,6 +264,8 @@ def read(self) -> Iterable[np.ndarray]:
 
 
 class RandomAccessReader(SequentialReader, ABC):
+    """Base class for readers that can access features by index."""
+
     @abstractmethod
     def __getitem__(
         self, index: Union[int, slice, List[int], np.ndarray]
@@ -258,6 +281,8 @@ def read(self) -> Iterator[np.ndarray]:
 
 
 class NetCDFReader(RandomAccessReader):
+    """Read features from NetCDF4 file."""
+
     _is_2d: bool = False
     _is_3d: bool = False
     _is_vlen: bool = False
@@ -321,6 +346,8 @@ def __getitem__(
 
 
 class RawAudioReader(RandomAccessReader):
+    """Read raw audio files."""
+
     def __init__(self, path: PathOrStr, sample_rate: int = 16000, **kwargs) -> None:
         super().__init__(path, **kwargs)
         self.sample_rate = sample_rate
@@ -344,6 +371,8 @@ def __getitem__(
 
 
 class CSVSequentialReader(SequentialReader):
+    """Read features from CSV file."""
+
     def open(self) -> None:
         with open(self.path) as fid:
             reader = csv.reader(fid)
@@ -365,6 +394,8 @@ def read(self) -> Iterator[np.ndarray]:
 
 
 class ARFFSequentialReader(SequentialReader):
+    """Read features from ARFF file."""
+
     def open(self) -> None:
         self.feature_names = []
         with open(self.path) as fid:
@@ -401,6 +432,18 @@ def read(self) -> Iterator[np.ndarray]:
 
 
 class SequentialWriter:
+    """Write features to file.
+
+    Parameters
+    ----------
+    names : List[str]
+        List of utterance names.
+    corpus : str
+        Corpus name.
+    feature_names : List[str]
+        List of feature names.
+    """
+
     def __init__(
         self,
         names: List[str],
@@ -409,6 +452,8 @@ def __init__(
     ) -> None:
         self.names = names
         self.corpus = corpus
+        if len(feature_names) == 0:
+            raise ValueError("Feature names must be provided.")
         self.feature_names = feature_names
 
     def write(self, path: PathOrStr, features: Iterable[np.ndarray], **kwargs):
@@ -420,6 +465,8 @@ def write(self, path: PathOrStr, features: Iterable[np.ndarray], **kwargs):
         meth(self, path, features, **kwargs)
 
     def write_arff(self, path: PathOrStr, features: Iterable[np.ndarray], **kwargs):
+        """Write features to ARFF file."""
+
         with open(path, "w") as fid:
             fid.write(f"@RELATION {self.corpus}\n\n")
             fid.write("@ATTRIBUTE name STRING\n")
@@ -435,6 +482,8 @@ def write_arff(self, path: PathOrStr, features: Iterable[np.ndarray], **kwargs):
                     writer.writerow([name] + list(inst))
 
     def write_csv(self, path, features: Iterable[np.ndarray], header=True, **kwargs):
+        """Write features to CSV file."""
+
         with open(path, "w") as fid:
             writer = csv.writer(fid, lineterminator="\n")
             if header:
@@ -453,6 +502,8 @@ def write_netcdf(
         chunksize: int = 1024,
         **kwargs,
     ):
+        """Write features to NetCDF4 file."""
+
         dataset = netCDF4.Dataset(path, "w")
         dataset.createDimension("name", len(self.names))
         dataset.createDimension("concat", 0)
@@ -495,6 +546,8 @@ def write_netcdf(
     def write_raw(
         self, path: PathOrStr, features: Iterable[np.ndarray], sr: int = 16000, **kwargs
     ):
+        """Write raw audio files."""
+
         output_dir = Path(path).with_suffix("")
         output_dir.mkdir(exist_ok=True, parents=True)
         for name, audio in zip(self.names, features):
@@ -510,11 +563,11 @@ def write_raw(
     }
 
 
-_READ_BACKENDS: Dict[str, Callable[..., FeaturesData]] = {
-    ".nc": read_netcdf,
-    ".arff": read_arff,
-    ".csv": read_csv,
-    ".txt": read_raw,
+_READ_BACKENDS_MEM: Dict[str, Callable[..., FeaturesData]] = {
+    ".nc": read_netcdf_mem,
+    ".arff": read_arff_mem,
+    ".csv": read_csv_mem,
+    ".txt": read_raw_mem,
 }
 
 _READ_BACKENDS_SEQUENTIAL: Dict[str, Type[SequentialReader]] = {
@@ -535,14 +588,14 @@ def register_format(
     if suffix[0] != ".":
         suffix = "." + suffix
     if read is not None:
-        _READ_BACKENDS[suffix] = read
+        _READ_BACKENDS_MEM[suffix] = read
     if write is not None:
         SequentialWriter._write_backends[suffix] = write
 
 
 def find_features_file(files: Iterable[Path]) -> Path:
     files = list(files)
-    for suffix in _READ_BACKENDS:
+    for suffix in _READ_BACKENDS_MEM:
         file = next((x for x in files if x.suffix == suffix), None)
         if file is not None:
             return file
@@ -550,10 +603,10 @@ def find_features_file(files: Iterable[Path]) -> Path:
 
 
 def read_features(path: PathOrStr, **kwargs) -> FeaturesData:
-    """Read features from given path."""
+    """Read features from given path into memory."""
 
     path = Path(path)
-    meth = _READ_BACKENDS.get(path.suffix)
+    meth = _READ_BACKENDS_MEM.get(path.suffix)
     if meth is None:
         raise ValueError(f"File format {path.suffix} not supported.")
     logger.debug(f"Reading features from {path}")
@@ -561,7 +614,7 @@ def read_features(path: PathOrStr, **kwargs) -> FeaturesData:
 
 
 def read_features_iterable(path: PathOrStr, **kwargs) -> SequentialReader:
-    """Read features from given path as an iterable."""
+    """Read features from given path either sequentially or by index."""
 
     path = Path(path)
     meth = _READ_BACKENDS_SEQUENTIAL.get(path.suffix)
@@ -577,11 +630,31 @@ def write_features(
     names: List[str],
     corpus: str = "",
     feature_names: List[str] = [],
-    slices: Union[np.ndarray, List[int]] = None,
+    slices: Union[np.ndarray, List[int], None] = None,
     **kwargs,
 ) -> None:
-    """Convenience function to write features to given path. Arguments
-    are the same as those passed to IterableWriter.
+    """Convenience function to write features to given path.
+
+    Parameters
+    ----------
+    path : PathOrStr
+        Path to write features to.
+    features : Union[Iterable[np.ndarray], np.ndarray]
+        Features to write. If a 3D array is given, the dimensions are
+        assumed to be (instances, frames, features). If a 2D array is
+        given, the dimensions are assumed to be (frames, features), and
+        the `slices` parameter is required. If an iterable of arrays is
+        given, each 1D or 2D array is assumed to be a single instance.
+    names : List[str]
+        Instance names.
+    corpus : str
+        Name of corpus. Default is empty string.
+    feature_names : List[str]
+        Feature names.
+    slices : Union[np.ndarray, List[int]], optional
+        Number of frames per instance. Only required if a 2D matrix is
+        given, and the first dimension is different than the number of
+        instances.
     """
     if slices is not None:
         if not isinstance(features, np.ndarray):
diff --git a/src/ertk/dataset/utils.py b/src/ertk/dataset/utils.py
index e9038b0..1eb7527 100644
--- a/src/ertk/dataset/utils.py
+++ b/src/ertk/dataset/utils.py
@@ -1,3 +1,5 @@
+"""Utilities for working with audio datasets."""
+
 import logging
 import subprocess
 from pathlib import Path
@@ -7,8 +9,12 @@
 
 from ertk.utils import PathOrStr, TqdmParallel
 
-# The below is because mypy does not yet support type narrowing upon redefinition
-# mypy: allow-redefinition
+__all__ = [
+    "get_audio_paths",
+    "resample_audio",
+    "resample_rename_clips",
+    "write_filelist",
+]
 
 
 def get_audio_paths(path: PathOrStr, absolute: bool = True) -> List[Path]:
diff --git a/src/ertk/preprocessing/__init__.py b/src/ertk/preprocessing/__init__.py
index 4f8c296..2fc2ce5 100644
--- a/src/ertk/preprocessing/__init__.py
+++ b/src/ertk/preprocessing/__init__.py
@@ -1,1 +1,40 @@
+"""
+Preprocessor classes and plugins
+================================
+
+This module contains base classes for audio clip processors, feature
+extractors and instance processors, as well as a number of plugins for
+each of these.
+
+
+Base classes
+------------
+.. autosummary::
+    :toctree: generated/
+
+    InstanceProcessor
+    AudioClipProcessor
+    FeatureExtractor
+
+
+Plugins
+-------
+.. autosummary::
+    :toctree: generated/
+
+    audioset
+    encodec
+    fairseq
+    huggingface
+    keras_apps
+    kmeans
+    opensmile
+    openxbow
+    phonemize
+    resample
+    spectrogram
+    speechbrain
+    vad_trim
+"""
+
 from ._base import AudioClipProcessor, FeatureExtractor, InstanceProcessor
diff --git a/src/ertk/preprocessing/_base.py b/src/ertk/preprocessing/_base.py
index 22c8758..1ac9de7 100644
--- a/src/ertk/preprocessing/_base.py
+++ b/src/ertk/preprocessing/_base.py
@@ -127,12 +127,24 @@ def get_config_type(cls) -> Type[ERTKConfig]:
 
     @classmethod
     def valid_processors(cls) -> List[str]:
-        """Get a list of all registered processor names."""
+        """Get a list of all registered processor names.
+
+        Returns
+        -------
+        names: list of str
+            The names of all registered processors.
+        """
         return sorted(cls._registry.keys() | cls._plugins.keys())
 
     @classmethod
     def get_default_config(cls) -> ERTKConfig:
-        """Get the default configuration for this processor."""
+        """Get the default configuration for this processor.
+
+        Returns
+        -------
+        config: ERTKConfig
+            The default configuration.
+        """
         return OmegaConf.structured(cls._config_type)
 
     @abstractmethod
@@ -207,6 +219,7 @@ def finish(self) -> None:
     @property
     @abstractmethod
     def feature_names(self) -> List[str]:
+        """The names of the features produced by this processor."""
         raise NotImplementedError()
 
 
@@ -268,6 +281,8 @@ def process_files(
 
 
 class FeatureExtractor(InstanceProcessor):
+    """Extracts features from instances."""
+
     @property
     def dim(self) -> int:
         """The dimensionality of the extracted features."""
diff --git a/src/ertk/preprocessing/_pipeline.py b/src/ertk/preprocessing/_pipeline.py
new file mode 100644
index 0000000..1ea2b3e
--- /dev/null
+++ b/src/ertk/preprocessing/_pipeline.py
@@ -0,0 +1,134 @@
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, Iterable, List, Union
+
+import numpy as np
+from omegaconf import MISSING
+
+from ertk.config import ERTKConfig
+from ertk.utils import PathOrStr
+
+from ._base import InstanceProcessor
+
+
+@dataclass
+class ProcessorConfig:
+    """Configuration for a processor."""
+
+    type: str = MISSING
+    """The type of processor to use (friendly name)."""
+    batch_size: int = 1
+    """The batch size to use for processing."""
+    config: Dict[str, Any] = field(default_factory=lambda: {})
+    """The processor configuration."""
+
+
+class ChainingType(Enum):
+    """The type of chaining to use for a processing pipeline."""
+
+    IN_MEMORY = "in_memory"
+    """Results from each processor are stored in memory."""
+    GENERATOR = "generator"
+    """Results from each processor are yielded as they are produced."""
+
+
+@dataclass
+class ProcessingPipelineConfig(ERTKConfig):
+    """Configuration for a processing pipeline."""
+
+    input: str = MISSING
+    """The input data to process."""
+    output: str = MISSING
+    """Where to write the output data."""
+    chaining: ChainingType = ChainingType.IN_MEMORY
+    """The type of chaining to use for the pipeline"""
+    pipeline: List[str] = MISSING
+    """The processors to use in the pipeline, in order"""
+    processors: Dict[str, ProcessorConfig] = MISSING
+    """The processors to use in the pipeline, and their configurations"""
+
+
+class ProcessingPipeline(InstanceProcessor):
+    """A processing pipeline.
+
+    Parameters
+    ----------
+    config: ProcessingPipelineConfig
+        The configuration for the pipeline.
+    """
+
+    pipeline: List[InstanceProcessor]
+    """The processors in the pipeline."""
+    chaining: ChainingType = ChainingType.IN_MEMORY
+    """The type of chaining to use for the pipeline"""
+    output: str
+    """Where to write the output data."""
+    config: ProcessingPipelineConfig
+    """The pipeline configuration."""
+
+    def __init__(self, config: ProcessingPipelineConfig) -> None:
+        from ertk.dataset.features import read_features_iterable
+
+        data = read_features_iterable(config.input)
+        self.data = data
+        self.config = config
+        self.chaining = config.chaining
+        self.output = config.output
+        self.pipeline = []
+        for name in config.pipeline:
+            # This implicitly checks that the processor exists, and that
+            # the config is valid
+            proc_conf = (
+                InstanceProcessor.get_processor_class(name)
+                .get_config_type()
+                .from_config(config.processors[name].config)
+            )
+            processor = InstanceProcessor.make_processor(name, proc_conf)
+            self.pipeline.append(processor)
+
+    def process_instance(self, x: np.ndarray, **kwargs) -> np.ndarray:
+        for processor in self.pipeline:
+            x = processor.process_instance(x, **kwargs)
+        return x
+
+    def process_batch(
+        self, batch: Union[Iterable[np.ndarray], np.ndarray], **kwargs
+    ) -> List[np.ndarray]:
+        for processor in self.pipeline:
+            batch = processor.process_batch(batch, **kwargs)
+        return list(batch)
+
+    def process_all(
+        self, xs: Union[Iterable[np.ndarray], np.ndarray], batch_size: int, **kwargs
+    ) -> Iterable[np.ndarray]:
+        for processor in self.pipeline:
+            xs = processor.process_all(xs, batch_size, **kwargs)
+            if self.chaining == ChainingType.IN_MEMORY:
+                xs = list(xs)
+        return xs
+
+    def finish(self) -> None:
+        for processor in self.pipeline:
+            processor.finish()
+
+    @property
+    def feature_names(self) -> List[str]:
+        return self.pipeline[-1].feature_names
+
+
+def load_pipeline(file: PathOrStr) -> ProcessingPipeline:
+    """Load a processing pipeline from a YAML file.
+
+    Parameters
+    ----------
+    file: pathlike or str
+        The path to the YAML file.
+
+    Returns
+    -------
+    pipeline: ProcessingPipeline
+        The loaded processing pipeline.
+    """
+
+    config = ProcessingPipelineConfig.from_file(file)
+    return ProcessingPipeline(config)
diff --git a/src/ertk/preprocessing/audioset/__init__.py b/src/ertk/preprocessing/audioset/__init__.py
index cfe9a04..c3796f6 100644
--- a/src/ertk/preprocessing/audioset/__init__.py
+++ b/src/ertk/preprocessing/audioset/__init__.py
@@ -1,3 +1,18 @@
+"""
+Audioset feature extractors
+===========================
+
+This module contains feature extractors for the Audioset dataset.
+
+.. autosummary::
+    :toctree:
+
+    VGGishExtractor
+    VGGishExtractorConfig
+    YamnetExtractor
+    YamnetExtractorConfig
+"""
+
 from .audioset import (
     VGGishExtractor,
     VGGishExtractorConfig,
diff --git a/src/ertk/preprocessing/audioset/audioset.py b/src/ertk/preprocessing/audioset/audioset.py
index bd787af..4388550 100644
--- a/src/ertk/preprocessing/audioset/audioset.py
+++ b/src/ertk/preprocessing/audioset/audioset.py
@@ -1,3 +1,5 @@
+"""Audioset feature extractors."""
+
 import logging
 import os
 from dataclasses import dataclass
@@ -13,6 +15,13 @@
 
 from .._base import AudioClipProcessor, FeatureExtractor
 
+__all__ = [
+    "YamnetExtractor",
+    "YamnetExtractorConfig",
+    "VGGishExtractor",
+    "VGGishExtractorConfig",
+]
+
 
 def _spectrogram(x: np.ndarray, sr: float) -> np.ndarray:
     return spectrogram(
@@ -32,12 +41,17 @@ def _spectrogram(x: np.ndarray, sr: float) -> np.ndarray:
 
 @dataclass
 class YamnetExtractorConfig(ERTKConfig):
+    """Configuration for a Yamnet extractor."""
+
     model_dir: str = MISSING
+    """The directory containing the YAMNet model files."""
 
 
 class YamnetExtractor(
     FeatureExtractor, AudioClipProcessor, fname="yamnet", config=YamnetExtractorConfig
 ):
+    """A YAMNet extractor."""
+
     config: YamnetExtractorConfig
 
     def __init__(self, config: YamnetExtractorConfig) -> None:
@@ -108,13 +122,19 @@ def feature_names(self) -> List[str]:
 
 @dataclass
 class VGGishExtractorConfig(ERTKConfig):
+    """Configuration for a VGGish extractor."""
+
     model_dir: str = MISSING
+    """The directory containing the VGGish model files."""
     postprocess: bool = True
+    """Whether to apply the VGGish postprocessing."""
 
 
 class VGGishExtractor(
     FeatureExtractor, AudioClipProcessor, fname="vggish", config=VGGishExtractorConfig
 ):
+    """A VGGish extractor."""
+
     config: VGGishExtractorConfig
 
     def __init__(self, config: VGGishExtractorConfig) -> None:
diff --git a/src/ertk/preprocessing/encodec.py b/src/ertk/preprocessing/encodec.py
index 560b6e6..1f40c8f 100644
--- a/src/ertk/preprocessing/encodec.py
+++ b/src/ertk/preprocessing/encodec.py
@@ -1,3 +1,5 @@
+"""Encodec feature extractor."""
+
 from dataclasses import dataclass
 from enum import Enum
 from typing import List
@@ -10,6 +12,8 @@
 
 from ._base import AudioClipProcessor, FeatureExtractor
 
+__all__ = ["EncodecExtractorConfig", "EncodecExtractor"]
+
 
 class Agg(Enum):
     MEAN = "mean"
@@ -24,17 +28,27 @@ class Model(Enum):
 
 @dataclass
 class EncodecExtractorConfig(ERTKConfig):
+    """Configuration for the Encodec feature extractor."""
+
     model: Model = Model.ENCODEC_48kHz
+    """The Encodec model to use."""
     aggregate: Agg = Agg.MEAN
+    """The aggregation method to use."""
     device: str = "cuda"
+    """The device to use for inference."""
     vq_ids: bool = False
+    """Whether to return the VQ-VAE IDs."""
     vq_ids_as_string: bool = True
+    """Whether to return the VQ-VAE IDs as a string."""
     max_input_len: int = 1500000
+    """The maximum input length for the model."""
 
 
 class EncodecExtractor(
     FeatureExtractor, AudioClipProcessor, fname="encodec", config=EncodecExtractorConfig
 ):
+    """Encodec feature extractor."""
+
     config: EncodecExtractorConfig
 
     def __init__(self, config: EncodecExtractorConfig) -> None:
diff --git a/src/ertk/preprocessing/fairseq.py b/src/ertk/preprocessing/fairseq.py
index abcc0b2..805d987 100644
--- a/src/ertk/preprocessing/fairseq.py
+++ b/src/ertk/preprocessing/fairseq.py
@@ -1,3 +1,5 @@
+"""Fairseq processor."""
+
 from dataclasses import dataclass, field
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union
@@ -13,31 +15,55 @@
 
 from ._base import AudioClipProcessor, FeatureExtractor
 
+__all__ = ["FairseqExtractorConfig", "FairseqExtractor"]
+
 
 class Agg(Enum):
+    """Aggregation method."""
+
     MEAN = "mean"
+    """Mean of the output sequence."""
     MAX = "max"
+    """Max of the output sequence."""
     NONE = "none"
+    """Return entire sequence."""
 
 
 @dataclass
 class FairseqExtractorConfig(ERTKConfig):
+    """Fairseq feature extractor configuration."""
+
     model_type: str = MISSING
+    """Model type."""
     checkpoint: str = MISSING
+    """Path to model checkpoint."""
     layer: str = "context"
+    """Layer to extract features from."""
     aggregate: Agg = Agg.MEAN
+    """Aggregation method."""
     device: str = "cuda"
+    """Device to run model on."""
     arg_overrides: Dict[str, Any] = field(default_factory=dict)
+    """Overrides for model arguments."""
     vq_path: Optional[str] = None
+    """Path to vector quantiser."""
     vq_ids: bool = False
+    """Whether to return VQ cluster ids."""
     vq_ids_as_string: bool = True
+    """Whether to return VQ cluster ids as a single string of integers
+    separated by spaces.
+    """
     max_input_len: int = 1500000
+    """Maximum input length."""
     mms_lang: str = "eng"
+    """Language to use for MMS models."""
 
 
 class FairseqExtractor(
     FeatureExtractor, AudioClipProcessor, fname="fairseq", config=FairseqExtractorConfig
 ):
+    """Fairseq processor."""
+
     config: FairseqExtractorConfig
 
     def __init__(self, config: FairseqExtractorConfig) -> None:
diff --git a/src/ertk/preprocessing/huggingface.py b/src/ertk/preprocessing/huggingface.py
index bbafd9e..dec15ff 100644
--- a/src/ertk/preprocessing/huggingface.py
+++ b/src/ertk/preprocessing/huggingface.py
@@ -1,3 +1,5 @@
+"""Processor using HuggingFace models."""
+
 from dataclasses import dataclass
 from enum import Enum
 from typing import Iterable, List, Optional, Union
@@ -11,17 +13,29 @@
 
 from ._base import AudioClipProcessor, FeatureExtractor
 
+__all__ = ["HuggingFaceExtractorConfig", "HuggingFaceExtractor"]
+
 
 class Task(Enum):
+    """Task to perform."""
+
     CTC = "ctc"
+    """CTC speech recognition."""
+    W2V2 = "w2v2"
+    """Wav2Vec2 speech recognition. Use if CTC does not work."""
     EMBEDDINGS = "embeddings"
+    """Extract embeddings."""
     S2T = "s2t"
+    """Speech2Text speech recognition. Use for speech2text models."""
     WHISPER = "whisper"
-    W2V2 = "w2v2"
+    """Whisper speech recognition. Use for Whisper models."""
     AUTO = "auto"
+    """Automatically detect model and task."""
 
 
 class Agg(Enum):
+    """Aggregation method for extracting embeddings."""
+
     MEAN = "mean"
     MAX = "max"
     NONE = "none"
@@ -29,14 +43,26 @@ class Agg(Enum):
 
 @dataclass
 class HuggingFaceExtractorConfig(ERTKConfig):
+    """HuggingFace feature extractor configuration."""
+
     model: str = MISSING
+    """Model name or path."""
     device: str = "cuda"
+    """Device to run model on."""
     task: Task = Task.CTC
+    """Task to perform."""
     agg: Optional[Agg] = Agg.MEAN
+    """Aggregation method for embeddings."""
     layer: Optional[str] = "context"
+    """Layer to extract embeddings from."""
     max_input_len: int = 1500000
+    """Maximum input length."""
     whisper_lang: Optional[str] = None
+    """Language to use for Whisper models."""
     max_new_tokens: int = 448
+    """Maximum number of generated tokens for speech2text and Whisper
+    models.
+    """
 
 
 class HuggingFaceExtractor(
@@ -45,6 +71,8 @@ class HuggingFaceExtractor(
     fname="huggingface",
     config=HuggingFaceExtractorConfig,
 ):
+    """Processor using HuggingFace models."""
+
     config: HuggingFaceExtractorConfig
 
     def __init__(self, config: HuggingFaceExtractorConfig) -> None:
@@ -67,14 +95,14 @@ def __init__(self, config: HuggingFaceExtractorConfig) -> None:
         if config.task == Task.CTC:
             self.model = AutoModelForCTC.from_pretrained(config.model)
             self.processor = AutoProcessor.from_pretrained(config.model)
+        elif config.task == Task.W2V2:
+            self.model = Wav2Vec2ForCTC.from_pretrained(config.model)
+            self.processor = Wav2Vec2Processor.from_pretrained(config.model)
         elif config.task == Task.S2T:
             self.model = Speech2TextForConditionalGeneration.from_pretrained(
                 config.model
             )
             self.processor = Speech2Text2Processor.from_pretrained(config.model)
-        elif config.task == Task.W2V2:
-            self.model = Wav2Vec2ForCTC.from_pretrained(config.model)
-            self.processor = Wav2Vec2Processor.from_pretrained(config.model)
         elif config.task == Task.WHISPER:
             self.model = WhisperForConditionalGeneration.from_pretrained(config.model)
             self.processor = WhisperProcessor.from_pretrained(config.model)
@@ -103,6 +131,7 @@ def process_instance(self, x: np.ndarray, **kwargs) -> np.ndarray:
             wav = processed.input_values.to(self.config.device)
         except AttributeError:
             wav = processed.input_features.to(self.config.device)
+
         if self.config.task in [Task.CTC, Task.W2V2]:
             output = self.model(wav).logits.argmax(-1)
             output = self.processor.batch_decode(output)
@@ -123,6 +152,7 @@ def process_instance(self, x: np.ndarray, **kwargs) -> np.ndarray:
             )
             output = self.processor.batch_decode(output, skip_special_tokens=True)
             return np.array([output], dtype=object)
+
         # Embeddings
         if self.config.layer == "context":
             output = self.model(wav).last_hidden_state.squeeze()
@@ -137,16 +167,38 @@ def process_instance(self, x: np.ndarray, **kwargs) -> np.ndarray:
     def process_batch(
         self, batch: Union[Iterable[np.ndarray], np.ndarray], **kwargs
     ) -> List[np.ndarray]:
-        wav = self.processor(
+        processed = self.processor(
             [x.squeeze()[: self.config.max_input_len] for x in batch],
             sampling_rate=kwargs.pop("sr"),
             return_tensors="pt",
             padding="longest",
-        ).input_values.to(self.config.device)
-        if self.config.task == Task.CTC:
+        )
+        try:
+            wav = processed.input_values.to(self.config.device)
+        except AttributeError:
+            wav = processed.input_features.to(self.config.device)
+
+        if self.config.task in [Task.CTC, Task.W2V2]:
             output = self.model(wav).logits.argmax(-1)
             output = self.processor.batch_decode(output)
             return [np.array([x], dtype=object) for x in output]
+        elif self.config.task in [Task.S2T, Task.WHISPER]:
+            forced_decoder_ids = None
+            if (
+                self.config.task == Task.WHISPER
+                and self.config.whisper_lang is not None
+            ):
+                forced_decoder_ids = self.processor.get_decoder_prompt_ids(
+                    task="transcribe", language=self.config.whisper_lang
+                )
+            output = self.model.generate(
+                wav,
+                forced_decoder_ids=forced_decoder_ids,
+                max_new_tokens=self.config.max_new_tokens,
+            )
+            output = self.processor.batch_decode(output, skip_special_tokens=True)
+            return [np.array([x], dtype=object) for x in output]
+
         # Embeddings
         raise RuntimeError("Cannot use batching for extracting embeddings.")
 
diff --git a/src/ertk/preprocessing/keras_apps.py b/src/ertk/preprocessing/keras_apps.py
index 637e228..b9225ff 100644
--- a/src/ertk/preprocessing/keras_apps.py
+++ b/src/ertk/preprocessing/keras_apps.py
@@ -1,3 +1,5 @@
+"""Feature extractor using Keras applications."""
+
 from dataclasses import dataclass
 from types import ModuleType
 from typing import Iterable, List, Optional, Union
@@ -12,9 +14,11 @@
 
 from ._base import AudioClipProcessor, FeatureExtractor
 
+__all__ = ["KerasAppsExtractorConfig", "KerasAppsExtractor"]
+
 
 @dataclass
-class ModelArgs:
+class _ModelArgs:
     module: ModuleType
     funcname: str
     size: int
@@ -41,44 +45,57 @@ def _init_models():
 
     global MODEL_MAP
     MODEL_MAP = {
-        "densenet121": ModelArgs(densenet, "DenseNet121", 224),
-        "densenet169": ModelArgs(densenet, "DenseNet169", 224),
-        "densenet201": ModelArgs(densenet, "DenseNet201", 224),
-        "efficientnet_b0": ModelArgs(efficientnet, "EfficientNetB0", 224),
-        "efficientnet_b1": ModelArgs(efficientnet, "EfficientNetB1", 240),
-        "efficientnet_b2": ModelArgs(efficientnet, "EfficientNetB2", 260),
-        "efficientnet_b3": ModelArgs(efficientnet, "EfficientNetB3", 300),
-        "efficientnet_b4": ModelArgs(efficientnet, "EfficientNetB4", 380),
-        "efficientnet_b5": ModelArgs(efficientnet, "EfficientNetB5", 456),
-        "efficientnet_b6": ModelArgs(efficientnet, "EfficientNetB6", 528),
-        "efficientnet_b7": ModelArgs(efficientnet, "EfficientNetB7", 600),
-        "inception_resnet_v2": ModelArgs(inception_resnet_v2, "InceptionResNetV2", 299),
-        "inception_v3": ModelArgs(inception_v3, "InceptionV3", 299),
-        "mobilenet": ModelArgs(mobilenet, "MobileNet", 224),
-        "mobilenet_v2": ModelArgs(mobilenet_v2, "MobileNetV2", 224),
-        "mobilenet_v3_large": ModelArgs(mobilenet_v3, "MobileNetV3Large", 224),
-        "mobilenet_v3_small": ModelArgs(mobilenet_v3, "MobileNetV3Small", 224),
-        "nasnet_large": ModelArgs(nasnet, "NASNetLarge", 331),
-        "nasnet_mobile": ModelArgs(nasnet, "NASNetMobile", 224),
-        "resnet50": ModelArgs(resnet, "ResNet50", 224),
-        "resnet50_v2": ModelArgs(resnet, "ResNet50V2", 224),
-        "resnet101": ModelArgs(resnet, "ResNet101", 224),
-        "resnet101_v2": ModelArgs(resnet, "ResNet101V2", 224),
-        "resnet152": ModelArgs(resnet, "ResNet152", 224),
-        "resnet152_v2": ModelArgs(resnet, "ResNet152V2", 224),
-        "vgg16": ModelArgs(vgg16, "VGG16", 224),
-        "vgg19": ModelArgs(vgg19, "VGG19", 224),
-        "xception": ModelArgs(xception, "Xception", 299),
+        "densenet121": _ModelArgs(densenet, "DenseNet121", 224),
+        "densenet169": _ModelArgs(densenet, "DenseNet169", 224),
+        "densenet201": _ModelArgs(densenet, "DenseNet201", 224),
+        "efficientnet_b0": _ModelArgs(efficientnet, "EfficientNetB0", 224),
+        "efficientnet_b1": _ModelArgs(efficientnet, "EfficientNetB1", 240),
+        "efficientnet_b2": _ModelArgs(efficientnet, "EfficientNetB2", 260),
+        "efficientnet_b3": _ModelArgs(efficientnet, "EfficientNetB3", 300),
+        "efficientnet_b4": _ModelArgs(efficientnet, "EfficientNetB4", 380),
+        "efficientnet_b5": _ModelArgs(efficientnet, "EfficientNetB5", 456),
+        "efficientnet_b6": _ModelArgs(efficientnet, "EfficientNetB6", 528),
+        "efficientnet_b7": _ModelArgs(efficientnet, "EfficientNetB7", 600),
+        "inception_resnet_v2": _ModelArgs(
+            inception_resnet_v2, "InceptionResNetV2", 299
+        ),
+        "inception_v3": _ModelArgs(inception_v3, "InceptionV3", 299),
+        "mobilenet": _ModelArgs(mobilenet, "MobileNet", 224),
+        "mobilenet_v2": _ModelArgs(mobilenet_v2, "MobileNetV2", 224),
+        "mobilenet_v3_large": _ModelArgs(mobilenet_v3, "MobileNetV3Large", 224),
+        "mobilenet_v3_small": _ModelArgs(mobilenet_v3, "MobileNetV3Small", 224),
+        "nasnet_large": _ModelArgs(nasnet, "NASNetLarge", 331),
+        "nasnet_mobile": _ModelArgs(nasnet, "NASNetMobile", 224),
+        "resnet50": _ModelArgs(resnet, "ResNet50", 224),
+        "resnet50_v2": _ModelArgs(resnet, "ResNet50V2", 224),
+        "resnet101": _ModelArgs(resnet, "ResNet101", 224),
+        "resnet101_v2": _ModelArgs(resnet, "ResNet101V2", 224),
+        "resnet152": _ModelArgs(resnet, "ResNet152", 224),
+        "resnet152_v2": _ModelArgs(resnet, "ResNet152V2", 224),
+        "vgg16": _ModelArgs(vgg16, "VGG16", 224),
+        "vgg19": _ModelArgs(vgg19, "VGG19", 224),
+        "xception": _ModelArgs(xception, "Xception", 299),
     }
 
 
 @dataclass
 class KerasAppsExtractorConfig(ERTKConfig):
+    """Configuration for KerasAppsExtractor."""
+
     model: str = MISSING
+    """Name of the model to use. See https://keras.io/api/applications/
+    for the list of available models.
+    """
     size: Optional[int] = None
+    """Size of the input image. If not specified, the default size of the
+    model will be used.
+    """
     n_mels: int = 128
+    """Number of mel bands for spectrogram extraction."""
     spec_type: str = "mel"
+    """Type of spectrogram to extract. Can be either 'mel' or 'linear'."""
     fmax: int = 8000
+    """Maximum frequency for spectrogram extraction."""
 
 
 class KerasAppsExtractor(
@@ -87,6 +104,8 @@ class KerasAppsExtractor(
     fname="keras_apps",
     config=KerasAppsExtractorConfig,
 ):
+    """Feature extractor using Keras applications."""
+
     config: KerasAppsExtractorConfig
 
     def __init__(self, config: KerasAppsExtractorConfig) -> None:
diff --git a/src/ertk/preprocessing/kmeans.py b/src/ertk/preprocessing/kmeans.py
index 7ef75ca..a8e462a 100644
--- a/src/ertk/preprocessing/kmeans.py
+++ b/src/ertk/preprocessing/kmeans.py
@@ -1,3 +1,5 @@
+"""Kmeans vector quantiser."""
+
 from dataclasses import dataclass
 from typing import List
 
@@ -8,13 +10,20 @@
 
 from ._base import FeatureExtractor
 
+__all__ = ["KMeansExtractorConfig", "KMeansExtractor"]
+
 
 @dataclass
 class KMeansExtractorConfig(ERTKConfig):
+    """KMeans feature extractor configuration."""
+
     pickle: str = MISSING
+    """Path to pickle file."""
 
 
 class KMeansExtractor(FeatureExtractor, fname="kmeans", config=KMeansExtractorConfig):
+    """KMeans vector quantiser."""
+
     config: KMeansExtractorConfig
 
     def __init__(self, config: KMeansExtractorConfig) -> None:
diff --git a/src/ertk/preprocessing/opensmile.py b/src/ertk/preprocessing/opensmile.py
index 7b9f21f..128cc4a 100644
--- a/src/ertk/preprocessing/opensmile.py
+++ b/src/ertk/preprocessing/opensmile.py
@@ -1,3 +1,5 @@
+"""OpenSMILE feature extraction."""
+
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Dict, List
@@ -10,12 +12,19 @@
 
 from ._base import AudioClipProcessor, FeatureExtractor
 
+__all__ = ["OpenSMILEExtractorConfig", "OpenSMILEExtractor"]
+
 
 @dataclass
 class OpenSMILEExtractorConfig(ERTKConfig):
+    """OpenSMILE feature extractor configuration."""
+
     opensmile_config: str = MISSING
+    """OpenSMILE config file."""
     levels: List[str] = field(default_factory=lambda: ["func"])
+    """OpenSMILE levels to extract."""
     opensmile_opts: Dict[str, str] = field(default_factory=dict)
+    """Additional OpenSMILE options."""
 
 
 class OpenSMILEExtractor(
@@ -24,6 +33,8 @@ class OpenSMILEExtractor(
     fname="opensmile",
     config=OpenSMILEExtractorConfig,
 ):
+    """OpenSMILE feature extractor."""
+
     config: OpenSMILEExtractorConfig
 
     def __init__(self, config: OpenSMILEExtractorConfig) -> None:
diff --git a/src/ertk/preprocessing/openxbow/__init__.py b/src/ertk/preprocessing/openxbow/__init__.py
index 8e8e8f5..76807a8 100644
--- a/src/ertk/preprocessing/openxbow/__init__.py
+++ b/src/ertk/preprocessing/openxbow/__init__.py
@@ -1,1 +1,15 @@
+"""
+OpenXBOW extractor
+==================
+
+This module contains the OpenXBOW extractor, which is a wrapper around
+the OpenXBOW Java program.
+
+.. autosummary::
+    :toctree:
+
+    OpenXBOWExtractor
+    OpenXBOWExtractorConfig
+"""
+
 from .openxbow import OpenXBOWExtractor, OpenXBOWExtractorConfig
diff --git a/src/ertk/preprocessing/openxbow/openxbow.py b/src/ertk/preprocessing/openxbow/openxbow.py
index dcb6c1c..932142f 100644
--- a/src/ertk/preprocessing/openxbow/openxbow.py
+++ b/src/ertk/preprocessing/openxbow/openxbow.py
@@ -1,3 +1,5 @@
+"""OpenXBOW feature extractor."""
+
 import os
 import subprocess
 import tempfile
@@ -13,15 +15,22 @@
 from ertk.dataset import write_features
 from ertk.preprocessing._base import FeatureExtractor
 
+__all__ = ["OpenXBOWExtractorConfig", "OpenXBOWExtractor"]
+
 
 @dataclass
 class OpenXBOWExtractorConfig(ERTKConfig):
+    """Configuration for the OpenXBOW feature extractor."""
+
     xbowargs: List[str] = MISSING
+    """The arguments to pass to OpenXBOW."""
 
 
 class OpenXBOWExtractor(
     FeatureExtractor, fname="openxbow", config=OpenXBOWExtractorConfig
 ):
+    """OpenXBOW feature extractor."""
+
     config: OpenXBOWExtractorConfig
     _dim = None
 
diff --git a/src/ertk/preprocessing/phonemize.py b/src/ertk/preprocessing/phonemize.py
index cd60962..239cd32 100644
--- a/src/ertk/preprocessing/phonemize.py
+++ b/src/ertk/preprocessing/phonemize.py
@@ -1,3 +1,5 @@
+"""Phonemize text using the phonemizer library."""
+
 import os
 import re
 import warnings
@@ -10,18 +12,33 @@
 
 from ._base import FeatureExtractor
 
+__all__ = ["PhonemizeConfig", "Phonemizer"]
+
 
 @dataclass
 class PhonemizeConfig(ERTKConfig):
+    """Phonemizer configuration."""
+
     language: str = "en-us"
+    """Language to use."""
     backend: str = "festival"
+    """Backend to use."""
     ph_sep: Optional[str] = " "
+    """Phoneme separator."""
     syl_sep: Optional[str] = None
+    """Syllable separator."""
     word_sep: Optional[str] = None
+    """Word separator."""
     preserve_punctuation: bool = False
+    """Preserve punctuation."""
     preserve_empty_lines: bool = True
+    """Preserve empty lines."""
     strip: bool = False
+    """Strip whitespace."""
     language_switch: str = "keep-flags"
+    """Language switch flagging. One of "keep-flags", "remove-flags",
+    "remove-utterance".
+    """
 
 
 def _festival_punct(texts: List[str]) -> Tuple[List[str], int]:
@@ -37,6 +54,8 @@ def _festival_punct(texts: List[str]) -> Tuple[List[str], int]:
 
 
 class Phonemizer(FeatureExtractor, fname="phonemize", config=PhonemizeConfig):
+    """Phonemizer using the phonemizer library."""
+
     config: PhonemizeConfig
 
     def __init__(self, config: PhonemizeConfig) -> None:
diff --git a/src/ertk/preprocessing/resample.py b/src/ertk/preprocessing/resample.py
index 9869028..87fa2ab 100644
--- a/src/ertk/preprocessing/resample.py
+++ b/src/ertk/preprocessing/resample.py
@@ -1,3 +1,5 @@
+"""Audio resampling using resampy."""
+
 from dataclasses import dataclass
 from typing import List
 
@@ -7,14 +9,22 @@
 from ertk.config import ERTKConfig
 from ertk.preprocessing._base import AudioClipProcessor
 
+__all__ = ["ResampleConfig", "Resampler"]
+
 
 @dataclass
 class ResampleConfig(ERTKConfig):
+    """Resample configuration."""
+
     sample_rate: int = 16000
+    """Target sample rate."""
     filter: str = "kaiser_fast"
+    """Resampling filter."""
 
 
 class Resampler(AudioClipProcessor, fname="resample", config=ResampleConfig):
+    """Audio resampler using resampy."""
+
     config: ResampleConfig
 
     def process_instance(self, x: np.ndarray, **kwargs) -> np.ndarray:
diff --git a/src/ertk/preprocessing/spectrogram.py b/src/ertk/preprocessing/spectrogram.py
index 8a8c4a0..afa580b 100644
--- a/src/ertk/preprocessing/spectrogram.py
+++ b/src/ertk/preprocessing/spectrogram.py
@@ -1,3 +1,5 @@
+"""Spectrogram extraction."""
+
 import warnings
 from dataclasses import dataclass
 from typing import List, Optional
@@ -9,6 +11,8 @@
 
 from ._base import AudioClipProcessor, FeatureExtractor
 
+__all__ = ["SpectrogramExtractorConfig", "SpectrogramExtractor", "spectrogram"]
+
 
 def spectrogram(
     audio: np.ndarray,
@@ -143,22 +147,57 @@ def spectrogram(
 
 @dataclass
 class SpectrogramExtractorConfig(ERTKConfig):
+    """Configuration for SpectrogramExtractor."""
+
     kind: str = "mel"
+    """The kind of spectrogram to calculate. "mel" calculates a mel
+    spectrogram, whereas "stft" just uses linear frequency.
+    """
     pre_emphasis: float = 0
+    """Amount of pre-emphasis to apply. Default is 0 (no pre-emphasis)."""
     window_size: float = 0.025
+    """Size of window in seconds. Default is 0.025 (25 ms)."""
     window_shift: float = 0.01
+    """Amount window moves each frame, in seconds. Default is 0.01 (10 ms)."""
     win_length_samp: Optional[int] = None
+    """The window length in samples. This overrides `window_size` if
+    given.
+    """
     hop_length_samp: Optional[int] = None
+    """The hop size in samples. This overrides `window_shift` if given."""
     n_mels: int = 128
+    """Number of mel bands."""
     htk_mel: bool = False
+    """Use HTK formula for mel calculation."""
     n_fft: int = 2048
+    """Number of FFT bins. Default is to use the next power of 2
+    greater than `window_size * sr`.
+    """
     n_chroma: int = 12
+    """Number of chroma bands."""
     clip_db: Optional[float] = None
+    """Whether to clip noise floor at a given level in dB below
+    maximum. Default is `None` which does no clipping. This is only
+    used if `to_db` is `True`.
+    """
     fmin: float = 0
+    """Minimum frequency for mel bands. Default is 0."""
     fmax: Optional[float] = 8000
+    """Maximum frequency for mel bands. Default is 8000."""
     power: int = 2
+    """Raise spectrogram magnitudes to this power. Default is 2 which
+    is the usual power spectrogram.
+    """
     to_log: Optional[str] = "db"
+    """Whether to convert spectrogram to logarithmic domain. Default is
+    `db` which converts to dB units. If `to_log=='log'` then the
+    natual logarithm is taken. Note that this argument is mainly
+    only useful if `power == 2`.
+    """
     mel_norm: str = "slaney"
+    """Normalisation to apply to mel filters. Default is "slaney" as in
+    librosa.
+    """
 
 
 class SpectrogramExtractor(
@@ -167,6 +206,8 @@ class SpectrogramExtractor(
     fname="spectrogram",
     config=SpectrogramExtractorConfig,
 ):
+    """Extracts a spectrogram from an audio clip."""
+
     config: SpectrogramExtractorConfig
 
     def process_instance(self, x: np.ndarray, **kwargs) -> np.ndarray:
diff --git a/src/ertk/preprocessing/speechbrain.py b/src/ertk/preprocessing/speechbrain.py
index d283df1..a3a8066 100644
--- a/src/ertk/preprocessing/speechbrain.py
+++ b/src/ertk/preprocessing/speechbrain.py
@@ -1,3 +1,5 @@
+"""Processing using SpeechBrain models."""
+
 from dataclasses import dataclass
 from enum import Enum
 from typing import List, Optional
@@ -11,14 +13,23 @@
 
 from ._base import AudioClipProcessor, FeatureExtractor
 
+__all__ = ["SpeechBrainExtractorConfig", "SpeechBrainExtractor"]
+
 
 class Task(Enum):
+    """Task to perform."""
+
     WHISPER = "whisper"
+    """Whisper speech recognition. Use for Whisper models."""
     EMBEDDINGS = "embeddings"
+    """Extract embeddings."""
     ASR = "asr"
+    """Automatic speech recognition."""
 
 
 class Agg(Enum):
+    """Aggregation method for extracting embeddings."""
+
     MEAN = "mean"
     MAX = "max"
     NONE = "none"
@@ -26,12 +37,20 @@ class Agg(Enum):
 
 @dataclass
 class SpeechBrainExtractorConfig(ERTKConfig):
+    """SpeechBrain feature extractor configuration."""
+
     model: str = MISSING
+    """Model name or path."""
     device: str = "cuda"
+    """Device to run model on."""
     task: Task = Task.EMBEDDINGS
+    """Task to perform."""
     agg: Optional[Agg] = Agg.MEAN
+    """Aggregation method for embeddings."""
     max_input_len: int = 1500000
+    """Maximum input length."""
     normalise: bool = True
+    """Normalise audio."""
 
 
 class SpeechBrainExtractor(
@@ -40,6 +59,8 @@ class SpeechBrainExtractor(
     fname="speechbrain",
     config=SpeechBrainExtractorConfig,
 ):
+    """Processing using SpeechBrain models."""
+
     config: SpeechBrainExtractorConfig
 
     def __init__(self, config: SpeechBrainExtractorConfig) -> None:
diff --git a/src/ertk/preprocessing/vad_trim.py b/src/ertk/preprocessing/vad_trim.py
index 4333bcf..884b0c3 100644
--- a/src/ertk/preprocessing/vad_trim.py
+++ b/src/ertk/preprocessing/vad_trim.py
@@ -1,3 +1,5 @@
+"""Voice activity detection (VAD) trimming."""
+
 from dataclasses import dataclass
 from enum import Enum
 from typing import List
@@ -8,26 +10,43 @@
 from ertk.config import ERTKConfig
 from ertk.preprocessing._base import AudioClipProcessor
 
+__all__ = ["VADTrimmerConfig", "VADTrimmer"]
+
 
 class Method(Enum):
+    """VAD trimming method."""
+
     mh2009 = "MH2009"
+    """Moattar and Homayounpour (2009)"""
     librosa = "librosa"
+    """librosa.effects.trim"""
 
 
 @dataclass
 class VADTrimmerConfig(ERTKConfig):
+    """VAD trimming configuration."""
+
     method: Method = Method.librosa
+    """VAD trimming method."""
     top_db: int = 60
+    """Top decibel threshold for librosa VAD."""
     energy_thresh: float = 2.5
+    """Energy threshold for MH2009 VAD."""
     freq_thresh: float = 150
+    """Frequency threshold for MH2009 VAD."""
     sf_thresh: float = 10
+    """Spectral flatness threshold for MH2009 VAD."""
     window: float = 0.01
+    """Window size in seconds for MH2009 VAD."""
     debug: bool = False
+    """Whether to plot debug figures."""
     min_speech: int = 5
+    """Minimum number of consecutive speech frames to keep."""
     min_silence: int = 10
+    """Minimum number of consecutive silence frames to keep."""
 
 
-def filter_silence_speech(x: np.ndarray, min_speech: int = 5, min_silence: int = 10):
+def _filter_silence_speech(x: np.ndarray, min_speech: int = 5, min_silence: int = 10):
     """Smooths a boolean array indicating segments of speech based on
     the given speech and silence thresholds.
     """
@@ -64,12 +83,43 @@ def mh2009_vad(
     min_speech: int = 5,
     min_silence: int = 10,
 ):
-    """Trim audio clips based on a VAD algorithm adapted from [1].
-
-    [1] M. H. Moattar and M. M. Homayounpour, "A simple but efficient
-    real-time Voice Activity Detection algorithm," 2009 17th European
-    Signal Processing Conference, Glasgow, 2009, pp. 2549-2553.
-    https://ieeexplore.ieee.org/abstract/document/7077834
+    """Trim audio clips based on a VAD algorithm adapted from [1]_.
+
+    Parameters
+    ----------
+    x : np.ndarray
+        Audio signal.
+    sr : float
+        Sample rate.
+    energy_thresh : float, optional
+        Energy threshold, by default 0.05.
+    freq_thresh : float, optional
+        Frequency threshold, by default 150.
+    sf_thresh : float, optional
+        Spectral flatness threshold, by default 10.
+    window : float, optional
+        Window size in seconds, by default 0.01.
+    debug : bool, optional
+        Whether to plot debug figures, by default False.
+    min_speech : int, optional
+        Minimum number of consecutive speech frames to keep, by default
+        5.
+    min_silence : int, optional
+        Minimum number of consecutive silence frames to keep, by default
+        10.
+
+    Returns
+    -------
+    np.ndarray
+        Boolean array indicating segments of voicing.
+
+    References
+    ----------
+    .. [1] M. H. Moattar and M. M. Homayounpour, "A simple but efficient
+           real-time Voice Activity Detection algorithm," 2009 17th
+           European Signal Processing Conference, Glasgow, 2009, pp.
+           2549-2553.
+           https://ieeexplore.ieee.org/abstract/document/7077834
     """
     window_samples = int(sr * window)
     n_fft = 1024
@@ -104,7 +154,7 @@ def mh2009_vad(
     count = cond1 + cond2 + cond3
     speech = count > 1
 
-    filter_silence_speech(speech, min_speech=min_speech, min_silence=min_silence)
+    _filter_silence_speech(speech, min_speech=min_speech, min_silence=min_silence)
 
     if not np.any(speech):
         if debug:
@@ -125,6 +175,8 @@ def mh2009_vad(
 
 
 class VADTrimmer(AudioClipProcessor, fname="vad_trim", config=VADTrimmerConfig):
+    """Voice activity detection (VAD) trimmer."""
+
     config: VADTrimmerConfig
 
     def process_instance(self, x: np.ndarray, **kwargs) -> np.ndarray:
diff --git a/src/ertk/pytorch/__init__.py b/src/ertk/pytorch/__init__.py
index 7b5abc2..a4e7356 100644
--- a/src/ertk/pytorch/__init__.py
+++ b/src/ertk/pytorch/__init__.py
@@ -1,1 +1,17 @@
+"""
+PyTorch backend
+===============
+
+This package contains the PyTorch backend for the ERTK.
+
+.. autosummary::
+    :toctree: generated/
+
+    models
+    train
+    dataset
+    classification
+    utils
+"""
+
 from .models import get_pt_model, get_pt_model_fn
diff --git a/src/ertk/pytorch/classification.py b/src/ertk/pytorch/classification.py
index 28ebe37..0ce4623 100644
--- a/src/ertk/pytorch/classification.py
+++ b/src/ertk/pytorch/classification.py
@@ -1,3 +1,5 @@
+"""PyTorch classification utilities."""
+
 import logging
 import warnings
 from collections import defaultdict
@@ -20,6 +22,12 @@
 from ertk.train import ExperimentResult, get_scores
 from ertk.utils import ScoreFunction
 
+__all__ = [
+    "pt_cross_validate",
+    "pt_train_val_test",
+    "train_mtl_model",
+]
+
 
 def pt_cross_validate(
     model_fn: Callable[..., Union[nn.Module, pl.LightningModule]],
diff --git a/src/ertk/pytorch/dataset.py b/src/ertk/pytorch/dataset.py
index 5fbf6c2..c5e15ac 100644
--- a/src/ertk/pytorch/dataset.py
+++ b/src/ertk/pytorch/dataset.py
@@ -1,3 +1,5 @@
+"""PyTorch dataset utilities."""
+
 from typing import Any, Callable, Collection, Dict, List, Optional, Tuple, Union
 
 import numpy as np
@@ -10,6 +12,14 @@
 from ertk.dataset import Dataset, load_datasets_config
 from ertk.train import ExperimentConfig
 
+__all__ = [
+    "DatasetFuncWrapper",
+    "DataModuleAdapter",
+    "MTLDataModule",
+    "MTLDataset",
+    "create_mtl_dataloader",
+]
+
 
 class DatasetFuncWrapper(TorchDataset[Tuple[torch.Tensor, ...]]):
     funcs: List[Callable[[torch.Tensor], torch.Tensor]]
@@ -42,6 +52,7 @@ class DataModuleAdapter(pl.LightningDataModule):
 
     def __init__(
         self,
+        *,
         load_config: Optional[ExperimentConfig] = None,
         dataset: Optional[Dataset] = None,
         train_select: Union[str, Dict[str, Collection[str]], np.ndarray, None] = None,
@@ -57,21 +68,15 @@ def __init__(
         self.batch_size = batch_size
         self.train_select = train_select
         self.val_select = val_select
-        self.test_select = test_select or val_select
+        self.test_select = test_select if test_select is not None else val_select
         self.dl_num_workers = dl_num_workers
 
     def setup(self, stage: Optional[str] = None) -> None:
         if self.load_config:
             data = load_datasets_config(self.load_config)
-            self.train_idx = data.get_idx_for_split(
-                self.load_config.eval.train_val_test.train
-            )
-            self.val_idx = data.get_idx_for_split(
-                self.load_config.eval.train_val_test.valid
-            )
-            self.test_idx = data.get_idx_for_split(
-                self.load_config.eval.train_val_test.test
-            )
+            self.train_idx = data.get_idx_for_split(self.load_config.eval.train)
+            self.val_idx = data.get_idx_for_split(self.load_config.eval.valid)
+            self.test_idx = data.get_idx_for_split(self.load_config.eval.test)
 
         if isinstance(self.train_select, np.ndarray):
             self.train_idx = self.train_select
@@ -135,16 +140,17 @@ def predict_dataloader(self) -> DataLoader:
 class MTLDataModule(DataModuleAdapter):
     def __init__(
         self,
+        *,
         dataset: Dataset,
         tasks: List[str],
-        train_select: Union[str, Dict[str, Collection[str]]],
-        val_select: Union[str, Dict[str, Collection[str]]],
-        test_select: Union[str, Dict[str, Collection[str]], None] = None,
+        train_select: Union[str, Dict[str, Collection[str]], np.ndarray, None] = None,
+        val_select: Union[str, Dict[str, Collection[str]], np.ndarray, None] = None,
+        test_select: Union[str, Dict[str, Collection[str]], np.ndarray, None] = None,
         batch_size: int = 32,
         dl_num_workers: int = 0,
     ):
         super().__init__(
-            dataset,
+            dataset=dataset,
             train_select=train_select,
             val_select=val_select,
             test_select=test_select,
diff --git a/src/ertk/pytorch/models/__init__.py b/src/ertk/pytorch/models/__init__.py
index df3f02c..feb4a74 100644
--- a/src/ertk/pytorch/models/__init__.py
+++ b/src/ertk/pytorch/models/__init__.py
@@ -1,8 +1,41 @@
+"""PyTorch implementations of various models.
+
+.. autosummary::
+    :toctree:
+
+    get_pt_model
+    get_pt_model_fn
+    PyTorchModelConfig
+    ERTKPyTorchModel
+    SimpleModel
+    SimpleClassificationModel
+    LightningWrapper
+
+
+Models
+------
+.. autosummary::
+    :toctree:
+
+    aldeneh2017
+    mlp
+    zhang2019
+
+
+Layers
+------
+.. autosummary::
+    :toctree:
+
+    layers
+"""
+
 from functools import partial
 from typing import Callable
 
 from torch.nn import Module
 
+# Needed for registration
 from . import aldeneh2017, mlp, wav2vec_ft, zhang2019
 from ._base import (
     ERTKPyTorchModel,
@@ -11,7 +44,16 @@
     SimpleClassificationModel,
     SimpleModel,
 )
-from .layers import Attention1D, make_fc
+
+__all__ = [
+    "ERTKPyTorchModel",
+    "LightningWrapper",
+    "PyTorchModelConfig",
+    "SimpleClassificationModel",
+    "SimpleModel",
+    "get_pt_model",
+    "get_pt_model_fn",
+]
 
 
 def get_pt_model_fn(name: str, **kwargs) -> Callable[..., Module]:
diff --git a/src/ertk/pytorch/models/_base.py b/src/ertk/pytorch/models/_base.py
index e8a83a8..f9a3c48 100644
--- a/src/ertk/pytorch/models/_base.py
+++ b/src/ertk/pytorch/models/_base.py
@@ -13,17 +13,40 @@
 
 @dataclass
 class PyTorchModelConfig(ERTKConfig):
+    """Base class for PyTorch model configuration dataclasses."""
+
     optimiser: str = "adam"
+    """The optimiser to use."""
     opt_params: Dict[str, Any] = field(default_factory=dict)
+    """Parameters to pass to the optimiser."""
     learning_rate: float = 1e-3
+    """The learning rate to use."""
     n_features: int = -1
+    """The number of features in the input data."""
     n_classes: int = -1
+    """The number of classes in the output data."""
     loss: str = "cross_entropy"
+    """The loss function to use."""
     loss_args: Dict[str, Any] = field(default_factory=dict)
+    """Arguments to pass to the loss function."""
 
 
 class ERTKPyTorchModel(pl.LightningModule, ABC):
+    """Base class for PyTorch models.
+
+    This class is a subclass of
+    :class:`pytorch_lightning.LightningModule` and implements the
+    :meth:`training_step` method. It also provides a
+    :meth:`get_model_class` method for retrieving a model class by name.
+
+    Parameters
+    ----------
+    config : PyTorchModelConfig
+        The model configuration.
+    """
+
     config: PyTorchModelConfig
+    """The model configuration."""
 
     _config_type: ClassVar[Type[PyTorchModelConfig]]
     _friendly_name: ClassVar[str]
@@ -65,6 +88,23 @@ def __init_subclass__(
 
     @classmethod
     def get_model_class(cls, name: str) -> Type["ERTKPyTorchModel"]:
+        """Get a model class by name.
+
+        Parameters
+        ----------
+        name : str
+            The name of the model class to retrieve.
+
+        Returns
+        -------
+        Type[ERTKPyTorchModel]
+            The model class.
+
+        Raises
+        ------
+        ValueError
+            If no model class with the given name is registered.
+        """
         try:
             return cls._registry[name]
         except KeyError as e:
@@ -72,18 +112,58 @@ def get_model_class(cls, name: str) -> Type["ERTKPyTorchModel"]:
 
     @classmethod
     def make_model(cls, name: str, config: PyTorchModelConfig) -> "ERTKPyTorchModel":
+        """Make a model by name.
+
+        Parameters
+        ----------
+        name : str
+            The name of the model class to retrieve.
+        config : PyTorchModelConfig
+            The model configuration.
+
+        Returns
+        -------
+        ERTKPyTorchModel
+            The model created with the given config.
+        """
         return cls.get_model_class(name)(config)
 
     @classmethod
     def get_config_type(cls) -> Type[PyTorchModelConfig]:
+        """Get the configuration dataclass for this model.
+
+        Returns
+        -------
+        Type[PyTorchModelConfig]
+            The configuration dataclass.
+
+        Notes
+        -----
+        This is a class method rather than a property because it is
+        needed before an instance of the model is created.
+        """
         return cls._config_type
 
     @classmethod
     def valid_models(cls) -> List[str]:
+        """Get a list of valid model names.
+
+        Returns
+        -------
+        List[str]
+            A list of valid model names.
+        """
         return list(cls._registry)
 
     @classmethod
     def get_default_config(cls) -> PyTorchModelConfig:
+        """Get the default configuration for this model.
+
+        Returns
+        -------
+        PyTorchModelConfig
+            The default configuration.
+        """
         return OmegaConf.structured(cls._config_type)
 
     @abstractmethod
@@ -94,6 +174,20 @@ def preprocess_input(self, x: torch.Tensor) -> torch.Tensor:
         """Perform any modifications to the input tensor `x`, which is
         extracted from a batch. This will typically have shape
         (batch_size, features) or (batch_size, timesteps, features).
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            The input tensor.
+
+        Returns
+        -------
+        torch.Tensor
+            The preprocessed input tensor.
+
+        Notes
+        -----
+        The default implementation returns `x` unchanged.
         """
         return x
 
diff --git a/src/ertk/pytorch/models/aldeneh2017.py b/src/ertk/pytorch/models/aldeneh2017.py
index 6e71480..2ba7bd1 100644
--- a/src/ertk/pytorch/models/aldeneh2017.py
+++ b/src/ertk/pytorch/models/aldeneh2017.py
@@ -1,3 +1,13 @@
+"""Implementation of the model architecture from [1]_.
+
+References
+----------
+.. [1] Z. Aldeneh and E. Mower Provost, 'Using regional saliency for
+       speech emotion recognition', in 2017 IEEE International
+       Conference on Acoustics, Speech and Signal Processing (ICASSP),
+       Mar. 2017, pp. 2741-2745, doi: 10.1109/ICASSP.2017.7952655.
+"""
+
 from dataclasses import dataclass, field
 from typing import List
 
@@ -7,6 +17,8 @@
 
 from ._base import PyTorchModelConfig, SimpleClassificationModel
 
+__all__ = ["Aldeneh2017Config", "Aldeneh2017Model"]
+
 
 @dataclass
 class Aldeneh2017Config(PyTorchModelConfig):
@@ -15,7 +27,9 @@ class Aldeneh2017Config(PyTorchModelConfig):
     dense_dims: List[int] = field(default_factory=lambda: [1024, 1024])
 
 
-class Model(SimpleClassificationModel, fname="aldeneh2017", config=Aldeneh2017Config):
+class Aldeneh2017Model(
+    SimpleClassificationModel, fname="aldeneh2017", config=Aldeneh2017Config
+):
     def __init__(self, config: Aldeneh2017Config):
         super().__init__(config)
 
diff --git a/src/ertk/pytorch/models/layers.py b/src/ertk/pytorch/models/layers.py
index ed53131..6f936f5 100644
--- a/src/ertk/pytorch/models/layers.py
+++ b/src/ertk/pytorch/models/layers.py
@@ -1,3 +1,5 @@
+"""PyTorch layers for use in models."""
+
 import math
 from typing import Any, Callable, List, Sequence
 
@@ -7,8 +9,18 @@
 
 from ertk.pytorch.utils import get_activation
 
+__all__ = ["Attention1D", "make_fc"]
+
 
 class Attention1D(nn.Module):
+    """1D attention layer.
+
+    Parameters
+    ----------
+    in_channels: int
+        The number of input channels.
+    """
+
     def __init__(self, in_channels: int) -> None:
         super().__init__()
 
diff --git a/src/ertk/pytorch/models/mlp.py b/src/ertk/pytorch/models/mlp.py
index e5fdb6e..63525a5 100644
--- a/src/ertk/pytorch/models/mlp.py
+++ b/src/ertk/pytorch/models/mlp.py
@@ -1,3 +1,5 @@
+"""Simple multi-layer perceptron (MLP) model."""
+
 from dataclasses import dataclass, field
 from typing import List
 
@@ -6,6 +8,8 @@
 from ._base import PyTorchModelConfig, SimpleClassificationModel
 from .layers import make_fc
 
+__all__ = ["MLPModel", "MLPConfig"]
+
 
 @dataclass
 class MLPConfig(PyTorchModelConfig):
@@ -14,7 +18,7 @@ class MLPConfig(PyTorchModelConfig):
     activation: str = "relu"
 
 
-class Model(SimpleClassificationModel, fname="mlp", config=MLPConfig):
+class MLPModel(SimpleClassificationModel, fname="mlp", config=MLPConfig):
     config: MLPConfig
 
     def __init__(self, config: MLPConfig):
diff --git a/src/ertk/pytorch/models/mtl.py b/src/ertk/pytorch/models/mtl.py
index b75ebcd..9fd2950 100644
--- a/src/ertk/pytorch/models/mtl.py
+++ b/src/ertk/pytorch/models/mtl.py
@@ -1,16 +1,19 @@
+"""Multi-task learning MTL utilities."""
+
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from typing import Callable, Dict, List
 
 import omegaconf
 import torch
-import torch.nn as nn
 
 from ertk.config import ERTKConfig
 from ertk.pytorch.utils import get_loss
 
 from ._base import ERTKPyTorchModel, PyTorchModelConfig
 
+__all__ = ["MTLModelConfig", "MTLModel", "MTLTaskConfig"]
+
 
 @dataclass
 class MTLTaskConfig(ERTKConfig):
diff --git a/src/ertk/pytorch/models/simple_mtl.py b/src/ertk/pytorch/models/simple_mtl.py
index 3fac6e2..cd9ed75 100644
--- a/src/ertk/pytorch/models/simple_mtl.py
+++ b/src/ertk/pytorch/models/simple_mtl.py
@@ -1,3 +1,7 @@
+"""Simple MTL model consisting of shared layers followed by
+task-specific layers.
+"""
+
 from collections import OrderedDict
 from dataclasses import dataclass, field
 from typing import Dict, List
@@ -9,6 +13,8 @@
 from .layers import make_fc
 from .mtl import MTLModel, MTLModelConfig
 
+__all__ = ["SimpleMTLModel", "SimpleMTLModelConfig"]
+
 
 @dataclass
 class SimpleMTLModelConfig(MTLModelConfig):
diff --git a/src/ertk/pytorch/models/zhang2019.py b/src/ertk/pytorch/models/zhang2019.py
index bacd89c..a873174 100644
--- a/src/ertk/pytorch/models/zhang2019.py
+++ b/src/ertk/pytorch/models/zhang2019.py
@@ -1,3 +1,14 @@
+"""Implementation of the model from [1]_.
+
+References
+----------
+.. [1] Z. Zhang, B. Wu, and B. Schuller, 'Attention-augmented End-to-end
+       Multi-task Learning for Emotion Prediction from Speech', in 2019
+       IEEE International Conference on Acoustics, Speech and Signal
+       Processing, May 2019, pp. 6705-6709, doi:
+       10.1109/ICASSP.2019.8682896.
+"""
+
 import torch
 from torch import nn
 
diff --git a/src/ertk/pytorch/train.py b/src/ertk/pytorch/train.py
index 8803656..7772458 100644
--- a/src/ertk/pytorch/train.py
+++ b/src/ertk/pytorch/train.py
@@ -1,3 +1,5 @@
+"""Training configuration for PyTorch models."""
+
 from dataclasses import dataclass, field
 from typing import Any, Dict, Optional
 
@@ -6,6 +8,13 @@
 from ertk.config import ERTKConfig
 from ertk.pytorch.models import PyTorchModelConfig
 
+__all__ = [
+    "PyTorchTrainConfig",
+    "PyTorchLoggingConfig",
+    "PyTorchDataConfig",
+    "PyTorchDataAugConfig",
+]
+
 
 @dataclass
 class PyTorchLoggingConfig(ERTKConfig):
diff --git a/src/ertk/pytorch/utils.py b/src/ertk/pytorch/utils.py
index 9fc3f32..730573f 100644
--- a/src/ertk/pytorch/utils.py
+++ b/src/ertk/pytorch/utils.py
@@ -1,3 +1,5 @@
+"""PyTorch utilities."""
+
 from typing import Optional
 
 import torch
@@ -5,6 +7,16 @@
 import torch.nn.functional as F
 import torchaudio
 
+__all__ = [
+    "get_activation",
+    "get_loss",
+    "load_tensorboard",
+    "frame_tensor",
+    "add_signal_awgn",
+    "add_gaussian",
+    "spectrogram",
+]
+
 
 def get_loss(loss: str, *args, **kwargs) -> nn.Module:
     """Get torch Module for activation function.
@@ -167,10 +179,7 @@ def add_signal_awgn(x: torch.Tensor, snr: float = 20) -> torch.Tensor:
     """
     x = x.squeeze(-1)
     noise = torch.randn(x.size(), dtype=x.dtype, device=x.device)
-    x_power = torch.mean(x ** 2, -1)
-    n_power = torch.mean(noise ** 2, -1)
-    K = torch.sqrt(x_power / n_power * 10 ** (-snr / 10))
-    return x + K.unsqueeze(-1) * noise
+    return torchaudio.functional.add_noise(x, noise, snr)
 
 
 def add_gaussian(x: torch.Tensor, sigma: float = 1, multiplicative: bool = True):
@@ -236,4 +245,8 @@ def spectrogram(
             power=power,
         )
     with torch.no_grad():
-        return module(x).transpose(-1, -2)
+        spec = module(x).transpose(-1, -2)
+        if to_db:
+            spec = torchaudio.transforms.AmplitudeToDB(
+                stype="power" if power == 2 else "magnitude", top_db=clip_db
+            )(spec)
diff --git a/src/ertk/sklearn/__init__.py b/src/ertk/sklearn/__init__.py
index e69de29..efc2e10 100644
--- a/src/ertk/sklearn/__init__.py
+++ b/src/ertk/sklearn/__init__.py
@@ -0,0 +1,17 @@
+"""
+Scikit-learn backend
+====================
+
+This module contains the scikit-learn backend for ERTK. It is used to
+provide a consistent interface for scikit-learn models, and to provide
+additional functionality.
+
+.. autosummary::
+    :toctree: generated/
+
+    models
+    classification
+    utils
+"""
+
+from .models import get_sk_model, get_sk_model_fn
diff --git a/src/ertk/sklearn/classification.py b/src/ertk/sklearn/classification.py
index 2d8bfce..c541303 100644
--- a/src/ertk/sklearn/classification.py
+++ b/src/ertk/sklearn/classification.py
@@ -1,3 +1,5 @@
+"""Scikit-learn classification functions."""
+
 import json
 import logging
 import time
@@ -14,6 +16,8 @@
 from ertk.train import ExperimentResult, get_pipeline_params, get_scores
 from ertk.utils import ScoreFunction
 
+__all__ = ["sk_cross_validate", "sk_train_val_test"]
+
 logger = logging.getLogger("ertk.sklearn.classification")
 
 
diff --git a/src/ertk/sklearn/models/__init__.py b/src/ertk/sklearn/models/__init__.py
index e90aa5f..73d2b8d 100644
--- a/src/ertk/sklearn/models/__init__.py
+++ b/src/ertk/sklearn/models/__init__.py
@@ -1,4 +1,21 @@
-"""Scikit-learn implementations of various models"""
+"""Scikit-learn implementations of various models.
+
+.. autosummary::
+    :toctree:
+
+    get_sk_model
+    get_sk_model_fn
+    SkModelConfig
+
+
+Models
+------
+.. autosummary::
+    :toctree:
+
+    svm
+    mtl
+"""
 
 from typing import Callable
 
@@ -12,6 +29,18 @@
 from .mtl import MTFLClassifier, MultiDimensionalCC
 from .svm import PrecomputedSVC
 
+__all__ = [
+    "ERTKSkModel",
+    "SkModelConfig",
+    "SkWrapperConfig",
+    "SkWrapperModel",
+    "MTFLClassifier",
+    "MultiDimensionalCC",
+    "PrecomputedSVC",
+    "get_sk_model",
+    "get_sk_model_fn",
+]
+
 CLASSIFIER_MAP = {
     "svm": PrecomputedSVC,
     "rf": RandomForestClassifier,
diff --git a/src/ertk/sklearn/models/_base.py b/src/ertk/sklearn/models/_base.py
index c83a592..4bf7781 100644
--- a/src/ertk/sklearn/models/_base.py
+++ b/src/ertk/sklearn/models/_base.py
@@ -32,15 +32,33 @@
 
 @dataclass
 class SkModelConfig(ERTKConfig):
+    """Configuration for a scikit-learn model."""
+
     kwargs: Dict[str, Any] = field(default_factory=dict)
+    """Keyword arguments to pass to the model constructor."""
     transform: Optional[str] = "std"
+    """The name of the transform to use."""
     transform_kwargs: Dict[str, Any] = field(default_factory=dict)
+    """Keyword arguments to pass to the transform constructor."""
     param_grid: Dict[str, Any] = field(default_factory=dict)
+    """The parameter grid to use for grid search."""
     param_grid_path: Optional[str] = None
+    """The path to a YAML file containing the parameter grid to use for
+    grid search.
+    """
 
 
 class ERTKSkModel(ABC):
+    """Base class for scikit-learn models.
+
+    Parameters
+    ----------
+    config: SkModelConfig
+        The model configuration.
+    """
+
     config: SkModelConfig
+    """The model configuration."""
 
     _config_type: ClassVar[Type[SkModelConfig]]
     _friendly_name: ClassVar[str]
@@ -98,10 +116,15 @@ def get_default_config(cls) -> SkModelConfig:
 
 
 class SkWrapperConfig(SkModelConfig):
+    """Configuration for a scikit-learn model wrapper."""
+
     kind: str = MISSING
+    """The name of the model to wrap."""
 
 
 class SkWrapperModel(ERTKSkModel, fname="wrapper", config=SkWrapperConfig):
+    """A wrapper for scikit-learn models."""
+
     config: SkWrapperConfig
 
     def __init__(self, config: SkWrapperConfig) -> None:
@@ -127,7 +150,15 @@ def __init__(self, config: SkWrapperConfig) -> None:
         self.clf = clf
 
 
-def gen_wrapper_for_clf(kind):
+def gen_wrapper_for_clf(kind: str) -> None:
+    """Generate a wrapper for a scikit-learn classifier.
+
+    Parameters
+    ----------
+    kind: str
+        The name of the classifier.
+    """
+
     class _(ERTKSkModel, fname=kind, config=SkModelConfig):
         def __init__(self, config: SkModelConfig) -> None:
             super().__init__(config)
diff --git a/src/ertk/sklearn/models/mtl.py b/src/ertk/sklearn/models/mtl.py
index aadaef5..7b029c5 100644
--- a/src/ertk/sklearn/models/mtl.py
+++ b/src/ertk/sklearn/models/mtl.py
@@ -1,3 +1,5 @@
+"""Multi-task learning models."""
+
 import warnings
 
 import liblinear.liblinearutil as liblinearutil
@@ -12,6 +14,8 @@
 from sklearn.utils import check_random_state
 from sklearn.utils.validation import check_is_fitted
 
+__all__ = ["BaseMTFL", "MTFLClassifier", "MultiDimensionalCC"]
+
 _EPS = 1e-15
 
 
@@ -83,7 +87,7 @@ def Dmin(self, a):
         return a / a.sum()
 
     def Dmin_eps(self, a, eps):
-        return self.Dmin(np.sqrt(a ** 2 + eps))
+        return self.Dmin(np.sqrt(a**2 + eps))
 
     def fit(self, X, y, tasks, sample_weight=None):
         self._fit_feat(X, y, tasks)
diff --git a/src/ertk/sklearn/models/svm.py b/src/ertk/sklearn/models/svm.py
index 8db86d2..4820d58 100644
--- a/src/ertk/sklearn/models/svm.py
+++ b/src/ertk/sklearn/models/svm.py
@@ -1,3 +1,5 @@
+"""Support Vector Machine (SVM) models for scikit-learn."""
+
 from functools import partial
 from typing import Any, Dict, Union
 
@@ -5,6 +7,8 @@
 from sklearn.metrics.pairwise import kernel_metrics
 from sklearn.svm import SVC
 
+__all__ = ["PrecomputedSVC"]
+
 
 def _get_kernel_func(
     kernel: str,
@@ -30,8 +34,8 @@ class PrecomputedSVC(SVC):
     """Class that wraps scikit-learn's SVC to precompute the kernel
     values in order to speed up training. The kernel parameter is a
     string which is transparently mapped to and from the corresponding
-    callable function with the relevant parameters (degree, gamma,
-    coef0). All other parameters are passed directly to SVC.
+    callable function with the relevant parameters (`degree`, `gamma`,
+    `coef0`). All other parameters are passed directly to SVC.
     """
 
     def __init__(
diff --git a/src/ertk/sklearn/utils.py b/src/ertk/sklearn/utils.py
index ea96487..1f8f840 100644
--- a/src/ertk/sklearn/utils.py
+++ b/src/ertk/sklearn/utils.py
@@ -1,3 +1,5 @@
+"""Utilities for scikit-learn estimators."""
+
 import warnings
 from typing import List
 
@@ -10,6 +12,13 @@
 from sklearn.pipeline import Pipeline
 from sklearn.preprocessing import LabelBinarizer
 
+__all__ = [
+    "get_estimator_tree",
+    "get_base_estimator",
+    "OneVsRestClassifier",
+    "GridSearchVal",
+]
+
 
 # Workaround for OneVsRest
 # https://github.com/scikit-learn/scikit-learn/issues/10882#issuecomment-376912896
diff --git a/src/ertk/tensorflow/__init__.py b/src/ertk/tensorflow/__init__.py
index 12a41fc..b990a4e 100644
--- a/src/ertk/tensorflow/__init__.py
+++ b/src/ertk/tensorflow/__init__.py
@@ -1,2 +1,18 @@
+"""
+TensorFlow backend
+==================
+
+This package contains the TensorFlow backend for the ERTK.
+
+.. autosummary::
+    :toctree: generated/
+
+    models
+    train
+    dataset
+    classification
+    utils
+"""
+
 from .models import get_tf_model, get_tf_model_fn
 from .utils import compile_wrap, init_gpu_memory_growth, test_fit
diff --git a/src/ertk/tensorflow/classification.py b/src/ertk/tensorflow/classification.py
index 7762f83..af7ca46 100644
--- a/src/ertk/tensorflow/classification.py
+++ b/src/ertk/tensorflow/classification.py
@@ -1,3 +1,5 @@
+"""TensorFlow classification utils."""
+
 import logging
 import time
 from collections import defaultdict
@@ -11,18 +13,25 @@
 from sklearn.model_selection import BaseCrossValidator
 from sklearn.pipeline import Pipeline
 
-from ertk.tensorflow.utils import (
+from ertk.tensorflow.dataset import (
     DataFunction,
-    TFModelFunction,
     get_data_fn,
-    init_gpu_memory_growth,
     tf_dataset_gen,
     tf_dataset_mem,
     tf_dataset_mem_ragged,
 )
+from ertk.tensorflow.utils import TFModelFunction, init_gpu_memory_growth
 from ertk.train import ExperimentResult, get_scores
 from ertk.utils import ScoreFunction, filter_kwargs
 
+__all__ = [
+    "tf_train_val_test",
+    "tf_cross_validate",
+    "tf_classification_metrics",
+    "TFClassifierWrapper",
+    "BalancedSparseCategoricalAccuracy",
+]
+
 logger = logging.getLogger(__name__)
 
 
diff --git a/src/ertk/tensorflow/dataset.py b/src/ertk/tensorflow/dataset.py
new file mode 100644
index 0000000..c8cf335
--- /dev/null
+++ b/src/ertk/tensorflow/dataset.py
@@ -0,0 +1,286 @@
+"""Dataset utilities for TensorFlow."""
+
+from typing import Callable, List, Optional, Tuple, Union
+
+import numpy as np
+import tensorflow as tf
+from keras.utils.data_utils import Sequence
+
+from ertk.utils import batch_arrays_by_length, shuffle_multiple
+
+__all__ = [
+    "tf_dataset_gen",
+    "tf_dataset_mem_ragged",
+    "tf_dataset_mem",
+    "BatchedSequence",
+    "BatchedFrameSequence",
+    "DataFunction",
+]
+
+
+DataFunction = Callable[..., tf.data.Dataset]
+
+
+def tf_dataset_gen(
+    x: np.ndarray,
+    y: np.ndarray,
+    sample_weight: Optional[np.ndarray] = None,
+    *,
+    batch_size: int = 64,
+    shuffle: bool = True,
+):
+    """Returns a TensorFlow generator Dataset instance with the given
+    data.
+
+    Parameters
+    ----------
+    x: numpy.ndarray
+        A 2- or 3-D data matrix of shape (n_instances, n_features) or
+        (n_instances, seq_len, n_features).
+    y: numpy.ndarray
+        A 1-D array of length n_instances containing numeric class
+        labels.
+    sample_weight: numpy.ndarray, optional
+        A 1-D array of length n_instances containing sample weights.
+        Added as third item in dataset if present.
+    batch_size: int
+        The batch size to use.
+    shuffle: boolean
+        Whether or not to shuffle the dataset. Note that shuffling is
+        done *before* batching, unlike in `create_tf_dataset_ragged()`.
+    """
+
+    def gen_inst():
+        if shuffle:
+            perm = np.random.permutation(len(x))
+        else:
+            perm = np.arange(len(x))
+
+        if sample_weight is None:
+            for i in perm:
+                yield x[i], y[i]
+        else:
+            for i in perm:
+                yield x[i], y[i], sample_weight[i]
+
+    sig: Tuple[tf.TensorSpec, ...] = (
+        tf.TensorSpec(shape=x[0].shape, dtype=tf.float32),
+        tf.TensorSpec(shape=(), dtype=tf.int64),
+    )
+    if sample_weight is not None:
+        sig += (tf.TensorSpec(shape=(), dtype=tf.float32),)
+    data = tf.data.Dataset.from_generator(gen_inst, output_signature=sig)
+    return data.batch(batch_size).prefetch(2)
+
+
+def tf_dataset_mem(
+    x: np.ndarray,
+    y: np.ndarray,
+    sample_weight: Optional[np.ndarray] = None,
+    *,
+    batch_size: int = 64,
+    shuffle: bool = True,
+) -> tf.data.Dataset:
+    """Returns a TensorFlow in-memory Dataset instance with the given
+    data.
+
+    Parameters
+    ----------
+    x: numpy.ndarray
+        A 2- or 3-D data matrix of shape (n_instances, n_features) or
+        (n_instances, seq_len, n_features).
+    y: numpy.ndarray
+        A 1-D array of length n_instances containing numeric class
+        labels.
+    sample_weight: numpy.ndarray, optional
+        A 1-D array of length n_instances containing sample weights.
+        Added as third item in dataset if present.
+    batch_size: int
+        The batch size to use.
+    shuffle: boolean
+        Whether or not to shuffle the dataset. Note that shuffling is
+        done *before* batching, unlike in `create_tf_dataset_ragged()`.
+    """
+    with tf.device("CPU"):
+        if sample_weight is None:
+            data = tf.data.Dataset.from_tensor_slices((x, y))
+        else:
+            data = tf.data.Dataset.from_tensor_slices((x, y, sample_weight))
+
+    if shuffle:
+        data = data.shuffle(len(x))
+    return data.batch(batch_size).prefetch(2)
+
+
+def tf_dataset_mem_ragged(
+    x: np.ndarray,
+    y: np.ndarray,
+    sample_weight: Optional[np.ndarray] = None,
+    *,
+    batch_size: int = 64,
+    shuffle: bool = True,
+) -> tf.data.Dataset:
+    """Returns a TensorFlow in-memory Dataset instance from
+    variable-length features.
+
+    Parameters
+    ----------
+    x: numpy.ndarray
+        A 3-D data matrix of shape (n_instances, length[i], n_features)
+        with variable length axis 1.
+    y: numpy.ndarray
+        A 1-D array of length n_instances containing numeric class
+        labels.
+    sample_weight: numpy.ndarray, optional
+        A 1-D array of length n_instances containing sample weights.
+        Added as third item in dataset if present.
+    batch_size: int
+        The batch size to use.
+    shuffle: boolean
+        Whether or not to shuffle the dataset. Note that shuffling is
+        done *after* batching, because sequences are sorted by length,
+        then batched in similar lengths.
+    """
+
+    def ragged_to_dense(x: tf.RaggedTensor, y):
+        return x.to_tensor(), y
+
+    def ragged_to_dense_weighted(x: tf.RaggedTensor, y, sample_weight):
+        return x.to_tensor(), y, sample_weight
+
+    # Sort according to length
+    perm = np.argsort([len(a) for a in x])
+    x = x[perm]
+    y = y[perm]
+    if sample_weight is not None:
+        sample_weight = sample_weight[perm]
+
+    ragged = tf.RaggedTensor.from_row_lengths(
+        np.concatenate(list(x)), [len(a) for a in x]
+    )
+    with tf.device("CPU"):
+        if sample_weight is None:
+            data = tf.data.Dataset.from_tensor_slices((ragged, y))
+        else:
+            data = tf.data.Dataset.from_tensor_slices((ragged, y, sample_weight))
+
+    # Group similar lengths in batches, then shuffle batches
+    data = data.batch(batch_size)
+    if shuffle:
+        data = data.shuffle(len(x) // batch_size + 1)
+
+    if sample_weight is None:
+        data = data.map(ragged_to_dense)
+    else:
+        data = data.map(ragged_to_dense_weighted)
+    return data.prefetch(2)
+
+
+class BatchedFrameSequence(Sequence):
+    """Creates a sequence of batches of frames to process.
+
+    Parameters
+    ----------
+    x: ndarray or list of ndarray
+        Sequences of vectors.
+    y: ndarray
+        Labels corresponding to sequences in x.
+    prebatched: bool, default = False
+        Whether or not x has already been grouped into batches.
+    batch_size: int, default = 32
+        Batch size to use. Each generated batch will be at most this size.
+    shuffle: bool, default = True
+        Whether to shuffle the order of the batches.
+    """
+
+    def __init__(
+        self,
+        x: Union[np.ndarray, List[np.ndarray]],
+        y: np.ndarray,
+        prebatched: bool = False,
+        batch_size: int = 32,
+        shuffle: bool = True,
+    ):
+        self.x = x
+        self.y = y
+        if not prebatched:
+            self.x, self.y = batch_arrays_by_length(
+                self.x, self.y, batch_size=batch_size, shuffle=shuffle
+            )
+        if shuffle:
+            self.x, self.y = shuffle_multiple(self.x, self.y, numpy_indexing=True)
+
+    def __len__(self):
+        return len(self.x)
+
+    def __getitem__(self, idx: int):
+        return self.x[idx], self.y[idx]
+
+
+class BatchedSequence(Sequence):
+    """Creates a sequence of batches to process.
+
+    Parameters
+    ----------
+    x: ndarray or list of ndarray
+        Instance feature vectors. Each vector is assumed to be for a different
+        instance.
+    y: ndarray
+        Labels corresponding to sequences in x.
+    prebatched: bool, default = False
+        Whether or not x has already been grouped into batches.
+    batch_size: int, default = 32
+        Batch size to use. Each generated batch will be at most this size.
+    shuffle: bool, default = True
+        Whether to shuffle the instances.
+    """
+
+    def __init__(
+        self, x: np.ndarray, y: np.ndarray, batch_size: int = 32, shuffle: bool = True
+    ):
+        self.x = x
+        self.y = y
+        self.batch_size = batch_size
+        if shuffle:
+            self.x, self.y = shuffle_multiple(self.x, self.y, numpy_indexing=True)
+
+    def __len__(self):
+        return int(np.ceil(len(self.x) / self.batch_size))
+
+    def __getitem__(self, idx: int):
+        sl = slice(idx * self.batch_size, (idx + 1) * self.batch_size)
+        return self.x[sl], self.y[sl]
+
+
+_DATA_FNS = {
+    "mem": tf_dataset_mem,
+    "gen": tf_dataset_gen,
+    "mem_ragged": tf_dataset_mem_ragged,
+}
+
+
+def get_data_fn(name: str) -> Callable[..., tf.data.Dataset]:
+    """Returns a data function by name.
+
+    Parameters
+    ----------
+    name: str
+        The name of the data function to return. Can be one of:
+        * "``mem``": In-memory data function.
+        * "``gen``": Generator data function.
+        * "``mem_ragged``": In-memory data function for ragged tensors.
+
+    Returns
+    -------
+    Callable[..., tf.data.Dataset]
+        The data function.
+    """
+    if ":" in name:
+        import importlib
+
+        modname, fname = name.split(":")
+        module = importlib.import_module(modname)
+        func = getattr(module, fname)
+    else:
+        func = _DATA_FNS[name]
+    return func
diff --git a/src/ertk/tensorflow/models/__init__.py b/src/ertk/tensorflow/models/__init__.py
index ea52472..242fa76 100644
--- a/src/ertk/tensorflow/models/__init__.py
+++ b/src/ertk/tensorflow/models/__init__.py
@@ -1,4 +1,36 @@
-"""TensorFlow implementations of various models."""
+"""TensorFlow implementations of various models.
+
+.. autosummary::
+    :toctree:
+
+    get_tf_model
+    get_tf_model_fn
+    TFModelConfig
+
+
+Models
+------
+.. autosummary::
+    :toctree:
+
+    aldeneh2017
+    audeep
+    depinto2020
+    iskhakova2020
+    latif2019
+    mlp
+    rbm
+    zhang2019
+    zhao2019
+
+
+Layers
+------
+.. autosummary::
+    :toctree:
+
+    layers
+"""
 
 import importlib
 from functools import partial
@@ -9,6 +41,8 @@
 
 from ._base import TFModelConfig
 
+__all__ = ["get_tf_model", "get_tf_model_fn", "TFModelConfig"]
+
 
 def get_tf_model_fn(name: str, **kwargs) -> TFModelFunction:
     """Get a TensorFlow model creation function.
diff --git a/src/ertk/tensorflow/models/_base.py b/src/ertk/tensorflow/models/_base.py
index 9a99b11..fea3170 100644
--- a/src/ertk/tensorflow/models/_base.py
+++ b/src/ertk/tensorflow/models/_base.py
@@ -3,12 +3,22 @@
 
 from ertk.config import ERTKConfig
 
+__all__ = ["TFModelConfig"]
+
 
 @dataclass
 class TFModelConfig(ERTKConfig):
+    """Configuration for a TensorFlow model."""
+
     optimiser: str = "adam"
+    """The optimiser to use."""
     opt_params: Dict[str, Any] = field(default_factory=dict)
+    """The parameters to pass to the optimiser."""
     learning_rate: float = 1e-3
+    """The learning rate to use."""
     n_features: int = -1
+    """The number of features in the input data."""
     n_classes: int = -1
+    """The number of classes in the output data."""
     loss: str = "sparse_categorical_crossentropy"
+    """The loss function to use."""
diff --git a/src/ertk/tensorflow/models/aldeneh2017.py b/src/ertk/tensorflow/models/aldeneh2017.py
index 4f0a144..1a23c7c 100644
--- a/src/ertk/tensorflow/models/aldeneh2017.py
+++ b/src/ertk/tensorflow/models/aldeneh2017.py
@@ -1,13 +1,15 @@
-"""Implementation of the final full model architecture from [1].
-
-[1] Z. Aldeneh and E. Mower Provost, 'Using regional saliency for speech
-emotion recognition', in 2017 IEEE International Conference on
-Acoustics, Speech and Signal Processing (ICASSP), Mar. 2017, pp.
-2741–2745, doi: 10.1109/ICASSP.2017.7952655.
+"""Implementation of the model architecture from [1]_.
+
+References
+----------
+.. [1] Z. Aldeneh and E. Mower Provost, 'Using regional saliency for
+       speech emotion recognition', in 2017 IEEE International
+       Conference on Acoustics, Speech and Signal Processing (ICASSP),
+       Mar. 2017, pp. 2741-2745, doi: 10.1109/ICASSP.2017.7952655.
 """
 
-from tensorflow.keras import Model
-from tensorflow.keras.layers import Conv1D, Dense, GlobalMaxPool1D, Input, concatenate
+from keras import Model
+from keras.layers import Conv1D, Dense, GlobalMaxPool1D, Input, concatenate
 
 __all__ = ["model"]
 
diff --git a/src/ertk/tensorflow/models/audeep.py b/src/ertk/tensorflow/models/audeep.py
index 043b3e5..6286cc1 100644
--- a/src/ertk/tensorflow/models/audeep.py
+++ b/src/ertk/tensorflow/models/audeep.py
@@ -1,3 +1,13 @@
+"""Implementation of auDeep model from [1]_.
+
+References
+----------
+.. [1] M. Freitag, S. Amiriparian, S. Pugachevskiy, N. Cummins, and B.
+       Schuller, “auDeep: Unsupervised learning of representations from
+       audio with deep recurrent neural networks,” The Journal of
+       Machine Learning Research, vol. 18, no. 1, pp. 6340-6344, 2017.
+"""
+
 import tensorflow as tf
 from tensorflow.keras import Model
 from tensorflow.keras.layers import (
diff --git a/src/ertk/tensorflow/models/depinto2020.py b/src/ertk/tensorflow/models/depinto2020.py
index 57a3341..cafacd4 100644
--- a/src/ertk/tensorflow/models/depinto2020.py
+++ b/src/ertk/tensorflow/models/depinto2020.py
@@ -1,23 +1,18 @@
-"""Implementation of short model from [1].
+"""Implementation of short model from [1]_.
 
-[1] M. G. de Pinto, M. Polignano, P. Lops, and G. Semeraro, 'Emotions
-Understanding Model from Spoken Language using Deep Neural Networks and
-Mel-Frequency Cepstral Coefficients', in 2020 IEEE Conference on
-Evolving and Adaptive Intelligent Systems (EAIS), May 2020, pp. 1–5,
-doi: 10.1109/EAIS48028.2020.9122698.
+References
+----------
+.. [1] M. G. de Pinto, M. Polignano, P. Lops, and G. Semeraro, 'Emotions
+       Understanding Model from Spoken Language using Deep Neural
+       Networks and Mel-Frequency Cepstral Coefficients', in 2020 IEEE
+       Conference on Evolving and Adaptive Intelligent Systems (EAIS),
+       May 2020, pp. 1-5, doi: 10.1109/EAIS48028.2020.9122698.
 """
 
+from keras import Model
+from keras.layers import Conv1D, Dense, Dropout, Flatten, Input, MaxPool1D, Reshape
 
-from tensorflow.keras import Model
-from tensorflow.keras.layers import (
-    Conv1D,
-    Dense,
-    Dropout,
-    Flatten,
-    Input,
-    MaxPool1D,
-    Reshape,
-)
+__all__ = ["model"]
 
 
 def model(n_features: int, n_classes: int) -> Model:
diff --git a/src/ertk/tensorflow/models/iskhakova2020.py b/src/ertk/tensorflow/models/iskhakova2020.py
index e0d179c..d037bb2 100644
--- a/src/ertk/tensorflow/models/iskhakova2020.py
+++ b/src/ertk/tensorflow/models/iskhakova2020.py
@@ -1,13 +1,16 @@
-"""Implementation of model from [1].
+"""Implementation of model from [1]_.
 
-[1] A. O. Iskhakova, D. A. Wolf, R. R. Galin, and M. V. Mamchenko, ‘1-D
-convolutional neural network based on the inner ear principle to
-automatically assess human’s emotional state’, E3S Web Conf., vol. 224,
-p. 01023, 2020, doi: 10.1051/e3sconf/202022401023.
+References
+----------
+.. [1] A. O. Iskhakova, D. A. Wolf, R. R. Galin, and M. V. Mamchenko,
+       '1-D convolutional neural network based on the inner ear
+       principle to automatically assess human's emotional state', E3S
+       Web Conf., vol. 224, p. 01023, 2020, doi:
+       10.1051/e3sconf/202022401023.
 """
 
-from tensorflow.keras import Model
-from tensorflow.keras.layers import (
+from keras import Model
+from keras.layers import (
     BatchNormalization,
     Conv1D,
     Dense,
diff --git a/src/ertk/tensorflow/models/latif2019.py b/src/ertk/tensorflow/models/latif2019.py
index 6cfa48d..e71f860 100644
--- a/src/ertk/tensorflow/models/latif2019.py
+++ b/src/ertk/tensorflow/models/latif2019.py
@@ -1,9 +1,11 @@
-"""Implementation of the model from [1].
+"""Implementation of the model from [1]_.
 
-[1] S. Latif, R. Rana, S. Khalifa, R. Jurdak, and J. Epps, 'Direct
-Modelling of Speech Emotion from Raw Speech', arXiv:1904.03833 [cs,
-eess], Jul. 2019, Accessed: Feb. 10, 2020. [Online]. Available:
-http://arxiv.org/abs/1904.03833.
+References
+----------
+.. [1] S. Latif, R. Rana, S. Khalifa, R. Jurdak, and J. Epps, 'Direct
+       Modelling of Speech Emotion from Raw Speech', arXiv:1904.03833
+       [cs, eess], Jul. 2019, Accessed: Feb. 10, 2020. [Online].
+       Available: http://arxiv.org/abs/1904.03833.
 """
 
 import tensorflow as tf
diff --git a/src/ertk/tensorflow/models/layers.py b/src/ertk/tensorflow/models/layers.py
index 81f746d..0168f43 100644
--- a/src/ertk/tensorflow/models/layers.py
+++ b/src/ertk/tensorflow/models/layers.py
@@ -1,5 +1,9 @@
+"""Custom Keras layers for use in TensorFlow models."""
+
 import tensorflow as tf
-from tensorflow.keras.layers import Layer
+from keras.layers import Layer
+
+__all__ = ["Attention1D"]
 
 
 class Attention1D(Layer):
diff --git a/src/ertk/tensorflow/models/mlp.py b/src/ertk/tensorflow/models/mlp.py
index 5bb597c..a8947ab 100644
--- a/src/ertk/tensorflow/models/mlp.py
+++ b/src/ertk/tensorflow/models/mlp.py
@@ -1,9 +1,11 @@
-"""Basic MLP model."""
+"""Basic multi-layer perceptron (MLP) model."""
 
 from typing import Sequence
 
-from tensorflow.keras.layers import Dense, Dropout, Input
-from tensorflow.keras.models import Model
+from keras.layers import Dense, Dropout, Input
+from keras.models import Model
+
+__all__ = ["model"]
 
 
 def model(
diff --git a/src/ertk/tensorflow/models/rbm.py b/src/ertk/tensorflow/models/rbm.py
index 1e70ba9..5c05594 100644
--- a/src/ertk/tensorflow/models/rbm.py
+++ b/src/ertk/tensorflow/models/rbm.py
@@ -1,3 +1,16 @@
+"""Implementation of Restricted Boltzmann Machines (RBMs) [1]_ in
+TensorFlow.
+
+References
+----------
+.. [1] G. E. Hinton, “A Practical Guide to Training Restricted Boltzmann
+       Machines,” in Neural Networks: Tricks of the Trade: Second
+       Edition, G. Montavon, G. B. Orr, and K.-R. Müller, Eds., in
+       Lecture Notes in Computer Science. Berlin, Heidelberg: Springer
+       Berlin Heidelberg, 2012, pp. 599-619. doi:
+       10.1007/978-3-642-35289-8_32.
+"""
+
 from enum import Enum
 from math import pi
 from pathlib import Path
@@ -35,11 +48,11 @@ class RBM:
 
     References
     ----------
-    [1] G. E. Hinton and R. R. Salakhutdinov, 'Reducing the
+    .. [1] G. E. Hinton and R. R. Salakhutdinov, 'Reducing the
     Dimensionality of Data with Neural Networks', Science, vol. 313, no.
     5786, pp. 504-507, Jul. 2006, doi: 10.1126/science.1127647.
 
-    [2] G. E. Hinton, 'A Practical Guide to Training Restricted
+    .. [2] G. E. Hinton, 'A Practical Guide to Training Restricted
     Boltzmann Machines', in Neural Networks: Tricks of the Trade: Second
     Edition, G. Montavon, G. B. Orr, and K.-R. Müller, Eds. Berlin,
     Heidelberg: Springer Berlin Heidelberg, 2012, pp. 599-619.
@@ -285,7 +298,7 @@ def train(
 
                 tf.summary.scalar("learining_rate", self.learning_rate, step=epoch)
                 tf.summary.scalar("momentum", self.momentum, step=epoch)
-                weights_l2 = tf.reduce_sum(self.W ** 2)
+                weights_l2 = tf.reduce_sum(self.W**2)
                 tf.summary.scalar("weights_l2", weights_l2, step=epoch)
 
                 if self.output_histograms:
diff --git a/src/ertk/tensorflow/models/transformer.py b/src/ertk/tensorflow/models/transformer.py
new file mode 100644
index 0000000..28cf09f
--- /dev/null
+++ b/src/ertk/tensorflow/models/transformer.py
@@ -0,0 +1,61 @@
+"""Basic Transformer model for time series classification."""
+
+import numpy as np
+import tensorflow as tf
+from keras import Model
+from keras.layers import (
+    Dense,
+    GlobalAveragePooling1D,
+    GlobalMaxPooling1D,
+    Input,
+    LayerNormalization,
+    MultiHeadAttention,
+    ReLU,
+    add,
+)
+
+__all__ = ["model"]
+
+
+def _encoder_layer(x, num_heads: int, model_dim: int):
+    att = MultiHeadAttention(num_heads=num_heads, key_dim=model_dim)(x, x)
+    x = LayerNormalization()(x + att)
+    ff = Dense(2048)(x)
+    ff = ReLU()(ff)
+    ff = Dense(model_dim)(ff)
+    return LayerNormalization()(x + ff)
+
+
+def _positional_encoding(seq_len: int, model_dim: int):
+    if model_dim % 2 != 0:
+        raise ValueError("model_dim must be divisible by 2.")
+    pos = np.arange(seq_len, dtype=np.float32)[:, None]
+    dims = np.arange(model_dim // 2, dtype=np.float32)[None, :]
+    angles = pos / np.power(10000, 2 * dims / model_dim)
+    # Need to include batch dimension
+    pe = np.empty((1, seq_len, model_dim))
+    pe[..., 0::2] = np.sin(angles)
+    pe[..., 1::2] = np.cos(angles)
+    return tf.constant(pe, dtype=tf.float32)
+
+
+def model(
+    n_features: int,
+    n_classes: int,
+    enc_layers: int = 2,
+    num_heads: int = 4,
+    seq_len: int = 512,
+    model_dim: int = 512,
+    pool: str = "max",
+) -> Model:
+    inputs = Input((seq_len, n_features))
+    x = Dense(model_dim, name="projection")(inputs)
+    x = add([x, _positional_encoding(seq_len, model_dim)], name="positional_encoding")
+    for _ in range(enc_layers):
+        x = _encoder_layer(x, num_heads=num_heads, model_dim=model_dim)
+    if pool == "max":
+        x = GlobalMaxPooling1D()(x)
+    else:
+        x = GlobalAveragePooling1D()(x)
+    x = Dense(n_classes, activation="softmax")(x)
+    return Model(inputs=inputs, outputs=x)
diff --git a/src/ertk/tensorflow/models/zhang2019.py b/src/ertk/tensorflow/models/zhang2019.py
index 0e00a80..7e3e558 100644
--- a/src/ertk/tensorflow/models/zhang2019.py
+++ b/src/ertk/tensorflow/models/zhang2019.py
@@ -1,10 +1,12 @@
-"""Implementation of the model from [1].
+"""Implementation of the model from [1]_.
 
-[1] Z. Zhang, B. Wu, and B. Schuller, 'Attention-augmented End-to-end
-Multi-task Learning for Emotion Prediction from Speech', in ICASSP 2019
-- 2019 IEEE International Conference on Acoustics, Speech and Signal
-Processing (ICASSP), May 2019, pp. 6705-6709, doi:
-10.1109/ICASSP.2019.8682896.
+References
+----------
+.. [1] Z. Zhang, B. Wu, and B. Schuller, 'Attention-augmented End-to-end
+       Multi-task Learning for Emotion Prediction from Speech', in 2019
+       IEEE International Conference on Acoustics, Speech and Signal
+       Processing, May 2019, pp. 6705-6709, doi:
+       10.1109/ICASSP.2019.8682896.
 """
 
 from typing import Optional
@@ -25,7 +27,7 @@
 )
 from keras.models import Model, Sequential
 
-from ertk.tensorflow.utils import tf_dataset_mem
+from ertk.tensorflow.dataset import tf_dataset_mem
 
 from .layers import Attention1D
 
diff --git a/src/ertk/tensorflow/models/zhao2019.py b/src/ertk/tensorflow/models/zhao2019.py
index 1e0054f..535c787 100644
--- a/src/ertk/tensorflow/models/zhao2019.py
+++ b/src/ertk/tensorflow/models/zhao2019.py
@@ -1,13 +1,13 @@
-"""Implementation of the model from [1].
+"""Implementation of the model from [1]_.
 
-[1] Z. Zhao et al., "Exploring Deep Spectrum Representations via
-Attention-Based Recurrent and Convolutional Neural Networks for Speech
-Emotion Recognition," IEEE Access, vol. 7, pp. 97515–97525, 2019, doi:
-10.1109/ACCESS.2019.2928625.
+.. [1] Z. Zhao et al., "Exploring Deep Spectrum Representations via
+       Attention-Based Recurrent and Convolutional Neural Networks for
+       Speech Emotion Recognition," IEEE Access, vol. 7, pp.
+       97515-97525, 2019, doi: 10.1109/ACCESS.2019.2928625.
 """
 
 import tensorflow as tf
-from tensorflow.keras.layers import (
+from keras.layers import (
     LSTM,
     RNN,
     BatchNormalization,
@@ -22,7 +22,7 @@
     Reshape,
     concatenate,
 )
-from tensorflow.keras.models import Model
+from keras.models import Model
 
 from .layers import Attention1D
 
diff --git a/src/ertk/tensorflow/train.py b/src/ertk/tensorflow/train.py
index 9f90816..0940f9d 100644
--- a/src/ertk/tensorflow/train.py
+++ b/src/ertk/tensorflow/train.py
@@ -1,8 +1,12 @@
+"""TensorFlow training configuration."""
+
 from dataclasses import dataclass
 from typing import Optional
 
 from ertk.config import ERTKConfig
 
+__all__ = ["TFLoggingConfig", "TFTrainConfig"]
+
 
 @dataclass
 class TFLoggingConfig(ERTKConfig):
diff --git a/src/ertk/tensorflow/utils.py b/src/ertk/tensorflow/utils.py
index f1b27c4..d459bed 100644
--- a/src/ertk/tensorflow/utils.py
+++ b/src/ertk/tensorflow/utils.py
@@ -1,3 +1,5 @@
+"""Utility functions for TensorFlow models."""
+
 from functools import wraps
 from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 
@@ -9,13 +11,17 @@
 from keras.models import Model
 from keras.optimizers import Optimizer
 from keras.optimizers import get as get_optimiser
-from keras.utils.data_utils import Sequence
 from sklearn.pipeline import Pipeline
 
-from ertk.utils import batch_arrays_by_length, shuffle_multiple
+__all__ = [
+    "compile_wrap",
+    "test_fit",
+    "init_gpu_memory_growth",
+    "print_linear_model_structure",
+    "TFModelFunction",
+]
 
 TFModelFunction = Callable[..., Union[Model, Pipeline]]
-DataFunction = Callable[..., tf.data.Dataset]
 
 
 def compile_wrap(
@@ -112,257 +118,7 @@ def test_fit(
     model.fit(train_data, validation_data=valid_data, epochs=2, verbose=1)
 
 
-def tf_dataset_gen(
-    x: np.ndarray,
-    y: np.ndarray,
-    sample_weight: Optional[np.ndarray] = None,
-    *,
-    batch_size: int = 64,
-    shuffle: bool = True,
-):
-    """Returns a TensorFlow generator Dataset instance with the given
-    data.
-
-    Parameters
-    ----------
-    x: numpy.ndarray
-        A 2- or 3-D data matrix of shape (n_instances, n_features) or
-        (n_instances, seq_len, n_features).
-    y: numpy.ndarray
-        A 1-D array of length n_instances containing numeric class
-        labels.
-    sample_weight: numpy.ndarray, optional
-        A 1-D array of length n_instances containing sample weights.
-        Added as third item in dataset if present.
-    batch_size: int
-        The batch size to use.
-    shuffle: boolean
-        Whether or not to shuffle the dataset. Note that shuffling is
-        done *before* batching, unlike in `create_tf_dataset_ragged()`.
-    """
-
-    def gen_inst():
-        if shuffle:
-            perm = np.random.permutation(len(x))
-        else:
-            perm = np.arange(len(x))
-
-        if sample_weight is None:
-            for i in perm:
-                yield x[i], y[i]
-        else:
-            for i in perm:
-                yield x[i], y[i], sample_weight[i]
-
-    sig: Tuple[tf.TensorSpec, ...] = (
-        tf.TensorSpec(shape=x[0].shape, dtype=tf.float32),
-        tf.TensorSpec(shape=(), dtype=tf.int64),
-    )
-    if sample_weight is not None:
-        sig += (tf.TensorSpec(shape=(), dtype=tf.float32),)
-    data = tf.data.Dataset.from_generator(gen_inst, output_signature=sig)
-    return data.batch(batch_size).prefetch(2)
-
-
-def tf_dataset_mem(
-    x: np.ndarray,
-    y: np.ndarray,
-    sample_weight: Optional[np.ndarray] = None,
-    *,
-    batch_size: int = 64,
-    shuffle: bool = True,
-) -> tf.data.Dataset:
-    """Returns a TensorFlow in-memory Dataset instance with the given
-    data.
-
-    Parameters
-    ----------
-    x: numpy.ndarray
-        A 2- or 3-D data matrix of shape (n_instances, n_features) or
-        (n_instances, seq_len, n_features).
-    y: numpy.ndarray
-        A 1-D array of length n_instances containing numeric class
-        labels.
-    sample_weight: numpy.ndarray, optional
-        A 1-D array of length n_instances containing sample weights.
-        Added as third item in dataset if present.
-    batch_size: int
-        The batch size to use.
-    shuffle: boolean
-        Whether or not to shuffle the dataset. Note that shuffling is
-        done *before* batching, unlike in `create_tf_dataset_ragged()`.
-    """
-    with tf.device("CPU"):
-        if sample_weight is None:
-            data = tf.data.Dataset.from_tensor_slices((x, y))
-        else:
-            data = tf.data.Dataset.from_tensor_slices((x, y, sample_weight))
-
-    if shuffle:
-        data = data.shuffle(len(x))
-    return data.batch(batch_size).prefetch(2)
-
-
-def tf_dataset_mem_ragged(
-    x: np.ndarray,
-    y: np.ndarray,
-    sample_weight: Optional[np.ndarray] = None,
-    *,
-    batch_size: int = 64,
-    shuffle: bool = True,
-) -> tf.data.Dataset:
-    """Returns a TensorFlow in-memory Dataset instance from
-    variable-length features.
-
-    Parameters
-    ----------
-    x: numpy.ndarray
-        A 3-D data matrix of shape (n_instances, length[i], n_features)
-        with variable length axis 1.
-    y: numpy.ndarray
-        A 1-D array of length n_instances containing numeric class
-        labels.
-    sample_weight: numpy.ndarray, optional
-        A 1-D array of length n_instances containing sample weights.
-        Added as third item in dataset if present.
-    batch_size: int
-        The batch size to use.
-    shuffle: boolean
-        Whether or not to shuffle the dataset. Note that shuffling is
-        done *after* batching, because sequences are sorted by length,
-        then batched in similar lengths.
-    """
-
-    def ragged_to_dense(x: tf.RaggedTensor, y):
-        return x.to_tensor(), y
-
-    def ragged_to_dense_weighted(x: tf.RaggedTensor, y, sample_weight):
-        return x.to_tensor(), y, sample_weight
-
-    # Sort according to length
-    perm = np.argsort([len(a) for a in x])
-    x = x[perm]
-    y = y[perm]
-    if sample_weight is not None:
-        sample_weight = sample_weight[perm]
-
-    ragged = tf.RaggedTensor.from_row_lengths(
-        np.concatenate(list(x)), [len(a) for a in x]
-    )
-    with tf.device("CPU"):
-        if sample_weight is None:
-            data = tf.data.Dataset.from_tensor_slices((ragged, y))
-        else:
-            data = tf.data.Dataset.from_tensor_slices((ragged, y, sample_weight))
-
-    # Group similar lengths in batches, then shuffle batches
-    data = data.batch(batch_size)
-    if shuffle:
-        data = data.shuffle(len(x) // batch_size + 1)
-
-    if sample_weight is None:
-        data = data.map(ragged_to_dense)
-    else:
-        data = data.map(ragged_to_dense_weighted)
-    return data.prefetch(2)
-
-
-class BatchedFrameSequence(Sequence):
-    """Creates a sequence of batches of frames to process.
-
-    Parameters
-    ----------
-    x: ndarray or list of ndarray
-        Sequences of vectors.
-    y: ndarray
-        Labels corresponding to sequences in x.
-    prebatched: bool, default = False
-        Whether or not x has already been grouped into batches.
-    batch_size: int, default = 32
-        Batch size to use. Each generated batch will be at most this size.
-    shuffle: bool, default = True
-        Whether to shuffle the order of the batches.
-    """
-
-    def __init__(
-        self,
-        x: Union[np.ndarray, List[np.ndarray]],
-        y: np.ndarray,
-        prebatched: bool = False,
-        batch_size: int = 32,
-        shuffle: bool = True,
-    ):
-        self.x = x
-        self.y = y
-        if not prebatched:
-            self.x, self.y = batch_arrays_by_length(
-                self.x, self.y, batch_size=batch_size, shuffle=shuffle
-            )
-        if shuffle:
-            self.x, self.y = shuffle_multiple(self.x, self.y, numpy_indexing=True)
-
-    def __len__(self):
-        return len(self.x)
-
-    def __getitem__(self, idx: int):
-        return self.x[idx], self.y[idx]
-
-
-class BatchedSequence(Sequence):
-    """Creates a sequence of batches to process.
-
-    Parameters
-    ----------
-    x: ndarray or list of ndarray
-        Instance feature vectors. Each vector is assumed to be for a different
-        instance.
-    y: ndarray
-        Labels corresponding to sequences in x.
-    prebatched: bool, default = False
-        Whether or not x has already been grouped into batches.
-    batch_size: int, default = 32
-        Batch size to use. Each generated batch will be at most this size.
-    shuffle: bool, default = True
-        Whether to shuffle the instances.
-    """
-
-    def __init__(
-        self, x: np.ndarray, y: np.ndarray, batch_size: int = 32, shuffle: bool = True
-    ):
-        self.x = x
-        self.y = y
-        self.batch_size = batch_size
-        if shuffle:
-            self.x, self.y = shuffle_multiple(self.x, self.y, numpy_indexing=True)
-
-    def __len__(self):
-        return int(np.ceil(len(self.x) / self.batch_size))
-
-    def __getitem__(self, idx: int):
-        sl = slice(idx * self.batch_size, (idx + 1) * self.batch_size)
-        return self.x[sl], self.y[sl]
-
-
-DATA_FNS = {
-    "mem": tf_dataset_mem,
-    "gen": tf_dataset_gen,
-    "mem_ragged": tf_dataset_mem_ragged,
-}
-
-
-def get_data_fn(name: str) -> Callable[..., tf.data.Dataset]:
-    if ":" in name:
-        import importlib
-
-        modname, fname = name.split(":")
-        module = importlib.import_module(modname)
-        func = getattr(module, fname)
-    else:
-        func = DATA_FNS[name]
-    return func
-
-
-def print_linear_model_structure(model: Model):
+def print_linear_model_structure(model: Model) -> None:
     """Prints the structure of a "sequential" model by listing the layer
     types and shapes in order.
 
@@ -392,7 +148,7 @@ def print_inner(model: Layer, depth: int = 0):
     print_inner(model)
 
 
-def init_gpu_memory_growth():
+def init_gpu_memory_growth() -> None:
     """Sets TensorFlow to allocate memory on GPU as needed instead of
     all at once.
     """
diff --git a/src/ertk/train.py b/src/ertk/train.py
index c2c925a..cb4b8f8 100644
--- a/src/ertk/train.py
+++ b/src/ertk/train.py
@@ -1,3 +1,41 @@
+"""
+Training and evaluation classes and functions
+=============================================
+
+Config classes
+--------------
+.. autosummary::
+    :toctree: generated/
+
+    TrainConfig
+    ExperimentResult
+    ExperimentConfig
+    ModelConfig
+    CrossValidationConfig
+    EvalConfig
+
+
+Splitting classes and functions
+-------------------------------
+.. autosummary::
+    :toctree: generated/
+
+    TrainValidation
+    ShuffleGroupKFold
+    ValidationSplit
+    get_cv_splitter
+
+
+Miscellaneous functions
+-----------------------
+.. autosummary::
+    :toctree: generated/
+
+    get_pipeline_params
+    get_scores
+    scores_to_df
+"""
+
 from dataclasses import dataclass, field
 from enum import Enum
 from pathlib import Path
@@ -26,6 +64,22 @@
 from ertk.dataset.dataset import DataLoadConfig, DataSelector
 from ertk.utils import PathOrStr, ScoreFunction, filter_kwargs
 
+__all__ = [
+    "TrainConfig",
+    "ExperimentResult",
+    "ExperimentConfig",
+    "ModelConfig",
+    "CrossValidationConfig",
+    "EvalConfig",
+    "TrainValidation",
+    "ShuffleGroupKFold",
+    "ValidationSplit",
+    "get_cv_splitter",
+    "get_pipeline_params",
+    "get_scores",
+    "scores_to_df",
+]
+
 
 class TrainValidation(BaseCrossValidator):
     """Validation method that uses the training set as validation set."""
@@ -88,7 +142,15 @@ def _iter_test_indices(self, X, y, groups):
 
 
 class ValidationSplit(BaseCrossValidator):
-    """Validation method that uses a pre-defined validation set."""
+    """Validation method that uses a pre-defined validation set.
+
+    Parameters
+    ----------
+    train_idx: list of int or np.ndarray
+        Indices of the training set.
+    valid_idx: list of int or np.ndarray
+        Indices of the validation set.
+    """
 
     def __init__(
         self,
@@ -277,36 +339,89 @@ def scores_to_df(
 
 @dataclass
 class ExperimentResult:
+    """Class to hold results of an experiment."""
+
     scores_dict: Dict[str, np.ndarray] = field(default_factory=dict)
+    """Dictionary of scores. Each key is a string and all values are
+    either scalar or arrays.
+    """
     scores_df: pd.DataFrame = field(default_factory=pd.DataFrame)
+    """Scores dataframe, with a column per key in `scores_dict` and row
+    per evaluation.
+    """
     predictions: Optional[np.ndarray] = None
+    """Array of predictions. If not given, this is None."""
 
 
 @dataclass
 class ModelConfig(ERTKConfig):
+    """Class to hold model configuration."""
+
     type: str = omegaconf.MISSING
+    """Model type. This is used to determine which model class to
+    instantiate.
+    """
     config: Any = None
+    """Model configuration. This is passed to the model class
+    constructor.
+    """
     args: Dict[str, Any] = field(default_factory=dict)
+    """Model arguments. These are passed to the model class constructor
+    as keyword arguments.
+    """
     args_path: Optional[str] = None
+    """Path to a YAML file containing model arguments. These are passed
+    to the model class constructor as keyword arguments.
+    """
     param_grid: Dict[str, Any] = field(default_factory=dict)
+    """Parameter grid for hyperparameter search. This is passed to
+    scikit-learn's `GridSearchCV` class.
+    """
     param_grid_path: Optional[str] = None
+    """Path to a YAML file containing a parameter grid for
+    hyperparameter search. This is passed to scikit-learn's
+    `GridSearchCV` class.
+    """
 
 
 @dataclass
 class CrossValidationConfig(ERTKConfig):
+    """Class to hold cross-validation configuration."""
+
     part: Optional[str] = None
+    """Partition to use for cross-validation. If not given, then
+    cross-validation is performed over instances instead of groups.
+    """
     kfold: int = omegaconf.MISSING
+    """Number of folds for cross-validation."""
     test_size: float = 0.2
+    """Fraction of data to use for testing. This is only used if `kfold`
+    is 1.
+    """
 
 
 @dataclass
 class EvalConfig(ERTKConfig):
+    """Class to hold evaluation configuration."""
+
     cv: Optional[CrossValidationConfig] = None
+    """Cross-validation configuration. If not given, then no
+    cross-validation is performed.
+    """
     train: Optional[DataSelector] = None
+    """Data selector for training data."""
     valid: Optional[DataSelector] = None
+    """Data selector for validation data."""
     test: Optional[DataSelector] = None
+    """Data selector for testing data."""
     inner_kfold: Optional[int] = None
+    """Number of folds for inner cross-validation, if using
+    hyperparameter search.
+    """
     inner_part: Optional[str] = None
+    """Partition to use for inner cross-validation, if using
+    hyperparameter search.
+    """
 
 
 class TransformClass(Enum):
@@ -316,17 +431,32 @@ class TransformClass(Enum):
 
 @dataclass
 class TrainConfig(ERTKConfig):
+    """Class to hold training configuration."""
+
     balanced: bool = True
+    """Whether to use class-balanced sample weights."""
     reps: int = 1
+    """Number of repetitions."""
     normalise: str = "online"
+    """Normalisation method. Can be one of "online", or "none"."""
     transform: TransformClass = TransformClass.std
+    """Transform method for normalisation. Can be one of "std",
+    "minmax", or "none".
+    """
     seq_transform: str = "global"
+    """Transform method for sequence normalisation."""
     n_jobs: int = 1
+    """Number of jobs to run in parallel."""
     verbose: int = 0
+    """Verbosity level."""
     label: str = "label"
+    """Name of label column."""
     sklearn: Any = None
+    """Scikit-learn configuration."""
     pytorch: Any = None
+    """PyTorch configuration."""
     tensorflow: Any = None
+    """TensorFlow configuration."""
 
 
 T = TypeVar("T", bound="ExperimentConfig")
@@ -334,13 +464,22 @@ class TrainConfig(ERTKConfig):
 
 @dataclass
 class ExperimentConfig(ERTKConfig):
+    """Class to hold experiment configuration."""
+
     name: str = "default"
+    """Experiment name."""
     data: DataLoadConfig = omegaconf.MISSING
+    """Data loading configuration."""
     model: ModelConfig = omegaconf.MISSING
+    """Model configuration."""
     eval: Optional[EvalConfig] = None
+    """Evaluation configuration."""
     evals: Dict[str, EvalConfig] = field(default_factory=dict)
+    """Dictionary of evaluation configurations."""
     training: TrainConfig = field(default_factory=TrainConfig)
+    """Training configuration."""
     results: str = ""
+    """Path to output results."""
 
     @classmethod
     def from_file(
diff --git a/src/ertk/transform.py b/src/ertk/transform.py
index 1d42dce..490c11c 100644
--- a/src/ertk/transform.py
+++ b/src/ertk/transform.py
@@ -1,8 +1,30 @@
+"""Transforms to use in estimators.
+
+.. autosummary::
+    :toctree: generated/
+
+    group_transform
+    instance_transform
+    GroupTransformWrapper
+    InstanceTransformWrapper
+    SequenceTransform
+    SequenceTransformWrapper
+"""
+
 import numpy as np
 from sklearn.base import BaseEstimator, TransformerMixin
 
 from ertk.utils import flat_to_inst, inst_to_flat
 
+__all__ = [
+    "group_transform",
+    "instance_transform",
+    "GroupTransformWrapper",
+    "InstanceTransformWrapper",
+    "SequenceTransform",
+    "SequenceTransformWrapper",
+]
+
 
 def group_transform(
     x: np.ndarray,
diff --git a/src/ertk/utils/__init__.py b/src/ertk/utils/__init__.py
index b360531..0c5b31b 100644
--- a/src/ertk/utils/__init__.py
+++ b/src/ertk/utils/__init__.py
@@ -1,3 +1,59 @@
+"""
+Utility functions
+=================
+
+This module contains miscellaneous utility functions.
+
+
+Array utilities
+---------------
+.. autosummary::
+    :toctree: generated/
+
+    batch_arrays_by_length
+    clip_arrays
+    flat_to_inst
+    frame_array
+    frame_arrays
+    inst_to_flat
+    is_mono_audio
+    make_array_array
+    pad_arrays
+    shuffle_multiple
+    transpose_time
+
+
+Generic utilities
+-----------------
+.. autosummary::
+    :toctree: generated/
+
+    batch_iterable
+    filter_kwargs
+    itmap
+    ordered_intersect
+    ordered_subsets
+    subsets
+
+
+Miscellaneous utilities
+-----------------------
+.. autosummary::
+    :toctree: generated/
+
+    TqdmMultiprocessing
+    TqdmParallel
+
+
+Type definitions
+----------------
+.. autosummary::
+    :toctree: generated/
+
+    PathOrStr
+    ScoreFunction
+"""
+
 from .array import (
     batch_arrays_by_length,
     clip_arrays,
diff --git a/src/ertk/utils/array.py b/src/ertk/utils/array.py
index baf23a3..b54b66f 100644
--- a/src/ertk/utils/array.py
+++ b/src/ertk/utils/array.py
@@ -1,11 +1,39 @@
+"""Utility functions for working with numpy arrays."""
+
 from typing import List, Optional, Sequence, Tuple, Union, overload
 
 import numpy as np
 
+__all__ = [
+    "batch_arrays_by_length",
+    "make_array_array",
+    "check_3d",
+    "is_mono_audio",
+    "frame_array",
+    "frame_arrays",
+    "pad_array",
+    "pad_arrays",
+    "clip_arrays",
+    "inst_to_flat",
+    "flat_to_inst",
+    "transpose_time",
+    "shuffle_multiple",
+]
+
 
 def make_array_array(x: List[np.ndarray]) -> np.ndarray:
-    """Helper function to make an array of arrays. The returned array
-    has `shape==(len(x),)` and `dtype==object`.
+    """Helper function to make an array of arrays.
+
+    Parameters
+    ----------
+    x: list of np.ndarray
+        The arrays to put in the array of arrays.
+
+    Returns
+    -------
+    arr: np.ndarray
+        The array of arrays, with ``shape==(len(x),)``, and
+        ``dtype==object``.
     """
     arr = np.empty(len(x), dtype=object)
     for i, a in enumerate(x):
@@ -16,6 +44,16 @@ def make_array_array(x: List[np.ndarray]) -> np.ndarray:
 def check_3d(arrays: Union[Sequence[np.ndarray], np.ndarray]) -> None:
     """Checks if an array is 3D or each array in a list is 2D. Raises an
     exception if this isn't the case.
+
+    Parameters
+    ----------
+    arrays: np.ndarray or list of np.ndarray
+        The arrays to check.
+
+    Raises
+    ------
+    ValueError
+        If any array is not 3D.
     """
     if any(len(x.shape) != 2 for x in arrays):
         raise ValueError("arrays must be 3D (contiguous or vlen).")
@@ -27,6 +65,16 @@ def is_mono_audio(x: np.ndarray) -> bool:
 
     Note that this method will also return True for '2D' sequences of
     length 1, so care should be taken for edge cases.
+
+    Parameters
+    ----------
+    x: np.ndarray
+        The array to check.
+
+    Returns
+    -------
+    is_mono: bool
+        Whether the array represents mono audio data.
     """
     return np.squeeze(x).ndim == 1
 
@@ -235,10 +283,22 @@ def pad_arrays(arrays: np.ndarray, pad: int) -> np.ndarray:
 
 
 def pad_arrays(arrays: Union[List[np.ndarray], np.ndarray], pad: int = 32):
-    """Pads each array to the nearest multiple of `pad` greater than the
-    array size. Assumes axis 0 of each sub-array, or axis 1 of x, is
-    the time axis. This is mainly a wrapper around `pad_array()` for
-    instance arrays.
+    """Pads each array to the nearest multiple of ``pad`` greater than the
+    array size.  This is mainly a wrapper around `pad_array()` that
+    takes care of padding ragged or non-contiguous arrays.
+
+    Parameters
+    ----------
+    arrays: list of np.ndarray
+        The arrays to process. Assumes axis 0 of each sub-array, or axis
+        1 of ``arrays``, is the time axis.
+    pad: int
+        The multiple to pad to.
+
+    Returns
+    -------
+    padded_arrays:
+        Padded arrays.
     """
     if isinstance(arrays, np.ndarray) and len(arrays.shape) > 1:
         # Contiguous 2D/3D
@@ -267,7 +327,21 @@ def clip_arrays(arrays: np.ndarray, length: int, copy: bool) -> np.ndarray:
 def clip_arrays(
     arrays: Union[List[np.ndarray], np.ndarray], length: int, copy: bool = True
 ):
-    """Clips each array to the specified maximum length."""
+    """Clips each array to the specified maximum sequence length.
+
+    Parameters
+    ----------
+    arrays: list of np.ndarray
+        The arrays to process. Assumes axis 0 of each sub-array, or axis
+        1 of ``arrays``, is the time axis.
+    length: int
+        The maximum length.
+
+    Returns
+    -------
+    clipped_arrays:
+        Clipped arrays.
+    """
     if isinstance(arrays, np.ndarray):
         if len(arrays.shape) > 1:
             return arrays[:, :length, ...].copy() if copy else arrays[:, :length, ...]
@@ -294,6 +368,17 @@ def transpose_time(arrays: Union[List[np.ndarray], np.ndarray]):
     array be 2-D.
 
     Note: This function modifies the arrays in-place.
+
+    Parameters
+    ----------
+    arrays: list of np.ndarray
+        The arrays to process. Assumes axis 0 of each sub-array, or axis
+        1 of ``arrays``, is the time axis.
+
+    Returns
+    -------
+    transposed_arrays:
+        Transposed arrays.
     """
     check_3d(arrays)
     if isinstance(arrays, np.ndarray) and len(arrays.shape) == 3:
@@ -322,6 +407,7 @@ def shuffle_multiple(
         Whether to use NumPy-style indexing or list comprehension.
 
     Returns:
+    --------
     shuffled_arrays: iterable of array-like
         The shuffled arrays.
     """
diff --git a/src/ertk/utils/generic.py b/src/ertk/utils/generic.py
index 71a58b2..b836c0b 100644
--- a/src/ertk/utils/generic.py
+++ b/src/ertk/utils/generic.py
@@ -1,3 +1,5 @@
+"""General utilities."""
+
 from itertools import chain, combinations, permutations, zip_longest
 from typing import (
     Any,
@@ -12,6 +14,15 @@
     overload,
 )
 
+__all__ = [
+    "itmap",
+    "ordered_intersect",
+    "batch_iterable",
+    "filter_kwargs",
+    "subsets",
+    "ordered_subsets",
+]
+
 T = TypeVar("T")
 T1 = TypeVar("T1")
 T2 = TypeVar("T2")
@@ -125,7 +136,7 @@ def batch_iterable(
         `batch_size` doesn't divide the iterable length.)
 
     Yields
-    -------
+    ------
     tuple
         A tuple of length `batch_size` with successive elements from the
         original iterable. If `return_last == False` then the last batch
diff --git a/src/ertk/utils/misc.py b/src/ertk/utils/misc.py
index 9c4082c..c8ef298 100644
--- a/src/ertk/utils/misc.py
+++ b/src/ertk/utils/misc.py
@@ -1,3 +1,5 @@
+"""Miscellaneous utilities."""
+
 import multiprocessing
 import os
 from functools import partial
@@ -6,6 +8,8 @@
 import joblib
 import tqdm
 
+__all__ = ["TqdmParallel", "TqdmMultiprocessing"]
+
 
 # Class adapted from user394430's answer here:
 # https://stackoverflow.com/a/61900501/10044861
diff --git a/src/ertk/utils/types.py b/src/ertk/utils/types.py
index 6571507..a31dbb8 100644
--- a/src/ertk/utils/types.py
+++ b/src/ertk/utils/types.py
@@ -1,8 +1,12 @@
+"""Type definitions."""
+
 from os import PathLike
 from typing import Callable, Union
 
 import numpy as np
 
+__all__ = ["PathOrStr", "ScoreFunction"]
+
 
 PathOrStr = Union[PathLike, str]
 ScoreFunction = Callable[[np.ndarray, np.ndarray], float]


Test patch:
diff --git a/src/ertk/tensorflow/models/test_fit.py b/src/ertk/tensorflow/models/test_fit.py
index 127e1a0..44a89ae 100644
--- a/src/ertk/tensorflow/models/test_fit.py
+++ b/src/ertk/tensorflow/models/test_fit.py
@@ -5,7 +5,8 @@
 
 from ertk.config import get_arg_mapping
 from ertk.tensorflow import get_tf_model
-from ertk.tensorflow.utils import init_gpu_memory_growth, tf_dataset_gen, tf_dataset_mem
+from ertk.tensorflow.dataset import tf_dataset_gen, tf_dataset_mem
+from ertk.tensorflow.utils import init_gpu_memory_growth
 
 
 @click.command()