Repository: CouncilDataProject/speakerbox (language: python)

Issue:
Use CML for CI/CD
https://cml.dev/doc/usage
feature/dataset-expansion-and-model-training
<!--
  Thank you for submitting a pull request!

  ⚠️⚠️ Please do the following before submitting: ⚠️⚠️

  - Read the CONTRIBUTING.md guide and make sure you've followed all the steps given.
  - Ensure that the code is up-to-date with the `main` branch.
  - Provide or update documentation for any feature added by your pull request.
  - Provide relevant tests for your feature or bug fix.

  ❗️ Also: ❗️

  Please name your pull request {development-type}/{short-description}.
  For example: feature/read-tiff-files
-->

### Description of Changes

<!-- Include a description of the proposed changes. -->

This adds the whole process of expanding a dataset from gecko:

```csv
gecko_annotation_file_1, audio_file_1,
gecko_annotation_file_2, audio_file_2,
...
```

into an annotated speaker id training set
```csv
speaker_label_1, audio_clip_1,
speaker_label_1, audio_clip_2,
speaker_label_2, audio_clip_3,
speaker_label_3, audio_clip_4,
...
```

and then training a new speaker identification model by fine tuning an existing large model from transformers / huggingface.

I'm now on my laptop and didn't save the confusion matrix but last time I trained on my desktop the model finished training with ~99% accuracy. I am a bit worried about overfitting but honestly idk, I am even evaluating on a holdout validation set 🤷 

The README is updated and if you have an nvidia GPU feel free to give the quickstart a go, if you don't have a GPU (like me right now on my laptop), I wouldn't try training because it takes ~4 hours to complete on CPU.


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/build-docs.yml b/.github/workflows/build-docs.yml
index 7fd1d8b..bd4e6f9 100644
--- a/.github/workflows/build-docs.yml
+++ b/.github/workflows/build-docs.yml
@@ -9,16 +9,18 @@ jobs:
   docs:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2.3.1
-        with:
-          persist-credentials: false
+      - uses: actions/checkout@v2
       - name: Set up Python
         uses: actions/setup-python@v2
         with:
-          python-version: 3.9
+          python-version: '3.10'
+      - name: Install System Dependencies
+        run: |
+          sudo apt-get install libsndfile1
+          sudo apt-get install ffmpeg
       - name: Install Dependencies
         run: |
-          pip install --upgrade pip
+          python -m pip install --upgrade pip
           pip install .[dev]
       - name: Generate Docs
         run: |
diff --git a/.github/workflows/build-main.yml b/.github/workflows/build-main.yml
index 7ca8378..ee9c2ed 100644
--- a/.github/workflows/build-main.yml
+++ b/.github/workflows/build-main.yml
@@ -15,12 +15,11 @@ on:
 
 jobs:
   test:
-    runs-on: ${{ matrix.os }}
+    runs-on: ubuntu-latest
     strategy:
       fail-fast: false
       matrix:
-        python-version: [3.8, 3.9]
-        os: [ubuntu-latest, macOS-latest, windows-latest]
+        python-version: [3.8, 3.9, '3.10']
 
     steps:
     - uses: actions/checkout@v2
@@ -28,12 +27,16 @@ jobs:
       uses: actions/setup-python@v2
       with:
         python-version: ${{ matrix.python-version }}
+    - name: Install System Dependencies
+      run: |
+        sudo apt-get install libsndfile1
+        sudo apt-get install ffmpeg
     - name: Install Dependencies
       run: |
         python -m pip install --upgrade pip
         pip install .[test]
     - name: Run tests with Tox
-      run: tox -e py
+      run: tox -e py -- --cpu
     - name: Upload Codecov
       uses: codecov/codecov-action@v1
 
@@ -45,7 +48,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v2
       with:
-        python-version: 3.9
+        python-version: '3.10'
     - name: Install Dependencies
       run: |
         python -m pip install --upgrade pip
@@ -63,7 +66,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v2
       with:
-        python-version: 3.9
+        python-version: '3.10'
     - name: Install Dependencies
       run: |
         python -m pip install --upgrade pip
diff --git a/.github/workflows/check-pr.yml b/.github/workflows/check-pr.yml
index 810ceb5..d38a012 100644
--- a/.github/workflows/check-pr.yml
+++ b/.github/workflows/check-pr.yml
@@ -4,12 +4,11 @@ on: pull_request
 
 jobs:
   test:
-    runs-on: ${{ matrix.os }}
+    runs-on: ubuntu-latest
     strategy:
       fail-fast: false
       matrix:
-        python-version: [3.8, 3.9]
-        os: [ubuntu-latest, macOS-latest, windows-latest]
+        python-version: [3.8, 3.9, '3.10']
 
     steps:
     - uses: actions/checkout@v2
@@ -17,15 +16,29 @@ jobs:
       uses: actions/setup-python@v2
       with:
         python-version: ${{ matrix.python-version }}
+    - name: Install System Dependencies
+      run: |
+        sudo apt-get install libsndfile1
+        sudo apt-get install ffmpeg
     - name: Install Dependencies
       run: |
         python -m pip install --upgrade pip
         pip install .[test]
     - name: Run tests with Tox
-      run: tox -e py
+      run: tox -e py -- --cpu
     - name: Upload Codecov
       uses: codecov/codecov-action@v1
-
+    - uses: iterative/setup-cml@v1
+      if: matrix.python-version == '3.10'
+    - name: Create CML report
+      if: matrix.python-version == '3.10'
+      env:
+        REPO_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+      run: |
+        cml publish test-outputs/trained-speakerbox/validation-confusion.png --md >> test-outputs/trained-speakerbox/results.md
+        currentHash=$(git rev-parse --short HEAD)
+        echo "**Results from commit:** $currentHash" >> test-outputs/trained-speakerbox/results.md
+        cml send-comment test-outputs/trained-speakerbox/results.md
   lint:
     runs-on: ubuntu-latest
 
@@ -34,7 +47,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v2
       with:
-        python-version: 3.9
+        python-version: '3.10'
     - name: Install Dependencies
       run: |
         python -m pip install --upgrade pip
@@ -49,7 +62,11 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v2
       with:
-        python-version: 3.9
+        python-version: '3.10'
+    - name: Install System Dependencies
+      run: |
+        sudo apt-get install libsndfile1
+        sudo apt-get install ffmpeg
     - name: Install Dependencies
       run: |
         pip install --upgrade pip
diff --git a/.gitignore b/.gitignore
index e1922e3..c00a4bd 100644
--- a/.gitignore
+++ b/.gitignore
@@ -111,5 +111,15 @@ ENV/
 # VSCode
 .vscode
 
+# Project specific
+chunked-audio/
+trained-speakerbox/
+test-outputs/
+test-output-trained-speakerbox/
+speakerbox/tests/data/diarized/
+
 # Datasets
-seattle-2021-proto/
\ No newline at end of file
+seattle-2021-proto/
+training-data/
+unannotated/
+cdp-datasets/
\ No newline at end of file
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 2752391..32dfa95 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -16,7 +16,7 @@ repos:
       hooks:
           - id: isort
     - repo: https://github.com/psf/black
-      rev: 21.9b0
+      rev: 22.3.0
       hooks:
           - id: black
     - repo: https://github.com/pre-commit/mirrors-mypy
diff --git a/MANIFEST.in b/MANIFEST.in
index 7b27fc1..d4c4407 100644
--- a/MANIFEST.in
+++ b/MANIFEST.in
@@ -8,4 +8,4 @@ recursive-exclude * __pycache__
 recursive-exclude * *.py[co]
 
 recursive-include docs *.rst conf.py Makefile make.bat *.jpg *.png *.gif
-graft speakerbox/data
+graft speakerbox/datasets/data
diff --git a/Makefile b/Makefile
index 5c3e5b1..8219891 100644
--- a/Makefile
+++ b/Makefile
@@ -33,6 +33,7 @@ clean:  ## Clean all build, python, and testing files
 	rm -fr build/
 	rm -fr dist/
 	rm -fr .eggs/
+	rm -fr test-outputs/
 	find . -name '*.egg-info' -exec rm -fr {} +
 	find . -name '*.egg' -exec rm -f {} +
 	find . -name '*.pyc' -exec rm -f {} +
@@ -46,9 +47,12 @@ clean:  ## Clean all build, python, and testing files
 	rm -fr .pytest_cache
 	rm -fr .mypy_cache
 
-build: ## run tox / run tests and lint
+build: ## run tox / run tests and lint / use GPU for training test
 	tox
 
+build-cpu: ## run tox / run tests and lint / use CPU to training test
+	tox -- --cpu
+
 gen-docs: ## generate Sphinx HTML documentation, including API docs
 	rm -f docs/speakerbox*.rst
 	rm -f docs/modules.rst
diff --git a/README.md b/README.md
index 3821fc1..a19bf26 100644
--- a/README.md
+++ b/README.md
@@ -4,7 +4,7 @@
 [![Documentation](https://github.com/CouncilDataProject/speakerbox/workflows/Documentation/badge.svg)](https://CouncilDataProject.github.io/speakerbox)
 [![Code Coverage](https://codecov.io/gh/CouncilDataProject/speakerbox/branch/main/graph/badge.svg)](https://codecov.io/gh/CouncilDataProject/speakerbox)
 
-Speaker Annotation for Transcripts using Audio Classification
+Workflow for Annotation and Speaker Classification Model Training from Audio
 
 ---
 
@@ -17,20 +17,176 @@ Speaker Annotation for Transcripts using Audio Classification
 
 For full package documentation please visit [councildataproject.github.io/speakerbox](https://councildataproject.github.io/speakerbox).
 
-## Quickstart
+![Speakerbox example workflow](https://raw.githubusercontent.com/CouncilDataProject/speakerbox/main/docs/_static/images/workflow.png)
 
-Load the 2021 Seattle Prototype Dataset, get summary statistics
-about speaker time, finally pull the matching audio file for each annotation file
-and store annotation file matched to audio as a `pandas.DataFrame`.
+## Problem
+
+Given a set of recordings of multi-speaker conversations or meetings:
+
+```
+example/
+├── 0.wav
+├── 1.wav
+├── 2.wav
+├── 3.wav
+├── 4.wav
+└── 5.wav
+```
+
+Where each conversation has some or all of a consistent set of speakers, for example:
+
+-   0.wav -- contains speakers: A, B, C, D, E
+-   1.wav -- contains speakers: B, D, E
+-   2.wav -- contains speakers: A, B, C
+-   3.wav -- contains speakers: A, B, C, D, E
+-   4.wav -- contains speakers: A, C, D
+-   5.wav -- contains speakers: A, B, C, D, E
+
+You want to train a model to classify portions of audio as one of the N known speakers
+in future conversations not included in your original training set.
+
+`f(audio) -> [(start_time, end_time, speaker), (start_time, end_time, speaker), ...]`
+
+i.e. `f(audio) -> [(2.4, 10.5, "A"), (10.8, 14.1, "D"), (14.8, 22.7, "B"), ...]`
+
+The `speakerbox` library contains methods for both generating datasets for annotation
+and for utilizing multiple audio annotation schemes to train such a model.
+
+## Workflow
+
+⚠️ `speakerbox` currently only works on Ubuntu due to upstream dependencies ⚠️
+
+### Diarization
+
+We quickly generate an annotated dataset by first diarizing (or clustering based
+on the features of speaker audio) portions of larger audio files and splitting each the
+of the clusters into their own directories that you can then manually clean up
+(by removing incorrectly clustered audio segments).
+
+#### Notes
+
+-   It is recommended to have each larger audio file named with a unique id that
+    can be used to act as a "conversation id".
+-   Diarization time depends on machine resources and make take a long time -- one
+    potential recommendation is to run a diarization script overnight and clean up the
+    produced annotations the following day.
+-   During this process audio will be duplicated in the form of smaller audio clips --
+    ensure you have enough space on your machine to complete this process before
+    you begin.
+-   Clustering accuracy depends on how many speakers there are, how distinct their
+    voices are, and how much speech is talking over one-another.
+-   If possible, try to find meetings where speakers have a roughly uniform distribution
+    of speaking durations.
 
 ```python
-from speakerbox import datasets
+from speakerbox import preprocess
+
+diarized_and_split_audio_dir = preprocess.diarize_and_split_audio("0.wav")
+```
+
+### Cleaning
+
+Diarization will produce a directory structure organized by unlabeled speakers with
+the audio clips that were clustered together.
+
+For example, if `"0.wav"` had three speakers, the produced directory structure may look
+like the following tree:
+
+```
+0/
+├── SPEAKER_00
+│   ├── 567-12928.wav
+│   ├── ...
+│   └── 76192-82901.wav
+├── SPEAKER_01
+│   ├── 34123-38918.wav
+│   ├── ...
+│   └── 88212-89111.wav
+└── SPEAKER_02
+    ├── ...
+    └── 53998-62821.wav
+```
 
-seattle_2021_ds_dir = datasets.unpack_seattle_2021_proto(clean=True)
-seattle_2021_ds_summary_stats = datasets.summarize_annotation_statistics(
-    seattle_2021_ds_dir / "annotations"
+We leave it to you as a user to then go through these directories and remove any audio
+clips that were incorrectly clustered together as well as renaming the sub-directories
+to their correct speaker labels
+(from "SPEAKER_00", "SPEAKER_01", "SPEAKER_02", etc. to "A", "B", "C", etc.).
+
+#### Notes
+
+-   Most operating systems have an audio playback application to queue an entire directory
+    of audio files as a playlist for playback. This makes it easy to listen to a whole
+    unlabeled sub-directory (i.e. "SPEAKER_00") at a time and pause playback and remove
+    files from the directory which were incorrectly clustered.
+
+### Training Preparation
+
+Once you have annotated what you think is enough conversations, you can try preparing
+a dataset for training.
+
+The following functions will prepare the audio for training by:
+
+1. Finding all labeled audio clips in the provided directories
+2. Chunk all found audio clips into smaller duration clips _(parametrizable)_
+3. Check that the provided annotated dataset meets the following conditions:
+    1. There is enough data such that the training, test, and validation subsets all
+       contain different conversation ids.
+    2. There is enough data such that the training, test, and validation subsets each
+       contain all labels present in the whole dataset.
+
+#### Notes
+
+-   During this process audio will be duplicated in the form of smaller audio clips --
+    ensure you have enough space on your machine to complete this process before
+    you begin.
+-   Directory names are used as conversation ids during dataset construction.
+
+```python
+from speakerbox import preprocess
+
+dataset = preprocess.expand_labeled_diarized_audio_dir_to_dataset(
+    labeled_diarized_audio_dir=[
+        "0/",  # The cleaned and checked audio clips for conversation id 0
+        "1/",  # ... conversation id 1
+        "2/",  # ... conversation id 2
+        "3/",  # ... conversation id 3
+        "4/",  # ... conversation id 4
+        "5/",  # ... conversation id 5
+    ]
 )
-seattle_2021_ds = datasets.pull_seattle_2021_proto_audio()
+
+dataset_dict, value_counts = preprocess.prepare_dataset(dataset)
+
+# You can print the value_counts dataframe to see how many audio clips of each label
+# (speaker) is present in each data subset.
+print(value_counts)
+```
+
+### Model Training and Evaluation
+
+Once you have your dataset prepared and available, you can provide it directly to the
+training function to begin training a new model.
+
+The `eval_model` function will store a filed called `results.md` with the accuracy,
+precision, and recall of the model and additionally store a file called
+`validation-confusion.png` which is a
+[confusion matrix](https://en.wikipedia.org/wiki/Confusion_matrix).
+
+#### Notes
+
+-   The model (and evaluation metrics) will be stored in a new directory called
+    `trained-speakerbox` _(parametrizable)_.
+-   Training time depends on how much data you have annotated and provided.
+-   It is recommended to train with an NVidia GPU with CUDA available to speed up
+    the training process.
+
+```python
+from speakerbox import train, eval_model
+
+# dataset_dict comes from previous preparation step
+train(dataset_dict)
+
+eval_model(dataset_dict["valid"])
 ```
 
 ## Development
diff --git a/docs/_static/images/workflow.png b/docs/_static/images/workflow.png
new file mode 100644
index 0000000..1350f35
Binary files /dev/null and b/docs/_static/images/workflow.png differ
diff --git a/docs/conf.py b/docs/conf.py
index 49a46ab..86cc6fc 100755
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -21,8 +21,6 @@
 import os
 import sys
 
-import sphinx_rtd_theme
-
 import speakerbox
 
 sys.path.insert(0, os.path.abspath(".."))
@@ -60,18 +58,18 @@
 # You can specify multiple suffix as a list of string:
 #
 source_suffix = {
-        ".rst": "restructuredtext",
-        ".txt": "markdown",
-        ".md": "markdown",
+    ".rst": "restructuredtext",
+    ".txt": "markdown",
+    ".md": "markdown",
 }
 
 # The master toctree document.
 master_doc = "index"
 
 # General information about the project.
-project = u"speakerbox"
-copyright = u'2022, Jackson Maxfield Brown'
-author = u"Jackson Maxfield Brown"
+project = "speakerbox"
+copyright = "2022, Jackson Maxfield Brown"
+author = "Jackson Maxfield Brown"
 
 # The version info for the project you"re documenting, acts as replacement
 # for |version| and |release|, also used in various other places throughout
@@ -106,16 +104,13 @@
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes.
 #
-html_theme = "sphinx_rtd_theme"
+html_theme = "furo"
 
 # Theme options are theme-specific and customize the look and feel of a
 # theme further.  For a list of options available for each theme, see the
 # documentation.
 #
-html_theme_options = {
-    "collapse_navigation": False,
-    "prev_next_buttons_location": "top",
-}
+html_theme_options = {}
 
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
@@ -135,15 +130,12 @@
     # The paper size ("letterpaper" or "a4paper").
     #
     # "papersize": "letterpaper",
-
     # The font size ("10pt", "11pt" or "12pt").
     #
     # "pointsize": "10pt",
-
     # Additional stuff for the LaTeX preamble.
     #
     # "preamble": "",
-
     # Latex figure (float) alignment
     #
     # "figure_align": "htbp",
@@ -153,9 +145,13 @@
 # (source start file, target name, title, author, documentclass
 # [howto, manual, or own class]).
 latex_documents = [
-    (master_doc, "speakerbox.tex",
-     u"speakerbox Documentation",
-     u"Jackson Maxfield Brown", "manual"),
+    (
+        master_doc,
+        "speakerbox.tex",
+        "speakerbox Documentation",
+        "Jackson Maxfield Brown",
+        "manual",
+    ),
 ]
 
 
@@ -163,11 +159,7 @@
 
 # One entry per manual page. List of tuples
 # (source start file, name, description, authors, manual section).
-man_pages = [
-    (master_doc, "speakerbox",
-     u"speakerbox Documentation",
-     [author], 1)
-]
+man_pages = [(master_doc, "speakerbox", "speakerbox Documentation", [author], 1)]
 
 
 # -- Options for Texinfo output ----------------------------------------
@@ -176,10 +168,13 @@
 # (source start file, target name, title, author,
 #  dir menu entry, description, category)
 texinfo_documents = [
-    (master_doc, "speakerbox",
-     u"speakerbox Documentation",
-     author,
-     "speakerbox",
-     "One line description of project.",
-     "Miscellaneous"),
+    (
+        master_doc,
+        "speakerbox",
+        "speakerbox Documentation",
+        author,
+        "speakerbox",
+        "One line description of project.",
+        "Miscellaneous",
+    ),
 ]
diff --git a/docs/index.rst b/docs/index.rst
index ddf49a8..5359c68 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -10,7 +10,6 @@ Welcome to speakerbox's documentation!
    installation
    Package modules <modules>
    contributing
-   math
 
 .. mdinclude:: ../README.md
 
diff --git a/docs/modules.rst b/docs/modules.rst
new file mode 100644
index 0000000..40b4265
--- /dev/null
+++ b/docs/modules.rst
@@ -0,0 +1,7 @@
+speakerbox
+==========
+
+.. toctree::
+   :maxdepth: 4
+
+   speakerbox
diff --git a/scripts/pull-and-diarize-seattle-2021.py b/scripts/pull-and-diarize-seattle-2021.py
new file mode 100644
index 0000000..c106a19
--- /dev/null
+++ b/scripts/pull-and-diarize-seattle-2021.py
@@ -0,0 +1,47 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import logging
+from pathlib import Path
+
+import torch
+from cdp_data import CDPInstances, datasets
+
+from speakerbox.preprocess import diarize_and_split_audio
+
+###############################################################################
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="[%(levelname)4s: %(module)s:%(lineno)4s %(asctime)s] %(message)s",
+)
+log = logging.getLogger(__name__)
+
+###############################################################################
+
+# Pull specific meetings
+for start_date, end_date in [
+    ("2021-05-24", "2021-05-25"),
+    ("2021-06-07", "2021-06-08"),
+    ("2021-09-20", "2021-09-21"),
+    ("2021-06-28", "2021-06-29"),
+    ("2021-07-12", "2021-07-13"),
+]:
+    datasets.get_session_dataset(
+        CDPInstances.Seattle,
+        start_datetime=start_date,
+        end_datetime=end_date,
+        store_audio=True,
+    )
+
+
+dataset_dir = Path(f"cdp-datasets/{CDPInstances.Seattle}")
+for audio_file in dataset_dir.glob("event-*/session-*/audio.wav"):
+    storage = audio_file.parent.parent.name
+    if not (Path(storage).exists() or Path(f"ANNOTATED-{storage}").exists()):
+        print("working on file:", audio_file)
+        print("storing:", storage)
+        torch.cuda.empty_cache()
+        diarize_and_split_audio(audio_file, storage_dir=storage)
+    else:
+        print("skipping", storage)
diff --git a/scripts/train-seattle-2021.py b/scripts/train-seattle-2021.py
new file mode 100644
index 0000000..03f37aa
--- /dev/null
+++ b/scripts/train-seattle-2021.py
@@ -0,0 +1,54 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import logging
+
+import pandas as pd
+
+from speakerbox import eval_model, preprocess, train
+from speakerbox.datasets import seattle_2021_proto
+
+###############################################################################
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="[%(levelname)4s: %(module)s:%(lineno)4s %(asctime)s] %(message)s",
+)
+log = logging.getLogger(__name__)
+
+###############################################################################
+
+# Pull / prep original Seattle data
+seattle_2021_ds_items = seattle_2021_proto.pull_all_files(
+    annotations_dir="training-data/seattle-2021-proto/annotations/",
+    transcript_output_dir="training-data/seattle-2021-proto/unlabeled_transcripts/",
+    audio_output_dir="training-data/seattle-2021-proto/audio/",
+)
+seattle_2021_ds = preprocess.expand_gecko_annotations_to_dataset(
+    seattle_2021_ds_items,
+    overwrite=True,
+)
+
+# Expand diarized data
+diarized_ds = ds = preprocess.expand_labeled_diarized_audio_dir_to_dataset(
+    [
+        "training-data/diarized/01e7f8bb1c03/",
+        "training-data/diarized/2cdf68ae3c2c/",
+        "training-data/diarized/6d6702d7b820/",
+        "training-data/diarized/9f55f22d8e61/",
+        "training-data/diarized/9f581faa5ece/",
+    ],
+    overwrite=True,
+)
+
+# Combine into single
+combined_ds = pd.concat([seattle_2021_ds, diarized_ds], ignore_index=True)
+
+# Generate train test validate splits
+dataset, _ = preprocess.prepare_dataset(combined_ds, equalize_data=False)
+
+# dataset.save_to_disk(SOME_PATH)
+
+# Train a model
+model_path = train(dataset)
+eval_model(dataset["valid"])
diff --git a/setup.py b/setup.py
index edbe72a..f86ddb9 100644
--- a/setup.py
+++ b/setup.py
@@ -12,8 +12,12 @@
     "pytest-runner>=5.2",
 ]
 
+seattle_reqs = [
+    "cdp-backend~=3.0.3",
+]
+
 test_requirements = [
-    "black>=19.10b0",
+    "black>=22.3.0",
     "codecov>=2.1.4",
     "flake8>=3.8.3",
     "flake8-debugger>=3.2.1",
@@ -23,6 +27,7 @@
     "pytest-cov>=2.9.0",
     "pytest-raises>=0.11",
     "tox>=3.15.2",
+    *seattle_reqs,
 ]
 
 dev_requirements = [
@@ -32,22 +37,33 @@
     "coverage>=5.1",
     "jupyterlab>=3.2.8",
     "m2r2>=0.2.7",
-    "pytest-runner>=5.2",
     "Sphinx>=3.4.3",
-    "sphinx_rtd_theme>=0.5.1",
+    "furo>=2022.4.7",
     "twine>=3.1.1",
     "wheel>=0.34.2",
 ]
 
 requirements = [
-    "cdp-backend>=3.0.3",
-    "pandas>=1.4.0",
+    "dataclasses_json~=0.5",
+    "datasets[audio]~=1.18",
+    "librosa~=0.8",
+    "matplotlib~=3.5",
+    "pandas~=1.0",
+    "pydub~=0.25",
+    "scikit-learn~=1.0",
+    "speechbrain~=0.5.11",
+    "torch~=1.10",
+    "torchaudio~=0.10",
+    "transformers~=4.16",
+    # until v2 release
+    "pyannote.audio @ https://github.com/pyannote/pyannote-audio/archive/develop.zip",
 ]
 
 extra_requirements = {
     "setup": setup_requirements,
     "test": test_requirements,
     "dev": dev_requirements,
+    "seattle": seattle_reqs,
     "all": [
         *requirements,
         *dev_requirements,
@@ -58,12 +74,13 @@
     author="Jackson Maxfield Brown",
     author_email="jmaxfieldbrown@gmail.com",
     classifiers=[
-        "Development Status :: 2 - Pre-Alpha",
+        "Development Status :: 3 - Alpha",
         "Intended Audience :: Developers",
         "License :: OSI Approved :: MIT License",
         "Natural Language :: English",
         "Programming Language :: Python :: 3.8",
         "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
     ],
     description="Speaker Annotation for Transcripts using Audio Classification",
     entry_points={
diff --git a/speakerbox/__init__.py b/speakerbox/__init__.py
index 29d2817..87b7dab 100644
--- a/speakerbox/__init__.py
+++ b/speakerbox/__init__.py
@@ -11,3 +11,6 @@
 
 def get_module_version() -> str:
     return __version__
+
+
+from .main import eval_model, train  # noqa: F401
diff --git a/speakerbox/datasets.py b/speakerbox/datasets.py
deleted file mode 100644
index c66a1a7..0000000
--- a/speakerbox/datasets.py
+++ /dev/null
@@ -1,378 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-import json
-import shutil
-import statistics
-from concurrent.futures import ThreadPoolExecutor
-from dataclasses import dataclass
-from functools import partial
-from pathlib import Path
-from typing import Dict, List, NamedTuple, Optional, Union
-
-import fireo
-import pandas as pd
-from cdp_backend.database import models as db_models
-from gcsfs import GCSFileSystem
-from google.auth.credentials import AnonymousCredentials
-from google.cloud.firestore import Client
-
-###############################################################################
-
-DATASETS_DIR = (Path(__file__).parent / "data").resolve(strict=True)
-SEATTLE_2021_PROTOTYPE_DATASET = "seattle-2021-proto.zip"
-SEATTLE_2021_PROTOTYPE_DATASET_DIR = Path(SEATTLE_2021_PROTOTYPE_DATASET).with_suffix(
-    ""
-)
-SEATTLE_2021_PROTOTYPE_ANNOTATIONS_DIR = (
-    SEATTLE_2021_PROTOTYPE_DATASET_DIR / "annotations"
-)
-SEATTLE_2021_PROTOTYPE_AUDIO_DIR = SEATTLE_2021_PROTOTYPE_DATASET_DIR / "audio"
-SEATTLE_2021_INFRA_STR = "cdp-seattle-staging-dbengvtn"
-
-###############################################################################
-# Dataset Utilities
-
-
-def _unpack_zip(
-    zipfile: Union[str, Path],
-    dest: Union[str, Path],
-    clean: bool = False,
-) -> Path:
-    """
-    Unzips the zipfile to the destination location.
-
-    Parameters
-    ----------
-    zipfile: Union[str, Path]
-        The zipfile to unpack.
-    dest: Union[str, Path]
-        The destination to unpack to.
-    clean: bool
-        If a directory already exists at the destination location, should the directory
-        be removed entirely before unpacking again.
-        Default: False
-
-    Returns
-    -------
-    dataset_path: Path
-        The path to the unpacked data.
-
-    Raises
-    ------
-    NotADirectoryError
-        A file exists at the specified destination.
-    FileExistsError
-        A directory exists at the specified destination and is not empty.
-    """
-    zipfile = Path(zipfile).resolve(strict=True)
-    dest = Path(dest)
-
-    # Check dest is a dir
-    if dest.is_file():
-        raise NotADirectoryError(dest)
-
-    # Clean dest if required
-    if dest.is_dir():
-        dest_is_empty = not any(dest.iterdir())
-        if not dest_is_empty:
-            if clean:
-                shutil.rmtree(dest)
-            else:
-                raise FileExistsError(
-                    f"Files found in target destination directory ({dest}) "
-                    f"but parameter `clean` set to False."
-                )
-
-    # Make new dir
-    dest.mkdir(parents=True, exist_ok=True)
-
-    # Extract
-    shutil.unpack_archive(zipfile, dest)
-
-    # Return extracted data dir
-    return dest.resolve(strict=True)
-
-
-def unpack_seattle_2021_proto(
-    dest: Union[str, Path] = SEATTLE_2021_PROTOTYPE_DATASET_DIR,
-    clean: bool = False,
-) -> Path:
-    """
-    Unpacks the Seattle 2021 Prototype Dataset from a zipfile to the provided
-    destination location.
-
-    Parameters
-    ----------
-    dest: Union[str, Path]
-        The destination to unpack to.
-        Default: A new directory in your current working directory with the standard
-        Seattle 2021 Prototype Dataset name.
-    clean: bool
-        If a directory already exists at the destination location, should the directory
-        be removed entirely before unpacking again.
-        Default: False
-
-    Returns
-    -------
-    dataset_path: Path
-        The path to the unpacked data.
-
-    Raises
-    ------
-    NotADirectoryError
-        A file exists at the specified destination.
-    FileExistsError
-        A directory exists at the specified destination and is not empty.
-    """
-    return _unpack_zip(
-        zipfile=DATASETS_DIR / SEATTLE_2021_PROTOTYPE_DATASET,
-        dest=dest,
-        clean=clean,
-    )
-
-
-class AnnotationAndAudio(NamedTuple):
-    annotation_file: Path
-    audio_file: Path
-
-
-def _get_matching_audio_for_seattle_2021_proto_annotation_file(
-    annotation_file: Path,
-    audio_output_dir: Path,
-    fs: GCSFileSystem,
-    overwrite: bool,
-) -> Optional[AnnotationAndAudio]:
-    # Remove suffix and use just the file name for the event id
-    event_id = annotation_file.with_suffix("").name
-
-    # Pull sessions with matching event id
-    matching_sessions = list(
-        db_models.Session.collection.filter(
-            "event_ref", "==", f"{db_models.Event.collection_name}/{event_id}"
-        ).fetch()
-    )
-
-    # There should only be a single session so drop this event if more than one
-    if len(matching_sessions) > 1:
-        return None
-
-    # Find any transcript name for this session and pull the content hash
-    matching_session = matching_sessions[0]
-    matching_transcript = db_models.Transcript.collection.filter(
-        "session_ref", "==", matching_session.key
-    ).get()
-    matching_transcript_file = matching_transcript.file_ref.get()
-    content_hash = matching_transcript_file.name.split("-")[0]
-
-    # Check overwrite or skip
-    audio_save_path = audio_output_dir / f"{event_id}.wav"
-    if audio_save_path.exists():
-        if overwrite:
-            fs.get(
-                f"gs://{SEATTLE_2021_INFRA_STR}.appspot.com/{content_hash}-audio.wav",
-                str(audio_save_path),
-            )
-    else:
-        fs.get(
-            f"gs://{SEATTLE_2021_INFRA_STR}.appspot.com/{content_hash}-audio.wav",
-            str(audio_save_path),
-        )
-
-    return AnnotationAndAudio(
-        annotation_file=annotation_file,
-        audio_file=audio_save_path,
-    )
-
-
-def pull_seattle_2021_proto_audio(
-    annotations_dir: Union[str, Path] = SEATTLE_2021_PROTOTYPE_ANNOTATIONS_DIR,
-    audio_output_dir: Union[str, Path] = SEATTLE_2021_PROTOTYPE_AUDIO_DIR,
-    overwrite: bool = False,
-) -> pd.DataFrame:
-    """
-    Using the filenames for the files found in the annotations directory as event
-    ids, pull the matching audio files and store them in the dataset directory.
-
-    Parameters
-    ----------
-    annotations_dir: Union[str, Path]
-        The path to the directory which contains the annotation files for a dataset.
-        Default: Default Seattle 2021 Prototype Dataset Annotations Directory
-    audio_output_dir: Union[str, Path]
-        The path to where to store the pulled audio files
-        Default: Default Seattle 2021 Prototype Dataset Audio Directory
-    overwrite: bool
-        If an audio file with a filename matching the event id already exists in the
-        output directory exists, should it be overwritten (safer) or skipped (faster).
-        Default: False (skipped)
-
-    Returns
-    -------
-    full_dataset: pd.DataFrame
-        A minimal dataframe with columns for the annotation file and the
-        matching audio file.
-
-    Raises
-    ------
-    NotADirectoryError
-        Provided annotations_dir is not a directory.
-    NotADirectoryError
-        A file exists at the specified audio output directory.
-
-    Examples
-    --------
-    >>> from speakerbox.datasets import (
-    ...     unpack_seattle_2021_proto,
-    ...     pull_seattle_2021_proto_audio,
-    ... )
-    ... ds_dir = unpack_seattle_2021_proto()
-    ... ds = pull_seattle_2021_proto_audio()
-    """
-    # Ensure dataset dir is valid
-    annotations_dir = Path(annotations_dir).resolve(strict=True)
-    audio_output_dir = Path(audio_output_dir).resolve()
-    if annotations_dir.is_file():
-        raise NotADirectoryError(annotations_dir)
-    if audio_output_dir.exists() and audio_output_dir.is_file():
-        raise NotADirectoryError(audio_output_dir)
-
-    # Make output dir
-    audio_output_dir.mkdir(parents=True, exist_ok=True)
-
-    # Create connection to DB and FS
-    fireo.connection(
-        client=Client(
-            project=SEATTLE_2021_INFRA_STR, credentials=AnonymousCredentials()
-        )
-    )
-    fs = GCSFileSystem(project=SEATTLE_2021_INFRA_STR, token="anon")
-
-    # Create partial function for threading
-    downloader = partial(
-        _get_matching_audio_for_seattle_2021_proto_annotation_file,
-        audio_output_dir=audio_output_dir,
-        fs=fs,
-        overwrite=overwrite,
-    )
-
-    # Thread audio file downloading
-    with ThreadPoolExecutor() as exe:
-        downloads = list(exe.map(downloader, annotations_dir.glob("*.json")))
-
-    # Unpack results
-    results: List[Dict[str, str]] = []
-    for annotation_and_audio in downloads:
-        if annotation_and_audio is not None:
-            results.append(
-                {
-                    "annotations": str(annotation_and_audio.annotation_file),
-                    "audio": str(annotation_and_audio.audio_file),
-                }
-            )
-
-    return pd.DataFrame(results)
-
-
-###############################################################################
-# Dataset Summarization
-
-
-@dataclass
-class SpeakerTime:
-    label: str
-    duration: float
-    percent_of_total: float
-
-
-@dataclass
-class DatasetSummary:
-    n_speakers: int
-    n_events: int
-    speaker_times: List[SpeakerTime]
-    mean_speaker_time: float
-    min_speaker_time: SpeakerTime
-    max_speaker_time: SpeakerTime
-    std_speaker_time: float
-    total_speaker_time: float
-
-
-def summarize_annotation_statistics(
-    annotations_dir: Union[str, Path],
-) -> DatasetSummary:
-    """
-    Summarize all annotation files into a single pandas DataFrame.
-
-    Parameters
-    ----------
-    annotations_dir: Union[str, Path]
-        The path to the directory which contains the annotation files for a dataset.
-
-    Returns
-    -------
-    summary: DatasetSummary
-        The summary stats for the dataset.
-
-    Examples
-    --------
-    >>> from speakerbox.datasets import (
-    ...     unpack_seattle_2021_proto,
-    ...     summarize_annotation_statistics,
-    ... )
-    ... ds_dir = unpack_seattle_2021_proto()
-    ... summary_stats = summarize_annotation_statistics(ds_dir / "annotations")
-
-    Raises
-    ------
-    FileNotFoundError
-        The provided dataset directory was not found.
-    FileNotFoundError
-        The provided dataset directory did not contain any annotation files.
-    NotADirectoryError
-        THe provided dataset directory was a file and not a directory.
-    """
-    # Ensure dataset dir is valid
-    annotations_dir = Path(annotations_dir).resolve(strict=True)
-    if annotations_dir.is_file():
-        raise NotADirectoryError(annotations_dir)
-
-    # Add up the speaking time for all speakers found in every JSON
-    # annotation file in the provided dataset directory
-    speaker_lut: Dict[str, float] = {}
-    n_events = 0
-    for annotation_file in annotations_dir.glob("*.json"):
-        # Open new annotation file
-        with open(annotation_file, "r") as open_f:
-            annotations = json.load(open_f)
-
-        # Iter through each "monologue" section and add up times
-        for monologue in annotations["monologues"]:
-            speaker = monologue["speaker"]["id"]
-            duration = monologue["end"] - monologue["start"]
-            if speaker not in speaker_lut:
-                speaker_lut[speaker] = duration
-            else:
-                speaker_lut[speaker] += duration
-
-        n_events += 1
-
-    # Compute summary stats
-    total_speaker_time = sum([duration for speaker, duration in speaker_lut.items()])
-    speaker_times: List[SpeakerTime] = [
-        SpeakerTime(
-            label=speaker,
-            duration=duration,
-            percent_of_total=duration / total_speaker_time,
-        )
-        for speaker, duration in speaker_lut.items()
-    ]
-    return DatasetSummary(
-        n_speakers=len(speaker_times),
-        n_events=n_events,
-        speaker_times=speaker_times,
-        mean_speaker_time=total_speaker_time / len(speaker_times),
-        min_speaker_time=min(speaker_times, key=lambda st: st.duration),
-        max_speaker_time=max(speaker_times, key=lambda st: st.duration),
-        std_speaker_time=statistics.stdev([st.duration for st in speaker_times]),
-        total_speaker_time=total_speaker_time,
-    )
diff --git a/speakerbox/datasets/__init__.py b/speakerbox/datasets/__init__.py
new file mode 100644
index 0000000..cf88464
--- /dev/null
+++ b/speakerbox/datasets/__init__.py
@@ -0,0 +1,3 @@
+# -*- coding: utf-8 -*-
+
+"""Datasets package for speakerbox."""
diff --git a/speakerbox/data/seattle-2021-proto.zip b/speakerbox/datasets/data/seattle-2021-proto.zip
similarity index 100%
rename from speakerbox/data/seattle-2021-proto.zip
rename to speakerbox/datasets/data/seattle-2021-proto.zip
diff --git a/speakerbox/datasets/seattle_2021_proto.py b/speakerbox/datasets/seattle_2021_proto.py
new file mode 100644
index 0000000..10d8bd4
--- /dev/null
+++ b/speakerbox/datasets/seattle_2021_proto.py
@@ -0,0 +1,248 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+
+import logging
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+from pathlib import Path
+from typing import List, Optional, Union
+
+import fireo
+from cdp_backend.database import models as db_models
+from gcsfs import GCSFileSystem
+from google.auth.credentials import AnonymousCredentials
+from google.cloud.firestore import Client
+from tqdm import tqdm
+
+from ..types import GeckoAnnotationAndAudio
+from ..utils import _unpack_zip
+
+###############################################################################
+
+DATASETS_DIR = (Path(__file__).parent / "data").resolve(strict=True)
+SEATTLE_2021_PROTOTYPE_DATASET = "seattle-2021-proto.zip"
+SEATTLE_2021_PROTOTYPE_DATASET_DIR = Path(SEATTLE_2021_PROTOTYPE_DATASET).with_suffix(
+    ""
+)
+SEATTLE_2021_PROTOTYPE_ANNOTATIONS_DIR = (
+    SEATTLE_2021_PROTOTYPE_DATASET_DIR / "annotations"
+)
+SEATTLE_2021_PROTOTYPE_TRANSCRIPT_DIR = (
+    SEATTLE_2021_PROTOTYPE_DATASET_DIR / "unlabeled_transcripts"
+)
+SEATTLE_2021_PROTOTYPE_AUDIO_DIR = SEATTLE_2021_PROTOTYPE_DATASET_DIR / "audio"
+SEATTLE_2021_INFRA_STR = "cdp-seattle-staging-dbengvtn"
+
+###############################################################################
+
+log = logging.getLogger(__name__)
+
+###############################################################################
+
+
+def unpack(
+    dest: Union[str, Path] = SEATTLE_2021_PROTOTYPE_DATASET_DIR,
+    clean: bool = False,
+) -> Path:
+    """
+    Unpacks the Seattle 2021 Prototype Dataset from a zipfile to the provided
+    destination location.
+
+    Parameters
+    ----------
+    dest: Union[str, Path]
+        The destination to unpack to.
+        Default: A new directory in your current working directory with the standard
+        Seattle 2021 Prototype Dataset name.
+    clean: bool
+        If a directory already exists at the destination location, should the directory
+        be removed entirely before unpacking again.
+        Default: False
+
+    Returns
+    -------
+    dataset_path: Path
+        The path to the unpacked data.
+
+    Raises
+    ------
+    NotADirectoryError
+        A file exists at the specified destination.
+    FileExistsError
+        A directory exists at the specified destination and is not empty.
+    """
+    return _unpack_zip(
+        zipfile=DATASETS_DIR / SEATTLE_2021_PROTOTYPE_DATASET,
+        dest=dest,
+        clean=clean,
+    )
+
+
+def _get_matching_unlabelled_transcript_and_audio(
+    annotation_file: Path,
+    transcript_output_dir: Path,
+    audio_output_dir: Path,
+    fs: GCSFileSystem,
+    overwrite: bool,
+    progress_bar: Optional[tqdm] = None,
+) -> Optional[GeckoAnnotationAndAudio]:
+    # Remove suffix and use just the file name for the event id
+    event_id = annotation_file.with_suffix("").name
+
+    # Pull sessions with matching event id
+    matching_sessions = list(
+        db_models.Session.collection.filter(
+            "event_ref", "==", f"{db_models.Event.collection_name}/{event_id}"
+        ).fetch()
+    )
+
+    # There should only be a single session so drop this event if more than one
+    if len(matching_sessions) > 1:
+        log.debug(
+            f"Multiple sessions found for event: {event_id}, "
+            f"dropping annotations from dataset."
+        )
+        return None
+
+    # Find any transcript name for this session and pull the content hash
+    matching_session = matching_sessions[0]
+    matching_transcript = db_models.Transcript.collection.filter(
+        "session_ref", "==", matching_session.key
+    ).get()
+    matching_transcript_file = matching_transcript.file_ref.get()
+    content_hash = matching_transcript_file.name.split("-")[0]
+
+    # Check overwrite or skip
+    transcript_save_path = transcript_output_dir / f"{event_id}.json"
+    audio_save_path = audio_output_dir / f"{event_id}.wav"
+    remote_audio_uri = (
+        f"gs://{SEATTLE_2021_INFRA_STR}.appspot.com/{content_hash}-audio.wav"
+    )
+
+    # Transcript
+    if transcript_save_path.exists():
+        if overwrite:
+            fs.get(matching_transcript_file.uri, str(transcript_save_path))
+        else:
+            log.debug(
+                f"Using existing transcript file found for event: {event_id}, "
+                f"({transcript_save_path})."
+            )
+    else:
+        fs.get(matching_transcript_file.uri, str(transcript_save_path))
+
+    # Audio
+    if audio_save_path.exists():
+        if overwrite:
+            fs.get(remote_audio_uri, str(audio_save_path))
+        else:
+            log.debug(
+                f"Using existing audio file found for event: {event_id}, "
+                f"({audio_save_path})."
+            )
+    else:
+        fs.get(remote_audio_uri, str(audio_save_path))
+
+    # Update progress
+    if progress_bar:
+        progress_bar.update()
+
+    return GeckoAnnotationAndAudio(
+        annotation_file=annotation_file,
+        audio_file=audio_save_path,
+    )
+
+
+def pull_all_files(
+    annotations_dir: Union[str, Path] = SEATTLE_2021_PROTOTYPE_ANNOTATIONS_DIR,
+    transcript_output_dir: Union[str, Path] = SEATTLE_2021_PROTOTYPE_TRANSCRIPT_DIR,
+    audio_output_dir: Union[str, Path] = SEATTLE_2021_PROTOTYPE_AUDIO_DIR,
+    overwrite: bool = False,
+) -> List[GeckoAnnotationAndAudio]:
+    """
+    Using the filenames for the files found in the annotations directory as event
+    ids, pull the matching audio files and store them in the dataset directory.
+
+    Parameters
+    ----------
+    annotations_dir: Union[str, Path]
+        The path to the directory which contains the annotation files for a dataset.
+        Default: Default Seattle 2021 Prototype Dataset Annotations Directory
+    transcript_output_dir: Union[str, Path]
+        The path to where to store the pulled unlabeled transcript files.
+        Default: Default Seattle 2021 Prototype Dataset Transcript Directory
+    audio_output_dir: Union[str, Path]
+        The path to where to store the pulled audio files.
+        Default: Default Seattle 2021 Prototype Dataset Audio Directory
+    overwrite: bool
+        If an audio file with a filename matching the event id already exists in the
+        output directory exists, should it be overwritten (safer) or skipped (faster).
+        Default: False (skipped)
+
+    Returns
+    -------
+    annotations_and_audios: List[GeckoAnnotationAndAudio]
+        A list of matching GeckoAnnotationAndAudio objects.
+
+    Raises
+    ------
+    NotADirectoryError
+        Provided annotations directory is not a directory.
+    NotADirectoryError
+        A file exists at the specified audio output directory.
+
+    Examples
+    --------
+    >>> from speakerbox.datasets import (
+    ...     unpack_seattle_2021_proto,
+    ...     pull_seattle_2021_proto_audio,
+    ... )
+    ... ds_dir = unpack_seattle_2021_proto()
+    ... ds = pull_seattle_2021_proto_audio()
+    """
+    # Ensure dataset dir is valid
+    annotations_dir = Path(annotations_dir).resolve(strict=True)
+    transcript_output_dir = Path(transcript_output_dir).resolve()
+    audio_output_dir = Path(audio_output_dir).resolve()
+    if annotations_dir.is_file():
+        raise NotADirectoryError(annotations_dir)
+    if transcript_output_dir.exists() and transcript_output_dir.is_file():
+        raise NotADirectoryError(transcript_output_dir)
+    if audio_output_dir.exists() and audio_output_dir.is_file():
+        raise NotADirectoryError(audio_output_dir)
+
+    # Make output dirs
+    transcript_output_dir.mkdir(parents=True, exist_ok=True)
+    audio_output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Create connection to DB and FS
+    fireo.connection(
+        client=Client(
+            project=SEATTLE_2021_INFRA_STR, credentials=AnonymousCredentials()
+        )
+    )
+    fs = GCSFileSystem(project=SEATTLE_2021_INFRA_STR, token="anon")
+
+    # Get filepaths used for pulling extra files
+    annotation_files = list(annotations_dir.glob("*.json"))
+
+    # Create progress bar
+    pbar = tqdm(total=len(annotation_files))
+
+    # Create partial function for threading
+    downloader = partial(
+        _get_matching_unlabelled_transcript_and_audio,
+        transcript_output_dir=transcript_output_dir,
+        audio_output_dir=audio_output_dir,
+        fs=fs,
+        overwrite=overwrite,
+        progress_bar=pbar,
+    )
+
+    # Thread audio file downloading
+    with ThreadPoolExecutor() as exe:
+        results = list(exe.map(downloader, annotation_files))
+
+    # Filter out bad events
+    return [r for r in results if isinstance(r, GeckoAnnotationAndAudio)]
diff --git a/speakerbox/main.py b/speakerbox/main.py
new file mode 100644
index 0000000..1c797c6
--- /dev/null
+++ b/speakerbox/main.py
@@ -0,0 +1,291 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import logging
+from pathlib import Path
+from typing import Any, Dict, Optional, Tuple
+
+import datasets
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import transformers
+from datasets import Audio, Dataset, DatasetDict, arrow_dataset, load_metric
+from sklearn.metrics import (
+    ConfusionMatrixDisplay,
+    accuracy_score,
+    log_loss,
+    precision_score,
+    recall_score,
+)
+from transformers import (
+    EvalPrediction,
+    Trainer,
+    TrainingArguments,
+    Wav2Vec2FeatureExtractor,
+    Wav2Vec2ForSequenceClassification,
+    feature_extraction_utils,
+    pipeline,
+)
+
+from .utils import set_global_seed
+
+###############################################################################
+
+log = logging.getLogger(__name__)
+
+###############################################################################
+
+DEFAULT_BASE_MODEL = "superb/wav2vec2-base-superb-sid"
+
+EVAL_RESULTS_TEMPLATE = """
+## Results
+
+* **Accuracy:** {accuracy}
+* **Precision:** {precision}
+* **Recall:** {recall}
+* **Validation Loss:** {loss}
+
+### Confusion
+"""
+
+DEFAULT_TRAINER_ARGUMENTS_ARGS = dict(
+    evaluation_strategy="epoch",
+    save_strategy="epoch",
+    learning_rate=3e-5,
+    per_device_train_batch_size=8,
+    gradient_accumulation_steps=1,
+    eval_accumulation_steps=40,
+    per_device_eval_batch_size=8,
+    num_train_epochs=5,
+    warmup_ratio=0.1,
+    logging_steps=10,
+    load_best_model_at_end=True,
+    metric_for_best_model="accuracy",
+    gradient_checkpointing=True,
+)
+
+
+###############################################################################
+
+
+def eval_model(
+    validation_dataset: Dataset,
+    model_name: str = "trained-speakerbox",
+) -> Tuple[float, float, float, float]:
+    """
+    Evaluate a trained model.
+
+    This will store two files in the model directory, one for the accuracy, precision,
+    and recall in a markdown file and the other is the generated top one confusion
+    matrix as a PNG file.
+
+    Parameters
+    ----------
+    validation_dataset: Dataset
+        The dataset to validate the model against.
+    model_name: str
+        A name for the model. This will also create a directory with the same name
+        to store the produced model in.
+        Default: "trained-speakerbox"
+
+    Returns
+    -------
+    accuracy: float
+        The model accuracy as returned by sklearn.metrics.accuracy_score.
+    precision: float
+        The model (weighted) precision as returned by sklearn.metrics.precision_score.
+    recall: float
+        The model (weighted) recall as returned by sklearn.metrics.recall_score.
+    loss: float
+        The model log loss as returned by sklearn.metrics.log_loss.
+    """
+    log.info("Setting up evaluation pipeline")
+    classifier = pipeline(
+        "audio-classification",
+        model=model_name,
+    )
+    log.info("Running eval")
+
+    def predict(example: datasets.arrow_dataset.Example) -> Dict[str, Any]:
+        pred = classifier(example["audio"], top_k=1000)
+        pred_as_dict = {i["label"]: i["score"] for i in pred}
+        top_pred = max(pred_as_dict, key=pred_as_dict.get)  # type: ignore
+        return {
+            "pred_label": top_pred,
+            "true_label": classifier.model.config.id2label[example["label"]],
+            "pred_scores": [i["score"] for i in pred],
+        }
+
+    validation_dataset = validation_dataset.map(predict)
+
+    # Create confusion
+    ConfusionMatrixDisplay.from_predictions(
+        validation_dataset["true_label"],
+        validation_dataset["pred_label"],
+    )
+    plt.xticks(rotation=45)
+    plt.yticks(rotation=45)
+    plt.savefig(f"{model_name}/validation-confusion.png", bbox_inches="tight")
+
+    # Compute metrics
+    accuracy = accuracy_score(
+        y_true=validation_dataset["true_label"],
+        y_pred=validation_dataset["pred_label"],
+    )
+    precision = precision_score(
+        y_true=validation_dataset["true_label"],
+        y_pred=validation_dataset["pred_label"],
+        average="weighted",
+    )
+    recall = recall_score(
+        y_true=validation_dataset["true_label"],
+        y_pred=validation_dataset["pred_label"],
+        average="weighted",
+    )
+    loss = log_loss(
+        y_true=validation_dataset["true_label"],
+        y_pred=validation_dataset["pred_scores"],
+    )
+
+    # Store metrics
+    with open(f"{model_name}/results.md", "w") as open_f:
+        open_f.write(
+            EVAL_RESULTS_TEMPLATE.format(
+                accuracy=accuracy,
+                precision=precision,
+                recall=recall,
+                loss=loss,
+            )
+        )
+
+    return (accuracy, precision, recall, loss)
+
+
+def train(
+    dataset: DatasetDict,
+    model_name: str = "trained-speakerbox",
+    model_base: str = DEFAULT_BASE_MODEL,
+    max_duration: float = 2.0,
+    seed: Optional[int] = None,
+    use_cpu: bool = False,
+    trainer_arguments_kws: Dict[str, Any] = DEFAULT_TRAINER_ARGUMENTS_ARGS,
+) -> Path:
+    """
+    Train a speaker classification model.
+
+    Parameters
+    ----------
+    dataset: DatasetDict
+        The datasets to use for training, testing, and validation.
+        Should only contain the columns/features: "label" and "audio".
+        The values in the "audio" column should be paths to the audio files.
+    model_name: str
+        A name for the model. This will also create a directory with the same name
+        to store the produced model in.
+        Default: "trained-speakerbox"
+    model_base: str
+        The model base to use before fine tuning.
+    max_duration: float
+        The maximum duration to use for each audio clip.
+        Any clips longer than this will be trimmed.
+        Default: 2.0
+    seed: Optional[int]
+        Seed to pass to torch, numpy, and Python RNGs.
+        Default: None (do not set a seed)
+    use_cpu: bool
+        Should the model be trained using CPU.
+        This also sets `no_cuda=True` on TrainerArguments.
+        Default: False (use GPU if available)
+    trainer_arguments_kws: Dict[Any]
+        Any additional keyword arguments to be passed to the HuggingFace
+        TrainerArguments object.
+        Default: DEFAULT_TRAINER_ARGUMENTS_ARGS
+
+    Returns
+    -------
+    model_storage_path: Path
+        The path to the directory where the model is stored.
+    """
+    # Handle seed
+    if seed:
+        set_global_seed(seed)
+
+    # Handle cpu
+    if use_cpu:
+        trainer_arguments_kws["no_cuda"] = True
+
+    # Load feature extractor
+    feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_base)
+
+    # Convert dataset audios
+    log.info("Casting all audio paths to torch Audio")
+    dataset = dataset.cast_column("audio", Audio(feature_extractor.sampling_rate))
+
+    # Construct label to id and vice-versa LUTs
+    label2id, id2label = dict(), dict()
+    for i, label in enumerate(dataset["train"].features["label"].names):
+        label2id[label] = str(i)
+        id2label[str(i)] = label
+
+    # Construct preprocessing function
+    def preprocess(
+        examples: arrow_dataset.Batch,
+    ) -> feature_extraction_utils.BatchFeature:
+        audio_arrays = [x["array"] for x in examples["audio"]]
+        inputs = feature_extractor(
+            audio_arrays,
+            sampling_rate=feature_extractor.sampling_rate,
+            max_length=int(feature_extractor.sampling_rate * max_duration),
+            do_normalize=True,
+            truncation=True,
+            padding=True,
+        )
+        return inputs
+
+    # Encode the dataset
+    log.info("Extracting features from audio")
+    dataset = dataset.map(preprocess, batched=True)
+
+    # Create AutoModel
+    log.info("Setting up Trainer")
+    model = Wav2Vec2ForSequenceClassification.from_pretrained(
+        model_base,
+        num_labels=len(id2label),
+        label2id=label2id,
+        id2label=id2label,
+        ignore_mismatched_sizes=True,
+    )
+
+    # Create fine tuning Trainer
+    args = TrainingArguments(
+        model_name,
+        **trainer_arguments_kws,
+    )
+
+    # Compute accuracy metrics
+    metric = load_metric("accuracy")
+
+    def compute_metrics(eval_pred: EvalPrediction) -> Optional[Dict]:
+        # Eval pred comes with both the predictions and the attention mask
+        # grab just the predictions
+        predictions = np.argmax(eval_pred.predictions[0], axis=-1)
+        return metric.compute(predictions=predictions, references=eval_pred.label_ids)
+
+    # Trainer and train!
+    trainer = Trainer(
+        model,
+        args,
+        train_dataset=dataset["train"],
+        eval_dataset=dataset["test"],
+        tokenizer=feature_extractor,
+        compute_metrics=compute_metrics,
+    )
+    torch.cuda.empty_cache()
+    transformers.logging.set_verbosity_info()
+    trainer.train()
+
+    # Save model
+    trainer.save_model()
+
+    return Path(model_name).resolve()
diff --git a/speakerbox/preprocess.py b/speakerbox/preprocess.py
new file mode 100644
index 0000000..5eb618e
--- /dev/null
+++ b/speakerbox/preprocess.py
@@ -0,0 +1,580 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import json
+import logging
+from pathlib import Path
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import pandas as pd
+from datasets import Dataset, DatasetDict
+from pyannote.audio import Pipeline
+from pydub import AudioSegment
+from sklearn.model_selection import train_test_split
+
+from .types import AnnotatedAudio, GeckoAnnotationAndAudio
+from .utils import set_global_seed
+
+###############################################################################
+
+log = logging.getLogger(__name__)
+
+###############################################################################
+# Annotation
+
+
+def diarize_and_split_audio(
+    audio_file: Union[str, Path],
+    storage_dir: Optional[Union[str, Path]] = None,
+    min_audio_chunk_duration: float = 0.5,
+    diarization_pipeline: Optional[Pipeline] = None,
+    seed: Optional[int] = None,
+) -> Path:
+    """
+    Diarize a single audio file and split the file into smaller chunks stored into
+    directories with the unlabeled speaker annotation.
+
+    Parameters
+    ----------
+    audio_file: Union[str, Path]
+        The audio file to diarize and split.
+    storage_dir: Optional[Union[str, Path]]
+        A specific directory to store the produced chunks to.
+        Default: None (use the audio file name to create a new directory)
+    min_audio_chunk_duration: float
+        Length of the minimum audio duration to allow through after chunking.
+        default: 0.5 seconds
+    diarization_pipeline: Optional[Pipeline]
+        A preloaded PyAnnote Pipeline.
+        Default: None (load default)
+    seed: Optional[int]
+        Seed to pass to torch, numpy, and Python RNGs.
+        Default: None (do not set a seed)
+
+    Returns
+    -------
+    storage_dir: Path
+        The path to where all the chunked audio was stored.
+
+    See Also
+    --------
+    expand_labeled_diarized_audio_dir_to_dataset
+        After labeling the audio in the produced diarized audio directory, expand
+        the labeled data into a dataset ready for training.
+
+    Notes
+    -----
+    The directory structure of the produced chunks will follow the pattern::
+
+        {storage_dir}/
+        ├── SPEAKER_00
+        │   ├── {start_time_millis}-{start_end_millis}.wav
+        │   └── {start_time_millis}-{start_end_millis}.wav
+        ├── SPEAKER_01
+        │   ├── {start_time_millis}-{start_end_millis}.wav
+        │   └── {start_time_millis}-{start_end_millis}.wav
+    """
+    if seed:
+        set_global_seed(seed)
+
+    # Handle audio file
+    if isinstance(audio_file, str):
+        audio_file = Path(audio_file)
+
+    # Ensure file exists
+    audio_file = audio_file.resolve(strict=True)
+
+    # Handle storage dir
+    if isinstance(storage_dir, str):
+        storage_dir = Path(storage_dir).resolve()
+    elif storage_dir is None:
+        # No storage dir provided
+        # Make one using the audio file name
+        audio_file_name = audio_file.with_suffix("").name
+        storage_dir = Path(audio_file_name)
+
+    # Make storage dir
+    storage_dir.mkdir(parents=True, exist_ok=True)
+
+    # Load pipeline
+    if diarization_pipeline is None:
+        diarization_pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization")
+
+    # Diarize
+    dia = diarization_pipeline(audio_file)
+
+    # Load audio once to pull and store chunks
+    audio = AudioSegment.from_file(audio_file)
+
+    # Iter labeled segments and store
+    for turn, _, speaker in dia.itertracks(yield_label=True):
+        # Make speaker dir if needed
+        speaker_dir = storage_dir / speaker.replace(" ", "_")
+        speaker_dir.mkdir(exist_ok=True)
+
+        # Only store chunk if duration is longer than param
+        if (turn.end - turn.start) >= min_audio_chunk_duration:
+            # Get chunk
+            turn_start_millis = turn.start * 1000
+            turn_end_millis = turn.end * 1000
+            chunk = audio[turn_start_millis:turn_end_millis]
+
+            # Save chunk
+            chunk_save_path = (
+                speaker_dir / f"{int(turn_start_millis)}-{int(turn_end_millis)}.wav"
+            )
+            chunk.export(chunk_save_path, format="wav")
+
+    return storage_dir
+
+
+###############################################################################
+# Dataset expansions
+
+
+def _chunk_audio(
+    monologue_start_time: float,
+    monologue_end_time: float,
+    max_audio_chunk_duration: float,
+    min_audio_chunk_duration: float,
+    audio: AudioSegment,
+    audio_output_dir: Path,
+    chunk_filename_base: str,
+    overwrite: bool,
+    annotated_audios: List[AnnotatedAudio],
+    conversation_id: str,
+    label: str,
+) -> List[AnnotatedAudio]:
+    # Split monologue into chunks
+    for chunk_i, chunk_start_time in enumerate(
+        np.arange(
+            monologue_start_time,
+            monologue_end_time,
+            max_audio_chunk_duration,
+        )
+    ):
+        # Get chunk start and end in millis
+        chunk_start_millis = chunk_start_time * 1000
+        chunk_end = chunk_start_time + max_audio_chunk_duration
+
+        # Determine chunk end time
+        # If start + chunk duration is longer than monologue
+        # Chunk needs to be cut at monologue end
+        if monologue_end_time > chunk_end:
+            chunk_end_millis = chunk_end * 1000
+        else:
+            chunk_end_millis = monologue_end_time * 1000
+
+        # Only allow chunk through if duration is greater than
+        # min chunk length
+        duration = (chunk_end_millis - chunk_start_millis) / 1000
+        if duration >= min_audio_chunk_duration:
+            # Get chunk
+            chunk = audio[chunk_start_millis:chunk_end_millis]
+
+            # Determine save pattern
+            chunk_save_path = audio_output_dir / (
+                f"{chunk_filename_base}-chunk_{chunk_i}.wav"
+            )
+            if chunk_save_path.exists() and not overwrite:
+                raise FileExistsError(chunk_save_path)
+
+            # Save audio chunk
+            chunk.export(chunk_save_path, format="wav")
+
+            # Append new chunk to list
+            annotated_audios.append(
+                AnnotatedAudio(
+                    conversation_id=conversation_id,
+                    label=label,
+                    audio=str(chunk_save_path),
+                    duration=duration,
+                )
+            )
+
+    return annotated_audios
+
+
+def expand_gecko_annotations_to_dataset(
+    annotations_and_audios: List[GeckoAnnotationAndAudio],
+    audio_output_dir: Union[str, Path] = Path("chunked-audio/"),
+    overwrite: bool = False,
+    min_audio_chunk_duration: float = 0.5,
+    max_audio_chunk_duration: float = 2.0,
+) -> pd.DataFrame:
+    """
+    Expand a list of annotation and audio files into a full dataset to be used for
+    training and testing a speaker classification model.
+
+    Parameters
+    ----------
+    annotations_and_audios: List[GeckoAnnotationAndAudio]
+        A list of annotation and their matching audio files to expand into a speaker,
+        audio file path, start and end times.
+    audio_output_dir: Union[str, Path]
+        A directory path to store the chunked audio files in.
+        Default: "chunked-audio" directory in the current working directory.
+    overwrite: bool
+        When writting out an audio chunk, should existing files be overwritten.
+        Default: False (do not overwrite)
+    min_audio_chunk_duration: float
+        Length of the minimum audio duration to allow through after chunking.
+        default: 0.5 seconds
+    max_audio_chunk_duration: float
+        Length of the maximum audio duration to split larger audio files into.
+        Default: 2.0 seconds
+
+    Returns
+    -------
+    dataset: pd.DataFrame
+        The expanded dataset with columns: conversation_id, label, audio, duration
+
+    Raises
+    ------
+    NotADirectoryError
+        A file exists at the specified destination.
+    FileExistsError
+        A file exists at the target chunk audio location but overwrite is False.
+
+    Notes
+    -----
+    Generated and attached conversation ids are pulled from the annotation file name.
+    """
+    # Ensure dataset dir is valid
+    audio_output_dir = Path(audio_output_dir).resolve()
+    if audio_output_dir.is_file():
+        raise NotADirectoryError(audio_output_dir)
+    audio_output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Iter annotations and audios
+    annotated_audios: List[AnnotatedAudio] = []
+    for aaa in annotations_and_audios:
+        audio = AudioSegment.from_file(aaa.audio_file)
+        with open(aaa.annotation_file, "r") as open_f:
+            annotations = json.load(open_f)
+
+        # Iter through each "monologue" section
+        for i, monologue in enumerate(annotations["monologues"]):
+            monologue_speaker = monologue["speaker"]["id"].lower()
+            monologue_start_time = monologue["start"]
+            monologue_end_time = monologue["end"]
+
+            # Construct basic filename
+            filename_base = (
+                f"{aaa.audio_file.with_suffix('').name}"
+                f"-monologue_{i}-{monologue_speaker}"
+            )
+
+            # Chunk audios for this monologue
+            annotated_audios = _chunk_audio(
+                monologue_start_time=monologue_start_time,
+                monologue_end_time=monologue_end_time,
+                max_audio_chunk_duration=max_audio_chunk_duration,
+                min_audio_chunk_duration=min_audio_chunk_duration,
+                audio=audio,
+                audio_output_dir=audio_output_dir,
+                chunk_filename_base=filename_base,
+                overwrite=overwrite,
+                annotated_audios=annotated_audios,
+                conversation_id=aaa.annotation_file.with_suffix("").name,
+                label=monologue_speaker,
+            )
+
+        log.debug(f"Completed expansion for annotation file: {aaa.annotation_file}")
+
+    # Merge all into a single dataframe
+    return pd.DataFrame([aa.to_dict() for aa in annotated_audios])  # type: ignore
+
+
+def expand_labeled_diarized_audio_dir_to_dataset(
+    labeled_diarized_audio_dir: Union[
+        str,
+        Path,
+        List[str],
+        List[Path],
+        List[Union[str, Path]],
+    ],
+    audio_output_dir: Union[str, Path] = Path("chunked-audio/"),
+    overwrite: bool = False,
+    min_audio_chunk_duration: float = 0.5,
+    max_audio_chunk_duration: float = 2.0,
+) -> pd.DataFrame:
+    """
+    Expand the provided labeled diarized audio into a dataset ready for training.
+
+    Parameters
+    ----------
+    labeled_diarized_audio_dir: Union[Union[str, Path], List[Union[str, Path]]]
+        A path or list of paths to diarization results directories. These directories
+        should no longer have the "SPEAKER_00", "SPEAKER_01", default labeling but
+        expert annotated labels.
+    audio_output_dir: Union[str, Path]
+        A directory path to store the chunked audio files in.
+        Default: "chunked-audio" directory in the current working directory.
+    overwrite: bool
+        When writting out an audio chunk, should existing files be overwritten.
+        Default: False (do not overwrite)
+    min_audio_chunk_duration: float
+        Length of the minimum audio duration to allow through after chunking.
+        default: 0.5 seconds
+    max_audio_chunk_duration: float
+        Length of the maximum audio duration to split larger audio files into.
+        Default: 2.0 seconds
+
+    Returns
+    -------
+    dataset: pd.DataFrame
+        The expanded dataset with columns: conversation_id, label, audio, duration
+
+    Raises
+    ------
+    NotADirectoryError
+        A file exists at the specified destination.
+    FileExistsError
+        A file exists at the target chunk audio location but overwrite is False.
+
+    See Also
+    --------
+    diarize_and_split_audio
+        Function to diarize an audio file and split into annotation directories.
+
+    Notes
+    -----
+    The provided labeled diarized audio directory(s) should have the
+    following structure::
+
+        {labeled_diarized_audio_dir}/
+        ├── label
+        │   ├── 1.wav
+        │   └── 2.wav
+        ├── second_label
+        │   ├── 1.wav
+        │   └── 2.wav
+
+    Generated and attached conversation ids are pulled from the labeled diarized audio
+    directory names.
+    """
+    # Ensure dataset dir is valid
+    audio_output_dir = Path(audio_output_dir).resolve()
+    if audio_output_dir.is_file():
+        raise NotADirectoryError(audio_output_dir)
+    audio_output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Normalize annotation dir to list
+    if isinstance(labeled_diarized_audio_dir, Path):
+        labeled_diarized_audio_dir = [labeled_diarized_audio_dir]
+
+    # Iter over all dia dirs
+    annotated_audios: List[AnnotatedAudio] = []
+    for anno_dir in labeled_diarized_audio_dir:
+        # Convert to path
+        anno_dir = Path(anno_dir).resolve(strict=True)
+
+        # Get the conversation id
+        conversation_id = anno_dir.name
+
+        # Iter over all labels
+        for label_dir in anno_dir.iterdir():
+            if label_dir.is_dir():
+                for i, audio_file in enumerate(label_dir.glob("*.wav")):
+                    audio = AudioSegment.from_file(audio_file)
+                    monologue_speaker = label_dir.name
+
+                    # Construct basic filename
+                    filename_base = (
+                        f"{audio_file.with_suffix('').name}"
+                        f"-monologue_{i}-{monologue_speaker}"
+                    )
+
+                    # Chunk audios for this monologue
+                    annotated_audios = _chunk_audio(
+                        monologue_start_time=0,
+                        monologue_end_time=audio.duration_seconds,
+                        max_audio_chunk_duration=max_audio_chunk_duration,
+                        min_audio_chunk_duration=min_audio_chunk_duration,
+                        audio=audio,
+                        audio_output_dir=audio_output_dir,
+                        chunk_filename_base=filename_base,
+                        overwrite=overwrite,
+                        annotated_audios=annotated_audios,
+                        conversation_id=conversation_id,
+                        label=monologue_speaker.lower(),
+                    )
+
+                log.debug(f"Completed expansion for diarized audio dir: {label_dir}")
+
+    # Merge all into a single dataframe
+    return pd.DataFrame([aa.to_dict() for aa in annotated_audios])  # type: ignore
+
+
+###############################################################################
+# Data expections and cleaning
+
+
+def prepare_dataset(
+    dataset: pd.DataFrame,
+    test_and_valid_size: float = 0.4,
+    equalize_data_within_splits: bool = False,
+    n_iterations: int = 100,
+    seed: Optional[int] = None,
+) -> Tuple[DatasetDict, pd.DataFrame]:
+    """
+    Prepare a dataset for training a new speakerbox / audio-classification model.
+    This function attempts to randomly create train, test, and validation splits
+    from the provided dataframe that meet the following two conditions:
+
+    1. There is data holdout by conversation_id. I.e. if the dataset contains data from
+    nine unique conversation ids, the training, test, and validation sets should all
+    have different conversation ids (train has 0, 1, 2, 3; test has 4, 5, 6; validation
+    has 7, 8).
+
+    2. There is data stratification by label. I.e. if the dataset contains nine unique
+    labels, the training, test, and validation sets should each have all nine labels
+    present (train, test, and validation all have labels 0-8).
+
+    Parameters
+    ----------
+    dataset: pd.DataFrame
+        An expanded dataset with columns: conversation_id, label, audio, duration
+    test_and_valid_size: float
+        How much of the dataset to use for the combined test and validation sets
+        as a percent (i.e. 0.4 = 40% of the dataset).
+        The test and validation sets will further split this in half (i.e. 0.4 = 40%
+        which means 20% of the total data for testing and 20% of the total data for
+        validation).
+    equalize_data_within_splits: bool
+        After finding valid train, test, and validation splits, should the data within
+        each split be reduced to have an equal number of data for each label.
+        Default: False (Do not equalize labels within splits)
+    n_iterations: int
+        The number of iterations to attempt to find viable train, test, and validation
+        sets.
+        Default: 100
+    seed: Optional[int]
+        Seed to pass to torch, numpy, and Python RNGs.
+        Default: None (do not set a seed)
+
+    Returns
+    -------
+    dataset: DatasetDict
+        The prepared dataset split into train, test, and validation splits.
+    value_counts: pd.DataFrame
+        A value count table where each row is a different label and each column is the
+        count of that label in the matching train, test, or validation set.
+
+    See Also
+    --------
+    expand_labeled_diarized_audio_dir_to_dataset
+        Function to move from a directory of diarized audio (or multiple) into a
+        dataset to provide to this function.
+    expand_gecko_annotations_to_dataset
+        Function to move from a gecko annotation file (or multiple) into a dataset
+        to provide to this function.
+
+    Raises
+    ------
+    ValueError
+        Could not find train, test, and validation sets that meet the holdout and
+        stratification criteria after n iterations. Recommended to annotate more data.
+    """
+    # Handle random seed
+    if seed:
+        set_global_seed(seed)
+
+    # Holdout by conversation id
+    labels = set(dataset.label.unique())
+    conversation_ids = dataset.conversation_id.unique()
+
+    # Reroll holdouts until all labels are present in each split
+    all_labels_present = False
+    iters = 0
+    log.info("Finding data splits for provided dataset")
+    while not all_labels_present:
+        # Get conversation ids split randomly
+        train_ids, test_and_valid_ids = train_test_split(
+            conversation_ids, test_size=test_and_valid_size
+        )
+        test_ids, valid_ids = train_test_split(test_and_valid_ids, test_size=0.5)
+
+        # Check the produced subsets
+        train_ds = dataset.loc[dataset.conversation_id.isin(train_ids)]
+        test_ds = dataset.loc[dataset.conversation_id.isin(test_ids)]
+        valid_ds = dataset.loc[dataset.conversation_id.isin(valid_ids)]
+        if (
+            set(train_ds.label.unique()) == labels
+            and set(test_ds.label.unique()) == labels
+            and set(valid_ds.label.unique()) == labels
+        ):
+            all_labels_present = True
+
+        iters += 1
+        if iters == n_iterations:
+            raise ValueError(
+                f"Could not find train, test, and validation sets that meet the "
+                f"holdout and stratification criteria after {n_iterations}. "
+                f"Recommended to annotate more data."
+            )
+        log.debug(f"Attempted train test validation split construction {iters} times.")
+
+    log.debug(f"train_ds conversation_ids: {train_ds.conversation_id.unique()}")
+    log.debug(f"test_ds conversation_ids: {test_ds.conversation_id.unique()}")
+    log.debug(f"valid_ds conversation_ids: {valid_ds.conversation_id.unique()}")
+
+    # Drop extra columns
+    train_ds = train_ds.drop(["conversation_id", "duration"], axis=1)
+    test_ds = test_ds.drop(["conversation_id", "duration"], axis=1)
+    valid_ds = valid_ds.drop(["conversation_id", "duration"], axis=1)
+
+    # Handle equalization
+    if equalize_data_within_splits:
+        # Subsets should be balanced by each label
+        # Group by label and get random sample with the min of all labels
+        train_groups = train_ds.groupby("label")
+        train_ds = pd.DataFrame(
+            train_groups.apply(
+                lambda x: x.sample(train_groups.size().min()).reset_index(drop=True)
+            )
+        )
+        test_groups = test_ds.groupby("label")
+        test_ds = pd.DataFrame(
+            test_groups.apply(
+                lambda x: x.sample(test_groups.size().min()).reset_index(drop=True)
+            )
+        )
+        valid_groups = valid_ds.groupby("label")
+        valid_ds = pd.DataFrame(
+            valid_groups.apply(
+                lambda x: x.sample(valid_groups.size().min()).reset_index(drop=True)
+            )
+        )
+
+    # Construct summary table
+    value_counts = pd.DataFrame(
+        {
+            "train_counts": train_ds.label.value_counts(),
+            "test_counts": test_ds.label.value_counts(),
+            "valid_counts": valid_ds.label.value_counts(),
+        }
+    )
+    log.info(f"Constructed train, test, validation sets contain:\n{value_counts}")
+
+    # Construct DatasetDict
+    train_ds = Dataset.from_pandas(train_ds, preserve_index=False)
+    train_ds = train_ds.class_encode_column("label")
+    test_ds = Dataset.from_pandas(test_ds, preserve_index=False)
+    test_ds = test_ds.class_encode_column("label")
+    valid_ds = Dataset.from_pandas(valid_ds, preserve_index=False)
+    valid_ds = valid_ds.class_encode_column("label")
+
+    # Return both the dataset dict and the value counts
+    return (
+        DatasetDict(
+            {
+                "train": train_ds,
+                "test": test_ds,
+                "valid": valid_ds,
+            }
+        ),
+        value_counts,
+    )
diff --git a/speakerbox/types.py b/speakerbox/types.py
new file mode 100644
index 0000000..fb44c21
--- /dev/null
+++ b/speakerbox/types.py
@@ -0,0 +1,24 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import NamedTuple
+
+from dataclasses_json import dataclass_json
+
+###############################################################################
+
+
+class GeckoAnnotationAndAudio(NamedTuple):
+    annotation_file: Path
+    audio_file: Path
+
+
+@dataclass_json
+@dataclass
+class AnnotatedAudio:
+    conversation_id: str
+    label: str
+    audio: str
+    duration: float
diff --git a/speakerbox/utils.py b/speakerbox/utils.py
new file mode 100644
index 0000000..3de88bc
--- /dev/null
+++ b/speakerbox/utils.py
@@ -0,0 +1,82 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import shutil
+from pathlib import Path
+from typing import Union
+
+###############################################################################
+
+
+def _unpack_zip(
+    zipfile: Union[str, Path],
+    dest: Union[str, Path],
+    clean: bool = False,
+) -> Path:
+    """
+    Unzips the zipfile to the destination location.
+
+    Parameters
+    ----------
+    zipfile: Union[str, Path]
+        The zipfile to unpack.
+    dest: Union[str, Path]
+        The destination to unpack to.
+    clean: bool
+        If a directory already exists at the destination location, should the directory
+        be removed entirely before unpacking again.
+        Default: False
+
+    Returns
+    -------
+    dataset_path: Path
+        The path to the unpacked data.
+
+    Raises
+    ------
+    NotADirectoryError
+        A file exists at the specified destination.
+    FileExistsError
+        A directory exists at the specified destination and is not empty.
+    """
+    zipfile = Path(zipfile).resolve(strict=True)
+    dest = Path(dest)
+
+    # Check dest is a dir
+    if dest.is_file():
+        raise NotADirectoryError(dest)
+
+    # Clean dest if required
+    if dest.is_dir():
+        dest_is_empty = not any(dest.iterdir())
+        if not dest_is_empty:
+            if clean:
+                shutil.rmtree(dest)
+            else:
+                raise FileExistsError(
+                    f"Files found in target destination directory ({dest}) "
+                    f"but parameter `clean` set to False."
+                )
+
+    # Make new dir
+    dest.mkdir(parents=True, exist_ok=True)
+
+    # Extract
+    shutil.unpack_archive(zipfile, dest)
+
+    # Return extracted data dir
+    return dest.resolve(strict=True)
+
+
+def set_global_seed(seed: int) -> None:
+    """
+    Set the global RNG seed for torch, numpy, and Python.
+    """
+    import random
+
+    import numpy as np
+    import torch
+
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
diff --git a/tox.ini b/tox.ini
index a18db7f..470401e 100644
--- a/tox.ini
+++ b/tox.ini
@@ -1,5 +1,5 @@
 [tox]
-envlist = py38, py39, lint
+envlist = py38, py39, py310, lint
 skip_missing_interpreters = true
 toxworkdir={env:TOX_WORK_DIR:.tox}
 


Test patch:
diff --git a/speakerbox/tests/conftest.py b/speakerbox/tests/conftest.py
index 2165cf3..19122bd 100644
--- a/speakerbox/tests/conftest.py
+++ b/speakerbox/tests/conftest.py
@@ -19,8 +19,24 @@
 from pathlib import Path
 
 import pytest
+from pytest import Parser
 
+from speakerbox.utils import _unpack_zip
 
-@pytest.fixture
+###############################################################################
+
+
+@pytest.fixture(scope="session")
 def data_dir() -> Path:
-    return Path(__file__).parent / "data"
+    d_dir = Path(__file__).parent / "data"
+
+    # Check for diarized audio
+    diarized_dir = d_dir / "diarized"
+    if not diarized_dir.exists():
+        _unpack_zip(d_dir / "diarized.zip", diarized_dir)
+
+    return d_dir
+
+
+def pytest_addoption(parser: Parser) -> None:
+    parser.addoption("--cpu", action="store_true", dest="use_cpu")
diff --git a/speakerbox/tests/data/__init__.py b/speakerbox/tests/data/__init__.py
new file mode 100644
index 0000000..8d03c27
--- /dev/null
+++ b/speakerbox/tests/data/__init__.py
@@ -0,0 +1,3 @@
+# -*- coding: utf-8 -*-
+
+"""Data for unit tests for speakerbox"""
diff --git a/speakerbox/tests/data/diarized.zip b/speakerbox/tests/data/diarized.zip
new file mode 100644
index 0000000..59ffe48
Binary files /dev/null and b/speakerbox/tests/data/diarized.zip differ
diff --git a/speakerbox/tests/data/short_0.wav b/speakerbox/tests/data/short_0.wav
new file mode 100644
index 0000000..fe14a48
Binary files /dev/null and b/speakerbox/tests/data/short_0.wav differ
diff --git a/speakerbox/tests/data/short_1.wav b/speakerbox/tests/data/short_1.wav
new file mode 100644
index 0000000..d014e2d
Binary files /dev/null and b/speakerbox/tests/data/short_1.wav differ
diff --git a/speakerbox/tests/data/short_2.wav b/speakerbox/tests/data/short_2.wav
new file mode 100644
index 0000000..8eafeaa
Binary files /dev/null and b/speakerbox/tests/data/short_2.wav differ
diff --git a/speakerbox/tests/test_datasets.py b/speakerbox/tests/test_datasets.py
new file mode 100644
index 0000000..debb67a
--- /dev/null
+++ b/speakerbox/tests/test_datasets.py
@@ -0,0 +1,33 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+from speakerbox import preprocess
+from speakerbox.datasets import seattle_2021_proto
+
+###############################################################################
+
+
+def test_unpack_pull_and_expand_seattle() -> None:
+    # Unpack / unzip the Seattle city council 2021 prototype dataset
+    seattle_2021_ds_dir = seattle_2021_proto.unpack(
+        dest="test-outputs/seattle-2021-proto/",
+        clean=True,
+    )
+
+    # Pull matching audio files for each annotation file
+    seattle_2021_ds_items = seattle_2021_proto.pull_all_files(
+        annotations_dir=seattle_2021_ds_dir / "annotations",
+        transcript_output_dir=seattle_2021_ds_dir / "transcripts",
+        audio_output_dir=seattle_2021_ds_dir / "audios",
+        overwrite=True,
+    )
+
+    # Expand from multiple matching large gecko annotation files and large audio files
+    # into many small audio clips with speaker labels
+    seattle_2021_ds = preprocess.expand_gecko_annotations_to_dataset(
+        seattle_2021_ds_items,
+        audio_output_dir="test-outputs/seattle-2021-proto/chunked-audio/",
+        overwrite=True,
+    )
+
+    assert len(seattle_2021_ds) > 0
diff --git a/speakerbox/tests/test_empty.py b/speakerbox/tests/test_empty.py
deleted file mode 100644
index 1c9326c..0000000
--- a/speakerbox/tests/test_empty.py
+++ /dev/null
@@ -1,6 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-
-def test_empty() -> None:
-    pass
diff --git a/speakerbox/tests/test_preprocess.py b/speakerbox/tests/test_preprocess.py
new file mode 100644
index 0000000..a3ddf8f
--- /dev/null
+++ b/speakerbox/tests/test_preprocess.py
@@ -0,0 +1,33 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+from pathlib import Path
+
+import pytest
+
+from speakerbox import preprocess
+
+###############################################################################
+
+
+@pytest.mark.parametrize(
+    "audio_filename",
+    [
+        "short_0.wav",
+        "short_1.wav",
+        "short_2.wav",
+    ],
+)
+def test_diarize_and_split_audio(data_dir: Path, audio_filename: str) -> None:
+    """
+    This is a smoke test. We do not check the validity of the produced clusters,
+    we simply care that the function runs through without error.
+    """
+    # Use data dir with filename
+    audio_path = data_dir / audio_filename
+
+    # Diarize
+    preprocess.diarize_and_split_audio(
+        audio_file=audio_path,
+        storage_dir=f"test-outputs/diarization/{audio_path.with_suffix('').name}",
+    )
diff --git a/speakerbox/tests/test_train.py b/speakerbox/tests/test_train.py
new file mode 100644
index 0000000..78ce375
--- /dev/null
+++ b/speakerbox/tests/test_train.py
@@ -0,0 +1,63 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+from pathlib import Path
+
+from pytest import Config
+
+from speakerbox import eval_model, preprocess, train
+from speakerbox.utils import set_global_seed
+
+###############################################################################
+
+
+def test_train(data_dir: Path, pytestconfig: Config) -> None:
+    """
+    This is a smoke test. We do not check the validity of the produced model,
+    we simply care that the functions run through without error.
+
+    Run with CPU with pytest:
+    pytest speakerbox/tests/test_train.py --cpu
+
+    Or with tox:
+    tox -- --cpu
+    """
+    # Set global seed for some level of reproducibility across tests
+    set_global_seed(20220421)
+
+    # Get diarized dirs
+    diarized_audio = data_dir / "diarized"
+    diarized_audio_dirs = [
+        dia_convo_dir
+        for dia_convo_dir in diarized_audio.iterdir()
+        if dia_convo_dir.is_dir()
+    ]
+
+    # Expand diarizations to full dataset
+    expanded_dataset = preprocess.expand_labeled_diarized_audio_dir_to_dataset(
+        labeled_diarized_audio_dir=diarized_audio_dirs,
+        audio_output_dir="test-outputs/training/chunked-audio/",
+        overwrite=True,
+    )
+
+    # Prepare and check
+    dataset_dict, _ = preprocess.prepare_dataset(
+        expanded_dataset,
+        equalize_data_within_splits=True,
+    )
+
+    # Store prepared dataset to disk
+    dataset_dict.save_to_disk("test-outputs/training/prepared-dataset/")
+
+    # Train
+    train(
+        dataset_dict,
+        model_name="test-outputs/trained-speakerbox",
+        use_cpu=pytestconfig.getoption("use_cpu"),
+    )
+
+    # Eval
+    eval_model(
+        dataset_dict["valid"],
+        model_name="test-outputs/trained-speakerbox",
+    )