Repository: castorini/howl (language: python)

Issue:
restructuring/refactoring files for howl/training
some files under howl/data/dataset (labeller, phone) should be located outside as a common module that can be added separately

I have noticed some classes are incorrectly located and can be organized better.


Discussion / hints:
Writers must be utilizing audio_data instead of reading the same sample from disk

Fix patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 04070657..7742d45e 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -15,7 +15,7 @@ repos:
     -   id: trailing-whitespace
     -   id: check-json
     -   id: name-tests-test
--   repo: https://github.com/pycqa/isort
+-   repo: https://github.com/timothycrosley/isort
     rev: 5.5.4
     hooks:
     -   id: isort
diff --git a/howl/__init__.py b/howl/__init__.py
index 3c8e1796..e69de29b 100644
--- a/howl/__init__.py
+++ b/howl/__init__.py
@@ -1,1 +0,0 @@
-from howl.logger import *
diff --git a/howl/context.py b/howl/context.py
index 4224b508..91ac998a 100644
--- a/howl/context.py
+++ b/howl/context.py
@@ -3,19 +3,15 @@
 from pathlib import Path
 from typing import List
 
-from howl.data.dataset import (
-    PhonePhrase,
-    PhoneticFrameLabeler,
-    PronunciationDictionary,
-    WakeWordDataset,
-    WordFrameLabeler,
-)
-from howl.data.searcher import (
+from howl.data.common.labeler import PhoneticFrameLabeler, WordFrameLabeler
+from howl.data.common.phone import PhonePhrase, PronunciationDictionary
+from howl.data.common.searcher import (
     LabelColoring,
     PhoneticTranscriptSearcher,
     WordTranscriptSearcher,
 )
-from howl.data.tokenize import Vocab
+from howl.data.common.vocab import Vocab
+from howl.data.dataset import WakeWordDataset
 from howl.settings import SETTINGS
 
 
diff --git a/howl/data/common/__init__.py b/howl/data/common/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/howl/data/common/batch.py b/howl/data/common/batch.py
new file mode 100644
index 00000000..17616e7c
--- /dev/null
+++ b/howl/data/common/batch.py
@@ -0,0 +1,61 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+__all__ = [
+    "ClassificationBatch",
+    "SequenceBatch",
+]
+
+
+@dataclass
+class ClassificationBatch:
+    audio_data: torch.Tensor
+    labels: Optional[torch.Tensor]
+    lengths: torch.Tensor
+
+    @classmethod
+    def from_single(cls, audio_clip: torch.Tensor, label: int) -> "ClassificationBatch":
+        return cls(audio_clip.unsqueeze(0), torch.tensor([label]), torch.tensor([audio_clip.size(-1)]))
+
+    def pin_memory(self):
+        self.audio_data.pin_memory()
+        self.labels.pin_memory()
+        self.lengths.pin_memory()
+
+    def to(self, device: torch.device) -> "ClassificationBatch":
+        self.audio_data = self.audio_data.to(device)
+        if self.labels is not None:
+            self.labels = self.labels.to(device)
+        self.lengths = self.lengths.to(device)
+        return self
+
+
+@dataclass
+class SequenceBatch:
+    audio_data: torch.Tensor  # (batch size, max audio length)
+    labels: torch.Tensor  # (batch size, max label length)
+    audio_lengths: Optional[torch.Tensor]
+    label_lengths: Optional[torch.Tensor]
+
+    def __post_init__(self):
+        if self.audio_lengths is None:
+            self.audio_lengths = torch.tensor([self.audio_data.size(0) for _ in range(self.audio_data.size(1))])
+            self.audio_lengths = self.audio_lengths.to(self.audio_data.device)
+        if self.label_lengths is None:
+            self.label_lengths = torch.tensor([self.labels.size(0) for _ in range(self.labels.size(1))])
+            self.label_lengths = self.label_lengths.to(self.label_lengths.device)
+
+    def pin_memory(self):
+        self.audio_data.pin_memory()
+        self.labels.pin_memory()
+        self.audio_lengths.pin_memory()
+        self.label_lengths.pin_memory()
+
+    def to(self, device: torch.device) -> "SequenceBatch":
+        self.audio_data = self.audio_data.to(device)
+        self.labels = self.labels.to(device)
+        self.audio_lengths = self.audio_lengths.to(device)
+        self.label_lengths = self.label_lengths.to(device)
+        return self
diff --git a/howl/data/common/example.py b/howl/data/common/example.py
new file mode 100644
index 00000000..b60819cd
--- /dev/null
+++ b/howl/data/common/example.py
@@ -0,0 +1,75 @@
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+
+import torch
+
+from howl.data.common.frame import FrameLabelData
+from howl.data.common.metadata import AudioClipMetadata
+
+__all__ = [
+    "AudioClipExample",
+    "ClassificationClipExample",
+    "EmplacableExample",
+    "WakeWordClipExample",
+]
+
+
+class EmplacableExample:
+    audio_data: torch.Tensor
+
+    def emplaced_audio_data(
+        self, audio_data: torch.Tensor, scale: float = 1, bias: float = 0, new: bool = False
+    ) -> "EmplacableExample":
+        raise NotImplementedError
+
+
+T = TypeVar("T", bound=AudioClipMetadata)
+
+
+class AudioClipExample(EmplacableExample, Generic[T]):
+    def __init__(self, metadata: T, audio_data: torch.Tensor, sample_rate: int):
+        self.metadata = metadata
+        self.audio_data = audio_data
+        self.sample_rate = sample_rate
+
+    def pin_memory(self):
+        self.audio_data.pin_memory()
+
+    def emplaced_audio_data(
+        self, audio_data: torch.Tensor, scale: float = 1, bias: float = 0, new: bool = False
+    ) -> "AudioClipExample":
+        metadata = self.metadata
+        if new:
+            metadata = deepcopy(metadata)
+            metadata.transcription = ""
+        return AudioClipExample(metadata, audio_data, self.sample_rate)
+
+
+@dataclass
+class WakeWordClipExample(AudioClipExample[AudioClipMetadata]):
+    def __init__(self, label_data: FrameLabelData, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.label_data = label_data
+
+    def emplaced_audio_data(
+        self, audio_data: torch.Tensor, scale: float = 1, bias: float = 0, new: bool = False
+    ) -> "WakeWordClipExample":
+        ex = super().emplaced_audio_data(audio_data, scale, bias, new)
+        label_data = {} if new else {scale * k + bias: v for k, v in self.label_data.timestamp_label_map.items()}
+        return WakeWordClipExample(
+            FrameLabelData(label_data, self.label_data.start_timestamp, self.label_data.char_indices),
+            ex.metadata,
+            audio_data,
+            self.sample_rate,
+        )
+
+
+@dataclass
+class ClassificationClipExample(AudioClipExample[AudioClipMetadata]):
+    def __init__(self, label, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.label = label
+
+    def emplaced_audio_data(self, audio_data: torch.Tensor, **kwargs) -> "ClassificationClipExample":
+        return ClassificationClipExample(self.label, self.metadata, audio_data, self.sample_rate)
diff --git a/howl/data/common/frame.py b/howl/data/common/frame.py
new file mode 100644
index 00000000..2d3d3f97
--- /dev/null
+++ b/howl/data/common/frame.py
@@ -0,0 +1,11 @@
+from dataclasses import dataclass
+from typing import List, Mapping, Tuple
+
+__all__ = ["FrameLabelData"]
+
+
+@dataclass
+class FrameLabelData:
+    timestamp_label_map: Mapping[float, int]
+    start_timestamp: List[Tuple[int, float]]
+    char_indices: List[Tuple[int, List[int]]]
diff --git a/howl/data/dataset/labeller.py b/howl/data/common/labeler.py
similarity index 81%
rename from howl/data/dataset/labeller.py
rename to howl/data/common/labeler.py
index dea33db0..be482b2f 100644
--- a/howl/data/dataset/labeller.py
+++ b/howl/data/common/labeler.py
@@ -1,10 +1,10 @@
 import string
 from typing import List
 
-from howl.data.dataset.phone import PhoneEnum, PhonePhrase, PronunciationDictionary
-from howl.data.tokenize import Vocab
-
-from .base import AudioClipMetadata, FrameLabelData
+from howl.data.common.frame import FrameLabelData
+from howl.data.common.metadata import AudioClipMetadata
+from howl.data.common.phone import PhoneEnum, PhonePhrase, PronunciationDictionary
+from howl.data.common.vocab import Vocab
 
 __all__ = ["FrameLabeler", "WordFrameLabeler", "PhoneticFrameLabeler"]
 
@@ -15,15 +15,32 @@ def compute_frame_labels(self, metadata: AudioClipMetadata) -> FrameLabelData:
 
 
 class PhoneticFrameLabeler(FrameLabeler):
-    def __init__(self, pronounce_dict: PronunciationDictionary, phrases: List[PhonePhrase]):
+    def __init__(
+        self, pronounce_dict: PronunciationDictionary, phrases: List[PhonePhrase]
+    ):
         self.pronounce_dict = pronounce_dict
         self.phrases = phrases
         punctuation_to_replace = str.maketrans(
-            {"‘": "'", "’": "'", "”": '"', "“": '"', "—": "-", "ä": "a", "ö": "o", "ō": "o", "é": "e", "à": "a"}
+            {
+                "‘": "'",
+                "’": "'",
+                "”": '"',
+                "“": '"',
+                "—": "-",
+                "ä": "a",
+                "ö": "o",
+                "ō": "o",
+                "é": "e",
+                "à": "a",
+            }
         )
         punctuation_to_remove = str.maketrans({key: None for key in string.punctuation})
         # First transformation is None because we want to process the original word first
-        self.punctuation_transforms = [None, punctuation_to_replace, punctuation_to_remove]
+        self.punctuation_transforms = [
+            None,
+            punctuation_to_replace,
+            punctuation_to_remove,
+        ]
 
     def transform(self, original_word: str) -> PhonePhrase:
         """Transform the word into list of Phones
@@ -104,7 +121,9 @@ def compute_frame_labels(self, metadata: AudioClipMetadata) -> FrameLabelData:
             if phrase:
                 phonetic_transcription.extend(phrase)
             elif len(original_word) > 0:
-                print(f"Failed to find phonemes for {original_word} {[ord(c) for c in original_word]}")
+                print(
+                    f"Failed to find phonemes for {original_word} {[ord(c) for c in original_word]}"
+                )
 
         # TODO: idx might not be the correct label due to repeated phonemes
         for idx, phrase in enumerate(self.phrases):
@@ -127,7 +146,9 @@ def compute_frame_labels(self, metadata: AudioClipMetadata) -> FrameLabelData:
                 #       char_index = phoneme_mapping[start]
                 #       end_time = metadata.end_timestamps[char_index]
                 #       frame_labels[end_time] = idx
-                frame_labels[metadata.end_timestamps[start + len(str(phrase)) - 1]] = idx
+                frame_labels[
+                    metadata.end_timestamps[start + len(str(phrase)) - 1]
+                ] = idx
                 start += 1
 
         # TODO: process phonetic_transcription to compute valid FrameLabelData
@@ -154,8 +175,15 @@ def compute_frame_labels(self, metadata: AudioClipMetadata) -> FrameLabelData:
                 label = self.vocab[word]
                 end_timestamp = metadata.end_timestamps[char_idx + word_size - 1]
                 frame_labels[end_timestamp] = label
-                char_indices.append((label, list(range(char_idx, char_idx + word_size))))
-                start_timestamp.append((label, metadata.end_timestamps[char_idx - 1] if char_idx > 0 else 0.0))
+                char_indices.append(
+                    (label, list(range(char_idx, char_idx + word_size)))
+                )
+                start_timestamp.append(
+                    (
+                        label,
+                        metadata.end_timestamps[char_idx - 1] if char_idx > 0 else 0.0,
+                    )
+                )
 
             char_idx += word_size + 1  # space
 
diff --git a/howl/data/common/metadata.py b/howl/data/common/metadata.py
new file mode 100644
index 00000000..a88a0464
--- /dev/null
+++ b/howl/data/common/metadata.py
@@ -0,0 +1,31 @@
+from pathlib import Path
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+from howl.data.common.phone import Phone, PhonePhrase
+
+__all__ = ["AudioClipMetadata", "UNKNOWN_TRANSCRIPTION", "NEGATIVE_CLASS"]
+
+
+UNKNOWN_TRANSCRIPTION = "[UNKNOWN]"
+NEGATIVE_CLASS = "[NEGATIVE]"
+
+
+class AudioClipMetadata(BaseModel):
+    path: Optional[Path] = Path(".")
+    phone_strings: Optional[List[str]]
+    words: Optional[List[str]]
+    phone_end_timestamps: Optional[List[float]]
+    word_end_timestamps: Optional[List[float]]
+    end_timestamps: Optional[List[float]]  # TODO: remove, backwards compat right now
+    transcription: Optional[str] = ""
+
+    # TODO:: id should be an explicit variable in order to support datasets creation with the audio data in memory
+    @property
+    def audio_id(self) -> str:
+        return self.path.name.split(".", 1)[0]
+
+    @property
+    def phone_phrase(self) -> Optional[PhonePhrase]:
+        return PhonePhrase([Phone(x) for x in self.phone_strings]) if self.phone_strings else None
diff --git a/howl/data/dataset/phone.py b/howl/data/common/phone.py
similarity index 100%
rename from howl/data/dataset/phone.py
rename to howl/data/common/phone.py
diff --git a/howl/data/searcher.py b/howl/data/common/searcher.py
similarity index 83%
rename from howl/data/searcher.py
rename to howl/data/common/searcher.py
index a815e7ad..2c310448 100644
--- a/howl/data/searcher.py
+++ b/howl/data/common/searcher.py
@@ -4,14 +4,18 @@
 from typing import List
 
 import numpy as np
-from howl.data.dataset.phone import PhonePhrase
-from howl.data.tokenize import Vocab, WakeWordTokenizer
+
+from howl.data.common.phone import PhonePhrase
+from howl.data.common.tokenizer import WakeWordTokenizer
+from howl.data.common.vocab import Vocab
 from howl.settings import SETTINGS
 
-__all__ = ['LabelColoring',
-           'PhoneticTranscriptSearcher',
-           'TranscriptSearcher',
-           'WordTranscriptSearcher']
+__all__ = [
+    "LabelColoring",
+    "PhoneticTranscriptSearcher",
+    "TranscriptSearcher",
+    "WordTranscriptSearcher",
+]
 
 
 class LabelColoring:
@@ -63,7 +67,9 @@ def __init__(self, vocab: Vocab, **kwargs):
         super().__init__(**kwargs)
         self.vocab = vocab
         self.tokenizer = WakeWordTokenizer(self.vocab, False)
-        self.inference_sequence_str = ''.join(map(str, self.settings.inference_sequence))
+        self.inference_sequence_str = "".join(
+            map(str, self.settings.inference_sequence)
+        )
         # self.wakeword = self.vocab.decode(self.settings.inference_sequence)
 
     def search(self, item: str) -> bool:
@@ -74,7 +80,7 @@ def search(self, item: str) -> bool:
             bool: true if wakeword is in the item
         """
         encoded_output = self.tokenizer.encode(item)
-        encoded_str = ''.join(map(str, encoded_output))
+        encoded_str = "".join(map(str, encoded_output))
         return self.inference_sequence_str in encoded_str
 
     def contains_any(self, item: str) -> bool:
@@ -109,16 +115,19 @@ class PhoneticTranscriptSearcher(TranscriptSearcher):
     def __init__(self, phrases: List[PhonePhrase], coloring: LabelColoring, **kwargs):
         super().__init__(**kwargs)
         self.phrases = phrases
-        label_map = [(phrase.audible_transcript, coloring.color_map[idx]) for idx, phrase in enumerate(phrases)]
+        label_map = [
+            (phrase.audible_transcript, coloring.color_map[idx])
+            for idx, phrase in enumerate(phrases)
+        ]
         buckets = defaultdict(list)
         for transcript, color in label_map:
             buckets[color].append(transcript)
         pattern_strings = []
         for _, transcripts in sorted(buckets.items(), key=lambda x: x[0]):
-            pattern_strings.append('(' + '|'.join(f'({x})' for x in transcripts) + ')')
+            pattern_strings.append("(" + "|".join(f"({x})" for x in transcripts) + ")")
         pattern_strings = np.array(pattern_strings)[self.settings.inference_sequence]
-        pattern_str = '^.*' + ' '.join(pattern_strings) + '.*$'
-        logging.info(f'Using search pattern {pattern_str}')
+        pattern_str = "^.*" + " ".join(pattern_strings) + ".*$"
+        logging.info(f"Using search pattern {pattern_str}")
         self.pattern = re.compile(pattern_str)
 
     def search(self, item: str) -> bool:
diff --git a/howl/data/common/tokenizer.py b/howl/data/common/tokenizer.py
new file mode 100644
index 00000000..295271f0
--- /dev/null
+++ b/howl/data/common/tokenizer.py
@@ -0,0 +1,41 @@
+from typing import List
+
+from howl.data.common.vocab import Vocab
+
+__all__ = ["WakeWordTokenizer", "TranscriptTokenizer"]
+
+
+class TranscriptTokenizer:
+    def encode(self, transcript: str) -> List[int]:
+        raise NotImplementedError
+
+    def decode(self, ids: List[int]) -> str:
+        raise NotImplementedError
+
+
+class WakeWordTokenizer(TranscriptTokenizer):
+    # Only used for ctc objective
+    def __init__(self, vocab: Vocab, ignore_oov: bool = True):
+        self.vocab = vocab
+        self.ignore_oov = ignore_oov
+
+    def decode(self, ids: List[int]) -> str:
+        return " ".join(self.vocab[id] for id in ids)
+
+    def encode(self, transcript: str) -> List[int]:
+        encoded_output = []
+
+        for word in transcript.lower().split():
+            vocab_found, remaining_transcript = self.vocab.trie.max_split(word)
+
+            # append corresponding label
+            if vocab_found and remaining_transcript == "":
+                # word exists in the vocab
+                encoded_output.append(self.vocab[word])
+            elif not self.ignore_oov:
+                # label oov word
+                if self.vocab.oov_token_id is None:
+                    raise ValueError("label for oov word is not specified")
+                encoded_output.append(self.vocab.oov_token_id)
+
+        return encoded_output
diff --git a/howl/data/tokenize.py b/howl/data/common/vocab.py
similarity index 58%
rename from howl/data/tokenize.py
rename to howl/data/common/vocab.py
index cf2c43ba..2009f884 100644
--- a/howl/data/tokenize.py
+++ b/howl/data/common/vocab.py
@@ -1,14 +1,6 @@
 from typing import List, Mapping, Tuple, Union
 
-__all__ = ['WakeWordTokenizer', 'TranscriptTokenizer', 'VocabTrie']
-
-
-class TranscriptTokenizer:
-    def encode(self, transcript: str) -> List[int]:
-        raise NotImplementedError
-
-    def decode(self, ids: List[int]) -> str:
-        raise NotImplementedError
+__all__ = ["Vocab"]
 
 
 class VocabTrie:
@@ -55,14 +47,13 @@ def max_split(self, tokens: str) -> Tuple[str, str]:
             counter += 1
         if not node.terminal:
             counter = 0
-        return tokens[: counter], tokens[counter:]
+        return tokens[:counter], tokens[counter:]
 
 
 class Vocab:
-    def __init__(self,
-                 word2idx: Union[Mapping[str, int], List[str]],
-                 oov_token_id: int = None,
-                 oov_word_repr: str = '[OOV]'):
+    def __init__(
+        self, word2idx: Union[Mapping[str, int], List[str]], oov_token_id: int = None, oov_word_repr: str = "[OOV]"
+    ):
         if isinstance(word2idx, List):
             word2idx = {word: idx for idx, word in enumerate(word2idx)}
         self.word2idx = {k.lower(): v for k, v in word2idx.items()}
@@ -82,35 +73,5 @@ def __getitem__(self, item: Union[str, int]) -> Union[str, int]:
         else:
             ret = self.idx2word.get(item, self.oov_word_repr)
         if ret is None:
-            raise ValueError(f'couldn\'t find token for {item}')
+            raise ValueError(f"couldn't find token for {item}")
         return ret
-
-
-class WakeWordTokenizer(TranscriptTokenizer):
-    # Only used for ctc objective
-    def __init__(self,
-                 vocab: Vocab,
-                 ignore_oov: bool = True):
-        self.vocab = vocab
-        self.ignore_oov = ignore_oov
-
-    def decode(self, ids: List[int]) -> str:
-        return ' '.join(self.vocab[id] for id in ids)
-
-    def encode(self, transcript: str) -> List[int]:
-        encoded_output = []
-
-        for word in transcript.lower().split():
-            vocab_found, remaining_transcript = self.vocab.trie.max_split(word)
-
-            # append corresponding label
-            if vocab_found and remaining_transcript == "":
-                # word exists in the vocab
-                encoded_output.append(self.vocab[word])
-            elif not self.ignore_oov:
-                # label oov word
-                if self.vocab.oov_token_id is None:
-                    raise ValueError("label for oov word is not specified")
-                encoded_output.append(self.vocab.oov_token_id)
-
-        return encoded_output
diff --git a/howl/data/dataset/__init__.py b/howl/data/dataset/__init__.py
index af5e7c2d..e69de29b 100644
--- a/howl/data/dataset/__init__.py
+++ b/howl/data/dataset/__init__.py
@@ -1,5 +0,0 @@
-from .base import *
-from .dataset import *
-from .labeller import *
-from .phone import *
-from .serialize import *
diff --git a/howl/data/dataset/base.py b/howl/data/dataset/base.py
deleted file mode 100644
index f1ffc71d..00000000
--- a/howl/data/dataset/base.py
+++ /dev/null
@@ -1,182 +0,0 @@
-import enum
-from collections import Counter
-from copy import deepcopy
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Generic, List, Mapping, Optional, Tuple, TypeVar
-
-import torch
-from pydantic import BaseModel
-
-from .phone import Phone, PhonePhrase
-
-__all__ = [
-    "AudioClipExample",
-    "AudioClipMetadata",
-    "AudioDatasetStatistics",
-    "ClassificationBatch",
-    "ClassificationClipExample",
-    "DatasetType",
-    "EmplacableExample",
-    "WakeWordClipExample",
-    "SequenceBatch",
-    "FrameLabelData",
-    "UNKNOWN_TRANSCRIPTION",
-    "NEGATIVE_CLASS",
-]
-
-
-UNKNOWN_TRANSCRIPTION = "[UNKNOWN]"
-NEGATIVE_CLASS = "[NEGATIVE]"
-
-
-@dataclass
-class FrameLabelData:
-    timestamp_label_map: Mapping[float, int]
-    start_timestamp: List[Tuple[int, float]]
-    char_indices: List[Tuple[int, List[int]]]
-
-
-@dataclass
-class AudioDatasetStatistics:
-    num_examples: int
-    audio_length_seconds: int
-    vocab_counts: Counter
-
-
-class AudioClipMetadata(BaseModel):
-    path: Optional[Path] = Path(".")
-    phone_strings: Optional[List[str]]
-    words: Optional[List[str]]
-    phone_end_timestamps: Optional[List[float]]
-    word_end_timestamps: Optional[List[float]]
-    end_timestamps: Optional[List[float]]  # TODO: remove, backwards compat right now
-    transcription: Optional[str] = ""
-
-    # TODO:: id should be an explicit variable in order to support datasets creation with the audio data in memory
-    @property
-    def audio_id(self) -> str:
-        return self.path.name.split(".", 1)[0]
-
-    @property
-    def phone_phrase(self) -> Optional[PhonePhrase]:
-        return PhonePhrase([Phone(x) for x in self.phone_strings]) if self.phone_strings else None
-
-
-class EmplacableExample:
-    audio_data: torch.Tensor
-
-    def emplaced_audio_data(
-        self, audio_data: torch.Tensor, scale: float = 1, bias: float = 0, new: bool = False
-    ) -> "EmplacableExample":
-        raise NotImplementedError
-
-
-T = TypeVar("T", bound=AudioClipMetadata)
-
-
-class AudioClipExample(EmplacableExample, Generic[T]):
-    def __init__(self, metadata: T, audio_data: torch.Tensor, sample_rate: int):
-        self.metadata = metadata
-        self.audio_data = audio_data
-        self.sample_rate = sample_rate
-
-    def pin_memory(self):
-        self.audio_data.pin_memory()
-
-    def emplaced_audio_data(
-        self, audio_data: torch.Tensor, scale: float = 1, bias: float = 0, new: bool = False
-    ) -> "AudioClipExample":
-        metadata = self.metadata
-        if new:
-            metadata = deepcopy(metadata)
-            metadata.transcription = ""
-        return AudioClipExample(metadata, audio_data, self.sample_rate)
-
-
-@dataclass
-class ClassificationBatch:
-    audio_data: torch.Tensor
-    labels: Optional[torch.Tensor]
-    lengths: torch.Tensor
-
-    @classmethod
-    def from_single(cls, audio_clip: torch.Tensor, label: int) -> "ClassificationBatch":
-        return cls(audio_clip.unsqueeze(0), torch.tensor([label]), torch.tensor([audio_clip.size(-1)]))
-
-    def pin_memory(self):
-        self.audio_data.pin_memory()
-        self.labels.pin_memory()
-        self.lengths.pin_memory()
-
-    def to(self, device: torch.device) -> "ClassificationBatch":
-        self.audio_data = self.audio_data.to(device)
-        if self.labels is not None:
-            self.labels = self.labels.to(device)
-        self.lengths = self.lengths.to(device)
-        return self
-
-
-@dataclass
-class SequenceBatch:
-    audio_data: torch.Tensor  # (batch size, max audio length)
-    labels: torch.Tensor  # (batch size, max label length)
-    audio_lengths: Optional[torch.Tensor]
-    label_lengths: Optional[torch.Tensor]
-
-    def __post_init__(self):
-        if self.audio_lengths is None:
-            self.audio_lengths = torch.tensor([self.audio_data.size(0) for _ in range(self.audio_data.size(1))])
-            self.audio_lengths = self.audio_lengths.to(self.audio_data.device)
-        if self.label_lengths is None:
-            self.label_lengths = torch.tensor([self.labels.size(0) for _ in range(self.labels.size(1))])
-            self.label_lengths = self.label_lengths.to(self.label_lengths.device)
-
-    def pin_memory(self):
-        self.audio_data.pin_memory()
-        self.labels.pin_memory()
-        self.audio_lengths.pin_memory()
-        self.label_lengths.pin_memory()
-
-    def to(self, device: torch.device) -> "SequenceBatch":
-        self.audio_data = self.audio_data.to(device)
-        self.labels = self.labels.to(device)
-        self.audio_lengths = self.audio_lengths.to(device)
-        self.label_lengths = self.label_lengths.to(device)
-        return self
-
-
-@dataclass
-class WakeWordClipExample(AudioClipExample[AudioClipMetadata]):
-    def __init__(self, label_data: FrameLabelData, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.label_data = label_data
-
-    def emplaced_audio_data(
-        self, audio_data: torch.Tensor, scale: float = 1, bias: float = 0, new: bool = False
-    ) -> "WakeWordClipExample":
-        ex = super().emplaced_audio_data(audio_data, scale, bias, new)
-        label_data = {} if new else {scale * k + bias: v for k, v in self.label_data.timestamp_label_map.items()}
-        return WakeWordClipExample(
-            FrameLabelData(label_data, self.label_data.start_timestamp, self.label_data.char_indices),
-            ex.metadata,
-            audio_data,
-            self.sample_rate,
-        )
-
-
-@dataclass
-class ClassificationClipExample(AudioClipExample[AudioClipMetadata]):
-    def __init__(self, label, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.label = label
-
-    def emplaced_audio_data(self, audio_data: torch.Tensor, **kwargs) -> "ClassificationClipExample":
-        return ClassificationClipExample(self.label, self.metadata, audio_data, self.sample_rate)
-
-
-class DatasetType(enum.Enum):
-    TRAINING = enum.auto()
-    DEV = enum.auto()
-    TEST = enum.auto()
-    UNSPECIFIED = enum.auto()
diff --git a/howl/data/dataset/common_voice_dataset_loader.py b/howl/data/dataset/common_voice_dataset_loader.py
new file mode 100644
index 00000000..04c16d10
--- /dev/null
+++ b/howl/data/dataset/common_voice_dataset_loader.py
@@ -0,0 +1,114 @@
+import logging
+from functools import partial
+from pathlib import Path
+from typing import Tuple
+
+import pandas as pd
+
+from howl.data.common.metadata import UNKNOWN_TRANSCRIPTION, AudioClipMetadata
+from howl.data.dataset.dataset import AudioClipDataset, DatasetType
+from howl.data.dataset.dataset_loader import RegisteredPathDatasetLoader
+from howl.utils.hash import sha256_int
+
+__all__ = [
+    "MozillaCommonVoiceLoader",
+    "MozillaKeywordLoader",
+    "MozillaWakeWordLoader",
+    "SoundIdSplitMozillaWakeWordLoader",
+    "SpeakerSplitMozillaWakeWordLoader",
+]
+
+
+class MozillaCommonVoiceLoader(RegisteredPathDatasetLoader, name="mozilla-cv"):
+    def load_splits(self, path: Path, **dataset_kwargs) -> Tuple[AudioClipDataset, AudioClipDataset, AudioClipDataset]:
+        def load(filename, set_type):
+            logging.info(f"Loading split {filename}...")
+            df = pd.read_csv(str(path / filename), sep="\t", quoting=3, na_filter=False)
+            metadata_list = []
+            for tup in df.itertuples():
+                metadata_list.append(
+                    AudioClipMetadata(path=(path / "clips" / tup.path).absolute(), transcription=tup.sentence,)
+                )
+            return AudioClipDataset(metadata_list=metadata_list, set_type=set_type, **dataset_kwargs)
+
+        assert path.exists(), "dataset path doesn't exist"
+        filenames = ("train.tsv", "dev.tsv", "test.tsv")
+        assert all((path / x).exists() for x in filenames), "dataset missing metadata"
+        return (
+            load("train.tsv", DatasetType.TRAINING),
+            load("dev.tsv", DatasetType.DEV),
+            load("test.tsv", DatasetType.TEST),
+        )
+
+
+class MozillaKeywordLoader(RegisteredPathDatasetLoader, name="mozilla-kw"):
+    def load_splits(self, path: Path, **dataset_kwargs) -> Tuple[AudioClipDataset, AudioClipDataset, AudioClipDataset]:
+        logging.info("Loading Mozilla keyword dataset...")
+        df = pd.read_csv(str(path / "validated.tsv"), sep="\t", quoting=3, na_filter=False)
+        md_splits = ([], [], [])
+        for tup in df.itertuples():
+            md = AudioClipMetadata(path=(path / "clips" / tup.path).absolute(), transcription=tup.sentence)
+            bucket = sha256_int(tup.client_id) % 100
+            if bucket < 80:
+                md_splits[0].append(md)
+            elif bucket < 90:
+                md_splits[1].append(md)
+            else:
+                md_splits[2].append(md)
+        return (
+            AudioClipDataset(metadata_list=md_splits[0], set_type=DatasetType.TRAINING, **dataset_kwargs,),
+            AudioClipDataset(metadata_list=md_splits[1], set_type=DatasetType.DEV, **dataset_kwargs),
+            AudioClipDataset(metadata_list=md_splits[2], set_type=DatasetType.TEST, **dataset_kwargs),
+        )
+
+
+class MozillaWakeWordLoader(RegisteredPathDatasetLoader, name="mozilla-ww"):
+    def __init__(
+        self, training_pct=80, dev_pct=10, test_pct=10, split_by_speaker=True, split="verified",
+    ):
+        self.split_by_speaker = split_by_speaker
+        total = training_pct + dev_pct + test_pct
+        training_pct = 100 * training_pct / total
+        dev_pct = 100 * dev_pct / total
+        test_pct = 100 * test_pct / total
+        self.cutoffs = (
+            training_pct,
+            dev_pct + training_pct,
+            training_pct + dev_pct + test_pct,
+        )
+        self.split = split
+
+    def load_splits(self, path: Path, **dataset_kwargs) -> Tuple[AudioClipDataset, AudioClipDataset, AudioClipDataset]:
+        assert path.exists(), "dataset path doesn't exist"
+        if self.split in {"verified", "rejected"}:
+            audio_path = path / self.split
+        else:
+            raise ValueError("Split ill-defined.")
+        assert audio_path.exists(), "dataset malformed"
+
+        sound_file_paths = audio_path.glob("*/*.ogg")
+        metadatas = ([], [], [])
+        logging.info("Loading wake word dataset...")
+        using_verified = self.split == "verified"
+        for sound_fp in sound_file_paths:
+            sound_id = sound_fp.stem
+            speaker_id = sound_fp.parent.name
+            if using_verified:
+                with open(str((sound_fp.parent / sound_id).with_suffix(".txt"))) as f:
+                    transcription = f.read()
+            else:
+                transcription = UNKNOWN_TRANSCRIPTION
+            metadata = AudioClipMetadata(path=sound_fp.absolute(), transcription=transcription)
+            bucket = sha256_int(speaker_id) if self.split_by_speaker else sha256_int(sound_id)
+            bucket %= 100
+            bucket = next(idx for idx, cutoff in enumerate(self.cutoffs) if bucket < cutoff)
+            metadatas[bucket].append(metadata)
+        return (
+            AudioClipDataset(metadata_list=metadatas[0], set_type=DatasetType.TRAINING, **dataset_kwargs,),
+            AudioClipDataset(metadata_list=metadatas[1], set_type=DatasetType.DEV, **dataset_kwargs),
+            AudioClipDataset(metadata_list=metadatas[2], set_type=DatasetType.TEST, **dataset_kwargs),
+        )
+
+
+SoundIdSplitMozillaWakeWordLoader = partial(MozillaWakeWordLoader, split_by_speaker=False)
+SpeakerSplitMozillaWakeWordLoader = partial(MozillaWakeWordLoader, split_by_speaker=True)
diff --git a/howl/data/dataset/dataset.py b/howl/data/dataset/dataset.py
index 6b6b41a2..55843210 100644
--- a/howl/data/dataset/dataset.py
+++ b/howl/data/dataset/dataset.py
@@ -1,38 +1,60 @@
+import enum
 from collections import Counter, defaultdict
 from copy import deepcopy
+from dataclasses import dataclass
 from functools import lru_cache
 from typing import Any, Callable, Generic, List, TypeVar
 
 import torch
 import torch.utils.data as tud
-from howl.data.searcher import WordTranscriptSearcher
+
+from howl.data.common.example import (
+    AudioClipExample,
+    ClassificationClipExample,
+    WakeWordClipExample,
+)
+from howl.data.common.labeler import FrameLabeler
+from howl.data.common.metadata import NEGATIVE_CLASS, AudioClipMetadata
+from howl.data.common.searcher import WordTranscriptSearcher
 from howl.settings import SETTINGS
 from howl.utils.audio import silent_load
-from howl.utils.hash import sha256_int
 
-from .base import (NEGATIVE_CLASS, AudioClipExample, AudioClipMetadata,
-                   AudioDatasetStatistics, ClassificationClipExample,
-                   DatasetType, WakeWordClipExample)
-from .labeller import FrameLabeler
-from .phone import PronunciationDictionary
+__all__ = [
+    "DatasetType",
+    "AudioDataset",
+    "AudioDatasetStatistics",
+    "HonkSpeechCommandsDataset",
+    "AudioClipDataset",
+    "WakeWordDataset",
+    "AudioClassificationDataset",
+]
+
+
+@dataclass
+class AudioDatasetStatistics:
+    num_examples: int
+    audio_length_seconds: int
+    vocab_counts: Counter
+
 
-__all__ = ['AudioDataset',
-           'HonkSpeechCommandsDataset',
-           'AudioClipDataset',
-           'WakeWordDataset',
-           'AudioClassificationDataset',
-           'Sha256Splitter']
+class DatasetType(enum.Enum):
+    TRAINING = enum.auto()
+    DEV = enum.auto()
+    TEST = enum.auto()
+    UNSPECIFIED = enum.auto()
 
 
-T = TypeVar('T', bound=AudioClipMetadata)
+T = TypeVar("T", bound=AudioClipMetadata)
 
 
 class AudioDataset(tud.Dataset, Generic[T]):
-    def __init__(self,
-                 metadata_list: List[T],
-                 sr: int = 16000,
-                 mono: bool = True,
-                 set_type: DatasetType = DatasetType.UNSPECIFIED):
+    def __init__(
+        self,
+        metadata_list: List[T],
+        sr: int = 16000,
+        mono: bool = True,
+        set_type: DatasetType = DatasetType.UNSPECIFIED,
+    ):
         self.metadata_list = metadata_list
         self.set_type = set_type
         self.sr = sr
@@ -65,15 +87,18 @@ def split(self, predicate_fn: Callable[[Any], bool]):
         x2.metadata_list = data_list2
         return x1, x2
 
-    def extend(self, other_dataset: 'AudioDataset'):
+    def extend(self, other_dataset: "AudioDataset"):
         self.metadata_list.extend(other_dataset.metadata_list)
         return self
 
     def __len__(self):
         return len(self.metadata_list)
 
-    def compute_statistics(self, word_searcher: WordTranscriptSearcher = None, compute_length: bool = True, use_trim: bool = True) -> AudioDatasetStatistics:
-        from howl.data.transform import trim
+    def compute_statistics(
+        self, word_searcher: WordTranscriptSearcher = None, compute_length: bool = True, use_trim: bool = True
+    ) -> AudioDatasetStatistics:
+        from howl.data.transform.operator import trim
+
         seconds = 0
         total_vocab_count = Counter()
         for ex in self:
@@ -87,7 +112,7 @@ def compute_statistics(self, word_searcher: WordTranscriptSearcher = None, compu
 
 
 class AudioClipDataset(AudioDataset[AudioClipMetadata]):
-    @ lru_cache(maxsize=SETTINGS.cache.cache_size)
+    @lru_cache(maxsize=SETTINGS.cache.cache_size)
     def __getitem__(self, idx) -> AudioClipExample:
         metadata = self.metadata_list[idx]
         audio_data = silent_load(str(metadata.path), self.sr, self.mono)
@@ -95,65 +120,54 @@ def __getitem__(self, idx) -> AudioClipExample:
 
 
 class WakeWordDataset(AudioDataset[AudioClipMetadata]):
-    def __init__(self,
-                 frame_labeler: FrameLabeler,
-                 *args,
-                 **kwargs):
+    def __init__(self, frame_labeler: FrameLabeler, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.frame_labeler = frame_labeler
 
-    @ lru_cache(maxsize=SETTINGS.cache.cache_size)
+    @lru_cache(maxsize=SETTINGS.cache.cache_size)
     def __getitem__(self, idx) -> WakeWordClipExample:
         metadata = self.metadata_list[idx]
         audio_data = silent_load(str(metadata.path), self.sr, self.mono)
-        return WakeWordClipExample(metadata=metadata,
-                                   audio_data=torch.from_numpy(audio_data),
-                                   sample_rate=self.sr,
-                                   label_data=self.frame_labeler.compute_frame_labels(metadata))
+        return WakeWordClipExample(
+            metadata=metadata,
+            audio_data=torch.from_numpy(audio_data),
+            sample_rate=self.sr,
+            label_data=self.frame_labeler.compute_frame_labels(metadata),
+        )
 
 
 class AudioClassificationDataset(AudioDataset[AudioClipMetadata]):
-    def __init__(self,
-                 label_map: defaultdict,
-                 *args,
-                 **kwargs):
+    def __init__(self, label_map: defaultdict, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.label_map = label_map
         self.vocab = {v: k for k, v in label_map.items()}
         self.vocab[label_map.get(None)] = NEGATIVE_CLASS
 
-    @ lru_cache(maxsize=SETTINGS.cache.cache_size)
+    @lru_cache(maxsize=SETTINGS.cache.cache_size)
     def __getitem__(self, idx) -> ClassificationClipExample:
         metadata = self.metadata_list[idx]
         audio_data = silent_load(str(metadata.path), self.sr, self.mono)
-        return ClassificationClipExample(metadata=metadata,
-                                         audio_data=torch.from_numpy(audio_data),
-                                         sample_rate=self.sr,
-                                         label=self.label_map[metadata.transcription])
+        return ClassificationClipExample(
+            metadata=metadata,
+            audio_data=torch.from_numpy(audio_data),
+            sample_rate=self.sr,
+            label=self.label_map[metadata.transcription],
+        )
 
 
 class HonkSpeechCommandsDataset(AudioClassificationDataset):
     def __init__(self, *args, silence_proportion: float = 0.1, **kwargs):
         super().__init__(*args, **kwargs)
         self.silence_proportion = silence_proportion
-        self.silence_label = self.label_map['__silence__']
+        self.silence_label = self.label_map["__silence__"]
 
     def __getitem__(self, idx) -> ClassificationClipExample:
         if idx < super().__len__():
             return super().__getitem__(idx)
-        return ClassificationClipExample(metadata=AudioClipMetadata(),
-                                         audio_data=torch.zeros(16000),
-                                         sample_rate=self.sr,
-                                         label=self.silence_label)
+        return ClassificationClipExample(
+            metadata=AudioClipMetadata(), audio_data=torch.zeros(16000), sample_rate=self.sr, label=self.silence_label
+        )
 
     def __len__(self):
         orig_len = super().__len__()
         return orig_len + int(self.silence_proportion * orig_len)
-
-
-class Sha256Splitter:
-    def __init__(self, target_pct: int):
-        self.target_pct = target_pct
-
-    def __call__(self, x: AudioClipMetadata) -> bool:
-        return (sha256_int(str(x.path)) % 100) < self.target_pct
diff --git a/howl/data/dataset/dataset_loader.py b/howl/data/dataset/dataset_loader.py
new file mode 100644
index 00000000..11fc690c
--- /dev/null
+++ b/howl/data/dataset/dataset_loader.py
@@ -0,0 +1,77 @@
+import json
+import logging
+from pathlib import Path
+from typing import Tuple
+
+from howl.data.common.metadata import AudioClipMetadata
+from howl.data.dataset.dataset import (
+    AudioClipDataset,
+    AudioDataset,
+    DatasetType,
+    WakeWordDataset,
+)
+from howl.utils.class_registry import ClassRegistry
+
+__all__ = [
+    "PathDatasetLoader",
+    "RegisteredPathDatasetLoader",
+    "MetadataLoaderMixin",
+    "AudioClipDatasetLoader",
+    "WakeWordDatasetLoader",
+    "RecursiveNoiseDatasetLoader",
+]
+
+
+class PathDatasetLoader:
+    def load_splits(self, path: Path, **dataset_kwargs) -> Tuple[AudioDataset, AudioDataset, AudioDataset]:
+        raise NotImplementedError
+
+
+class RegisteredPathDatasetLoader(PathDatasetLoader, ClassRegistry):
+    registered_map = {}
+
+
+class MetadataLoaderMixin:
+    dataset_class = None
+    metadata_class = None
+    default_prefix = ""
+
+    def load_splits(self, path: Path, prefix: str = None, **dataset_kwargs):
+        def load(jsonl_name):
+            metadata_list = []
+            with open(jsonl_name) as f:
+                for json_str in iter(f.readline, ""):
+                    metadata = self.metadata_class(**json.loads(json_str))
+                    metadata.path = (path / "audio" / metadata.path).absolute()
+                    metadata_list.append(metadata)
+                return metadata_list
+
+        if prefix is None:
+            prefix = self.default_prefix
+        logging.info(f"Loading flat dataset from {path}...")
+        training_path = path / f"{prefix}metadata-{DatasetType.TRAINING.name.lower()}.jsonl"
+        dev_path = path / f"{prefix}metadata-{DatasetType.DEV.name.lower()}.jsonl"
+        test_path = path / f"{prefix}metadata-{DatasetType.TEST.name.lower()}.jsonl"
+        return (
+            self.dataset_class(metadata_list=load(training_path), set_type=DatasetType.TRAINING, **dataset_kwargs,),
+            self.dataset_class(metadata_list=load(dev_path), set_type=DatasetType.DEV, **dataset_kwargs),
+            self.dataset_class(metadata_list=load(test_path), set_type=DatasetType.TEST, **dataset_kwargs,),
+        )
+
+
+class AudioClipDatasetLoader(MetadataLoaderMixin, RegisteredPathDatasetLoader, name="clip"):
+    dataset_class = AudioClipDataset
+    metadata_class = AudioClipMetadata
+
+
+class WakeWordDatasetLoader(MetadataLoaderMixin, PathDatasetLoader):
+    default_prefix = "aligned-"
+    dataset_class = WakeWordDataset
+    metadata_class = AudioClipMetadata
+
+
+class RecursiveNoiseDatasetLoader:
+    def load(self, path: Path, **dataset_kwargs) -> AudioClipDataset:
+        wav_names = path.glob("**/*.wav")
+        metadata_list = [AudioClipMetadata(path=filename.absolute(), transcription="") for filename in wav_names]
+        return AudioClipDataset(metadata_list=metadata_list, set_type=DatasetType.TRAINING, **dataset_kwargs)
diff --git a/howl/data/dataset/dataset_writer.py b/howl/data/dataset/dataset_writer.py
new file mode 100644
index 00000000..8aecc047
--- /dev/null
+++ b/howl/data/dataset/dataset_writer.py
@@ -0,0 +1,74 @@
+import logging
+from copy import deepcopy
+from pathlib import Path
+from typing import TypeVar
+
+import soundfile
+from tqdm import tqdm
+
+from howl.data.common.metadata import AudioClipMetadata
+from howl.data.dataset.dataset import AudioClipDataset, DatasetType
+from howl.utils.audio import silent_load
+
+__all__ = [
+    "AudioDatasetWriter",
+    "AudioDatasetMetadataWriter",
+]
+
+
+class AudioDatasetMetadataWriter:
+    def __init__(
+        self, dataset_path: Path, set_type: DatasetType, prefix: str = "", mode: str = "w",
+    ):
+        self.filename = str(dataset_path / f"{prefix}metadata-{set_type.name.lower()}.jsonl")
+        self.mode = mode
+
+    def __enter__(self):
+        self.f = open(self.filename, self.mode)
+        return self
+
+    def write(self, metadata: AudioClipMetadata):
+        metadata = deepcopy(metadata)
+        with metadata.path.with_suffix(".lab").open("w") as f:
+            f.write(f"{metadata.transcription}\n")
+        metadata.path = metadata.path.name
+        self.f.write(metadata.json() + "\n")
+
+    def __exit__(self, *args):
+        self.f.close()
+
+
+T = TypeVar("T", bound=AudioClipDataset)
+
+
+class AudioDatasetWriter:
+    def __init__(
+        self, dataset: AudioClipDataset, prefix: str = "", mode: str = "w", print_progress: bool = True,
+    ):
+        self.dataset = dataset
+        self.print_progress = print_progress
+        self.mode = mode
+        self.prefix = prefix
+
+    def write(self, folder: Path):
+        def process(metadata: AudioClipMetadata):
+            new_path = (audio_folder / metadata.audio_id).with_suffix(".wav")
+            # TODO:: process function should also take in sample (AudioClipExample)
+            #        and use sample.audio_data when metadata.path does not exist
+            if not new_path.exists():
+                audio_data = silent_load(str(metadata.path), self.dataset.sr, self.dataset.mono)
+                soundfile.write(str(new_path), audio_data, self.dataset.sr)
+            metadata.path = new_path
+
+        logging.info(f"Writing flat dataset to {folder}...")
+        folder.mkdir(exist_ok=True)
+        audio_folder = folder / "audio"
+        audio_folder.mkdir(exist_ok=True)
+        with AudioDatasetMetadataWriter(folder, self.dataset.set_type, prefix=self.prefix, mode=self.mode) as writer:
+            for metadata in tqdm(self.dataset.metadata_list, disable=not self.print_progress, desc="Writing files",):
+                try:
+                    process(metadata)
+                except EOFError:
+                    logging.warning(f"Skipping bad file {metadata.path}")
+                    continue
+                writer.write(metadata)
diff --git a/howl/data/dataset/gsc_dataset_loader.py b/howl/data/dataset/gsc_dataset_loader.py
new file mode 100644
index 00000000..5322e504
--- /dev/null
+++ b/howl/data/dataset/gsc_dataset_loader.py
@@ -0,0 +1,47 @@
+from collections import defaultdict
+from pathlib import Path
+from typing import List, Tuple
+
+from howl.data.common.metadata import AudioClipMetadata
+from howl.data.dataset.dataset import AudioClassificationDataset, DatasetType
+from howl.data.dataset.dataset_loader import RegisteredPathDatasetLoader
+
+__all__ = [
+    "GoogleSpeechCommandsDatasetLoader",
+]
+
+
+class GoogleSpeechCommandsDatasetLoader(RegisteredPathDatasetLoader, name="gsc"):
+    def __init__(self, vocab: List[str] = None, use_bg_noise: bool = False):
+        self.vocab = vocab
+        self.use_bg_noise = use_bg_noise
+
+    def load_splits(
+        self, path: Path, **dataset_kwargs
+    ) -> Tuple[
+        AudioClassificationDataset, AudioClassificationDataset, AudioClassificationDataset,
+    ]:
+        def load(set_type):
+            metadata_list = []
+            for path in all_list:
+                key = str(Path(path.parent.name) / path.name)
+                if file_map[key] != set_type:
+                    continue
+                metadata_list.append(AudioClipMetadata(path=path.absolute(), transcription=path.parent.name))
+            return AudioClassificationDataset(
+                metadata_list=metadata_list, label_map=label_map, set_type=set_type, **dataset_kwargs
+            )
+
+        file_map = defaultdict(lambda: DatasetType.TRAINING)
+        with (path / "testing_list.txt").open() as f:
+            file_map.update({k: DatasetType.TEST for k in f.read().split("\n")})
+        with (path / "validation_list.txt").open() as f:
+            file_map.update({k: DatasetType.DEV for k in f.read().split("\n")})
+        all_list = list(path.glob("*/*.wav"))
+        if not self.use_bg_noise:
+            all_list = list(filter(lambda x: "noise" not in str(x), all_list))
+        folders = sorted(list(path.glob("*/")))
+        vocab = [x.name for x in folders] if self.vocab is None else self.vocab
+        label_map = defaultdict(lambda: len(vocab))
+        label_map.update({k: idx for idx, k in enumerate(vocab)})
+        return load(DatasetType.TRAINING), load(DatasetType.DEV), load(DatasetType.TEST)
diff --git a/howl/data/dataset/hey_snips_dataset_loader.py b/howl/data/dataset/hey_snips_dataset_loader.py
new file mode 100644
index 00000000..5276992d
--- /dev/null
+++ b/howl/data/dataset/hey_snips_dataset_loader.py
@@ -0,0 +1,80 @@
+import json
+import logging
+from multiprocessing import Pool
+from pathlib import Path
+from typing import Tuple
+
+from tqdm import tqdm
+
+from howl.data.common.metadata import AudioClipMetadata
+from howl.data.dataset.dataset import AudioClipDataset, DatasetType
+from howl.data.dataset.dataset_loader import RegisteredPathDatasetLoader
+from howl.utils.transcribe import SpeechToText
+
+__all__ = [
+    "HeySnipsWakeWordLoader",
+]
+
+
+def transcribe_hey_snips_audio(path, metadata):
+    stt = SpeechToText()
+    path = (path / metadata["audio_file_path"]).absolute()
+    transcription = "hey snips"
+    if metadata["is_hotword"] == 0:  # negative sample
+        transcription = stt.transcribe(path)
+
+    return path, transcription
+
+
+class HeySnipsWakeWordLoader(RegisteredPathDatasetLoader, name="hey-snips"):
+    def __init__(self, num_processes=8):
+        self.stt = SpeechToText()
+        self.num_processes = num_processes
+        self.pool = Pool(processes=self.num_processes)
+
+    def load_splits(self, path: Path, **dataset_kwargs) -> Tuple[AudioClipDataset, AudioClipDataset, AudioClipDataset]:
+        self.path = path
+
+        def load(filename, set_type):
+            logging.info(f"Loading split {filename}...")
+
+            metadata_list = []
+            with open(str(self.path / filename)) as f:
+                raw_metadata_list = json.load(f)
+
+            total = len(raw_metadata_list)
+
+            transcription_fail_count = 0
+            pbar = tqdm(range(0, total, self.num_processes), desc=f"{filename}")
+            for starting_idx in pbar:
+                transcription_results = []
+                for metadata in raw_metadata_list[starting_idx : starting_idx + self.num_processes]:
+                    transcription_results.append(
+                        self.pool.apply_async(transcribe_hey_snips_audio, (self.path, metadata,))
+                    )
+
+                processing_count = starting_idx + len(transcription_results)
+
+                for result in transcription_results:
+                    pbar.set_postfix(dict(transcription_fail_rate=f"{transcription_fail_count}/{processing_count}"))
+                    path, transcription = result.get()
+
+                    if transcription == "":
+                        transcription_fail_count += 1
+                        continue
+
+                    metadata_list.append(AudioClipMetadata(path=path, transcription=transcription))
+
+            logging.info(f"{transcription_fail_count}/{total} samples have empty transcription")
+
+            return AudioClipDataset(metadata_list=metadata_list, set_type=set_type, **dataset_kwargs)
+
+        assert path.exists(), "dataset path doesn't exist"
+        filenames = ("train.json", "dev.json", "test.json")
+        assert all((path / x).exists() for x in filenames), "dataset missing metadata"
+
+        return (
+            load("train.json", DatasetType.TRAINING),
+            load("dev.json", DatasetType.DEV),
+            load("test.json", DatasetType.TEST),
+        )
diff --git a/howl/data/dataset/serialize.py b/howl/data/dataset/serialize.py
deleted file mode 100644
index 61abbe7d..00000000
--- a/howl/data/dataset/serialize.py
+++ /dev/null
@@ -1,336 +0,0 @@
-import json
-import logging
-from collections import defaultdict
-from copy import deepcopy
-from functools import partial
-from multiprocessing import Pool
-from pathlib import Path
-from typing import List, Tuple, TypeVar
-
-import pandas as pd
-import soundfile
-from tqdm import tqdm
-
-from howl.registered import RegisteredObjectBase
-from howl.utils.audio import silent_load
-from howl.utils.hash import sha256_int
-from howl.utils.transcribe import SpeechToText
-
-from .base import UNKNOWN_TRANSCRIPTION, AudioClipMetadata, DatasetType
-from .dataset import (
-    AudioClassificationDataset,
-    AudioClipDataset,
-    AudioDataset,
-    WakeWordDataset,
-)
-
-__all__ = [
-    "AudioDatasetWriter",
-    "AudioClipDatasetLoader",
-    "MozillaWakeWordLoader",
-    "RegisteredPathDatasetLoader",
-    "MozillaCommonVoiceLoader",
-    "AudioDatasetMetadataWriter",
-    "WakeWordDatasetLoader",
-    "GoogleSpeechCommandsDatasetLoader",
-    "MozillaKeywordLoader",
-    "PathDatasetLoader",
-    "RecursiveNoiseDatasetLoader",
-    "HeySnipsWakeWordLoader",
-]
-
-
-class AudioDatasetMetadataWriter:
-    def __init__(self, dataset_path: Path, set_type: DatasetType, prefix: str = "", mode: str = "w"):
-        self.filename = str(dataset_path / f"{prefix}metadata-{set_type.name.lower()}.jsonl")
-        self.mode = mode
-
-    def __enter__(self):
-        self.f = open(self.filename, self.mode)
-        return self
-
-    def write(self, metadata: AudioClipMetadata):
-        metadata = deepcopy(metadata)
-        with metadata.path.with_suffix(".lab").open("w") as f:
-            f.write(f"{metadata.transcription}\n")
-        metadata.path = metadata.path.name
-        self.f.write(metadata.json() + "\n")
-
-    def __exit__(self, *args):
-        self.f.close()
-
-
-T = TypeVar("T", bound=AudioClipDataset)
-
-
-class AudioDatasetWriter:
-    def __init__(self, dataset: AudioClipDataset, prefix: str = "", mode: str = "w", print_progress: bool = True):
-        self.dataset = dataset
-        self.print_progress = print_progress
-        self.mode = mode
-        self.prefix = prefix
-
-    def write(self, folder: Path):
-        def process(metadata: AudioClipMetadata):
-            new_path = (audio_folder / metadata.audio_id).with_suffix(".wav")
-            # TODO:: process function should also take in sample (AudioClipExample)
-            #        and use sample.audio_data when metadata.path does not exist
-            if not new_path.exists():
-                audio_data = silent_load(str(metadata.path), self.dataset.sr, self.dataset.mono)
-                soundfile.write(str(new_path), audio_data, self.dataset.sr)
-            metadata.path = new_path
-
-        logging.info(f"Writing flat dataset to {folder}...")
-        folder.mkdir(exist_ok=True)
-        audio_folder = folder / "audio"
-        audio_folder.mkdir(exist_ok=True)
-        with AudioDatasetMetadataWriter(folder, self.dataset.set_type, prefix=self.prefix, mode=self.mode) as writer:
-            for metadata in tqdm(self.dataset.metadata_list, disable=not self.print_progress, desc="Writing files"):
-                try:
-                    process(metadata)
-                except EOFError:
-                    logging.warning(f"Skipping bad file {metadata.path}")
-                    continue
-                writer.write(metadata)
-
-
-class PathDatasetLoader:
-    def load_splits(self, path: Path, **dataset_kwargs) -> Tuple[AudioDataset, AudioDataset, AudioDataset]:
-        raise NotImplementedError
-
-
-class RegisteredPathDatasetLoader(PathDatasetLoader, RegisteredObjectBase):
-    registered_map = {}
-
-
-class MetadataLoaderMixin:
-    dataset_class = None
-    metadata_class = None
-    default_prefix = ""
-
-    def load_splits(self, path: Path, prefix: str = None, **dataset_kwargs):
-        def load(jsonl_name):
-            metadata_list = []
-            with open(jsonl_name) as f:
-                for json_str in iter(f.readline, ""):
-                    metadata = self.metadata_class(**json.loads(json_str))
-                    metadata.path = (path / "audio" / metadata.path).absolute()
-                    metadata_list.append(metadata)
-                return metadata_list
-
-        if prefix is None:
-            prefix = self.default_prefix
-        logging.info(f"Loading flat dataset from {path}...")
-        training_path = path / f"{prefix}metadata-{DatasetType.TRAINING.name.lower()}.jsonl"
-        dev_path = path / f"{prefix}metadata-{DatasetType.DEV.name.lower()}.jsonl"
-        test_path = path / f"{prefix}metadata-{DatasetType.TEST.name.lower()}.jsonl"
-        return (
-            self.dataset_class(metadata_list=load(training_path), set_type=DatasetType.TRAINING, **dataset_kwargs),
-            self.dataset_class(metadata_list=load(dev_path), set_type=DatasetType.DEV, **dataset_kwargs),
-            self.dataset_class(metadata_list=load(test_path), set_type=DatasetType.TEST, **dataset_kwargs),
-        )
-
-
-class AudioClipDatasetLoader(MetadataLoaderMixin, RegisteredPathDatasetLoader, name="clip"):
-    dataset_class = AudioClipDataset
-    metadata_class = AudioClipMetadata
-
-
-class WakeWordDatasetLoader(MetadataLoaderMixin, PathDatasetLoader):
-    default_prefix = "aligned-"
-    dataset_class = WakeWordDataset
-    metadata_class = AudioClipMetadata
-
-
-def transcribe_hey_snips_audio(path, metadata):
-    stt = SpeechToText()
-    path = (path / metadata["audio_file_path"]).absolute()
-    transcription = "hey snips"
-    if metadata["is_hotword"] == 0:  # negative sample
-        transcription = stt.transcribe(path)
-
-    return path, transcription
-
-
-class HeySnipsWakeWordLoader(RegisteredPathDatasetLoader, name="hey-snips"):
-    def __init__(self, num_processes=8):
-        self.stt = SpeechToText()
-        self.num_processes = num_processes
-        self.pool = Pool(processes=self.num_processes)
-
-    def load_splits(self, path: Path, **dataset_kwargs) -> Tuple[AudioClipDataset, AudioClipDataset, AudioClipDataset]:
-        self.path = path
-
-        def load(filename, set_type):
-            logging.info(f"Loading split {filename}...")
-
-            metadata_list = []
-            with open(str(self.path / filename)) as f:
-                raw_metadata_list = json.load(f)
-
-            total = len(raw_metadata_list)
-
-            transcription_fail_count = 0
-            pbar = tqdm(range(0, total, self.num_processes), desc=f"{filename}")
-            for starting_idx in pbar:
-                transcription_results = []
-                for metadata in raw_metadata_list[starting_idx : starting_idx + self.num_processes]:
-                    transcription_results.append(
-                        self.pool.apply_async(transcribe_hey_snips_audio, (self.path, metadata,))
-                    )
-
-                processing_count = starting_idx + len(transcription_results)
-
-                for result in transcription_results:
-                    pbar.set_postfix(dict(transcription_fail_rate=f"{transcription_fail_count}/{processing_count}"))
-                    path, transcription = result.get()
-
-                    if transcription == "":
-                        transcription_fail_count += 1
-                        continue
-
-                    metadata_list.append(AudioClipMetadata(path=path, transcription=transcription))
-
-            logging.info(f"{transcription_fail_count}/{total} samples have empty transcription")
-
-            return AudioClipDataset(metadata_list=metadata_list, set_type=set_type, **dataset_kwargs)
-
-        assert path.exists(), "dataset path doesn't exist"
-        filenames = ("train.json", "dev.json", "test.json")
-        assert all((path / x).exists() for x in filenames), "dataset missing metadata"
-
-        return (
-            load("train.json", DatasetType.TRAINING),
-            load("dev.json", DatasetType.DEV),
-            load("test.json", DatasetType.TEST),
-        )
-
-
-class GoogleSpeechCommandsDatasetLoader(RegisteredPathDatasetLoader, name="gsc"):
-    def __init__(self, vocab: List[str] = None, use_bg_noise: bool = False):
-        self.vocab = vocab
-        self.use_bg_noise = use_bg_noise
-
-    def load_splits(
-        self, path: Path, **dataset_kwargs
-    ) -> Tuple[AudioClassificationDataset, AudioClassificationDataset, AudioClassificationDataset]:
-        def load(set_type):
-            metadata_list = []
-            for path in all_list:
-                key = str(Path(path.parent.name) / path.name)
-                if file_map[key] != set_type:
-                    continue
-                metadata_list.append(AudioClipMetadata(path=path.absolute(), transcription=path.parent.name))
-            return AudioClassificationDataset(
-                metadata_list=metadata_list, label_map=label_map, set_type=set_type, **dataset_kwargs
-            )
-
-        file_map = defaultdict(lambda: DatasetType.TRAINING)
-        with (path / "testing_list.txt").open() as f:
-            file_map.update({k: DatasetType.TEST for k in f.read().split("\n")})
-        with (path / "validation_list.txt").open() as f:
-            file_map.update({k: DatasetType.DEV for k in f.read().split("\n")})
-        all_list = list(path.glob("*/*.wav"))
-        if not self.use_bg_noise:
-            all_list = list(filter(lambda x: "noise" not in str(x), all_list))
-        folders = sorted(list(path.glob("*/")))
-        vocab = [x.name for x in folders] if self.vocab is None else self.vocab
-        label_map = defaultdict(lambda: len(vocab))
-        label_map.update({k: idx for idx, k in enumerate(vocab)})
-        return load(DatasetType.TRAINING), load(DatasetType.DEV), load(DatasetType.TEST)
-
-
-class MozillaCommonVoiceLoader(RegisteredPathDatasetLoader, name="mozilla-cv"):
-    def load_splits(self, path: Path, **dataset_kwargs) -> Tuple[AudioClipDataset, AudioClipDataset, AudioClipDataset]:
-        def load(filename, set_type):
-            logging.info(f"Loading split {filename}...")
-            df = pd.read_csv(str(path / filename), sep="\t", quoting=3, na_filter=False)
-            metadata_list = []
-            for tup in df.itertuples():
-                metadata_list.append(
-                    AudioClipMetadata(path=(path / "clips" / tup.path).absolute(), transcription=tup.sentence)
-                )
-            return AudioClipDataset(metadata_list=metadata_list, set_type=set_type, **dataset_kwargs)
-
-        assert path.exists(), "dataset path doesn't exist"
-        filenames = ("train.tsv", "dev.tsv", "test.tsv")
-        assert all((path / x).exists() for x in filenames), "dataset missing metadata"
-        return (
-            load("train.tsv", DatasetType.TRAINING),
-            load("dev.tsv", DatasetType.DEV),
-            load("test.tsv", DatasetType.TEST),
-        )
-
-
-class MozillaKeywordLoader(RegisteredPathDatasetLoader, name="mozilla-kw"):
-    def load_splits(self, path: Path, **dataset_kwargs) -> Tuple[AudioClipDataset, AudioClipDataset, AudioClipDataset]:
-        logging.info("Loading Mozilla keyword dataset...")
-        df = pd.read_csv(str(path / "validated.tsv"), sep="\t", quoting=3, na_filter=False)
-        md_splits = ([], [], [])
-        for tup in df.itertuples():
-            md = AudioClipMetadata(path=(path / "clips" / tup.path).absolute(), transcription=tup.sentence)
-            bucket = sha256_int(tup.client_id) % 100
-            if bucket < 80:
-                md_splits[0].append(md)
-            elif bucket < 90:
-                md_splits[1].append(md)
-            else:
-                md_splits[2].append(md)
-        return (
-            AudioClipDataset(metadata_list=md_splits[0], set_type=DatasetType.TRAINING, **dataset_kwargs),
-            AudioClipDataset(metadata_list=md_splits[1], set_type=DatasetType.DEV, **dataset_kwargs),
-            AudioClipDataset(metadata_list=md_splits[2], set_type=DatasetType.TEST, **dataset_kwargs),
-        )
-
-
-class MozillaWakeWordLoader(RegisteredPathDatasetLoader, name="mozilla-ww"):
-    def __init__(self, training_pct=80, dev_pct=10, test_pct=10, split_by_speaker=True, split="verified"):
-        self.split_by_speaker = split_by_speaker
-        total = training_pct + dev_pct + test_pct
-        training_pct = 100 * training_pct / total
-        dev_pct = 100 * dev_pct / total
-        test_pct = 100 * test_pct / total
-        self.cutoffs = (training_pct, dev_pct + training_pct, training_pct + dev_pct + test_pct)
-        self.split = split
-
-    def load_splits(self, path: Path, **dataset_kwargs) -> Tuple[AudioClipDataset, AudioClipDataset, AudioClipDataset]:
-        assert path.exists(), "dataset path doesn't exist"
-        if self.split in {"verified", "rejected"}:
-            audio_path = path / self.split
-        else:
-            raise ValueError("Split ill-defined.")
-        assert audio_path.exists(), "dataset malformed"
-
-        sound_file_paths = audio_path.glob("*/*.ogg")
-        metadatas = ([], [], [])
-        logging.info("Loading wake word dataset...")
-        using_verified = self.split == "verified"
-        for sound_fp in sound_file_paths:
-            sound_id = sound_fp.stem
-            speaker_id = sound_fp.parent.name
-            if using_verified:
-                with open(str((sound_fp.parent / sound_id).with_suffix(".txt"))) as f:
-                    transcription = f.read()
-            else:
-                transcription = UNKNOWN_TRANSCRIPTION
-            metadata = AudioClipMetadata(path=sound_fp.absolute(), transcription=transcription)
-            bucket = sha256_int(speaker_id) if self.split_by_speaker else sha256_int(sound_id)
-            bucket %= 100
-            bucket = next(idx for idx, cutoff in enumerate(self.cutoffs) if bucket < cutoff)
-            metadatas[bucket].append(metadata)
-        return (
-            AudioClipDataset(metadata_list=metadatas[0], set_type=DatasetType.TRAINING, **dataset_kwargs),
-            AudioClipDataset(metadata_list=metadatas[1], set_type=DatasetType.DEV, **dataset_kwargs),
-            AudioClipDataset(metadata_list=metadatas[2], set_type=DatasetType.TEST, **dataset_kwargs),
-        )
-
-
-SoundIdSplitMozillaWakeWordLoader = partial(MozillaWakeWordLoader, split_by_speaker=False)
-SpeakerSplitMozillaWakeWordLoader = partial(MozillaWakeWordLoader, split_by_speaker=True)
-
-
-class RecursiveNoiseDatasetLoader:
-    def load(self, path: Path, **dataset_kwargs) -> AudioClipDataset:
-        wav_names = path.glob("**/*.wav")
-        metadata_list = [AudioClipMetadata(path=filename.absolute(), transcription="") for filename in wav_names]
-        return AudioClipDataset(metadata_list=metadata_list, set_type=DatasetType.TRAINING, **dataset_kwargs)
diff --git a/howl/data/stitcher.py b/howl/data/stitcher.py
index 1df0e431..0bc1823f 100644
--- a/howl/data/stitcher.py
+++ b/howl/data/stitcher.py
@@ -7,14 +7,10 @@
 import torch
 from tqdm import tqdm
 
-from howl.data.dataset import (
-    AudioClipDataset,
-    AudioClipExample,
-    AudioClipMetadata,
-    AudioDataset,
-    DatasetType,
-)
-from howl.data.tokenize import Vocab
+from howl.data.common.example import AudioClipExample
+from howl.data.common.metadata import AudioClipMetadata
+from howl.data.common.vocab import Vocab
+from howl.data.dataset.dataset import AudioClipDataset, AudioDataset, DatasetType
 from howl.settings import SETTINGS
 from howl.utils.sphinx_keyword_detector import SphinxKeywordDetector
 
diff --git a/howl/data/transform/__init__.py b/howl/data/transform/__init__.py
index bc8e3914..e69de29b 100644
--- a/howl/data/transform/__init__.py
+++ b/howl/data/transform/__init__.py
@@ -1,3 +0,0 @@
-from .base import *
-from .augment import *
-from .meyda import *
diff --git a/howl/data/transform/base.py b/howl/data/transform/base.py
deleted file mode 100644
index ad4e9e29..00000000
--- a/howl/data/transform/base.py
+++ /dev/null
@@ -1,245 +0,0 @@
-import random
-from typing import Iterable, List, Sequence
-
-import librosa.effects as effects
-import numpy as np
-import torch
-import torch.nn as nn
-from howl.data.dataset import (ClassificationBatch, EmplacableExample,
-                               SequenceBatch, WakeWordClipExample)
-from howl.data.tokenize import TranscriptTokenizer
-
-__all__ = ['Composition',
-           'compose',
-           'ZmuvTransform',
-           'random_slice',
-           'WakeWordFrameBatchifier',
-           'AudioSequenceBatchifier',
-           'batchify',
-           'identity',
-           'trim',
-           'truncate_length']
-
-
-class Composition(nn.Module):
-    def __init__(self, modules):
-        super().__init__()
-        self.modules = modules
-        self._module_list = nn.ModuleList(list(filter(lambda x: isinstance(x, nn.Module), modules)))
-
-    def forward(self, *args):
-        for mod in self.modules:
-            args = mod(*args)
-            args = (args,)
-        return args[0]
-
-
-class IdentityTransform(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, x):
-        return x
-
-
-def compose(*collate_modules):
-    return Composition(collate_modules)
-
-
-def trim(examples: Sequence[EmplacableExample], top_db: int = 40):
-    return [ex.emplaced_audio_data(torch.from_numpy(effects.trim(ex.audio_data.cpu().numpy(),
-                                                                 top_db=top_db)[0])) for ex in examples]
-
-
-def random_slice(examples: Sequence[WakeWordClipExample],
-                 max_window_size: int = 16000) -> Sequence[WakeWordClipExample]:
-    new_examples = []
-    for ex in examples:
-        if ex.audio_data.size(-1) < max_window_size:
-            new_examples.append(ex)
-            continue
-        a = random.randint(0, ex.audio_data.size(-1) - max_window_size)
-        new_examples.append(ex.emplaced_audio_data(ex.audio_data[..., a:a + max_window_size]))
-    return new_examples
-
-
-def truncate_length(examples: Sequence[EmplacableExample], length: int = None):
-    return [ex.emplaced_audio_data(ex.audio_data[..., :length]) for ex in examples]
-
-
-def batchify(examples: Sequence[EmplacableExample], label_provider=None):
-    examples = sorted(examples, key=lambda x: x.audio_data.size()[-1], reverse=True)
-    lengths = torch.tensor([ex.audio_data.size(-1) for ex in examples])
-    max_length = max(ex.audio_data.size(-1) for ex in examples)
-    audio_tensor = [torch.cat((ex.audio_data.squeeze(), torch.zeros(max_length - ex.audio_data.size(-1))), -1) for
-                    ex in examples]
-    audio_tensor = torch.stack(audio_tensor)
-    labels = torch.tensor(list(map(label_provider, examples))) if label_provider else None
-    return ClassificationBatch(audio_tensor, labels, lengths)
-
-
-def tensorize_audio_data(audio_data_lst: List[torch.Tensor],
-                         max_length: int = None,
-                         rand_append: bool = False,
-                         **extra_data_lists):
-    lengths = np.array([audio_data.size(-1) for audio_data in audio_data_lst])
-    sort_indices = np.argsort(-lengths)
-    audio_data_lst = np.array(audio_data_lst, dtype=object)[sort_indices].tolist()
-    extra_data_lists = {k: np.array(v, dtype=object)[sort_indices].tolist() for k, v in extra_data_lists.items()}
-
-    if max_length is None:
-        max_length = max(audio_data.size(-1) for audio_data in audio_data_lst)
-    audio_tensor = []
-    for audio_data in audio_data_lst:
-        squeezed_data = audio_data.squeeze()
-        if squeezed_data.dim() == 0:
-            squeezed_data = squeezed_data.unsqueeze(0)
-        if rand_append and random.random() < 0.5:
-            x = (torch.zeros(max_length - audio_data.size(-1)), squeezed_data)
-        else:
-            x = (squeezed_data, torch.zeros(max_length - audio_data.size(-1)))
-        audio_tensor.append(torch.cat(x, -1))
-    return torch.stack(audio_tensor), extra_data_lists
-
-
-def pad(data_list, element=0, max_length=None):
-    if max_length is None:
-        max_length = max(map(len, data_list))
-    data_list = [x + [element] * (max_length - len(x)) for x in data_list]
-    return data_list
-
-
-class AudioSequenceBatchifier:
-    def __init__(self,
-                 negative_label: int,
-                 tokenizer: TranscriptTokenizer,
-                 sample_rate: int = 16000):
-        self.sample_rate = sample_rate
-        self.tokenizer = tokenizer
-        self.negative_label = negative_label
-
-    def __call__(self, examples: Sequence[WakeWordClipExample]) -> SequenceBatch:
-        audio_data_lst = []
-        labels_lst = []
-        for ex in examples:
-            labels = self.tokenizer.encode(ex.metadata.transcription)
-            labels_lst.append(labels)
-            audio_data_lst.append(ex.audio_data)
-        audio_data_lengths = [audio_data.size(-1) for audio_data in audio_data_lst]
-        labels_lengths = list(map(len, labels_lst))
-        audio_tensor, data = tensorize_audio_data(audio_data_lst,
-                                                  labels_lst=labels_lst,
-                                                  labels_lengths=labels_lengths,
-                                                  input_lengths=audio_data_lengths)
-        labels_lst = torch.tensor(pad(data['labels_lst'], element=self.negative_label))
-        labels_lengths = torch.tensor(data['labels_lengths'])
-        return SequenceBatch(audio_tensor, labels_lst, torch.tensor(data['input_lengths']), labels_lengths)
-
-
-class WakeWordFrameBatchifier:
-    def __init__(self,
-                 negative_label: int,
-                 positive_sample_prob: float = 0.5,
-                 window_size_ms: int = 500,
-                 sample_rate: int = 16000,
-                 positive_delta_ms: int = 150,
-                 eps_ms: int = 20,
-                 pad_to_window: bool = True):
-        self.positive_sample_prob = positive_sample_prob
-        self.window_size_ms = window_size_ms
-        self.sample_rate = sample_rate
-        self.negative_label = negative_label
-        self.positive_delta_ms = positive_delta_ms
-        self.eps_ms = eps_ms
-        self.pad_to_window = pad_to_window
-
-    def __call__(self, examples: Sequence[WakeWordClipExample]) -> ClassificationBatch:
-        new_examples = []
-        for ex in examples:
-            # if the sample does not contain any positive label
-            if not ex.label_data.timestamp_label_map:
-                new_examples.append((self.negative_label,
-                                     random_slice([ex], int(self.sample_rate * self.window_size_ms / 1000))[0]))
-                continue
-
-            select_negative = random.random() > self.positive_sample_prob
-
-            # pick a random positive word to genrate a positive sample
-            if not select_negative:
-                end_ms, label = random.choice(list(ex.label_data.timestamp_label_map.items()))
-                end_ms_rand = end_ms + (random.random() * self.eps_ms)
-                b = int((end_ms_rand / 1000) * self.sample_rate)
-                a = max(b - int((self.window_size_ms / 1000) * self.sample_rate), 0)
-                if random.random() < 0:
-                    closest_ms = min(filter(lambda k: end_ms - k > 0, ex.label_data.timestamp_label_map.keys()),
-                                     key=lambda k: end_ms - k,
-                                     default=-1)
-                    if closest_ms >= 0:
-                        a = max(a, int((closest_ms / 1000) * self.sample_rate))
-                if b - a < 0:
-                    select_negative = True
-                else:
-                    new_examples.append((label, ex.emplaced_audio_data(ex.audio_data[..., a:b])))
-
-            # use the interval of negative labels to generate a negative sample from audio with positive label
-            if select_negative:
-                positive_intervals = [(v - self.positive_delta_ms, v + self.positive_delta_ms)
-                                      for v in ex.label_data.timestamp_label_map.values()]
-                positive_intervals = sorted(positive_intervals, key=lambda x: x[0])
-                negative_intervals = []
-                last_positive = 0
-                for a, b in positive_intervals:
-                    if last_positive < a:
-                        negative_intervals.append((last_positive, a))
-                    last_positive = b
-                negative_intervals.append((b, int(len(ex.audio_data) / 16000 * 1000)))
-                a, b = random.choice(negative_intervals)
-                if b - a > self.window_size_ms:
-                    a = random.randint(0, int(b - self.window_size_ms))
-                    b = a + self.window_size_ms
-                new_examples.append((self.negative_label, ex.emplaced_audio_data(ex.audio_data[..., a:b])))
-
-        labels_lst = [x[0] for x in new_examples]
-        max_length = int(self.window_size_ms / 1000 * self.sample_rate) if self.pad_to_window else None
-        audio_tensor, extra_data = tensorize_audio_data([x[1].audio_data for x in new_examples],
-                                                        rand_append=True,
-                                                        max_length=max_length,
-                                                        labels_lst=labels_lst,
-                                                        lengths=[x[1].audio_data.size(-1) for x in new_examples])
-        labels_tensor = torch.tensor(extra_data['labels_lst'])
-        lengths = torch.tensor(extra_data['lengths'])
-        return ClassificationBatch(audio_tensor, labels_tensor, lengths)
-
-
-def identity(x):
-    return x
-
-
-class ZmuvTransform(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.register_buffer('total', torch.zeros(1))
-        self.register_buffer('mean', torch.zeros(1))
-        self.register_buffer('mean2', torch.zeros(1))
-
-    def update(self, data, mask=None):
-        with torch.no_grad():
-            if mask is not None:
-                data = data * mask
-                mask_size = mask.sum().item()
-            else:
-                mask_size = data.numel()
-            self.mean = (data.sum() + self.mean * self.total) / (self.total + mask_size)
-            self.mean2 = ((data ** 2).sum() + self.mean2 * self.total) / (self.total + mask_size)
-            self.total += mask_size
-
-    def initialize(self, iterable: Iterable[torch.Tensor]):
-        for ex in iterable:
-            self.update(ex)
-
-    @property
-    def std(self):
-        return (self.mean2 - self.mean ** 2).sqrt()
-
-    def forward(self, x):
-        return (x - self.mean) / self.std
diff --git a/howl/data/transform/batchifier.py b/howl/data/transform/batchifier.py
new file mode 100644
index 00000000..07d69eb4
--- /dev/null
+++ b/howl/data/transform/batchifier.py
@@ -0,0 +1,118 @@
+import random
+from typing import Sequence
+
+import torch
+
+from howl.data.common.batch import ClassificationBatch, SequenceBatch
+from howl.data.common.example import WakeWordClipExample
+from howl.data.common.tokenizer import TranscriptTokenizer
+from howl.data.transform.operator import pad, random_slice, tensorize_audio_data
+
+__all__ = ["WakeWordFrameBatchifier", "AudioSequenceBatchifier"]
+
+
+class AudioSequenceBatchifier:
+    def __init__(self, negative_label: int, tokenizer: TranscriptTokenizer, sample_rate: int = 16000):
+        self.sample_rate = sample_rate
+        self.tokenizer = tokenizer
+        self.negative_label = negative_label
+
+    def __call__(self, examples: Sequence[WakeWordClipExample]) -> SequenceBatch:
+        audio_data_lst = []
+        labels_lst = []
+        for ex in examples:
+            labels = self.tokenizer.encode(ex.metadata.transcription)
+            labels_lst.append(labels)
+            audio_data_lst.append(ex.audio_data)
+        audio_data_lengths = [audio_data.size(-1) for audio_data in audio_data_lst]
+        labels_lengths = list(map(len, labels_lst))
+        audio_tensor, data = tensorize_audio_data(
+            audio_data_lst, labels_lst=labels_lst, labels_lengths=labels_lengths, input_lengths=audio_data_lengths
+        )
+        labels_lst = torch.tensor(pad(data["labels_lst"], element=self.negative_label))
+        labels_lengths = torch.tensor(data["labels_lengths"])
+        return SequenceBatch(audio_tensor, labels_lst, torch.tensor(data["input_lengths"]), labels_lengths)
+
+
+class WakeWordFrameBatchifier:
+    def __init__(
+        self,
+        negative_label: int,
+        positive_sample_prob: float = 0.5,
+        window_size_ms: int = 500,
+        sample_rate: int = 16000,
+        positive_delta_ms: int = 150,
+        eps_ms: int = 20,
+        pad_to_window: bool = True,
+    ):
+        self.positive_sample_prob = positive_sample_prob
+        self.window_size_ms = window_size_ms
+        self.sample_rate = sample_rate
+        self.negative_label = negative_label
+        self.positive_delta_ms = positive_delta_ms
+        self.eps_ms = eps_ms
+        self.pad_to_window = pad_to_window
+
+    def __call__(self, examples: Sequence[WakeWordClipExample]) -> ClassificationBatch:
+        new_examples = []
+        for ex in examples:
+            # if the sample does not contain any positive label
+            if not ex.label_data.timestamp_label_map:
+                new_examples.append(
+                    (self.negative_label, random_slice([ex], int(self.sample_rate * self.window_size_ms / 1000))[0])
+                )
+                continue
+
+            select_negative = random.random() > self.positive_sample_prob
+
+            # pick a random positive word to genrate a positive sample
+            if not select_negative:
+                end_ms, label = random.choice(list(ex.label_data.timestamp_label_map.items()))
+                end_ms_rand = end_ms + (random.random() * self.eps_ms)
+                b = int((end_ms_rand / 1000) * self.sample_rate)
+                a = max(b - int((self.window_size_ms / 1000) * self.sample_rate), 0)
+                if random.random() < 0:
+                    closest_ms = min(
+                        filter(lambda k: end_ms - k > 0, ex.label_data.timestamp_label_map.keys()),
+                        key=lambda k: end_ms - k,
+                        default=-1,
+                    )
+                    if closest_ms >= 0:
+                        a = max(a, int((closest_ms / 1000) * self.sample_rate))
+                if b - a < 0:
+                    select_negative = True
+                else:
+                    new_examples.append((label, ex.emplaced_audio_data(ex.audio_data[..., a:b])))
+
+            # use the interval of negative labels to generate a negative sample from audio with positive label
+            if select_negative:
+                positive_intervals = [
+                    (v - self.positive_delta_ms, v + self.positive_delta_ms)
+                    for v in ex.label_data.timestamp_label_map.values()
+                ]
+                positive_intervals = sorted(positive_intervals, key=lambda x: x[0])
+                negative_intervals = []
+                last_positive = 0
+                for a, b in positive_intervals:
+                    if last_positive < a:
+                        negative_intervals.append((last_positive, a))
+                    last_positive = b
+                negative_intervals.append((b, int(len(ex.audio_data) / 16000 * 1000)))
+                a, b = random.choice(negative_intervals)
+                if b - a > self.window_size_ms:
+                    a = random.randint(0, int(b - self.window_size_ms))
+                    b = a + self.window_size_ms
+                new_examples.append((self.negative_label, ex.emplaced_audio_data(ex.audio_data[..., a:b])))
+
+        labels_lst = [x[0] for x in new_examples]
+        max_length = int(self.window_size_ms / 1000 * self.sample_rate) if self.pad_to_window else None
+        audio_tensor, extra_data = tensorize_audio_data(
+            [x[1].audio_data for x in new_examples],
+            rand_append=True,
+            max_length=max_length,
+            labels_lst=labels_lst,
+            lengths=[x[1].audio_data.size(-1) for x in new_examples],
+        )
+        labels_tensor = torch.tensor(extra_data["labels_lst"])
+        lengths = torch.tensor(extra_data["lengths"])
+        return ClassificationBatch(audio_tensor, labels_tensor, lengths)
diff --git a/howl/data/transform/operator.py b/howl/data/transform/operator.py
new file mode 100644
index 00000000..ba700ca8
--- /dev/null
+++ b/howl/data/transform/operator.py
@@ -0,0 +1,146 @@
+import random
+from typing import Iterable, List, Sequence
+
+import librosa.effects as effects
+import numpy as np
+import torch
+import torch.nn as nn
+
+from howl.data.common.batch import ClassificationBatch
+from howl.data.common.example import EmplacableExample, WakeWordClipExample
+
+__all__ = [
+    "Composition",
+    "compose",
+    "ZmuvTransform",
+    "random_slice",
+    "batchify",
+    "identity",
+    "trim",
+    "truncate_length",
+]
+
+
+class Composition(nn.Module):
+    def __init__(self, modules):
+        super().__init__()
+        self.modules = modules
+        self._module_list = nn.ModuleList(list(filter(lambda x: isinstance(x, nn.Module), modules)))
+
+    def forward(self, *args):
+        for mod in self.modules:
+            args = mod(*args)
+            args = (args,)
+        return args[0]
+
+
+def compose(*collate_modules):
+    return Composition(collate_modules)
+
+
+class IdentityTransform(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x
+
+
+def identity(x):
+    return x
+
+
+def trim(examples: Sequence[EmplacableExample], top_db: int = 40):
+    return [
+        ex.emplaced_audio_data(torch.from_numpy(effects.trim(ex.audio_data.cpu().numpy(), top_db=top_db)[0]))
+        for ex in examples
+    ]
+
+
+def random_slice(
+    examples: Sequence[WakeWordClipExample], max_window_size: int = 16000
+) -> Sequence[WakeWordClipExample]:
+    new_examples = []
+    for ex in examples:
+        if ex.audio_data.size(-1) < max_window_size:
+            new_examples.append(ex)
+            continue
+        a = random.randint(0, ex.audio_data.size(-1) - max_window_size)
+        new_examples.append(ex.emplaced_audio_data(ex.audio_data[..., a : a + max_window_size]))
+    return new_examples
+
+
+def truncate_length(examples: Sequence[EmplacableExample], length: int = None):
+    return [ex.emplaced_audio_data(ex.audio_data[..., :length]) for ex in examples]
+
+
+def batchify(examples: Sequence[EmplacableExample], label_provider=None):
+    examples = sorted(examples, key=lambda x: x.audio_data.size()[-1], reverse=True)
+    lengths = torch.tensor([ex.audio_data.size(-1) for ex in examples])
+    max_length = max(ex.audio_data.size(-1) for ex in examples)
+    audio_tensor = [
+        torch.cat((ex.audio_data.squeeze(), torch.zeros(max_length - ex.audio_data.size(-1))), -1) for ex in examples
+    ]
+    audio_tensor = torch.stack(audio_tensor)
+    labels = torch.tensor(list(map(label_provider, examples))) if label_provider else None
+    return ClassificationBatch(audio_tensor, labels, lengths)
+
+
+def tensorize_audio_data(
+    audio_data_lst: List[torch.Tensor], max_length: int = None, rand_append: bool = False, **extra_data_lists
+):
+    lengths = np.array([audio_data.size(-1) for audio_data in audio_data_lst])
+    sort_indices = np.argsort(-lengths)
+    audio_data_lst = np.array(audio_data_lst, dtype=object)[sort_indices].tolist()
+    extra_data_lists = {k: np.array(v, dtype=object)[sort_indices].tolist() for k, v in extra_data_lists.items()}
+
+    if max_length is None:
+        max_length = max(audio_data.size(-1) for audio_data in audio_data_lst)
+    audio_tensor = []
+    for audio_data in audio_data_lst:
+        squeezed_data = audio_data.squeeze()
+        if squeezed_data.dim() == 0:
+            squeezed_data = squeezed_data.unsqueeze(0)
+        if rand_append and random.random() < 0.5:
+            x = (torch.zeros(max_length - audio_data.size(-1)), squeezed_data)
+        else:
+            x = (squeezed_data, torch.zeros(max_length - audio_data.size(-1)))
+        audio_tensor.append(torch.cat(x, -1))
+    return torch.stack(audio_tensor), extra_data_lists
+
+
+def pad(data_list, element=0, max_length=None):
+    if max_length is None:
+        max_length = max(map(len, data_list))
+    data_list = [x + [element] * (max_length - len(x)) for x in data_list]
+    return data_list
+
+
+class ZmuvTransform(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("total", torch.zeros(1))
+        self.register_buffer("mean", torch.zeros(1))
+        self.register_buffer("mean2", torch.zeros(1))
+
+    def update(self, data, mask=None):
+        with torch.no_grad():
+            if mask is not None:
+                data = data * mask
+                mask_size = mask.sum().item()
+            else:
+                mask_size = data.numel()
+            self.mean = (data.sum() + self.mean * self.total) / (self.total + mask_size)
+            self.mean2 = ((data ** 2).sum() + self.mean2 * self.total) / (self.total + mask_size)
+            self.total += mask_size
+
+    def initialize(self, iterable: Iterable[torch.Tensor]):
+        for ex in iterable:
+            self.update(ex)
+
+    @property
+    def std(self):
+        return (self.mean2 - self.mean ** 2).sqrt()
+
+    def forward(self, x):
+        return (x - self.mean) / self.std
diff --git a/howl/data/transform/augment.py b/howl/data/transform/transform.py
similarity index 72%
rename from howl/data/transform/augment.py
rename to howl/data/transform/transform.py
index 0690bf01..1b20bfff 100644
--- a/howl/data/transform/augment.py
+++ b/howl/data/transform/transform.py
@@ -1,27 +1,29 @@
-from dataclasses import dataclass
-from typing import Sequence
 import math
 import random
+from dataclasses import dataclass
+from typing import Sequence
 
-from torchaudio.transforms import MelSpectrogram, ComputeDeltas
 import librosa
 import torch
 import torch.nn as nn
+from torchaudio.transforms import ComputeDeltas, MelSpectrogram
 
-from howl.data.dataset import EmplacableExample, WakeWordClipExample, AudioClipDataset
+from howl.data.common.example import EmplacableExample, WakeWordClipExample
+from howl.data.dataset.dataset import AudioClipDataset
+from howl.data.transform.meyda import MeydaMelSpectrogram
 from howl.settings import SETTINGS
-from .meyda import MeydaMelSpectrogram
 
-
-__all__ = ['AugmentationParameter',
-           'AugmentModule',
-           'TimeshiftTransform',
-           'TimestretchTransform',
-           'NoiseTransform',
-           'DatasetMixer',
-           'StandardAudioTransform',
-           'SpecAugmentTransform',
-           'NegativeSampleTransform']
+__all__ = [
+    "AugmentationParameter",
+    "AugmentModule",
+    "TimeshiftTransform",
+    "TimestretchTransform",
+    "NoiseTransform",
+    "DatasetMixer",
+    "StandardAudioTransform",
+    "SpecAugmentTransform",
+    "NegativeSampleTransform",
+]
 
 
 @dataclass
@@ -32,7 +34,7 @@ class AugmentationParameter:
     prob: float = 0.75
     enabled: bool = True
 
-    def copy_from(self, op: 'AugmentationParameter'):
+    def copy_from(self, op: "AugmentationParameter"):
         self.current_value_idx = op.current_value_idx
         self.prob = op.prob
         self.enabled = op.enabled
@@ -43,7 +45,7 @@ def magnitude(self):
 
     @classmethod
     def from_dict(cls, data_dict):
-        return cls(data_dict['domain'], data_dict['name'], data_dict['current_value_idx'], data_dict['prob'])
+        return cls(data_dict["domain"], data_dict["name"], data_dict["current_value_idx"], data_dict["prob"])
 
 
 class AugmentModule(nn.Module):
@@ -79,13 +81,13 @@ def forward(self, x, **kwargs):
 class NegativeSampleTransform(AugmentModule):
     @property
     def default_params(self):
-        return AugmentationParameter([0.2, 0.3, 0.4, 0.5], 'chunk_size', 1, prob=0.3),
+        return (AugmentationParameter([0.2, 0.3, 0.4, 0.5], "chunk_size", 1, prob=0.3),)
 
     @torch.no_grad()
     def augment(self, param: AugmentationParameter, examples: Sequence[WakeWordClipExample], **kwargs):
         new_examples = []
         for example in examples:
-            audio_data = example.audio_data[..., :int(example.audio_data.size(-1) * param.magnitude)]
+            audio_data = example.audio_data[..., : int(example.audio_data.size(-1) * param.magnitude)]
             example = example.emplaced_audio_data(audio_data)
             example.contains_wake_word = False
             new_examples.append(example)
@@ -99,7 +101,7 @@ def __init__(self, sr=16000):
 
     @property
     def default_params(self):
-        return AugmentationParameter([0.25, 0.5, 0.75, 1], 'timeshift', 0),
+        return (AugmentationParameter([0.25, 0.5, 0.75, 1], "timeshift", 0),)
 
     @torch.no_grad()
     def augment(self, param: AugmentationParameter, examples: Sequence[EmplacableExample], **kwargs):
@@ -107,7 +109,9 @@ def augment(self, param: AugmentationParameter, examples: Sequence[EmplacableExa
         for example in examples:
             w = min(int(self.rand.random() * param.magnitude * self.sr), int(0.5 * example.audio_data.size(-1)))
             audio_data = example.audio_data
-            audio_data = audio_data[..., w:] if self.rand.random() < 0.5 else audio_data[..., :example.audio_data.size(-1) - w]
+            audio_data = (
+                audio_data[..., w:] if self.rand.random() < 0.5 else audio_data[..., : example.audio_data.size(-1) - w]
+            )
             new_examples.append(example.emplaced_audio_data(audio_data))
         return new_examples
 
@@ -115,7 +119,7 @@ def augment(self, param: AugmentationParameter, examples: Sequence[EmplacableExa
 class TimestretchTransform(AugmentModule):
     @property
     def default_params(self):
-        return AugmentationParameter([0.025, 0.05, 0.15, 0.2, 0.25], 'timestretch', 2, prob=0.3),
+        return (AugmentationParameter([0.025, 0.05, 0.15, 0.2, 0.25], "timestretch", 2, prob=0.3),)
 
     @torch.no_grad()
     def augment(self, param, examples: Sequence[EmplacableExample], **kwargs):
@@ -130,20 +134,24 @@ def augment(self, param, examples: Sequence[EmplacableExample], **kwargs):
 class NoiseTransform(AugmentModule):
     @property
     def default_params(self):
-        return AugmentationParameter([0.0001, 0.00025, 0.0005, 0.001, 0.002], 'white', 3),\
-               AugmentationParameter([1 / 20000, 1 / 15000, 1 / 10000, 1 / 5000, 1 / 2500], 'salt_pepper', 2)
+        return (
+            AugmentationParameter([0.0001, 0.00025, 0.0005, 0.001, 0.002], "white", 3),
+            AugmentationParameter([1 / 20000, 1 / 15000, 1 / 10000, 1 / 5000, 1 / 2500], "salt_pepper", 2),
+        )
 
     @torch.no_grad()
     def augment(self, param, examples: Sequence[EmplacableExample], **kwargs):
         new_examples = []
         for example in examples:
             waveform = example.audio_data
-            if param.name == 'white':
+            if param.name == "white":
                 strength = param.magnitude * self.rand.random()
                 noise_mask = torch.empty_like(waveform).normal_(0, strength)
             else:
                 prob = param.magnitude * self.rand.random()
-                noise_mask = torch.empty_like(waveform).bernoulli_(prob / 2) - torch.empty_like(waveform).bernoulli_(prob / 2)
+                noise_mask = torch.empty_like(waveform).bernoulli_(prob / 2) - torch.empty_like(waveform).bernoulli_(
+                    prob / 2
+                )
             noise_mask.clamp_(-1, 1)
             waveform = (waveform + noise_mask).clamp_(-1, 1)
             new_examples.append(example.emplaced_audio_data(waveform))
@@ -151,18 +159,17 @@ def augment(self, param, examples: Sequence[EmplacableExample], **kwargs):
 
 
 class DatasetMixer(AugmentModule):
-    def __init__(self,
-                 background_noise_dataset: AudioClipDataset,
-                 do_replace: bool = False,
-                 **kwargs):
+    def __init__(self, background_noise_dataset: AudioClipDataset, do_replace: bool = False, **kwargs):
         self.do_replace = do_replace
         super().__init__(**kwargs)
         self.dataset = background_noise_dataset
 
     @property
     def default_params(self):
-        return (AugmentationParameter([0.1, 0.2, 0.3, 0.4, 0.5], 'strength', 1),
-                AugmentationParameter([0], 'replace', 0, prob=0.1 if self.do_replace else 0))
+        return (
+            AugmentationParameter([0.1, 0.2, 0.3, 0.4, 0.5], "strength", 1),
+            AugmentationParameter([0], "replace", 0, prob=0.1 if self.do_replace else 0),
+        )
 
     @torch.no_grad()
     def augment(self, param, examples: Sequence[EmplacableExample], **kwargs):
@@ -175,7 +182,7 @@ def augment(self, param, examples: Sequence[EmplacableExample], **kwargs):
             b = self.rand.randint(waveform.size(-1), bg_ex.size(-1))
             a = b - waveform.size(-1)
             bg_audio = bg_ex[..., a:b]
-            alpha = 1 if param.name == 'replace' else self.rand.random() * param.magnitude
+            alpha = 1 if param.name == "replace" else self.rand.random() * param.magnitude
             mixed_wf = waveform * (1 - alpha) + bg_audio * alpha
             ex = example.emplaced_audio_data(mixed_wf, new=alpha == 1)
             new_examples.append(ex)
@@ -187,25 +194,33 @@ def __init__(self):
         super().__init__()
         settings = SETTINGS.audio_transform
         if settings.use_meyda_spectrogram:
-            self.spec_transform = MeydaMelSpectrogram(n_mels=settings.num_mels,
-                                                      sample_rate=settings.sample_rate,
-                                                      n_fft=settings.num_fft,
-                                                      hop_length=settings.hop_length)
+            self.spec_transform = MeydaMelSpectrogram(
+                n_mels=settings.num_mels,
+                sample_rate=settings.sample_rate,
+                n_fft=settings.num_fft,
+                hop_length=settings.hop_length,
+            )
         else:
-            self.spec_transform = MelSpectrogram(n_mels=settings.num_mels,
-                                                 sample_rate=settings.sample_rate,
-                                                 n_fft=settings.num_fft,
-                                                 hop_length=settings.hop_length)
-
-        self.vtlp_transform = apply_vtlp(MelSpectrogram(n_mels=settings.num_mels,
-                                                        sample_rate=settings.sample_rate,
-                                                        n_fft=settings.num_fft,
-                                                        hop_length=settings.hop_length))
+            self.spec_transform = MelSpectrogram(
+                n_mels=settings.num_mels,
+                sample_rate=settings.sample_rate,
+                n_fft=settings.num_fft,
+                hop_length=settings.hop_length,
+            )
+
+        self.vtlp_transform = apply_vtlp(
+            MelSpectrogram(
+                n_mels=settings.num_mels,
+                sample_rate=settings.sample_rate,
+                n_fft=settings.num_fft,
+                hop_length=settings.hop_length,
+            )
+        )
         self.delta_transform = ComputeDeltas()
 
     @property
     def default_params(self) -> Sequence[AugmentationParameter]:
-        return AugmentationParameter([0], 'vtlp', 0),
+        return (AugmentationParameter([0], "vtlp", 0),)
 
     @torch.no_grad()
     def _execute_op(self, op, audio, mels_only=False, deltas_only=False):
@@ -231,8 +246,10 @@ def compute_lengths(self, length: torch.Tensor):
 class SpecAugmentTransform(AugmentModule):
     @property
     def default_params(self) -> Sequence[AugmentationParameter]:
-        return AugmentationParameter([2, 5, 10, 20, 25], 'sa_freq', 2),\
-               AugmentationParameter([10, 50, 75, 125, 150], 'sa_time', 2)
+        return (
+            AugmentationParameter([2, 5, 10, 20, 25], "sa_freq", 2),
+            AugmentationParameter([10, 50, 75, 125, 150], "sa_time", 2),
+        )
 
     def tmask(self, x, T):
         for idx in range(x.size(0)):
@@ -241,22 +258,22 @@ def tmask(self, x, T):
                 t0 = self.rand.randrange(0, x.size(3) - t)
             except ValueError:
                 continue
-            x[idx, :, :, t0:t0 + t] = 0
+            x[idx, :, :, t0 : t0 + t] = 0
         return x
 
     def fmask(self, x, F):
         for idx in range(x.size(0)):
             f = self.rand.randrange(0, F)
             f0 = self.rand.randrange(0, x.size(2) - f)
-            x[idx, :, f0:f0 + f] = 0
+            x[idx, :, f0 : f0 + f] = 0
         return x
 
     @torch.no_grad()
     def augment(self, param, x, **kwargs):
         with torch.no_grad():
-            if param.name == 'sa_freq':
+            if param.name == "sa_freq":
                 return self.fmask(x, param.magnitude)
-            elif param.name == 'sa_time':
+            elif param.name == "sa_time":
                 return self.tmask(x, param.magnitude)
         return x
 
@@ -266,7 +283,7 @@ def augment(self, param, x, **kwargs):
 """
 BSD 2-Clause License
 
-Copyright (c) 2017 Facebook Inc. (Soumith Chintala), 
+Copyright (c) 2017 Facebook Inc. (Soumith Chintala),
 All rights reserved.
 
 Redistribution and use in source and binary forms, with or without
@@ -291,6 +308,7 @@ def augment(self, param, x, **kwargs):
 OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 """
 
+
 def create_vtlp_fb_matrix(n_freqs, f_min, f_max, n_mels, sample_rate, alpha, f_hi=4800, training=True):
     # type: (int, float, float, int, int, float, int, bool) -> torch.Tensor
     # freq bins
@@ -308,8 +326,9 @@ def create_vtlp_fb_matrix(n_freqs, f_min, f_max, n_mels, sample_rate, alpha, f_h
     if training:
         f_pts[f_pts <= f_hi * min(alpha, 1) / alpha] *= alpha
         f = f_pts[f_pts > f_hi * min(alpha, 1) / alpha]
-        f_pts[f_pts > f_hi * min(alpha, 1) / alpha] = S / 2 - ((S / 2 - f_hi * min(alpha, 1)) /
-                                                               (S / 2 - f_hi * min(alpha, 1) / alpha)) * (S / 2 - f)
+        f_pts[f_pts > f_hi * min(alpha, 1) / alpha] = S / 2 - (
+            (S / 2 - f_hi * min(alpha, 1)) / (S / 2 - f_hi * min(alpha, 1) / alpha)
+        ) * (S / 2 - f)
     # calculate the difference between each mel point and each stft freq point in hertz
     f_diff = f_pts[1:] - f_pts[:-1]  # (n_mels + 1)
     slopes = f_pts.unsqueeze(0) - all_freqs.unsqueeze(1)  # (n_freqs, n_mels + 2)
@@ -323,24 +342,31 @@ def create_vtlp_fb_matrix(n_freqs, f_min, f_max, n_mels, sample_rate, alpha, f_h
 
 class VtlpMelScale(nn.Module):
 
-    __constants__ = ['n_mels', 'sample_rate', 'f_min', 'f_max']
+    __constants__ = ["n_mels", "sample_rate", "f_min", "f_max"]
 
-    def __init__(self, n_mels=128, sample_rate=16000, f_min=0., f_max=None, n_stft=None):
+    def __init__(self, n_mels=128, sample_rate=16000, f_min=0.0, f_max=None, n_stft=None):
         super().__init__()
         self.n_mels = n_mels
         self.sample_rate = sample_rate
         self.f_max = f_max if f_max is not None else float(sample_rate // 2)
         self.f_min = f_min
 
-        assert f_min <= self.f_max, 'Require f_min: %f < f_max: %f' % (f_min, self.f_max)
+        assert f_min <= self.f_max, "Require f_min: %f < f_max: %f" % (f_min, self.f_max)
 
     def forward(self, specgram):
         # pack batch
         shape = specgram.size()
         specgram = specgram.reshape(-1, shape[-2], shape[-1])
 
-        fb = create_vtlp_fb_matrix(specgram.size(1), self.f_min, self.f_max, self.n_mels, self.sample_rate,
-                                   random.random() * 0.2 + 0.9, training=self.training).to(specgram.device)
+        fb = create_vtlp_fb_matrix(
+            specgram.size(1),
+            self.f_min,
+            self.f_max,
+            self.n_mels,
+            self.sample_rate,
+            random.random() * 0.2 + 0.9,
+            training=self.training,
+        ).to(specgram.device)
         # (channel, frequency, time).transpose(...) dot (frequency, n_mels)
         # -> (channel, time, n_mels).transpose(...)
         mel_specgram = torch.matmul(specgram.transpose(1, 2), fb).transpose(1, 2)
@@ -348,10 +374,11 @@ def forward(self, specgram):
         mel_specgram = mel_specgram.reshape(shape[:-2] + mel_specgram.shape[-2:])
         return mel_specgram
 
+
 # END LICENSED BLOCK ###########
 
 
 def apply_vtlp(mel_spectrogram: MelSpectrogram):
     s = mel_spectrogram
     s.mel_scale = VtlpMelScale(s.n_mels, s.sample_rate, s.f_min, s.f_max, s.n_fft // 2 + 1)
-    return s
\ No newline at end of file
+    return s
diff --git a/howl/dataset/__init__.py b/howl/dataset/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/howl/logger.py b/howl/logger.py
deleted file mode 100644
index 64a6d4c7..00000000
--- a/howl/logger.py
+++ /dev/null
@@ -1,10 +0,0 @@
-import os
-
-import coloredlogs
-
-
-__all__ = []
-
-
-coloredlogs.install(level=os.environ.get('HOWL_LOG_LEVEL', 'INFO'),
-                    fmt='%(asctime)s [%(levelname)s] %(module)s: %(message)s')
diff --git a/howl/model/base.py b/howl/model/base.py
index 3fde9ea0..8ae681e4 100644
--- a/howl/model/base.py
+++ b/howl/model/base.py
@@ -3,13 +3,12 @@
 import torch
 import torch.nn as nn
 
-from howl.registered import RegisteredObjectBase
+from howl.utils.class_registry import ClassRegistry
 
+__all__ = ["RegisteredModel", "ConvertedStaticModel"]
 
-__all__ = ['RegisteredModel', 'ConvertedStaticModel']
 
-
-class RegisteredModel(nn.Module, RegisteredObjectBase):
+class RegisteredModel(nn.Module, ClassRegistry):
     registered_map = {}
 
     def __init__(self, num_labels: int):
@@ -38,11 +37,8 @@ def streaming_state(self, x: Any):
         pass
 
 
-class ConvertedStaticModel(RegisteredModel, name='converted'):
-    def __init__(self,
-                 model: RegisteredModel,
-                 frame_window_size: int,
-                 frame_stride_size: int):
+class ConvertedStaticModel(RegisteredModel, name="converted"):
+    def __init__(self, model: RegisteredModel, frame_window_size: int, frame_stride_size: int):
         super().__init__(model.num_labels)
         self.model = model
         self.frame_window_size = frame_window_size
@@ -55,12 +51,12 @@ def compute_length(self, length: int):
 
     def forward(self, x, lengths):
         first = True
-        window = x[:, :, :, self.frame_window_size:]
+        window = x[:, :, :, self.frame_window_size :]
         idx = self.frame_stride_size
         outputs = []
         while first or window.size(3) == self.frame_window_size:
             first = False
             outputs.append(self.model(window, lengths))
-            window = x[:, :, :, idx:idx + self.frame_window_size]
+            window = x[:, :, :, idx : idx + self.frame_window_size]
             idx += self.frame_stride_size
         return torch.stack(outputs)
diff --git a/howl/model/rnn.py b/howl/model/rnn.py
index f9907c42..48ea44e6 100644
--- a/howl/model/rnn.py
+++ b/howl/model/rnn.py
@@ -1,19 +1,14 @@
 from typing import Any
 
-from pydantic import BaseSettings
-from torch.nn.utils.rnn import pad_packed_sequence, pack_padded_sequence
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from pydantic import BaseSettings
+from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
 
 from .base import RegisteredModel
 
-
-__all__ = ['LASEncoderConfig',
-           'FixedAttentionModuleConfig',
-           'LASClassifierConfig',
-           'LASEncoder',
-           'LASClassifier']
+__all__ = ["LASEncoderConfig", "FixedAttentionModuleConfig", "LASClassifierConfig", "LASEncoder", "LASClassifier"]
 
 
 class LASEncoderConfig(BaseSettings):
@@ -43,13 +38,15 @@ class LstmConfig(BaseSettings):
     num_labels: int = 2
 
 
-class SequentialLstm(RegisteredModel, name='seq-lstm'):
+class SequentialLstm(RegisteredModel, name="seq-lstm"):
     def __init__(self, num_labels: int, config: LstmConfig = LstmConfig()):
         super().__init__(num_labels)
         self.lstm = nn.LSTM(config.num_mels, config.hidden_size)
-        self.dnn = nn.Sequential(nn.Linear(config.hidden_size, int(2 * config.hidden_size)),
-                                 nn.ReLU(),
-                                 nn.Linear(int(2 * config.hidden_size), num_labels))
+        self.dnn = nn.Sequential(
+            nn.Linear(config.hidden_size, int(2 * config.hidden_size)),
+            nn.ReLU(),
+            nn.Linear(int(2 * config.hidden_size), num_labels),
+        )
         self.hc = None
 
     @property
@@ -74,13 +71,15 @@ def forward(self, x, lengths):
         return self.dnn(rnn_seq)
 
 
-class SimpleLstm(RegisteredModel, name='lstm'):
+class SimpleLstm(RegisteredModel, name="lstm"):
     def __init__(self, num_labels: int, config: LstmConfig = LstmConfig()):
         super().__init__(num_labels)
         self.lstm = nn.LSTM(config.num_mels, config.hidden_size)
-        self.dnn = nn.Sequential(nn.Linear(config.hidden_size, int(2 * config.hidden_size)),
-                                 nn.ReLU(),
-                                 nn.Linear(int(2 * config.hidden_size), num_labels))
+        self.dnn = nn.Sequential(
+            nn.Linear(config.hidden_size, int(2 * config.hidden_size)),
+            nn.ReLU(),
+            nn.Linear(int(2 * config.hidden_size), num_labels),
+        )
         self.hc = None
 
     def forward(self, x, lengths):
@@ -92,24 +91,28 @@ def forward(self, x, lengths):
         return self.dnn(hc[0].squeeze(0))
 
 
-class SimpleGru(RegisteredModel, name='gru'):
+class SimpleGru(RegisteredModel, name="gru"):
     def __init__(self, num_labels: int, config: LASEncoderConfig = LASEncoderConfig()):
         super().__init__(num_labels)
         conv1 = nn.Conv2d(1, config.num_latent_channels, 3, padding=(1, 3))
         conv2 = nn.Conv2d(config.num_latent_channels, 1, 3, padding=1)
-        self.conv_encoder = nn.Sequential(conv1,
-                                          nn.BatchNorm2d(config.num_latent_channels),
-                                          nn.ReLU(),
-                                          nn.MaxPool2d((1, 2 if config.use_maxpool else 1)),
-                                          conv2,
-                                          nn.ReLU(),
-                                          nn.BatchNorm2d(1))
+        self.conv_encoder = nn.Sequential(
+            conv1,
+            nn.BatchNorm2d(config.num_latent_channels),
+            nn.ReLU(),
+            nn.MaxPool2d((1, 2 if config.use_maxpool else 1)),
+            conv2,
+            nn.ReLU(),
+            nn.BatchNorm2d(1),
+        )
         self.use_maxpool = config.use_maxpool
         self.lstm_encoder = nn.GRU(config.num_mels, config.hidden_size, bidirectional=False)
-        self.dnn = nn.Sequential(nn.Linear(config.hidden_size, int(2 * config.hidden_size)),
-                                 nn.ReLU(),
-                                 nn.Dropout(0.2),
-                                 nn.Linear(int(2 * config.hidden_size), num_labels))
+        self.dnn = nn.Sequential(
+            nn.Linear(config.hidden_size, int(2 * config.hidden_size)),
+            nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(int(2 * config.hidden_size), num_labels),
+        )
 
     def forward(self, x, lengths):
         if lengths is None:
@@ -133,16 +136,20 @@ def __init__(self, config: LASEncoderConfig):
         out_channels = config.num_latent_channels
         hidden_size = config.hidden_size
         self.use_maxpool = config.use_maxpool
-        self.conv1 = conv1 = nn.Conv2d(config.num_spec_channels, stride=1, kernel_size=3, padding=2, out_channels=out_channels)
+        self.conv1 = conv1 = nn.Conv2d(
+            config.num_spec_channels, stride=1, kernel_size=3, padding=2, out_channels=out_channels
+        )
         self.conv2 = conv2 = nn.Conv2d(out_channels, stride=1, kernel_size=3, padding=2, out_channels=out_channels)
-        self.conv_encoder = nn.Sequential(conv1,
-                                          nn.BatchNorm2d(out_channels),
-                                          nn.ReLU(),
-                                          nn.MaxPool2d((1, 2 if config.use_maxpool else 1)),
-                                          conv2,
-                                          nn.BatchNorm2d(out_channels),
-                                          nn.ReLU(),
-                                          nn.MaxPool2d((1, 2 if config.use_maxpool else 1)))
+        self.conv_encoder = nn.Sequential(
+            conv1,
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(),
+            nn.MaxPool2d((1, 2 if config.use_maxpool else 1)),
+            conv2,
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(),
+            nn.MaxPool2d((1, 2 if config.use_maxpool else 1)),
+        )
         self.lstm_encoder = nn.LSTM(out_channels * 44, hidden_size, config.num_layers, bias=True, bidirectional=True)
 
     def forward(self, x, lengths):
@@ -175,24 +182,26 @@ def forward(self, rnn_seq, mask=None):
         rnn_seq = values.view(values.size(0), values.size(1), self.num_heads, values.size(2) // self.num_heads)
         keys = keys.view(values.size(0), keys.size(1), self.num_heads, keys.size(2) // self.num_heads)
         cvec = self.context_vec.view(-1, self.num_heads).unsqueeze(-1).expand(-1, -1, rnn_seq.size(0))
-        logits = torch.einsum('ijkl,lki->ijk', rnn_seq, cvec)
+        logits = torch.einsum("ijkl,lki->ijk", rnn_seq, cvec)
         if mask is not None:
             mask = (1 - mask) * -100
             logits += mask.unsqueeze(-1).expand_as(logits)
         scores = F.softmax(logits, 0)
-        vec = torch.einsum('ijk,ijkl->jkl', scores, keys)
+        vec = torch.einsum("ijk,ijkl->jkl", scores, keys)
         return vec.view(vec.size(0), -1)
 
 
-class LASClassifier(RegisteredModel, name='las'):
+class LASClassifier(RegisteredModel, name="las"):
     def __init__(self, num_labels: int, config: LASClassifierConfig = LASClassifierConfig()):
         super().__init__(num_labels)
         self.encoder = LASEncoder(config.las_config)
         self.attn = FixedAttentionModule(config.fixed_attn_config)
-        self.fc = nn.Sequential(nn.Linear(config.las_config.hidden_size * 2, config.dnn_size),
-                                nn.ReLU(),
-                                nn.Dropout(config.dropout),
-                                nn.Linear(config.dnn_size, num_labels))
+        self.fc = nn.Sequential(
+            nn.Linear(config.las_config.hidden_size * 2, config.dnn_size),
+            nn.ReLU(),
+            nn.Dropout(config.dropout),
+            nn.Linear(config.dnn_size, num_labels),
+        )
 
     def forward(self, x, lengths):
         rnn_seq, rnn_out = self.encoder(x, lengths)
diff --git a/howl/registered.py b/howl/utils/class_registry.py
similarity index 84%
rename from howl/registered.py
rename to howl/utils/class_registry.py
index 90dc56e5..fbea4653 100644
--- a/howl/registered.py
+++ b/howl/utils/class_registry.py
@@ -1,10 +1,9 @@
 from typing import List, Type, TypeVar
 
+T = TypeVar("T", bound="ClassRegistry")
 
-T = TypeVar('T', bound='RegisteredObjectBase')
 
-
-class RegisteredObjectBase:
+class ClassRegistry:
     registered_map = {}
 
     def __init_subclass__(cls, name: str = None):
diff --git a/howl/utils/hash.py b/howl/utils/hash.py
index 7823a510..15aaacc0 100644
--- a/howl/utils/hash.py
+++ b/howl/utils/hash.py
@@ -1,7 +1,17 @@
 from hashlib import sha256
 
+from howl.data.common.metadata import AudioClipMetadata
+
 
 def sha256_int(value: str):
     m = sha256()
     m.update(value.encode())
     return int(m.hexdigest(), 16)
+
+
+class Sha256Splitter:
+    def __init__(self, target_pct: int):
+        self.target_pct = target_pct
+
+    def __call__(self, x: AudioClipMetadata) -> bool:
+        return (sha256_int(str(x.path)) % 100) < self.target_pct
diff --git a/howl/utils/logger.py b/howl/utils/logger.py
new file mode 100644
index 00000000..8b7bccaa
--- /dev/null
+++ b/howl/utils/logger.py
@@ -0,0 +1,10 @@
+import os
+
+import coloredlogs
+
+__all__ = []
+
+
+coloredlogs.install(
+    level=os.environ.get("HOWL_LOG_LEVEL", "INFO"), fmt="%(asctime)s [%(levelname)s] %(module)s: %(message)s"
+)
diff --git a/training/__init__.py b/training/__init__.py
index 3c8e1796..e69de29b 100644
--- a/training/__init__.py
+++ b/training/__init__.py
@@ -1,1 +0,0 @@
-from howl.logger import *
diff --git a/training/run/attach_alignment.py b/training/run/attach_alignment.py
index baf505f7..2149520a 100644
--- a/training/run/attach_alignment.py
+++ b/training/run/attach_alignment.py
@@ -1,60 +1,65 @@
+import argparse
 from itertools import chain
 from pathlib import Path
-import argparse
 
 from textgrids import TextGrid
 from tqdm import tqdm
 
-from training.align import MfaTextGridConverter, StubAligner
-from howl.data.dataset import AudioClipDatasetLoader, AudioDatasetMetadataWriter
 from howl.data.dataset.base import AudioClipMetadata as AlignedAudioClipMetadata
+from howl.data.dataset.dataset_loader import AudioClipDatasetLoader
+from howl.data.dataset.dataset_writer import AudioDatasetMetadataWriter
 from howl.settings import SETTINGS
+from training.align import MfaTextGridConverter, StubAligner
 
 
 def main():
     def load_mfa_align():
-        converter = MfaTextGridConverter(use_phones=SETTINGS.training.token_type == 'phone')
+        converter = MfaTextGridConverter(use_phones=SETTINGS.training.token_type == "phone")
         id_align_map = {}
 
-        for tg_path in args.align_folder.glob('**/*.TextGrid'):
+        for tg_path in args.align_folder.glob("**/*.TextGrid"):
             tg = TextGrid(str(tg_path.absolute()))
-            audio_id = tg_path.name.split('.', 1)[0]
+            audio_id = tg_path.name.split(".", 1)[0]
             id_align_map[audio_id] = converter.convert(tg)
         return id_align_map
 
     def load_stub_align():
         id_align_map = {}
-        for ex in tqdm(chain(train_ds, dev_ds, test_ds), total=sum(map(len, (train_ds, dev_ds, test_ds)))):
+        for ex in tqdm(chain(train_ds, dev_ds, test_ds), total=sum(map(len, (train_ds, dev_ds, test_ds))),):
             id_align_map[ex.metadata.audio_id] = StubAligner().align(ex)
         return id_align_map
 
     parser = argparse.ArgumentParser()
-    parser.add_argument('--mfa-folder', '-i', dest='align_folder', type=Path)
-    parser.add_argument('--align-type', type=str, default='mfa', choices=('mfa', 'stub'))
+    parser.add_argument("--mfa-folder", "-i", dest="align_folder", type=Path)
+    parser.add_argument("--align-type", type=str, default="mfa", choices=("mfa", "stub"))
     args = parser.parse_args()
 
     ds_kwargs = dict(sr=SETTINGS.audio.sample_rate, mono=SETTINGS.audio.use_mono)
     ds_path = Path(SETTINGS.dataset.dataset_path)
     train_ds, dev_ds, test_ds = AudioClipDatasetLoader().load_splits(ds_path, **ds_kwargs)
 
-    if args.align_type == 'mfa':
+    if args.align_type == "mfa":
         id_align_map = load_mfa_align()
-    elif args.align_type == 'stub':
+    elif args.align_type == "stub":
         id_align_map = load_stub_align()
     else:
         raise ValueError
 
     for ds in (train_ds, dev_ds, test_ds):
-        with AudioDatasetMetadataWriter(ds_path, ds.set_type, 'aligned-', mode='w') as writer:
+        with AudioDatasetMetadataWriter(ds_path, ds.set_type, "aligned-", mode="w") as writer:
             for ex in tqdm(ds, total=len(ds)):
                 try:
                     transcription = id_align_map[ex.metadata.audio_id]
-                    writer.write(AlignedAudioClipMetadata(path=ex.metadata.path,
-                                                          transcription=transcription.transcription,
-                                                          end_timestamps=transcription.end_timestamps))
+                    writer.write(
+                        AlignedAudioClipMetadata(
+                            path=ex.metadata.path,
+                            transcription=transcription.transcription,
+                            end_timestamps=transcription.end_timestamps,
+                        )
+                    )
                 except KeyError:
                     pass
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()
diff --git a/training/run/create_raw_dataset.py b/training/run/create_raw_dataset.py
index b7f79407..58e7171a 100644
--- a/training/run/create_raw_dataset.py
+++ b/training/run/create_raw_dataset.py
@@ -2,8 +2,9 @@
 from pathlib import Path
 
 from howl.context import InferenceContext
-from howl.data.dataset import (AudioDataset, AudioDatasetWriter,
-                               RegisteredPathDatasetLoader)
+from howl.data.dataset.dataset import AudioDataset
+from howl.data.dataset.dataset_loader import RegisteredPathDatasetLoader
+from howl.data.dataset.dataset_writer import AudioDatasetWriter
 from howl.settings import SETTINGS
 from howl.utils.hash import sha256_int
 
@@ -19,11 +20,13 @@
 """
 
 
-def print_stats(header: str, context: InferenceContext, * datasets: AudioDataset, compute_length=True):
-    word_searcher = context.searcher if context.token_type == 'word' else None
+def print_stats(header: str, context: InferenceContext, *datasets: AudioDataset, compute_length=True):
+    word_searcher = context.searcher if context.token_type == "word" else None
     for ds in datasets:
         logging.info(
-            f'{header} ({ds.set_type}) statistics: {ds.compute_statistics(word_searcher=word_searcher, compute_length=compute_length)}')
+            f"{header} ({ds.set_type}) "
+            f"statistics: {ds.compute_statistics(word_searcher=word_searcher, compute_length=compute_length)}"
+        )
 
 
 def main():
@@ -38,44 +41,42 @@ def filter_fn(x):
         return False
 
     apb = ArgumentParserBuilder()
-    apb.add_options(opt('--negative-pct',
-                        type=int,
-                        default=1,
-                        help='The percentage of the dataset to check for negative examples.'),
-                    opt('--positive-pct', type=int, default=100,
-                        help='The percentage of the dataset to check for positive examples.'),
-                    opt('--input-path', '-i', type=str),
-                    opt('--dataset-type',
-                        type=str,
-                        default='mozilla-cv',
-                        choices=RegisteredPathDatasetLoader.registered_names()))
+    apb.add_options(
+        opt(
+            "--negative-pct", type=int, default=1, help="The percentage of the dataset to check for negative examples.",
+        ),
+        opt(
+            "--positive-pct",
+            type=int,
+            default=100,
+            help="The percentage of the dataset to check for positive examples.",
+        ),
+        opt("--input-path", "-i", type=str),
+        opt("--dataset-type", type=str, default="mozilla-cv", choices=RegisteredPathDatasetLoader.registered_names(),),
+    )
     args = apb.parser.parse_args()
     if args.input_path is None:
         args.input_path = SETTINGS.raw_dataset.common_voice_dataset_path
 
-    ctx = InferenceContext(SETTINGS.training.vocab,
-                           token_type=SETTINGS.training.token_type)
-    loader = RegisteredPathDatasetLoader.find_registered_class(
-        args.dataset_type)()
-    ds_kwargs = dict(sr=SETTINGS.audio.sample_rate,
-                     mono=SETTINGS.audio.use_mono)
-    train_ds, dev_ds, test_ds = loader.load_splits(
-        Path(args.input_path), **ds_kwargs)
+    ctx = InferenceContext(SETTINGS.training.vocab, token_type=SETTINGS.training.token_type)
+    loader = RegisteredPathDatasetLoader.find_registered_class(args.dataset_type)()
+    ds_kwargs = dict(sr=SETTINGS.audio.sample_rate, mono=SETTINGS.audio.use_mono)
+    train_ds, dev_ds, test_ds = loader.load_splits(Path(args.input_path), **ds_kwargs)
 
-    if args.dataset_type == 'mozilla-cv':
+    if args.dataset_type == "mozilla-cv":
         train_ds = train_ds.filter(filter_fn)
         dev_ds = dev_ds.filter(filter_fn)
         test_ds = test_ds.filter(filter_fn)
 
-    print_stats('Dataset', ctx, train_ds, dev_ds, test_ds, compute_length=True)
+    print_stats("Dataset", ctx, train_ds, dev_ds, test_ds, compute_length=True)
 
     for ds in train_ds, dev_ds, test_ds:
         try:
             AudioDatasetWriter(ds).write(Path(SETTINGS.dataset.dataset_path))
         except KeyboardInterrupt:
-            logging.info('Skipping...')
+            logging.info("Skipping...")
             pass
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()
diff --git a/training/run/generate_precise_dataset.py b/training/run/generate_precise_dataset.py
index 1c20a53f..29341ab4 100644
--- a/training/run/generate_precise_dataset.py
+++ b/training/run/generate_precise_dataset.py
@@ -4,19 +4,23 @@
 
 import librosa
 import numpy as np
+from tqdm import tqdm
+
 from howl.context import InferenceContext
-from howl.data.dataset import (DatasetType, RecursiveNoiseDatasetLoader,
-                               Sha256Splitter, WakeWordDataset,
-                               WakeWordDatasetLoader)
+from howl.data.dataset.dataset import DatasetType, WakeWordDataset
+from howl.data.dataset.dataset_loader import (
+    RecursiveNoiseDatasetLoader,
+    WakeWordDatasetLoader,
+)
 from howl.data.transform import DatasetMixer
 from howl.settings import SETTINGS
-from tqdm import tqdm
+from howl.utils.hash import Sha256Splitter
 
 from .args import ArgumentParserBuilder, opt
 from .create_raw_dataset import print_stats
 
 """
-This script is used to transform datasets for howl to dataset for Mycroft-precise 
+This script is used to transform datasets for howl to dataset for Mycroft-precise
 
 sample command:
 python -m training.run.generate_precise_dataset.py --i <datasets to convert> --o <location of the output dataset>
@@ -25,7 +29,7 @@
 
 def main():
     def copy_files(dataset, output_dir, deep_copy=False):
-        print('copying files to', output_dir)
+        print("copying files to", output_dir)
         output_dir.mkdir(parents=True, exist_ok=True)
         for item in tqdm(dataset):
             output_path = output_dir / item.metadata.path.name
@@ -38,37 +42,40 @@ def copy_files(dataset, output_dir, deep_copy=False):
                 pass
 
     def write_files(dataset, output_dir, mixer: DatasetMixer):
-        print('copying files to', output_dir)
+        print("copying files to", output_dir)
         output_dir.mkdir(parents=True, exist_ok=True)
         for item in tqdm(dataset):
-            item, = mixer([item])
+            (item,) = mixer([item])
             output_path = output_dir / item.metadata.path.name
             librosa.output.write_wav(output_path, item.audio_data.numpy(), 16000)
 
     apb = ArgumentParserBuilder()
     apb.add_options(
-        opt('--dataset-paths', '-i', type=str, nargs='+', default=[SETTINGS.dataset.dataset_path]),
-        opt('--output-paths', '-o', type=str, default='data/precise'),
-        opt('--deep-copy', action='store_true'))
+        opt("--dataset-paths", "-i", type=str, nargs="+", default=[SETTINGS.dataset.dataset_path],),
+        opt("--output-paths", "-o", type=str, default="data/precise"),
+        opt("--deep-copy", action="store_true"),
+    )
     args = apb.parser.parse_args()
 
     if args.deep_copy:
-        print('copying the audio files')
+        print("copying the audio files")
     else:
-        print('generating symlink files')
+        print("generating symlink files")
 
-    use_frame = SETTINGS.training.objective == 'frame'
-    ctx = InferenceContext(SETTINGS.training.vocab, token_type=SETTINGS.training.token_type, use_blank=not use_frame)
+    use_frame = SETTINGS.training.objective == "frame"
+    ctx = InferenceContext(SETTINGS.training.vocab, token_type=SETTINGS.training.token_type, use_blank=not use_frame,)
 
     loader = WakeWordDatasetLoader()
-    ds_kwargs = dict(sr=SETTINGS.audio.sample_rate, mono=SETTINGS.audio.use_mono, frame_labeler=ctx.labeler)
+    ds_kwargs = dict(sr=SETTINGS.audio.sample_rate, mono=SETTINGS.audio.use_mono, frame_labeler=ctx.labeler,)
 
     inference_settings = SETTINGS.inference_engine
-    wakeword = '_'.join(np.array(SETTINGS.training.vocab)[inference_settings.inference_sequence]).strip()
+    wakeword = "_".join(np.array(SETTINGS.training.vocab)[inference_settings.inference_sequence]).strip()
 
-    ww_train_ds, ww_dev_ds, ww_test_ds = WakeWordDataset(metadata_list=[], set_type=DatasetType.TRAINING, **ds_kwargs), \
-        WakeWordDataset(metadata_list=[], set_type=DatasetType.DEV, **ds_kwargs), \
-        WakeWordDataset(metadata_list=[], set_type=DatasetType.TEST, **ds_kwargs)
+    ww_train_ds, ww_dev_ds, ww_test_ds = (
+        WakeWordDataset(metadata_list=[], set_type=DatasetType.TRAINING, **ds_kwargs),
+        WakeWordDataset(metadata_list=[], set_type=DatasetType.DEV, **ds_kwargs),
+        WakeWordDataset(metadata_list=[], set_type=DatasetType.TEST, **ds_kwargs),
+    )
 
     for ds_path in args.dataset_paths:
         ds_path = Path(ds_path)
@@ -77,21 +84,21 @@ def write_files(dataset, output_dir, mixer: DatasetMixer):
         ww_dev_ds.extend(dev_ds)
         ww_test_ds.extend(test_ds)
 
-    print_stats(f'Wake word dataset', ctx, ww_train_ds, ww_dev_ds, ww_test_ds)
+    print_stats("Wake word dataset", ctx, ww_train_ds, ww_dev_ds, ww_test_ds)
 
     output_path = Path(args.output_paths) / wakeword
 
     ww_train_pos_ds = ww_train_ds.filter(lambda x: ctx.searcher.search(x.transcription), clone=True)
-    print_stats(f'train positive dataset', ctx, ww_train_pos_ds)
-    copy_files(ww_train_pos_ds, output_path / 'wake-word', args.deep_copy)
+    print_stats("train positive dataset", ctx, ww_train_pos_ds)
+    copy_files(ww_train_pos_ds, output_path / "wake-word", args.deep_copy)
 
     ww_train_neg_ds = ww_train_ds.filter(lambda x: not ctx.searcher.search(x.transcription), clone=True)
-    print_stats(f'train negative dataset', ctx, ww_train_neg_ds)
-    copy_files(ww_train_neg_ds, output_path / 'not-wake-word', args.deep_copy)
+    print_stats("train negative dataset", ctx, ww_train_neg_ds)
+    copy_files(ww_train_neg_ds, output_path / "not-wake-word", args.deep_copy)
 
-    noise_ds = RecursiveNoiseDatasetLoader().load(Path(SETTINGS.raw_dataset.noise_dataset_path),
-                                                  sr=SETTINGS.audio.sample_rate,
-                                                  mono=SETTINGS.audio.use_mono)
+    noise_ds = RecursiveNoiseDatasetLoader().load(
+        Path(SETTINGS.raw_dataset.noise_dataset_path), sr=SETTINGS.audio.sample_rate, mono=SETTINGS.audio.use_mono,
+    )
     noise_ds_train, noise_ds_dev = noise_ds.split(Sha256Splitter(80))
     noise_ds_dev, noise_ds_test = noise_ds_dev.split(Sha256Splitter(50))
 
@@ -99,25 +106,25 @@ def write_files(dataset, output_dir, mixer: DatasetMixer):
     test_mixer = DatasetMixer(noise_ds_test, seed=10, do_replace=False)
 
     ww_dev_pos_ds = ww_dev_ds.filter(lambda x: ctx.searcher.search(x.transcription), clone=True)
-    print_stats(f'dev positive dataset', ctx, ww_dev_pos_ds)
-    copy_files(ww_dev_pos_ds, output_path / 'dev/wake-word', args.deep_copy)
-    write_files(ww_dev_pos_ds, output_path / 'noisy-dev/wake-word', dev_mixer)
+    print_stats("dev positive dataset", ctx, ww_dev_pos_ds)
+    copy_files(ww_dev_pos_ds, output_path / "dev/wake-word", args.deep_copy)
+    write_files(ww_dev_pos_ds, output_path / "noisy-dev/wake-word", dev_mixer)
 
     ww_dev_neg_ds = ww_dev_ds.filter(lambda x: not ctx.searcher.search(x.transcription), clone=True)
-    print_stats(f'dev negative dataset', ctx, ww_dev_neg_ds)
-    copy_files(ww_dev_neg_ds, output_path / 'dev/not-wake-word', args.deep_copy)
-    write_files(ww_dev_neg_ds, output_path / 'noisy-dev/not-wake-word', dev_mixer)
+    print_stats("dev negative dataset", ctx, ww_dev_neg_ds)
+    copy_files(ww_dev_neg_ds, output_path / "dev/not-wake-word", args.deep_copy)
+    write_files(ww_dev_neg_ds, output_path / "noisy-dev/not-wake-word", dev_mixer)
 
     ww_test_pos_ds = ww_test_ds.filter(lambda x: ctx.searcher.search(x.transcription), clone=True)
-    print_stats(f'test positive dataset', ctx, ww_test_pos_ds)
-    copy_files(ww_test_pos_ds, output_path / 'test/wake-word', args.deep_copy)
-    write_files(ww_test_pos_ds, output_path / 'noisy-test/wake-word', test_mixer)
+    print_stats("test positive dataset", ctx, ww_test_pos_ds)
+    copy_files(ww_test_pos_ds, output_path / "test/wake-word", args.deep_copy)
+    write_files(ww_test_pos_ds, output_path / "noisy-test/wake-word", test_mixer)
 
     ww_test_neg_ds = ww_test_ds.filter(lambda x: not ctx.searcher.search(x.transcription), clone=True)
-    print_stats(f'test negative dataset', ctx, ww_test_neg_ds)
-    copy_files(ww_test_neg_ds, output_path / 'test/not-wake-word', args.deep_copy)
-    write_files(ww_test_neg_ds, output_path / 'noisy-test/not-wake-word', test_mixer)
+    print_stats("test negative dataset", ctx, ww_test_neg_ds)
+    copy_files(ww_test_neg_ds, output_path / "test/not-wake-word", args.deep_copy)
+    write_files(ww_test_neg_ds, output_path / "noisy-test/not-wake-word", test_mixer)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()
diff --git a/training/run/pretrain_gsc.py b/training/run/pretrain_gsc.py
index dd7fa3e7..27709af2 100644
--- a/training/run/pretrain_gsc.py
+++ b/training/run/pretrain_gsc.py
@@ -1,22 +1,30 @@
+import logging
 from collections import Counter
 from functools import partial
 from pathlib import Path
-import logging
 
-from tqdm import trange, tqdm
-from torch.optim.adamw import AdamW
 import torch
 import torch.nn as nn
+from torch.optim.adamw import AdamW
+from tqdm import tqdm, trange
 
-from .args import ArgumentParserBuilder, opt
-from howl.data.dataset import GoogleSpeechCommandsDatasetLoader
 from howl.data.dataloader import StandardAudioDataLoaderBuilder
-from howl.data.transform import compose, ZmuvTransform, StandardAudioTransform,\
-    NoiseTransform, batchify, TimeshiftTransform, truncate_length
-from howl.settings import SETTINGS
+from howl.data.dataset.gsc_dataset_loader import GoogleSpeechCommandsDatasetLoader
+from howl.data.transform import (
+    NoiseTransform,
+    StandardAudioTransform,
+    TimeshiftTransform,
+    ZmuvTransform,
+    batchify,
+    compose,
+    truncate_length,
+)
 from howl.model import RegisteredModel, Workspace
+from howl.settings import SETTINGS
 from howl.utils.random import set_seed
 
+from .args import ArgumentParserBuilder, opt
+
 
 def main():
     def evaluate_accuracy(data_loader, prefix: str, save: bool = False):
@@ -28,16 +36,17 @@ def evaluate_accuracy(data_loader, prefix: str, save: bool = False):
         counter = Counter()
         for idx, batch in enumerate(pbar):
             batch = batch.to(device)
-            scores = model(zmuv_transform(std_transform(batch.audio_data)),
-                           std_transform.compute_lengths(batch.lengths))
+            scores = model(
+                zmuv_transform(std_transform(batch.audio_data)), std_transform.compute_lengths(batch.lengths),
+            )
             num_tot += scores.size(0)
             labels = batch.labels.to(device)
             counter.update(labels.tolist())
             num_corr += (scores.max(1)[1] == labels).float().sum().item()
             acc = num_corr / num_tot
-            pbar.set_postfix(accuracy=f'{acc:.4}')
+            pbar.set_postfix(accuracy=f"{acc:.4}")
         if save and not args.eval:
-            writer.add_scalar(f'{prefix}/Metric/acc', acc, epoch_idx)
+            writer.add_scalar(f"{prefix}/Metric/acc", acc, epoch_idx)
             ws.increment_model(model, acc / 10)
         elif args.eval:
             tqdm.write(str(counter))
@@ -46,10 +55,12 @@ def evaluate_accuracy(data_loader, prefix: str, save: bool = False):
         return num_corr / num_tot
 
     apb = ArgumentParserBuilder()
-    apb.add_options(opt('--model', type=str, choices=RegisteredModel.registered_names(), default='las'),
-                    opt('--workspace', type=str, default=str(Path('workspaces') / 'default')),
-                    opt('--load-weights', action='store_true'),
-                    opt('--eval', action='store_true'))
+    apb.add_options(
+        opt("--model", type=str, choices=RegisteredModel.registered_names(), default="las",),
+        opt("--workspace", type=str, default=str(Path("workspaces") / "default")),
+        opt("--load-weights", action="store_true"),
+        opt("--eval", action="store_true"),
+    )
     args = apb.parser.parse_args()
 
     ws = Workspace(Path(args.workspace), delete_existing=not args.eval)
@@ -66,53 +77,50 @@ def evaluate_accuracy(data_loader, prefix: str, save: bool = False):
     zmuv_transform = ZmuvTransform().to(device)
     batchifier = partial(batchify, label_provider=lambda x: x.label)
     truncater = partial(truncate_length, length=int(SETTINGS.training.max_window_size_seconds * sr))
-    train_comp = compose(truncater,
-                         TimeshiftTransform().train(),
-                         NoiseTransform().train(),
-                         batchifier)
+    train_comp = compose(truncater, TimeshiftTransform().train(), NoiseTransform().train(), batchifier)
     prep_dl = StandardAudioDataLoaderBuilder(train_ds, collate_fn=batchifier).build(1)
     prep_dl.shuffle = True
     train_dl = StandardAudioDataLoaderBuilder(train_ds, collate_fn=train_comp).build(SETTINGS.training.batch_size)
-    dev_dl = StandardAudioDataLoaderBuilder(dev_ds, collate_fn=compose(truncater, batchifier)).build(SETTINGS.training.batch_size)
-    test_dl = StandardAudioDataLoaderBuilder(test_ds, collate_fn=compose(truncater, batchifier)).build(SETTINGS.training.batch_size)
+    dev_dl = StandardAudioDataLoaderBuilder(dev_ds, collate_fn=compose(truncater, batchifier)).build(
+        SETTINGS.training.batch_size
+    )
+    test_dl = StandardAudioDataLoaderBuilder(test_ds, collate_fn=compose(truncater, batchifier)).build(
+        SETTINGS.training.batch_size
+    )
 
     model = RegisteredModel.find_registered_class(args.model)(30).to(device)
     params = list(filter(lambda x: x.requires_grad, model.parameters()))
-    optimizer = AdamW(params, SETTINGS.training.learning_rate, weight_decay=SETTINGS.training.weight_decay)
-    logging.info(f'{sum(p.numel() for p in params)} parameters')
+    optimizer = AdamW(params, SETTINGS.training.learning_rate, weight_decay=SETTINGS.training.weight_decay,)
+    logging.info(f"{sum(p.numel() for p in params)} parameters")
     criterion = nn.CrossEntropyLoss()
 
-    if (ws.path / 'zmuv.pt.bin').exists():
-        zmuv_transform.load_state_dict(torch.load(str(ws.path / 'zmuv.pt.bin')))
+    if (ws.path / "zmuv.pt.bin").exists():
+        zmuv_transform.load_state_dict(torch.load(str(ws.path / "zmuv.pt.bin")))
     else:
-        for idx, batch in enumerate(tqdm(prep_dl, desc='Constructing ZMUV')):
+        for idx, batch in enumerate(tqdm(prep_dl, desc="Constructing ZMUV")):
             batch.to(device)
             zmuv_transform.update(std_transform(batch.audio_data))
             if idx == 2000:  # TODO: quick debugging, remove later
                 break
         logging.info(dict(zmuv_mean=zmuv_transform.mean, zmuv_std=zmuv_transform.std))
-    torch.save(zmuv_transform.state_dict(), str(ws.path / 'zmuv.pt.bin'))
+    torch.save(zmuv_transform.state_dict(), str(ws.path / "zmuv.pt.bin"))
 
     if args.load_weights:
         ws.load_model(model, best=True)
     if args.eval:
         ws.load_model(model, best=True)
-        evaluate_accuracy(dev_dl, 'Dev')
-        evaluate_accuracy(test_dl, 'Test')
+        evaluate_accuracy(dev_dl, "Dev")
+        evaluate_accuracy(test_dl, "Test")
         return
 
     ws.write_args(args)
     ws.write_settings(SETTINGS)
-    writer.add_scalar('Meta/Parameters', sum(p.numel() for p in params))
+    writer.add_scalar("Meta/Parameters", sum(p.numel() for p in params))
     dev_acc = 0
     for epoch_idx in trange(SETTINGS.training.num_epochs, position=0, leave=True):
         model.train()
         std_transform.train()
-        pbar = tqdm(train_dl,
-                    total=len(train_dl),
-                    position=1,
-                    desc='Training',
-                    leave=True)
+        pbar = tqdm(train_dl, total=len(train_dl), position=1, desc="Training", leave=True)
         for batch in pbar:
             batch.to(device)
             audio_data = zmuv_transform(std_transform(batch.audio_data))
@@ -123,18 +131,18 @@ def evaluate_accuracy(data_loader, prefix: str, save: bool = False):
             loss = criterion(scores, labels)
             loss.backward()
             optimizer.step()
-            pbar.set_postfix(dict(loss=f'{loss.item():.3}'))
-            writer.add_scalar('Training/Loss', loss.item(), epoch_idx)
+            pbar.set_postfix(dict(loss=f"{loss.item():.3}"))
+            writer.add_scalar("Training/Loss", loss.item(), epoch_idx)
 
         for group in optimizer.param_groups:
-            group['lr'] *= SETTINGS.training.lr_decay
-        dev_acc = evaluate_accuracy(dev_dl, 'Dev', save=True)
-    test_acc = evaluate_accuracy(test_dl, 'Test')
+            group["lr"] *= SETTINGS.training.lr_decay
+        dev_acc = evaluate_accuracy(dev_dl, "Dev", save=True)
+    test_acc = evaluate_accuracy(test_dl, "Test")
 
     print("model: ", args.model)
     print("dev_acc: ", dev_acc)
     print("test_acc: ", test_acc)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()
diff --git a/training/run/stitch_vocab_samples.py b/training/run/stitch_vocab_samples.py
index d64ea1db..e313b546 100644
--- a/training/run/stitch_vocab_samples.py
+++ b/training/run/stitch_vocab_samples.py
@@ -1,13 +1,11 @@
 import argparse
 from pathlib import Path
 
-from howl.data.dataset import (
-    AudioDatasetWriter,
-    WakeWordDatasetLoader,
-    WordFrameLabeler,
-)
+from howl.data.common.labeler import WordFrameLabeler
+from howl.data.common.vocab import Vocab
+from howl.data.dataset.dataset_loader import WakeWordDatasetLoader
+from howl.data.dataset.dataset_writer import AudioDatasetWriter
 from howl.data.stitcher import WordStitcher
-from howl.data.tokenize import Vocab
 from howl.settings import SETTINGS
 
 """Using aligned dataset, generate wakeword samples by stitching vocab samples
@@ -48,7 +46,7 @@ def main():
         help="train/dev/test pct for stitched dataset (default: [0.5, 0.25, 0.25])",
     )
     parser.add_argument(
-        "--disable-detect-keyword", action="store_false", help="disable keyword detection based verifcation"
+        "--disable-detect-keyword", action="store_false", help="disable keyword detection based verifcation",
     )
 
     args = parser.parse_args()
@@ -58,7 +56,7 @@ def main():
 
     vocab = Vocab(SETTINGS.training.vocab)
     labeler = WordFrameLabeler(vocab)
-    ds_kwargs = dict(sr=SETTINGS.audio.sample_rate, mono=SETTINGS.audio.use_mono, frame_labeler=labeler)
+    ds_kwargs = dict(sr=SETTINGS.audio.sample_rate, mono=SETTINGS.audio.use_mono, frame_labeler=labeler,)
 
     # load aligned datasets
     train_ds, dev_ds, test_ds = WakeWordDatasetLoader().load_splits(aligned_ds_path, **ds_kwargs)
diff --git a/training/run/train.py b/training/run/train.py
index 3de5171b..f8a9a0ac 100644
--- a/training/run/train.py
+++ b/training/run/train.py
@@ -8,15 +8,13 @@
 from tqdm import tqdm, trange
 
 from howl.context import InferenceContext
+from howl.data.common.tokenizer import WakeWordTokenizer
 from howl.data.dataloader import StandardAudioDataLoaderBuilder
-from howl.data.dataset import (
-    DatasetType,
+from howl.data.dataset.dataset import DatasetType, WakeWordDataset
+from howl.data.dataset.dataset_loader import (
     RecursiveNoiseDatasetLoader,
-    Sha256Splitter,
-    WakeWordDataset,
     WakeWordDatasetLoader,
 )
-from howl.data.tokenize import WakeWordTokenizer
 from howl.data.transform import (
     AudioSequenceBatchifier,
     DatasetMixer,
@@ -30,6 +28,7 @@
 from howl.model import ConfusionMatrix, ConvertedStaticModel, RegisteredModel, Workspace
 from howl.model.inference import FrameInferenceEngine, SequenceInferenceEngine
 from howl.settings import SETTINGS
+from howl.utils.hash import Sha256Splitter
 from howl.utils.random import set_seed
 
 from .args import ArgumentParserBuilder, opt
@@ -95,26 +94,30 @@ def do_evaluate():
         evaluate_engine(ww_test_pos_ds, "Test positive", positive_set=True)
         evaluate_engine(ww_test_neg_ds, "Test negative", positive_set=False)
         if SETTINGS.training.use_noise_dataset:
-            evaluate_engine(ww_test_pos_ds, "Test noisy positive", positive_set=True, mixer=test_mixer)
-            evaluate_engine(ww_test_neg_ds, "Test noisy negative", positive_set=False, mixer=test_mixer)
+            evaluate_engine(
+                ww_test_pos_ds, "Test noisy positive", positive_set=True, mixer=test_mixer,
+            )
+            evaluate_engine(
+                ww_test_neg_ds, "Test noisy negative", positive_set=False, mixer=test_mixer,
+            )
 
     apb = ArgumentParserBuilder()
     apb.add_options(
-        opt("--model", type=str, choices=RegisteredModel.registered_names(), default="las"),
+        opt("--model", type=str, choices=RegisteredModel.registered_names(), default="las",),
         opt("--workspace", type=str, default=str(Path("workspaces") / "default")),
         opt("--load-weights", action="store_true"),
         opt("--load-last", action="store_true"),
         opt("--no-dev-per-epoch", action="store_false", dest="dev_per_epoch"),
-        opt("--dataset-paths", "-i", type=str, nargs="+", default=[SETTINGS.dataset.dataset_path]),
+        opt("--dataset-paths", "-i", type=str, nargs="+", default=[SETTINGS.dataset.dataset_path],),
         opt("--eval", action="store_true"),
     )
     args = apb.parser.parse_args()
 
     use_frame = SETTINGS.training.objective == "frame"
-    ctx = InferenceContext(SETTINGS.training.vocab, token_type=SETTINGS.training.token_type, use_blank=not use_frame)
+    ctx = InferenceContext(SETTINGS.training.vocab, token_type=SETTINGS.training.token_type, use_blank=not use_frame,)
     if use_frame:
         batchifier = WakeWordFrameBatchifier(
-            ctx.negative_label, window_size_ms=int(SETTINGS.training.max_window_size_seconds * 1000)
+            ctx.negative_label, window_size_ms=int(SETTINGS.training.max_window_size_seconds * 1000),
         )
         criterion = nn.CrossEntropyLoss()
     else:
@@ -126,7 +129,7 @@ def do_evaluate():
     writer = ws.summary_writer
     set_seed(SETTINGS.training.seed)
     loader = WakeWordDatasetLoader()
-    ds_kwargs = dict(sr=SETTINGS.audio.sample_rate, mono=SETTINGS.audio.use_mono, frame_labeler=ctx.labeler)
+    ds_kwargs = dict(sr=SETTINGS.audio.sample_rate, mono=SETTINGS.audio.use_mono, frame_labeler=ctx.labeler,)
 
     ww_train_ds, ww_dev_ds, ww_test_ds = (
         WakeWordDataset(metadata_list=[], set_type=DatasetType.TRAINING, **ds_kwargs),
@@ -158,7 +161,7 @@ def do_evaluate():
 
     if SETTINGS.training.use_noise_dataset:
         noise_ds = RecursiveNoiseDatasetLoader().load(
-            Path(SETTINGS.raw_dataset.noise_dataset_path), sr=SETTINGS.audio.sample_rate, mono=SETTINGS.audio.use_mono
+            Path(SETTINGS.raw_dataset.noise_dataset_path), sr=SETTINGS.audio.sample_rate, mono=SETTINGS.audio.use_mono,
         )
         logging.info(f"Loaded {len(noise_ds.metadata_list)} noise files.")
         noise_ds_train, noise_ds_dev = noise_ds.split(Sha256Splitter(80))
@@ -176,7 +179,7 @@ def do_evaluate():
     if SETTINGS.training.convert_static:
         model = ConvertedStaticModel(model, 40, 10)
     params = list(filter(lambda x: x.requires_grad, model.parameters()))
-    optimizer = AdamW(params, SETTINGS.training.learning_rate, weight_decay=SETTINGS.training.weight_decay)
+    optimizer = AdamW(params, SETTINGS.training.learning_rate, weight_decay=SETTINGS.training.weight_decay,)
     logging.info(f"{sum(p.numel() for p in params)} parameters")
 
     if (ws.path / "zmuv.pt.bin").exists():
@@ -210,7 +213,7 @@ def do_evaluate():
             batch.to(device)
             if use_frame:
                 scores = model(
-                    zmuv_transform(std_transform(batch.audio_data)), std_transform.compute_lengths(batch.lengths)
+                    zmuv_transform(std_transform(batch.audio_data)), std_transform.compute_lengths(batch.lengths),
                 )
                 loss = criterion(scores, batch.labels)
             else:
@@ -235,7 +238,9 @@ def do_evaluate():
         writer.add_scalar("Training/LearningRate", group["lr"], epoch_idx)
 
         if args.dev_per_epoch:
-            evaluate_engine(ww_dev_pos_ds, "Dev positive", positive_set=True, save=True, write_errors=False)
+            evaluate_engine(
+                ww_dev_pos_ds, "Dev positive", positive_set=True, save=True, write_errors=False,
+            )
 
     do_evaluate()
 


Test patch:
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index e9ea8f12..17e43c72 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -19,6 +19,7 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         pip install setuptools wheel
+        sudo apt-get update
         sudo apt-get install libasound-dev portaudio19-dev libportaudio2 libportaudiocpp0 libpulse-dev
         pip install -r requirements.txt
         pip install -r requirements_training.txt
diff --git a/test/data/dataset/phone_test.py b/test/data/common/phone_test.py
similarity index 97%
rename from test/data/dataset/phone_test.py
rename to test/data/common/phone_test.py
index 76443448..db25c68b 100644
--- a/test/data/dataset/phone_test.py
+++ b/test/data/common/phone_test.py
@@ -1,7 +1,7 @@
 import unittest
 from pathlib import Path
 
-from howl.data.dataset import PhonePhrase, PronunciationDictionary
+from howl.data.common.phone import PhonePhrase, PronunciationDictionary
 
 
 class TestPhonePhrase(unittest.TestCase):
diff --git a/test/data/common/tokenizer_test.py b/test/data/common/tokenizer_test.py
new file mode 100644
index 00000000..df683deb
--- /dev/null
+++ b/test/data/common/tokenizer_test.py
@@ -0,0 +1,42 @@
+import unittest
+
+from howl.data.common.tokenizer import WakeWordTokenizer
+from howl.data.common.vocab import Vocab
+
+
+class TestWakeWordTokenizer(unittest.TestCase):
+    def test_decode(self):
+        """test decode
+        """
+        vocab = Vocab({"Cat": 0, "cap": 1}, oov_token_id=2, oov_word_repr="<OOV>")
+        tokenizer = WakeWordTokenizer(vocab)
+
+        decoded_string = tokenizer.decode([0, 1, 1])
+        self.assertEqual(decoded_string, "Cat cap cap")
+
+        decoded_string = tokenizer.decode([])
+        self.assertEqual(decoded_string, "")
+
+    def test_encode(self):
+        """test encode
+        """
+        oov_label = 2
+        vocab = Vocab({"Cat": 0, "can": 1}, oov_token_id=oov_label, oov_word_repr="<OOV>")
+        original_text = "Catch me if you CAN"
+
+        # ignore oov words
+        tokenizer = WakeWordTokenizer(vocab)
+        encoded_output = tokenizer.encode(original_text)
+        self.assertEqual(len(encoded_output), 1)
+        self.assertEqual(encoded_output[0], 1)
+
+        # consider oov words
+        tokenizer = WakeWordTokenizer(vocab, False)
+        encoded_output = tokenizer.encode(original_text)
+        self.assertEqual(len(encoded_output), len(original_text.split()))
+        self.assertEqual(encoded_output[0], oov_label)
+        self.assertEqual(encoded_output[-1], 1)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/data/tokenize_test.py b/test/data/common/vocab_test.py
similarity index 64%
rename from test/data/tokenize_test.py
rename to test/data/common/vocab_test.py
index 33e3f202..8ad0774d 100644
--- a/test/data/tokenize_test.py
+++ b/test/data/common/vocab_test.py
@@ -1,14 +1,13 @@
 import unittest
 
-from howl.data.tokenize import Vocab, VocabTrie, WakeWordTokenizer
+from howl.data.common.vocab import Vocab, VocabTrie
 
 
 class TestVocab(unittest.TestCase):
-
     def test_vocab(self):
         """test Vocab class
         """
-        vocab = Vocab({"Hello": 0, "World": 1}, oov_token_id=2, oov_word_repr='<OOV>')
+        vocab = Vocab({"Hello": 0, "World": 1}, oov_token_id=2, oov_word_repr="<OOV>")
         self.assertEqual(len(vocab), 2)
 
         self.assertEqual(vocab[0], "Hello")
@@ -46,7 +45,6 @@ def test_vocab(self):
 
 
 class TestVocabTrie(unittest.TestCase):
-
     def test_vocab_trie(self):
         """test different cases of max splits
         """
@@ -73,40 +71,5 @@ def test_vocab_trie(self):
         self.assertEqual(remaining_text, remaining_text)
 
 
-class TestWakeWordTokenizer(unittest.TestCase):
-
-    def test_decode(self):
-        """test decode
-        """
-        vocab = Vocab({"Cat": 0, "cap": 1}, oov_token_id=2, oov_word_repr='<OOV>')
-        tokenizer = WakeWordTokenizer(vocab)
-
-        decoded_string = tokenizer.decode([0, 1, 1])
-        self.assertEqual(decoded_string, "Cat cap cap")
-
-        decoded_string = tokenizer.decode([])
-        self.assertEqual(decoded_string, "")
-
-    def test_encode(self):
-        """test encode
-        """
-        oov_label = 2
-        vocab = Vocab({"Cat": 0, "can": 1}, oov_token_id=oov_label, oov_word_repr='<OOV>')
-        original_text = "Catch me if you CAN"
-
-        # ignore oov words
-        tokenizer = WakeWordTokenizer(vocab)
-        encoded_output = tokenizer.encode(original_text)
-        self.assertEqual(len(encoded_output), 1)
-        self.assertEqual(encoded_output[0], 1)
-
-        # consider oov words
-        tokenizer = WakeWordTokenizer(vocab, False)
-        encoded_output = tokenizer.encode(original_text)
-        self.assertEqual(len(encoded_output), len(original_text.split()))
-        self.assertEqual(encoded_output[0], oov_label)
-        self.assertEqual(encoded_output[-1], 1)
-
-
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()
diff --git a/test/data/dataset/dataset_test.py b/test/data/dataset/dataset_test.py
index a37029c2..89473b4e 100644
--- a/test/data/dataset/dataset_test.py
+++ b/test/data/dataset/dataset_test.py
@@ -2,14 +2,11 @@
 
 import torch
 
-from howl.data.dataset import (
-    AudioClipExample,
-    AudioClipMetadata,
-    AudioDataset,
-    DatasetType,
-)
-from howl.data.searcher import WordTranscriptSearcher
-from howl.data.tokenize import Vocab
+from howl.data.common.example import AudioClipExample
+from howl.data.common.metadata import AudioClipMetadata
+from howl.data.common.searcher import WordTranscriptSearcher
+from howl.data.common.vocab import Vocab
+from howl.data.dataset.dataset import AudioDataset, DatasetType
 from howl.settings import SETTINGS
 
 
diff --git a/test/data/dataset/serialize_test.py b/test/data/dataset/dataset_writer_test.py
similarity index 86%
rename from test/data/dataset/serialize_test.py
rename to test/data/dataset/dataset_writer_test.py
index 0a85c824..6443dacd 100644
--- a/test/data/dataset/serialize_test.py
+++ b/test/data/dataset/dataset_writer_test.py
@@ -5,13 +5,10 @@
 import pytest
 import torch
 
-from howl.data.dataset import (
-    AudioClipDataset,
-    AudioClipExample,
-    AudioClipMetadata,
-    AudioDatasetWriter,
-    DatasetType,
-)
+from howl.data.common.example import AudioClipExample
+from howl.data.common.metadata import AudioClipMetadata
+from howl.data.dataset.dataset import AudioClipDataset, DatasetType
+from howl.data.dataset.dataset_writer import AudioDatasetWriter
 from howl.settings import SETTINGS
 
 
diff --git a/test/data/stitcher_test.py b/test/data/stitcher_test.py
index c81e2fd1..d5589908 100644
--- a/test/data/stitcher_test.py
+++ b/test/data/stitcher_test.py
@@ -2,9 +2,10 @@
 import unittest
 from pathlib import Path
 
-from howl.data.dataset import WakeWordDatasetLoader, WordFrameLabeler
+from howl.data.common.labeler import WordFrameLabeler
+from howl.data.common.vocab import Vocab
+from howl.data.dataset.dataset_loader import WakeWordDatasetLoader
 from howl.data.stitcher import WordStitcher
-from howl.data.tokenize import Vocab
 from howl.settings import SETTINGS
 
 
@@ -22,7 +23,7 @@ def test_compute_statistics(self):
         labeler = WordFrameLabeler(vocab)
 
         loader = WakeWordDatasetLoader()
-        ds_kwargs = dict(sr=SETTINGS.audio.sample_rate, mono=SETTINGS.audio.use_mono, frame_labeler=labeler)
+        ds_kwargs = dict(sr=SETTINGS.audio.sample_rate, mono=SETTINGS.audio.use_mono, frame_labeler=labeler,)
 
         test_dataset_path = Path("test/test_data/stitcher")
         stitched_dataset_path = test_dataset_path / "stitched"
diff --git a/test/model/inference_test.py b/test/model/inference_test.py
index c513eb22..79448536 100644
--- a/test/model/inference_test.py
+++ b/test/model/inference_test.py
@@ -1,17 +1,16 @@
 import unittest
 
-from howl.data.tokenize import Vocab
-from howl.data.searcher import WordTranscriptSearcher
+from howl.data.common.searcher import WordTranscriptSearcher
+from howl.data.common.vocab import Vocab
 from howl.settings import SETTINGS
 
 
 class TestWordTranscriptSearcher(unittest.TestCase):
-
     def test_search(self):
         """test search
         """
         SETTINGS.inference_engine.inference_sequence = [0, 1]
-        vocab = Vocab({"hello": 0, "world": 1}, oov_token_id=2, oov_word_repr='<OOV>')
+        vocab = Vocab({"hello": 0, "world": 1}, oov_token_id=2, oov_word_repr="<OOV>")
         tokenizer = WordTranscriptSearcher(vocab)
 
         self.assertTrue(tokenizer.search("hello world"))
@@ -23,7 +22,7 @@ def test_contains_any(self):
         """test contains_any
         """
         SETTINGS.inference_engine.inference_sequence = [0, 1]
-        vocab = Vocab({"hello": 0, "world": 1}, oov_token_id=2, oov_word_repr='<OOV>')
+        vocab = Vocab({"hello": 0, "world": 1}, oov_token_id=2, oov_word_repr="<OOV>")
         tokenizer = WordTranscriptSearcher(vocab)
 
         self.assertTrue(tokenizer.contains_any("hello world"))
@@ -35,7 +34,7 @@ def test_count_vocab(self):
         """test count_vocab
         """
         SETTINGS.inference_engine.inference_sequence = [0, 1]
-        vocab = Vocab({"Hello": 0, "World": 1}, oov_token_id=2, oov_word_repr='<OOV>')
+        vocab = Vocab({"Hello": 0, "World": 1}, oov_token_id=2, oov_word_repr="<OOV>")
         tokenizer = WordTranscriptSearcher(vocab)
 
         # test string with valid vocabs
@@ -57,5 +56,5 @@ def test_count_vocab(self):
         self.assertEqual(counter["World"], 0)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()