Repository: wfondrie/depthcharge (language: python)

Issue:
Why there is a replacement from "I" to "L" in peptide sequence  during the tokenize function of _PeptideTransformer in transformer.py?
Hi, casanovo and depthcharge are excellent work. However, there is some confuse in transformer.py. Why there is a replacement from "I" to "L" in peptide sequence  during the tokenize function of _PeptideTransformer? Will this replacement influence the predict precusion of peptide sequence incuding "I" during inference?
Provide tokenization function
Have a separate tokenization function to unify this behavior across depthcharge (i.e. [to calculate masses from peptide sequences](https://github.com/wfondrie/depthcharge/blob/main/depthcharge/masses.py#L93)) and Casanovo (during [evaluation of the predictions](https://github.com/Noble-Lab/casanovo/blob/main/casanovo/denovo/evaluate.py#L222)).
Preprocess spectra before adding to index
I did some coding on my flight to Japan and moved the preprocessing code from Casanovo to depthcharge to fix [Casanovo#56](https://github.com/Noble-Lab/casanovo/issues/56). The upshot is that invalid spectra are now no longer included in the index on the depthcharge side, rather than later being replaced by a dummy spectrum by Casanovo that will get an incorrect prediction.

This needs some unit tests, but I'm unlikely to be able to add those within the next two weeks while I'm in Japan.

One remaining question is whether we also want to introduce a parameter that specifies the minimum number of peaks (after preprocessing) for spectra to be included?

Linked to [Casanovo#181](https://github.com/Noble-Lab/casanovo/pull/181).


Discussion / hints:
Hi @liangzhendong123 👋 - thanks for checking our depthcharge and casanovo! 

> Why there is a replacement from "I" to "L" in peptide sequence during the tokenize function of _PeptideTransformer? 

This choice was made because "I" is indistinguishable from "L" in mass spectrometry applications. Both I and L have the same mass and the small difference in their structures appear to have a negligible affect on the fragment ion intensities we observe in a mass spectrum. In hindsight however, I should probably make this an optional default instead.

> Will this replacement influence the predict precusion of peptide sequence incuding "I" during inference?

Yes, but this is true either implicitly or explicitly for all mass spectrometry-based *de novo* peptide sequencing approaches.
One more related thing is that Casanovo sometimes apparently predicts a double N-terminal acetylation, e.g. `+43.006+43.006PEPTIDEK`, in which case the `+43.006+43.006` is not found in the AA mass dictionary. How should we handle such weird predictions that Casanovo comes up with? Currently mass calculation crashes with a `KeyError`, which doesn't seem desired behavior.
@bittremieux, this also brings up a larger question in my mind: should we try and parse peptides sequences and their modifications into the same format, say the mass enclosed in `[]`? Such standardization could help prevent these edge cases and I think I could change the parser to specifically handale `[]`, `()`, and non-enclosed modifications.
Good point. General parsing functionality could be useful to harmonize different ways people encode peptide sequences. Now we've kinda been using the format that MassIVE uses, but there are various different flavors. A parsing function that translates sequences to a consistent internal representation would enable depthcharge/Casanovo to deal with different input formats. (The dev version of) spectrum_utils also supports the full ProForma specification to encode peptidoforms, including controlled vocabularies to encode modifications.

On the other hand, I think there will always be edge cases. Is it worth the effort and likely continuing trouble to have wide support, versus just telling people what they should use? This is only a tangential part of the functionality.
Good points. I think you're right that we shouldn't worry too much about parsing. However, I am inclined to let users specify a custom regex if they have a different style of modification.

>One more related thing is that Casanovo sometimes apparently predicts a double N-terminal acetylation, e.g. +43.006+43.006PEPTIDEK, in which case the +43.006+43.006 is not found in the AA mass dictionary. How should we handle such weird predictions that Casanovo comes up with? Currently mass calculation crashes with a KeyError, which doesn't seem desired behavior.

This would be easy to account for, except that carbamylation with an ammonia neutral loss `"+43.006-17.027"` is apparently a n-terminal modification that I found when looking at MassIVE-KB. Is that a mistake on my part?

More generally, it may be better to create an new token for every terminal modification + residue. 

Either way, ideally then we would restrict Casanovo to only predict the non-terminal tokens during inference. 

Thoughts?
        
> This would be easy to account for, except that carbamylation with an ammonia neutral loss "+43.006-17.027" is apparently a n-terminal modification that I found when looking at MassIVE-KB. Is that a mistake on my part?

No, I've seen it in the data as well and I never really worried about it because I figured it must be fine. Does it not make sense chemically?

> More generally, it may be better to create an new token for every terminal modification + residue.

You mean have `+43.006A`, `+43.006C`, `+43.006D`, ... ? That's going to expand our alphabet massively, and there likely won't be that many examples for each of those.
Although you indeed see only a single peak for the combination of the N-terminal modification and the first amino acid, rather than separate peaks for each, I think practically it's not the best solution.

> Either way, ideally then we would restrict Casanovo to only predict the non-terminal tokens during inference.

Indeed, I'd been thinking about that as well. It's as easy as masking out the indexes that correspond to N-terminal modifications after the first token has been predicted. But we should maybe think about how to represent that in the model. We can hardcode those indexes based on our current AA alphabet, but ideally the model should be able to do this dynamically. Some kind of flag in the AA alphabet config to indicate whether certain tokens are restricted to C- or N-term?
I actually had an idea to implement this better. Instead of just moving all of the preprocessing options from Casanovo to depthcharge, maybe it would be cleaner if the preprocessing happens on the depthcharge side, but it is done by giving it a callable with a specific signature (i.e. raw spectrum in, processed spectrum out, ValueError to discard the spectrum). That would make the preprocessing in depthcharge much more general, with each application that can fully define preprocessing based on its preferences. And it still solves the initial problem.

What do you think @wfondrie? The initial coding was done at the tail end of a 14h flight and I thought of this alternative while only sleeping 3 hours last night due to jet lag, so ymmv.
Nice! I think this is a good start. One of the things to consider is making even more flexible. Instead of having a fixed set of transformations, we could execute arbitrary functions, kind of like I setup for the Dataset classes:
https://github.com/wfondrie/depthcharge/blob/25d617f5ad421a5c0dcefd19628449293b414da3/depthcharge/data/datasets.py#L23-L26

If we did do that, I would think we'd actually want to move the `preprocessing_fn` to the parsing side of things, rather than as something the dataset does. The only reason I made this design choice originally is so that someone could re-use the same index and try out multiple preprocessing functions, obviously with the runtime cost of preprocessing. I'm open to arguments to the latter though!
Indeed, that was my proposal as well. I forgot that there's a `preprocessing_fn` in the dataloader already. Why did we reimplement basically the same thing in Casanovo then instead of using `preprocessing_fn` from depthcharge?

I think that moving `preprocessing_fn` to the parsers is a good solution. I agree that conceptually you might not want to recreate the index every time a simple preprocessing option changes, but in practice I doubt that's going to happen very often. (We don't change it in Casanovo, and for ANN-SoLo I've noticed the same and I'm making a similar change.) Whereas this change would avoid the issue with not being able to skip invalid items in a dataloader. So the upside outweighs the potential negative aspect imo.
> Why did we reimplement basically the same thing in Casanovo then instead of using preprocessing_fn from depthcharge?

To be honest, I don't know 🤷‍♂️ - I think originally Melih copied and pasted the dataset code from depthcharge because it needed some minor tweaking for casanovo, but I don't know if that's still the case. 

> I think that moving preprocessing_fn to the parsers is a good solution. I agree that conceptually you might not want to recreate the index every time a simple preprocessing option changes, but in practice I doubt that's going to happen very often. (We don't change it in Casanovo, and for ANN-SoLo I've noticed the same and I'm making a similar change.) Whereas this change would avoid the issue with not being able to skip invalid items in a dataloader. So the upside outweighs the potential negative aspect imo.

I'm fine with this change - I went back and forth debating with my original implementation anyway. 


Fix patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index dc6b724..cf2fcf8 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -8,12 +8,12 @@ repos:
   - id: trailing-whitespace
   - id: detect-private-key
 - repo: https://github.com/charliermarsh/ruff-pre-commit
-  rev: v0.0.257
+  rev: v0.0.269
   hooks:
     - id: ruff
-      args: ['--fix']
+      args: ['--show-fixes']
 - repo: https://github.com/psf/black
-  rev: 23.1.0
+  rev: 23.3.0
   hooks:
     - id: black
       language_version: python3.10
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 74cd617..6bf6860 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,19 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+### Added
+- Providing a proper tokenization class (also resolves #24 and #18)
+- First-class support for ProForma peptide annotations, thanks to `spectrum_utils` and `pyteomics`.
+- Adding primitive dataclasses for peptides, peptide ions, mass spectra ... and even small molecules 🚀
+- Adding type hints to everything and stricter linting with Ruff.
+- Adding a ton of tests.
+- Tight integration with `spectrum_utils` 💪
+
+### Changed
+- Moving preprocessing onto parsing instead of data loading (similar to @bittremieux's proposal in #31)
+- Combining the SpectrumIndex and SpectrumDataset classes into one.
+- Changing peak encodings. Instead of encoding the intensity using a linear projection and summing with the sinusoidal m/z encodings, now the intensity is also sinusoidally encoded and is combined with the sinusoidal m/z encodings using a linear layer.
+
 
 ## [v0.2.2] - 2023-05-15
 ### Fixed
diff --git a/codecov.yml b/codecov.yml
new file mode 100644
index 0000000..0d24ede
--- /dev/null
+++ b/codecov.yml
@@ -0,0 +1,6 @@
+# Only allow commits that increase or maintain test coverage.
+coverage:
+  status:
+    project:
+      default:
+        target: auto  # auto compares coverage to the previous base commit
diff --git a/depthcharge/__init__.py b/depthcharge/__init__.py
index 4a49993..295c778 100644
--- a/depthcharge/__init__.py
+++ b/depthcharge/__init__.py
@@ -1,7 +1,17 @@
-"""Initialize the depthcharge package"""
+"""Initialize the depthcharge package."""
+from . import (
+    data,
+    encoders,
+    feedforward,
+    tokenizers,
+    transformers,
+)
+from .primitives import (
+    MassSpectrum,
+    Molecule,
+    Peptide,
+    PeptideIons,
+)
 from .version import _get_version
-from . import components
-from . import data
-from . import utils
 
 __version__ = _get_version()
diff --git a/depthcharge/components/__init__.py b/depthcharge/components/__init__.py
deleted file mode 100644
index 45b2a99..0000000
--- a/depthcharge/components/__init__.py
+++ /dev/null
@@ -1,4 +0,0 @@
-"""Components for building cool models"""
-from .transformers import SpectrumEncoder, PeptideDecoder
-from .feedforward import FeedForward
-from .mixins import ModelMixin
diff --git a/depthcharge/components/mixins.py b/depthcharge/components/mixins.py
deleted file mode 100644
index d5c9335..0000000
--- a/depthcharge/components/mixins.py
+++ /dev/null
@@ -1,20 +0,0 @@
-"""Useful mixins for model classes"""
-import pandas as pd
-
-
-class ModelMixin:
-    """Add some useful methods for models."""
-
-    def init(self):
-        """Initialize the ModelMixin"""
-        self._history = []
-
-    @property
-    def history(self):
-        """The training history of a model."""
-        return pd.DataFrame(self._history)
-
-    @property
-    def n_parameters(self):
-        """The number of learnable parameters."""
-        return sum(p.numel() for p in self.parameters() if p.requires_grad)
diff --git a/depthcharge/components/transformers.py b/depthcharge/components/transformers.py
deleted file mode 100644
index a92c3dc..0000000
--- a/depthcharge/components/transformers.py
+++ /dev/null
@@ -1,459 +0,0 @@
-"""Base Transformer models for working with mass spectra and peptides"""
-import re
-
-import torch
-
-from .encoders import FloatEncoder, PeakEncoder, PositionalEncoder
-from ..masses import PeptideMass
-from .. import utils
-
-
-class SpectrumEncoder(torch.nn.Module):
-    """A Transformer encoder for input mass spectra.
-
-    Parameters
-    ----------
-    dim_model : int, optional
-        The latent dimensionality to represent peaks in the mass spectrum.
-    n_head : int, optional
-        The number of attention heads in each layer. ``dim_model`` must be
-        divisible by ``n_head``.
-    dim_feedforward : int, optional
-        The dimensionality of the fully connected layers in the Transformer
-        layers of the model.
-    n_layers : int, optional
-        The number of Transformer layers.
-    dropout : float, optional
-        The dropout probability for all layers.
-    peak_encoder : bool, optional
-        Use positional encodings m/z values of each peak.
-    dim_intensity: int or None, optional
-        The number of features to use for encoding peak intensity.
-        The remaining (``dim_model - dim_intensity``) are reserved for
-        encoding the m/z value.
-    """
-
-    def __init__(
-        self,
-        dim_model=128,
-        n_head=8,
-        dim_feedforward=1024,
-        n_layers=1,
-        dropout=0,
-        peak_encoder=True,
-        dim_intensity=None,
-    ):
-        """Initialize a SpectrumEncoder"""
-        super().__init__()
-
-        self.latent_spectrum = torch.nn.Parameter(torch.randn(1, 1, dim_model))
-
-        if peak_encoder:
-            self.peak_encoder = PeakEncoder(
-                dim_model,
-                dim_intensity=dim_intensity,
-            )
-        else:
-            self.peak_encoder = torch.nn.Linear(2, dim_model)
-
-        # The Transformer layers:
-        layer = torch.nn.TransformerEncoderLayer(
-            d_model=dim_model,
-            nhead=n_head,
-            dim_feedforward=dim_feedforward,
-            batch_first=True,
-            dropout=dropout,
-        )
-
-        self.transformer_encoder = torch.nn.TransformerEncoder(
-            layer,
-            num_layers=n_layers,
-        )
-
-    def forward(self, spectra):
-        """The forward pass.
-
-        Parameters
-        ----------
-        spectra : torch.Tensor of shape (n_spectra, n_peaks, 2)
-            The spectra to embed. Axis 0 represents a mass spectrum, axis 1
-            contains the peaks in the mass spectrum, and axis 2 is essentially
-            a 2-tuple specifying the m/z-intensity pair for each peak. These
-            should be zero-padded, such that all of the spectra in the batch
-            are the same length.
-
-        Returns
-        -------
-        latent : torch.Tensor of shape (n_spectra, n_peaks + 1, dim_model)
-            The latent representations for the spectrum and each of its
-            peaks.
-        mem_mask : torch.Tensor
-            The memory mask specifying which elements were padding in X.
-        """
-        zeros = ~spectra.sum(dim=2).bool()
-        mask = [
-            torch.tensor([[False]] * spectra.shape[0]).type_as(zeros),
-            zeros,
-        ]
-        mask = torch.cat(mask, dim=1)
-        peaks = self.peak_encoder(spectra)
-
-        # Add the spectrum representation to each input:
-        latent_spectra = self.latent_spectrum.expand(peaks.shape[0], -1, -1)
-
-        peaks = torch.cat([latent_spectra, peaks], dim=1)
-        return self.transformer_encoder(peaks, src_key_padding_mask=mask), mask
-
-    @property
-    def device(self):
-        """The current device for the model"""
-        return next(self.parameters()).device
-
-
-class _PeptideTransformer(torch.nn.Module):
-    """A transformer base class for peptide sequences.
-
-    Parameters
-    ----------
-    dim_model : int
-        The latent dimensionality to represent the amino acids in a peptide
-        sequence.
-    pos_encoder : bool
-        Use positional encodings for the amino acid sequence.
-    residues: Dict or str {"massivekb", "canonical"}, optional
-        The amino acid dictionary and their masses. By default this is only
-        the 20 canonical amino acids, with cysteine carbamidomethylated. If
-        "massivekb", this dictionary will include the modifications found in
-        MassIVE-KB. Additionally, a dictionary can be used to specify a custom
-        collection of amino acids and masses.
-    max_charge : int
-        The maximum charge to embed.
-    """
-
-    def __init__(
-        self,
-        dim_model,
-        pos_encoder,
-        residues,
-        max_charge,
-    ):
-        super().__init__()
-        self.reverse = False
-        self._peptide_mass = PeptideMass(residues=residues)
-        self._amino_acids = list(self._peptide_mass.masses.keys()) + ["$"]
-        self._idx2aa = {i + 1: aa for i, aa in enumerate(self._amino_acids)}
-        self._aa2idx = {aa: i for i, aa in self._idx2aa.items()}
-
-        if pos_encoder:
-            self.pos_encoder = PositionalEncoder(dim_model)
-        else:
-            self.pos_encoder = torch.nn.Identity()
-
-        self.charge_encoder = torch.nn.Embedding(max_charge, dim_model)
-        self.aa_encoder = torch.nn.Embedding(
-            len(self._amino_acids) + 1,
-            dim_model,
-            padding_idx=0,
-        )
-
-    def tokenize(self, sequence, partial=False):
-        """Transform a peptide sequence into tokens
-
-        Parameters
-        ----------
-        sequence : str
-            A peptide sequence.
-
-        Returns
-        -------
-        torch.Tensor
-            The token for each amino acid in the peptide sequence.
-        """
-        if not isinstance(sequence, str):
-            return sequence  # Assume it is already tokenized.
-
-        sequence = sequence.replace("I", "L")
-        sequence = re.split(r"(?<=.)(?=[A-Z])", sequence)
-        if self.reverse:
-            sequence = list(reversed(sequence))
-
-        if not partial:
-            sequence += ["$"]
-
-        tokens = [self._aa2idx[aa] for aa in sequence]
-        tokens = torch.tensor(tokens, device=self.device)
-        return tokens
-
-    def detokenize(self, tokens):
-        """Transform tokens back into a peptide sequence.
-
-        Parameters
-        ----------
-        tokens : torch.Tensor of shape (n_amino_acids,)
-            The token for each amino acid in the peptide sequence.
-
-        Returns
-        -------
-        list of str
-            The amino acids in the peptide sequence.
-        """
-        sequence = [self._idx2aa.get(i.item(), "") for i in tokens]
-        if "$" in sequence:
-            idx = sequence.index("$")
-            sequence = sequence[: idx + 1]
-
-        if self.reverse:
-            sequence = list(reversed(sequence))
-
-        return sequence
-
-    @property
-    def vocab_size(self):
-        """Return the number of amino acids"""
-        return len(self._aa2idx)
-
-    @property
-    def device(self):
-        """The current device for the model"""
-        return next(self.parameters()).device
-
-
-class PeptideEncoder(_PeptideTransformer):
-    """A transformer encoder for peptide sequences.
-
-    Parameters
-    ----------
-    dim_model : int
-        The latent dimensionality to represent the amino acids in a peptide
-        sequence.
-    n_head : int, optional
-        The number of attention heads in each layer. ``dim_model`` must be
-        divisible by ``n_head``.
-    dim_feedforward : int, optional
-        The dimensionality of the fully connected layers in the Transformer
-        layers of the model.
-    n_layers : int, optional
-        The number of Transformer layers.
-    dropout : float, optional
-        The dropout probability for all layers.
-    pos_encoder : bool, optional
-        Use positional encodings for the amino acid sequence.
-    residues: Dict or str {"massivekb", "canonical"}, optional
-        The amino acid dictionary and their masses. By default this is only
-        the 20 canonical amino acids, with cysteine carbamidomethylated. If
-        "massivekb", this dictionary will include the modifications found in
-        MassIVE-KB. Additionally, a dictionary can be used to specify a custom
-        collection of amino acids and masses.
-    max_charge : int, optional
-        The maximum charge state for peptide sequences.
-    """
-
-    def __init__(
-        self,
-        dim_model=128,
-        n_head=8,
-        dim_feedforward=1024,
-        n_layers=1,
-        dropout=0,
-        pos_encoder=True,
-        residues="canonical",
-        max_charge=5,
-    ):
-        """Initialize a PeptideEncoder"""
-        super().__init__(
-            dim_model=dim_model,
-            pos_encoder=pos_encoder,
-            residues=residues,
-            max_charge=max_charge,
-        )
-
-        # The Transformer layers:
-        layer = torch.nn.TransformerEncoderLayer(
-            d_model=dim_model,
-            nhead=n_head,
-            dim_feedforward=dim_feedforward,
-            batch_first=True,
-            dropout=dropout,
-        )
-
-        self.transformer_encoder = torch.nn.TransformerEncoder(
-            layer,
-            num_layers=n_layers,
-        )
-
-    def forward(self, sequences, charges):
-        """Predict the next amino acid for a collection of sequences.
-
-        Parameters
-        ----------
-        sequences : list of str or list of torch.Tensor of length batch_size
-            The partial peptide sequences for which to predict the next
-            amino acid. Optionally, these may be the token indices instead
-            of a string.
-        charges : torch.Tensor of size (batch_size,)
-            The charge state of the peptide
-
-        Returns
-        -------
-        latent : torch.Tensor of shape (n_sequences, len_sequence, dim_model)
-            The latent representations for the spectrum and each of its
-            peaks.
-        mem_mask : torch.Tensor
-            The memory mask specifying which elements were padding in X.
-        """
-        sequences = utils.listify(sequences)
-        tokens = [self.tokenize(s) for s in sequences]
-        tokens = torch.nn.utils.rnn.pad_sequence(tokens, batch_first=True)
-        encoded = self.aa_encoder(tokens)
-
-        # Encode charges
-        charges = self.charge_encoder(charges - 1)[:, None]
-        encoded = torch.cat([charges, encoded], dim=1)
-
-        # Create mask
-        mask = ~encoded.sum(dim=2).bool()
-
-        # Add positional encodings
-        encoded = self.pos_encoder(encoded)
-
-        # Run through the model:
-        latent = self.transformer_encoder(encoded, src_key_padding_mask=mask)
-        return latent, mask
-
-
-class PeptideDecoder(_PeptideTransformer):
-    """A transformer decoder for peptide sequences.
-
-    Parameters
-    ----------
-    dim_model : int, optional
-        The latent dimensionality to represent peaks in the mass spectrum.
-    n_head : int, optional
-        The number of attention heads in each layer. ``dim_model`` must be
-        divisible by ``n_head``.
-    dim_feedforward : int, optional
-        The dimensionality of the fully connected layers in the Transformer
-        layers of the model.
-    n_layers : int, optional
-        The number of Transformer layers.
-    dropout : float, optional
-        The dropout probability for all layers.
-    pos_encoder : bool, optional
-        Use positional encodings for the amino acid sequence.
-    reverse : bool, optional
-        Sequence peptides from c-terminus to n-terminus.
-    residues: Dict or str {"massivekb", "canonical"}, optional
-        The amino acid dictionary and their masses. By default this is only
-        the 20 canonical amino acids, with cysteine carbamidomethylated. If
-        "massivekb", this dictionary will include the modifications found in
-        MassIVE-KB. Additionally, a dictionary can be used to specify a custom
-        collection of amino acids and masses.
-    """
-
-    def __init__(
-        self,
-        dim_model=128,
-        n_head=8,
-        dim_feedforward=1024,
-        n_layers=1,
-        dropout=0,
-        pos_encoder=True,
-        reverse=True,
-        residues="canonical",
-        max_charge=5,
-    ):
-        """Initialize a PeptideDecoder"""
-        super().__init__(
-            dim_model=dim_model,
-            pos_encoder=pos_encoder,
-            residues=residues,
-            max_charge=max_charge,
-        )
-        self.reverse = reverse
-
-        # Additional model components
-        self.mass_encoder = FloatEncoder(dim_model)
-        layer = torch.nn.TransformerDecoderLayer(
-            d_model=dim_model,
-            nhead=n_head,
-            dim_feedforward=dim_feedforward,
-            batch_first=True,
-            dropout=dropout,
-        )
-
-        self.transformer_decoder = torch.nn.TransformerDecoder(
-            layer,
-            num_layers=n_layers,
-        )
-
-        self.final = torch.nn.Linear(dim_model, len(self._amino_acids) + 1)
-
-    def forward(self, sequences, precursors, memory, memory_key_padding_mask):
-        """Predict the next amino acid for a collection of sequences.
-
-        Parameters
-        ----------
-        sequences : list of str or list of torch.Tensor
-            The partial peptide sequences for which to predict the next
-            amino acid. Optionally, these may be the token indices instead
-            of a string.
-        precursors : torch.Tensor of size (batch_size, 2)
-            The measured precursor mass (axis 0) and charge (axis 1) of each
-            tandem mass spectrum
-        memory : torch.Tensor of shape (batch_size, n_peaks, dim_model)
-            The representations from a ``TransformerEncoder``, such as a
-           ``SpectrumEncoder``.
-        memory_key_padding_mask : torch.Tensor of shape (batch_size, n_peaks)
-            The mask that indicates which elements of ``memory`` are padding.
-
-        Returns
-        -------
-        scores : torch.Tensor of size (batch_size, len_sequence, n_amino_acids)
-            The raw output for the final linear layer. These can be Softmax
-            transformed to yield the probability of each amino acid for the
-            prediction.
-        tokens : torch.Tensor of size (batch_size, len_sequence)
-            The input padded tokens.
-
-        """
-        # Prepare sequences
-        if sequences is not None:
-            sequences = utils.listify(sequences)
-            tokens = [self.tokenize(s) for s in sequences]
-            tokens = torch.nn.utils.rnn.pad_sequence(tokens, batch_first=True)
-        else:
-            tokens = torch.tensor([[]]).to(self.device)
-
-        # Prepare mass and charge
-        masses = self.mass_encoder(precursors[:, None, 0])
-        charges = self.charge_encoder(precursors[:, 1].int() - 1)
-        precursors = masses + charges[:, None, :]
-
-        # Feed through model:
-        if sequences is None:
-            tgt = precursors
-        else:
-            tgt = torch.cat([precursors, self.aa_encoder(tokens)], dim=1)
-
-        tgt_key_padding_mask = tgt.sum(axis=2) == 0
-        tgt = self.pos_encoder(tgt)
-        tgt_mask = generate_tgt_mask(tgt.shape[1]).type_as(precursors)
-        preds = self.transformer_decoder(
-            tgt=tgt,
-            memory=memory,
-            tgt_mask=tgt_mask,
-            tgt_key_padding_mask=tgt_key_padding_mask,
-            memory_key_padding_mask=memory_key_padding_mask.to(self.device),
-        )
-        return self.final(preds), tokens
-
-
-def generate_tgt_mask(sz):
-    """Generate a square mask for the sequence.
-
-    Parameters
-    ----------
-    sz : int
-        The length of the target sequence.
-    """
-    return ~torch.triu(torch.ones(sz, sz, dtype=torch.bool)).transpose(0, 1)
diff --git a/depthcharge/constants.py b/depthcharge/constants.py
new file mode 100644
index 0000000..d9b6cf9
--- /dev/null
+++ b/depthcharge/constants.py
@@ -0,0 +1,6 @@
+"""Constants."""
+HYDROGEN = 1.007825035
+OXYGEN = 15.99491463
+H2O = 2 * HYDROGEN + OXYGEN
+PROTON = 1.00727646688
+C13 = 1.003355
diff --git a/depthcharge/data/__init__.py b/depthcharge/data/__init__.py
index 04ebe1f..304ed7a 100644
--- a/depthcharge/data/__init__.py
+++ b/depthcharge/data/__init__.py
@@ -1,9 +1,4 @@
-"""The Pytorch Datasets"""
+"""The Pytorch Datasets."""
 from . import preprocessing
-
-from .datasets import (
-    SpectrumDataset,
-    AnnotatedSpectrumDataset,
-)
-
-from .hdf5 import SpectrumIndex, AnnotatedSpectrumIndex
+from .peptide_datasets import PeptideDataset
+from .spectrum_datasets import AnnotatedSpectrumDataset, SpectrumDataset
diff --git a/depthcharge/data/datasets.py b/depthcharge/data/datasets.py
deleted file mode 100644
index fd69be0..0000000
--- a/depthcharge/data/datasets.py
+++ /dev/null
@@ -1,301 +0,0 @@
-"""A PyTorch Dataset class for annotated spectra."""
-import torch
-import numpy as np
-from torch.utils.data import Dataset
-
-from . import preprocessing
-from .. import utils
-
-
-class SpectrumDataset(Dataset):
-    """Parse and retrieve collections of mass spectra.
-
-    Parameters
-    ----------
-    spectrum_index : depthcharge.data.SpectrumIndex
-        The collection of spectra to use as a dataset.
-    n_peaks : int, optional
-        Keep only the top-n most intense peaks in any spectrum. ``None``
-        retains all of the peaks.
-    min_mz : float, optional
-        The minimum m/z to include. The default is 140 m/z, in order to
-        exclude TMT and iTRAQ reporter ions.
-    preprocessing_fn: Callable or list of Callable, optional
-        The function(s) used to preprocess the mass spectra. See the
-        preprocessing module for details. ``None``, the default, square root
-        transforms the intensities and scales them to unit norm.
-
-    Attributes
-    ----------
-    n_peaks : int
-        The maximum number of mass speak to consider for each mass spectrum.
-    min_mz : float
-        The minimum m/z to consider for each mass spectrum.
-    n_spectra : int
-    index : depthcharge.data.SpectrumIndex
-    """
-
-    def __init__(
-        self,
-        spectrum_index,
-        n_peaks=200,
-        min_mz=140,
-        preprocessing_fn=None,
-    ):
-        """Initialize a SpectrumDataset"""
-        super().__init__()
-        self.n_peaks = n_peaks
-        self.min_mz = min_mz
-        self.preprocessing_fn = preprocessing_fn
-        self._index = spectrum_index
-
-    def __len__(self):
-        """The number of spectra."""
-        return self.n_spectra
-
-    def __getitem__(self, idx):
-        """Return a single mass spectrum.
-
-        Parameters
-        ----------
-        idx : int
-            The index to return.
-
-        Returns
-        -------
-        spectrum : torch.Tensor of shape (n_peaks, 2)
-            The mass spectrum where ``spectrum[:, 0]`` are the m/z values and
-            ``spectrum[:, 1]`` are their associated intensities.
-        precursor_mz : float
-            The m/z of the precursor.
-        precursor_charge : int
-            The charge of the precursor.
-        """
-        batch = self.index[idx]
-        spec = self._process_peaks(*batch)
-        if not spec.sum():
-            spec = torch.tensor([[0, 1]]).float()
-
-        return spec, batch[2], batch[3]
-
-    def get_spectrum_id(self, idx):
-        """Return the identifier for a single mass spectrum.
-
-        Parameters
-        ----------
-        idx : int
-            The index of a mass spectrum within the SpectrumIndex.
-
-        Returns
-        -------
-        ms_data_file : str
-            The mass spectrometry data file from which the mass spectrum was
-            originally parsed.
-        identifier : str
-            The mass spectrum identifier, per PSI recommendations.
-        """
-        with self.index:
-            return self.index.get_spectrum_id(idx)
-
-    def _process_peaks(
-        self,
-        mz_array,
-        int_array,
-        precursor_mz,
-        precursor_charge,
-    ):
-        """Process each mass spectrum.
-
-        Parameters
-        ----------
-        mz_array : numpy.ndarray of shape (n_peaks,)
-            The m/z values of the peaks in the spectrum.
-        int_array : numpy.ndarray of shape (n_peaks,)
-            The intensity values of the peaks in the spectrum.
-        precursor_mz : float
-            The precursor m/z.
-        precursor_charge : int
-            The precursor charge.
-
-        Returns
-        -------
-        torch.Tensor of shape (n_peaks, 2)
-            The mass spectrum where column 0 are the m/z values and
-            column 1 are their associated intensities.
-        """
-        if self.min_mz is not None:
-            keep = mz_array >= self.min_mz
-            mz_array = mz_array[keep]
-            int_array = int_array[keep]
-
-        if len(int_array) > self.n_peaks:
-            top_p = np.argpartition(int_array, -self.n_peaks)[-self.n_peaks :]
-            top_p = np.sort(top_p)
-            mz_array = mz_array[top_p]
-            int_array = int_array[top_p]
-
-        mz_array = torch.tensor(mz_array)
-        int_array = torch.tensor(int_array)
-        for func in self.preprocessing_fn:
-            mz_array, int_array = func(
-                mz_array=mz_array,
-                int_array=int_array,
-                precursor_mz=precursor_mz,
-                precursor_charge=precursor_charge,
-            )
-
-        return torch.vstack([mz_array, int_array]).T.float()
-
-    @property
-    def n_spectra(self):
-        """The total number of spectra."""
-        return self.index.n_spectra
-
-    @property
-    def index(self):
-        """The underyling SpectrumIndex."""
-        return self._index
-
-    @property
-    def preprocessing_fn(self):
-        """The functions to preprocess each mass spectrum."""
-        return self._preprocessing_fn
-
-    @preprocessing_fn.setter
-    def preprocessing_fn(self, funcs):
-        """Set the preprocessing_fn"""
-        if funcs is None:
-            funcs = preprocessing.sqrt_and_norm
-
-        self._preprocessing_fn = utils.listify(funcs)
-
-    @staticmethod
-    def collate_fn(batch):
-        """This is the collate function for a SpectrumDataset.
-
-        The mass spectra must be padded so that they fit nicely as a tensor.
-        However, the padded elements are ignored during the subsequent steps.
-
-        Parameters
-        ----------
-        batch : tuple of tuple of torch.Tensor
-            A batch of data from an AnnotatedSpectrumDataset.
-
-        Returns
-        -------
-        spectra : torch.Tensor of shape (batch_size, n_peaks, 2)
-            The mass spectra to sequence, where ``X[:, :, 0]`` are the m/z
-            values and ``X[:, :, 1]`` are their associated intensities.
-        precursors : torch.Tensor of shape (batch_size, 2)
-            The precursor mass and charge state.
-        """
-        spec, mz, charge = list(zip(*batch))
-        charge = torch.tensor(charge)
-        mass = (torch.tensor(mz) - 1.007276) * charge
-        precursors = torch.vstack([mass, charge]).T.float()
-        spec = torch.nn.utils.rnn.pad_sequence(
-            spec,
-            batch_first=True,
-        )
-        return spec, precursors
-
-
-class AnnotatedSpectrumDataset(SpectrumDataset):
-    """Parse and retrieve collections of mass spectra
-
-    Parameters
-    ----------
-    annotated_spectrum_index : depthcharge.data.SpectrumIndex
-        The collection of annotated mass spectra to use as a dataset.
-    n_peaks : int, optional
-        Keep only the top-n most intense peaks in any spectrum. ``None``
-        retains all of the peaks.
-    min_mz : float, optional
-        The minimum m/z to include. The default is 140 m/z, in order to
-        exclude TMT and iTRAQ reporter ions.
-    preprocessing_fn: Callable or list of Callable, optional
-        The function(s) used to preprocess the mass spectra. See the
-        preprocessing module for details. ``None``, the default, square root
-        transforms the intensities and scales them to unit norm.
-
-    Attributes
-    ----------
-    n_peaks : int
-        The maximum number of mass speak to consider for each mass spectrum.
-    min_mz : float
-        The minimum m/z to consider for each mass spectrum.
-    n_spectra : int
-    index : depthcharge.data.AnnotatedSpectrumIndex
-    """
-
-    def __init__(
-        self,
-        annotated_spectrum_index,
-        n_peaks=200,
-        min_mz=140,
-        preprocessing_fn=None,
-    ):
-        """Initialize an AnnotatedSpectrumDataset"""
-        super().__init__(
-            annotated_spectrum_index,
-            n_peaks=n_peaks,
-            min_mz=min_mz,
-            preprocessing_fn=preprocessing_fn,
-        )
-
-    def __getitem__(self, idx):
-        """Return a single annotated mass spectrum.
-
-        Parameters
-        ----------
-        idx : int
-            The index to return.
-
-        Returns
-        -------
-        spectrum : torch.Tensor of shape (n_peaks, 2)
-            The mass spectrum where ``spectrum[:, 0]`` are the m/z values and
-            ``spectrum[:, 1]`` are their associated intensities.
-        precursor_mz : float
-            The m/z of the precursor.
-        precursor_charge : int
-            The charge of the precursor.
-        annotation : str
-            The annotation for the mass spectrum.
-        """
-        *batch, seq = self.index[idx]
-        spec = self._process_peaks(*batch)
-        if not spec.sum():
-            spec = torch.tensor([[0, 1]]).float()
-
-        return spec, batch[2], batch[3], seq
-
-    @staticmethod
-    def collate_fn(batch):
-        """This is the collate function for an AnnotatedSpectrumDataset.
-
-        The mass spectra must be padded so that they fit nicely as a tensor.
-        However, the padded elements are ignored during the subsequent steps.
-
-        Parameters
-        ----------
-        batch : tuple of tuple of torch.Tensor
-            A batch of data from an AnnotatedSpectrumDataset.
-
-        Returns
-        -------
-        spectra : torch.Tensor of shape (batch_size, n_peaks, 2)
-            The mass spectra to sequence, where ``X[:, :, 0]`` are the m/z
-            values and ``X[:, :, 1]`` are their associated intensities.
-        precursors : torch.Tensor of shape (batch_size, 2)
-            The precursor mass and charge state.
-        sequence : list of str
-            The peptide sequence annotations.
-
-        """
-        spec, mz, charge, seq = list(zip(*batch))
-        charge = torch.tensor(charge)
-        mass = (torch.tensor(mz) - 1.007276) * charge
-        precursors = torch.vstack([mass, charge]).T.float()
-        spec = torch.nn.utils.rnn.pad_sequence(spec, batch_first=True)
-        return spec, precursors, np.array(seq)
diff --git a/depthcharge/data/hdf5.py b/depthcharge/data/hdf5.py
deleted file mode 100644
index a3f3cbb..0000000
--- a/depthcharge/data/hdf5.py
+++ /dev/null
@@ -1,485 +0,0 @@
-"""Parse mass spectra into an HDF5 file format."""
-import logging
-from pathlib import Path
-
-import h5py
-import numpy as np
-
-from .. import utils
-from .parsers import MzmlParser, MzxmlParser, MgfParser
-
-LOGGER = logging.getLogger(__name__)
-
-
-class SpectrumIndex:
-    """Store and access a collection of mass spectra.
-
-    This class parses one or more mzML file, converts it to an HDF5 file
-    format. This allows depthcharge to access spectra from many different
-    files quickly and without loading them all into memory.
-
-    Parameters
-    ----------
-    index_path : str
-        The name and path of the HDF5 file index. If the path does
-        not contain the `.h5` or `.hdf5` extension, `.hdf5` will be added.
-    ms_data_files : str or list of str, optional
-        The mzML to include in this collection.
-    ms_level : int, optional
-        The level of tandem mass spectra to use.
-    valid_charge : Iterable[int], optional
-        Only consider spectra with the specified precursor charges. If `None`,
-        any precursor charge is accepted.
-    annotated : bool, optional
-        Whether or not the index contains spectrum annotations.
-    overwite : bool, optional
-        Overwrite previously indexed files? If ``False`` and new files are
-        provided, they will be appended to the collection.
-
-    Attributes
-    ----------
-    ms_files : list of str
-    path : Path
-    ms_level : int
-    valid_charge : Optional[Iterable[int]]
-    annotated : bool
-    overwrite : bool
-    n_spectra : int
-    n_peaks : int
-    """
-
-    def __init__(
-        self,
-        index_path,
-        ms_data_files=None,
-        ms_level=2,
-        valid_charge=None,
-        annotated=False,
-        overwrite=False,
-    ):
-        """Initialize a SpectrumIndex"""
-        index_path = Path(index_path)
-        if index_path.suffix not in [".h5", ".hdf5"]:
-            index_path = Path(str(index_path) + ".hdf5")
-
-        # Set attributes and check parameters:
-        self._path = index_path
-        self._ms_level = utils.check_positive_int(ms_level, "ms_level")
-        self._valid_charge = valid_charge
-        self._annotated = bool(annotated)
-        self._overwrite = bool(overwrite)
-        self._handle = None
-        self._file_offsets = np.array([0])
-        self._file_map = {}
-        self._locs = {}
-
-        # Create the file if it doesn't exist.
-        if not self.path.exists() or self.overwrite:
-            with h5py.File(self.path, "w") as index:
-                index.attrs["ms_level"] = self.ms_level
-                index.attrs["n_spectra"] = 0
-                index.attrs["n_peaks"] = 0
-                index.attrs["annotated"] = self.annotated
-
-        # Else, verify that the previous index uses the same ms_level.
-        else:
-            with self:
-                try:
-                    assert self._handle.attrs["ms_level"] == self.ms_level
-                    assert self._handle.attrs["annotated"] == self.annotated
-                except (KeyError, AssertionError):
-                    raise ValueError(
-                        f"'{self.path}' already exists, but was created with "
-                        "incompatible parameters. Use 'overwrite=True' to "
-                        "overwrite it."
-                    )
-
-                self._reindex()
-
-        # Now parse spectra.
-        if ms_data_files is not None:
-            ms_data_files = utils.listify(ms_data_files)
-            LOGGER.info("Reading %i files...", len(ms_data_files))
-            for ms_file in ms_data_files:
-                self.add_file(ms_file)
-
-    def _reindex(self):
-        """Update the file mappings and offsets."""
-        offsets = []
-        for idx in range(len(self._handle)):
-            grp = self._handle[str(idx)]
-            offsets.append(grp.attrs["n_spectra"])
-            self._file_map[grp.attrs["path"]] = idx
-
-        self._file_offsets = np.cumsum([0] + offsets)
-
-        # Build a map of 1D indices to 2D locations:
-        grp_idx = 0
-        for lin_idx in range(offsets[-1]):
-            grp_idx += lin_idx >= offsets[grp_idx]
-            row_idx = lin_idx - offsets[grp_idx - 1]
-            self._locs[lin_idx] = (grp_idx, row_idx)
-
-    def _get_parser(self, ms_data_file):
-        """Get the parser for the MS data file.
-
-        Parameters
-        ----------
-        ms_data_file : Path
-            The mass spectrometry data file to be parsed.
-
-        Returns
-        -------
-        MzmlParser, MzxmlParser, or MgfParser
-            The appropriate parser for the file.
-        """
-        kw_args = dict(ms_level=self.ms_level, valid_charge=self.valid_charge)
-
-        if ms_data_file.suffix.lower() == ".mzml":
-            return MzmlParser(ms_data_file, **kw_args)
-
-        if ms_data_file.suffix.lower() == ".mzxml":
-            return MzxmlParser(ms_data_file, **kw_args)
-
-        if ms_data_file.suffix.lower() == ".mgf":
-            return MgfParser(ms_data_file, **kw_args)
-
-        raise ValueError("Only mzML, mzXML, and MGF files are supported.")
-
-    def _assemble_metadata(self, parser):
-        """Assemble the metadata.
-
-        Parameters
-        ----------
-        parser : MzmlParser or MgfParser
-            The parser to use.
-
-        Returns
-        -------
-        numpy.ndarray of shape (n_spectra,)
-            The file metadata.
-        """
-        meta_types = [
-            ("precursor_mz", np.float32),
-            ("precursor_charge", np.uint8),
-            ("offset", np.uint64),
-            ("scan_id", np.uint32),
-        ]
-
-        metadata = np.empty(parser.n_spectra, dtype=meta_types)
-        metadata["precursor_mz"] = parser.precursor_mz
-        metadata["precursor_charge"] = parser.precursor_charge
-        metadata["offset"] = parser.offset
-        metadata["scan_id"] = parser.scan_id
-        return metadata
-
-    def add_file(self, ms_data_file):
-        """Add a mass spectrometry data file to the index.
-
-        Parameters
-        ----------
-        ms_data_file : str or Path
-            The mass spectrometry data file to add. It must be in an mzML or
-            MGF file format and use an ``.mzML``, ``.mzXML``, or ``.mgf`` file
-            extension.
-        """
-        ms_data_file = Path(ms_data_file)
-        if str(ms_data_file) in self._file_map:
-            return
-
-        # Read the file:
-        parser = self._get_parser(ms_data_file)
-        parser.read()
-
-        # Create the tables:
-        metadata = self._assemble_metadata(parser)
-
-        spectrum_types = [
-            ("mz_array", np.float64),
-            ("intensity_array", np.float32),
-        ]
-
-        spectra = np.zeros(parser.n_peaks, dtype=spectrum_types)
-        spectra["mz_array"] = parser.mz_arrays
-        spectra["intensity_array"] = parser.intensity_arrays
-
-        # Write the tables:
-        with h5py.File(self.path, "a") as index:
-            group_index = len(index)
-            group = index.create_group(str(group_index))
-            group.attrs["path"] = str(ms_data_file)
-            group.attrs["n_spectra"] = parser.n_spectra
-            group.attrs["n_peaks"] = parser.n_peaks
-            group.attrs["id_type"] = parser.id_type
-
-            # Update overall stats:
-            index.attrs["n_spectra"] += parser.n_spectra
-            index.attrs["n_peaks"] += parser.n_peaks
-
-            # Add the datasets:
-            group.create_dataset(
-                "metadata",
-                data=metadata,
-            )
-
-            group.create_dataset(
-                "spectra",
-                data=spectra,
-            )
-
-            try:
-                group.create_dataset(
-                    "annotations",
-                    data=parser.annotations,
-                    dtype=h5py.string_dtype(),
-                )
-            except (KeyError, AttributeError):
-                pass
-
-            self._file_map[str(ms_data_file)] = group_index
-            end_offset = self._file_offsets[-1] + parser.n_spectra
-            self._file_offsets = np.append(self._file_offsets, [end_offset])
-
-            # Update the locations:
-            grp_idx = len(self._file_offsets) - 2
-            for row_idx in range(parser.n_spectra):
-                lin_idx = row_idx + self._file_offsets[-2]
-                self._locs[lin_idx] = (grp_idx, row_idx)
-
-    def get_spectrum(self, idx):
-        """Access a mass spectrum.
-
-        Parameters
-        ----------
-        idx : int
-            The index of the index of the mass spectrum to look-up.
-
-        Returns
-        -------
-        tuple of numpy.ndarray
-            The m/z values, intensity values, precurosr m/z, precurosr charge,
-            and the spectrum annotation.
-        """
-        group_index, row_index = self._locs[idx]
-        if self._handle is None:
-            raise RuntimeError("Use the context manager for access.")
-
-        grp = self._handle[str(group_index)]
-        metadata = grp["metadata"]
-        spectra = grp["spectra"]
-        offsets = metadata["offset"][row_index : row_index + 2]
-
-        start_offset = offsets[0]
-        if offsets.shape[0] == 2:
-            stop_offset = offsets[1]
-        else:
-            stop_offset = spectra.shape[0]
-
-        spectrum = spectra[start_offset:stop_offset]
-        precursor = metadata[row_index]
-        out = (
-            spectrum["mz_array"],
-            spectrum["intensity_array"],
-            precursor["precursor_mz"],
-            precursor["precursor_charge"],
-        )
-
-        return out
-
-    def get_spectrum_id(self, idx):
-        """Get the identifier for a mass spectrum.
-
-        Parameters
-        ----------
-        idx : int
-            The index of the mass spectrum in the SpectrumIndex.
-
-        Returns
-        -------
-        ms_data_file : str
-            The mass spectrometry data file from which the mass spectrum was
-            originally parsed.
-        identifier : str
-            The mass spectrum identifier, per PSI recommendations.
-        """
-        group_index, row_index = self._locs[idx]
-        if self._handle is None:
-            raise RuntimeError("Use the context manager for access.")
-
-        grp = self._handle[str(group_index)]
-        ms_data_file = grp.attrs["path"]
-        identifier = grp["metadata"][row_index]["scan_id"]
-        prefix = grp.attrs["id_type"]
-        return ms_data_file, f"{prefix}={identifier}"
-
-    def __len__(self):
-        """The number of spectra in the index."""
-        return self._offsets[-1]
-
-    def __getitem__(self, idx):
-        """Access a mass spectrum.
-
-        Parameters
-        ----------
-        idx : int
-            The overall index of the mass spectrum to retrieve.
-
-        Returns
-        -------
-        tuple of numpy.ndarray
-            The m/z values, intensity values, precurosr m/z, precurosr charge,
-            and the annotation (if available).
-        """
-        if self._handle is None:
-            with self:
-                return self.get_spectrum(idx)
-
-        return self.get_spectrum(idx)
-
-    def __enter__(self):
-        """Open the index file for reading."""
-        self._handle = h5py.File(
-            self.path,
-            "r",
-            rdcc_nbytes=int(3e8),
-            rdcc_nslots=1024000,
-        )
-        return self
-
-    def __exit__(self, *args):
-        """Close the HDF5 file."""
-        self._handle.close()
-        self._handle = None
-
-    @property
-    def ms_files(self):
-        """The files currently in the index."""
-        return list(self._file_map.keys())
-
-    @property
-    def path(self):
-        """The path to the underyling HDF5 index file."""
-        return self._path
-
-    @property
-    def ms_level(self):
-        """The MS level of tandem mass spectra in the collection."""
-        return self._ms_level
-
-    @property
-    def valid_charge(self):
-        """Valid precursor charges for spectra to be included."""
-        return self._valid_charge
-
-    @property
-    def annotated(self):
-        """Whether or not the index contains spectrum annotations."""
-        return self._annotated
-
-    @property
-    def overwrite(self):
-        """Overwrite a previous index?"""
-        return self._overwrite
-
-    @property
-    def n_spectra(self):
-        """The total number of mass spectra in the index."""
-        if self._handle is None:
-            with self:
-                return self._handle.attrs["n_spectra"]
-
-        return self._handle.attrs["n_spectra"]
-
-    @property
-    def n_peaks(self):
-        """The total number of mass peaks in the index."""
-        if self._handle is None:
-            with self:
-                return self._handle.attrs["n_peaks"]
-
-        return self._handle.attrs["n_peaks"]
-
-
-class AnnotatedSpectrumIndex(SpectrumIndex):
-    """Store and access a collection of annotated mass spectra.
-
-    This class parses one or more mzML file, converts it to our HDF5 index
-    format. This allows us to access spectra from many different files quickly
-    and without loading them all into memory.
-
-    Parameters
-    ----------
-    index_path : str
-        The name and path of the HDF5 file index. If the path does
-        not contain the `.h5` or `.hdf5` extension, `.hdf5` will be added.
-    ms_data_files : str or list of str, optional
-        The MGF to include in this collection.
-    ms_level : int, optional
-        The level of tandem mass spectra to use.
-    valid_charge : Iterable[int], optional
-        Only consider spectra with the specified precursor charges. If `None`,
-        any precursor charge is accepted.
-    overwite : bool
-        Overwrite previously indexed files? If ``False`` and new files are
-        provided, they will be appended to the collection.
-
-    Attributes
-    ----------
-    ms_files : list of str
-    path : Path
-    ms_level : int
-    valid_charge : Optional[Iterable[int]]
-    overwrite : bool
-    n_spectra : int
-    n_peaks : int
-    """
-
-    def __init__(
-        self,
-        index_path,
-        ms_data_files=None,
-        ms_level=2,
-        valid_charge=None,
-        overwrite=False,
-    ):
-        """Initialize an AnnotatedSpectrumIndex"""
-        super().__init__(
-            index_path=index_path,
-            ms_data_files=ms_data_files,
-            ms_level=ms_level,
-            valid_charge=valid_charge,
-            annotated=True,
-            overwrite=overwrite,
-        )
-
-    def _get_parser(self, ms_data_file):
-        """Get the parser for the MS data file"""
-        if ms_data_file.suffix.lower() == ".mgf":
-            return MgfParser(
-                ms_data_file,
-                ms_level=self.ms_level,
-                annotations=True,
-            )
-
-        raise ValueError("Only MGF files are currently supported.")
-
-    def get_spectrum(self, idx):
-        """Access a mass spectrum.
-
-        Parameters
-        ----------
-        idx : int
-            The index of the mass spectrum in the AnnotatedSpectrumIndex.
-
-        Returns
-        -------
-        tuple of numpy.ndarray
-            The m/z values, intensity values, precurosr m/z, precurosr charge,
-            and the spectrum annotation.
-        """
-        spec_info = super().get_spectrum(idx)
-        group_index, row_index = self._locs[idx]
-        grp = self._handle[str(group_index)]
-        annotations = grp["annotations"]
-        spec_ann = annotations[row_index].decode()
-        return (*spec_info, spec_ann)
diff --git a/depthcharge/data/parsers.py b/depthcharge/data/parsers.py
index 9371e7e..ece6105 100644
--- a/depthcharge/data/parsers.py
+++ b/depthcharge/data/parsers.py
@@ -1,14 +1,20 @@
-"""Mass spectrometry data parsers"""
+"""Mass spectrometry data parsers."""
+from __future__ import annotations
+
 import logging
-from pathlib import Path
 from abc import ABC, abstractmethod
+from collections.abc import Callable, Iterable
+from os import PathLike
+from pathlib import Path
 
 import numpy as np
-from tqdm.auto import tqdm
+from pyteomics.mgf import MGF
 from pyteomics.mzml import MzML
 from pyteomics.mzxml import MzXML
-from pyteomics.mgf import MGF
+from tqdm.auto import tqdm
 
+from .. import utils
+from ..primitives import MassSpectrum
 
 LOGGER = logging.getLogger(__name__)
 
@@ -18,10 +24,12 @@ class BaseParser(ABC):
 
     Parameters
     ----------
-    ms_data_file : str or Path
+    ms_data_file : PathLike
         The mzML file to parse.
     ms_level : int
         The MS level of the spectra to parse.
+    preprocessing_fn : Callable or Iterable[Callable], optional
+        The function(s) used to preprocess the mass spectra.
     valid_charge : Iterable[int], optional
         Only consider spectra with the specified precursor charges. If `None`,
         any precursor charge is accepted.
@@ -31,14 +39,20 @@ class BaseParser(ABC):
 
     def __init__(
         self,
-        ms_data_file,
-        ms_level,
-        valid_charge=None,
-        id_type="scan",
-    ):
-        """Initialize the BaseParser"""
+        ms_data_file: PathLike,
+        ms_level: int,
+        preprocessing_fn: Callable | Iterable[Callable] | None = None,
+        valid_charge: Iterable[int] | None = None,
+        id_type: str = "scan",
+    ) -> None:
+        """Initialize the BaseParser."""
         self.path = Path(ms_data_file)
         self.ms_level = ms_level
+        if preprocessing_fn is None:
+            self.preprocessing_fn = []
+        else:
+            self.preprocessing_fn = utils.listify(preprocessing_fn)
+
         self.valid_charge = None if valid_charge is None else set(valid_charge)
         self.id_type = id_type
         self.offset = None
@@ -47,24 +61,28 @@ def __init__(
         self.scan_id = []
         self.mz_arrays = []
         self.intensity_arrays = []
+        self.annotations = None
 
     @abstractmethod
-    def open(self):
-        """Open the file as an iterable"""
-        pass
+    def open(self) -> Iterable:
+        """Open the file as an iterable."""
 
     @abstractmethod
-    def parse_spectrum(self, spectrum):
-        """Parse a single spectrum
+    def parse_spectrum(self, spectrum: dict) -> MassSpectrum | None:
+        """Parse a single spectrum.
 
         Parameters
         ----------
         spectrum : dict
             The dictionary defining the spectrum in a given format.
+
+        Returns
+        -------
+        MassSpectrum or None
+            The parsed mass spectrum or None if it is skipped.
         """
-        pass
 
-    def read(self):
+    def read(self) -> BaseParser:
         """Read the ms data file.
 
         Returns
@@ -75,7 +93,21 @@ def read(self):
         with self.open() as spectra:
             for spectrum in tqdm(spectra, desc=str(self.path), unit="spectra"):
                 try:
-                    self.parse_spectrum(spectrum)
+                    spectrum = self.parse_spectrum(spectrum)
+                    if spectrum is None:
+                        continue
+
+                    if self.preprocessing_fn is not None:
+                        for processor in self.preprocessing_fn:
+                            spectrum = processor(spectrum)
+
+                    self.mz_arrays.append(spectrum.mz)
+                    self.intensity_arrays.append(spectrum.intensity)
+                    self.precursor_mz.append(spectrum.precursor_mz)
+                    self.precursor_charge.append(spectrum.precursor_charge)
+                    self.scan_id.append(_parse_scan_id(spectrum.scan_id))
+                    if self.annotations is not None:
+                        self.annotations.append(spectrum.label)
                 except (IndexError, KeyError, ValueError):
                     n_skipped += 1
 
@@ -102,12 +134,12 @@ def read(self):
         return self
 
     @property
-    def n_spectra(self):
-        """The number of spectra"""
+    def n_spectra(self) -> int:
+        """The number of spectra."""
         return self.offset.shape[0]
 
     @property
-    def n_peaks(self):
+    def n_peaks(self) -> int:
         """The number of peaks in the file."""
         return self.mz_arrays.shape[0]
 
@@ -117,37 +149,36 @@ class MzmlParser(BaseParser):
 
     Parameters
     ----------
-    ms_data_file : str or Path
+    ms_data_file : PathLike
         The mzML file to parse.
-    ms_level : int
+    ms_level : int, optional
         The MS level of the spectra to parse.
+    preprocessing_fn : Callable or Iterable[Callable], optional
+        The function(s) used to preprocess the mass spectra.
     valid_charge : Iterable[int], optional
         Only consider spectra with the specified precursor charges. If `None`,
         any precursor charge is accepted.
     """
 
-    def __init__(self, ms_data_file, ms_level=2, valid_charge=None):
-        """Initialize the MzmlParser."""
-        super().__init__(
-            ms_data_file,
-            ms_level=ms_level,
-            valid_charge=valid_charge,
-        )
-
-    def open(self):
-        """Open the mzML file for reading"""
+    def open(self) -> Iterable[dict]:
+        """Open the mzML file for reading."""
         return MzML(str(self.path))
 
-    def parse_spectrum(self, spectrum):
+    def parse_spectrum(self, spectrum: dict) -> MassSpectrum | None:
         """Parse a single spectrum.
 
         Parameters
         ----------
         spectrum : dict
             The dictionary defining the spectrum in mzML format.
+
+        Returns
+        -------
+        MassSpectrum or None
+            The parsed mass spectrum or None if not at the correct MS level.
         """
         if spectrum["ms level"] != self.ms_level:
-            return
+            return None
 
         if self.ms_level > 1:
             precursor = spectrum["precursorList"]["precursor"][0]
@@ -163,13 +194,16 @@ def parse_spectrum(self, spectrum):
             precursor_mz, precursor_charge = None, 0
 
         if self.valid_charge is None or precursor_charge in self.valid_charge:
-            self.mz_arrays.append(spectrum["m/z array"])
-            self.intensity_arrays.append(spectrum["intensity array"])
-            self.precursor_mz.append(precursor_mz)
-            self.precursor_charge.append(precursor_charge)
-            self.scan_id.append(_parse_scan_id(spectrum["id"]))
-        else:
-            raise ValueError("Invalid precursor charge")
+            return MassSpectrum(
+                filename=str(self.path),
+                scan_id=spectrum["id"],
+                mz=spectrum["m/z array"],
+                intensity=spectrum["intensity array"],
+                precursor_mz=precursor_mz,
+                precursor_charge=precursor_charge,
+            )
+
+        raise ValueError("Invalid precursor charge")
 
 
 class MzxmlParser(BaseParser):
@@ -177,37 +211,36 @@ class MzxmlParser(BaseParser):
 
     Parameters
     ----------
-    ms_data_file : str or Path
+    ms_data_file : PathLike
         The mzXML file to parse.
-    ms_level : int
+    ms_level : int, optional
         The MS level of the spectra to parse.
+    preprocessing_fn : Callable or Iterable[Callable], optional
+        The function(s) used to preprocess the mass spectra.
     valid_charge : Iterable[int], optional
         Only consider spectra with the specified precursor charges. If `None`,
         any precursor charge is accepted.
     """
 
-    def __init__(self, ms_data_file, ms_level=2, valid_charge=None):
-        """Initialize the MzxmlParser."""
-        super().__init__(
-            ms_data_file,
-            ms_level=ms_level,
-            valid_charge=valid_charge,
-        )
-
-    def open(self):
-        """Open the mzXML file for reading"""
+    def open(self) -> Iterable[dict]:
+        """Open the mzXML file for reading."""
         return MzXML(str(self.path))
 
-    def parse_spectrum(self, spectrum):
+    def parse_spectrum(self, spectrum: dict) -> MassSpectrum | None:
         """Parse a single spectrum.
 
         Parameters
         ----------
         spectrum : dict
             The dictionary defining the spectrum in mzXML format.
+
+        Returns
+        -------
+        MassSpectrum
+            The parsed mass spectrum.
         """
         if spectrum["msLevel"] != self.ms_level:
-            return
+            return None
 
         if self.ms_level > 1:
             precursor = spectrum["precursorMz"][0]
@@ -217,13 +250,16 @@ def parse_spectrum(self, spectrum):
             precursor_mz, precursor_charge = None, 0
 
         if self.valid_charge is None or precursor_charge in self.valid_charge:
-            self.mz_arrays.append(spectrum["m/z array"])
-            self.intensity_arrays.append(spectrum["intensity array"])
-            self.precursor_mz.append(precursor_mz)
-            self.precursor_charge.append(precursor_charge)
-            self.scan_id.append(_parse_scan_id(spectrum["id"]))
-        else:
-            raise ValueError("Invalid precursor charge")
+            return MassSpectrum(
+                filename=str(self.path),
+                scan_id=spectrum["id"],
+                mz=spectrum["m/z array"],
+                intensity=spectrum["intensity array"],
+                precursor_mz=precursor_mz,
+                precursor_charge=precursor_charge,
+            )
+
+        raise ValueError("Invalid precursor charge")
 
 
 class MgfParser(BaseParser):
@@ -231,10 +267,12 @@ class MgfParser(BaseParser):
 
     Parameters
     ----------
-    ms_data_file : str or Path
+    ms_data_file : PathLike
         The MGF file to parse.
-    ms_level : int
+    ms_level : int, optional
         The MS level of the spectra to parse.
+    preprocessing_fn : Callable or Iterable[Callable], optional
+        The function(s) used to preprocess the mass spectra.
     valid_charge : Iterable[int], optional
         Only consider spectra with the specified precursor charges. If `None`,
         any precursor charge is accepted.
@@ -244,26 +282,30 @@ class MgfParser(BaseParser):
 
     def __init__(
         self,
-        ms_data_file,
-        ms_level=2,
-        valid_charge=None,
-        annotations=False,
-    ):
+        ms_data_file: PathLike,
+        ms_level: int = 2,
+        preprocessing_fn: Callable | Iterable[Callable] | None = None,
+        valid_charge: Iterable[int] | None = None,
+        annotations: bool = False,
+    ) -> None:
         """Initialize the MgfParser."""
         super().__init__(
             ms_data_file,
             ms_level=ms_level,
+            preprocessing_fn=preprocessing_fn,
             valid_charge=valid_charge,
             id_type="index",
         )
-        self.annotations = [] if annotations else None
+        if annotations:
+            self.annotations = []
+
         self._counter = -1
 
-    def open(self):
-        """Open the MGF file for reading"""
+    def open(self) -> Iterable[dict]:
+        """Open the MGF file for reading."""
         return MGF(str(self.path))
 
-    def parse_spectrum(self, spectrum):
+    def parse_spectrum(self, spectrum: dict) -> MassSpectrum:
         """Parse a single spectrum.
 
         Parameters
@@ -280,19 +322,25 @@ def parse_spectrum(self, spectrum):
             precursor_mz, precursor_charge = None, 0
 
         if self.annotations is not None:
-            self.annotations.append(spectrum["params"].get("seq"))
+            label = spectrum["params"].get("seq")
+        else:
+            label = None
 
         if self.valid_charge is None or precursor_charge in self.valid_charge:
-            self.mz_arrays.append(spectrum["m/z array"])
-            self.intensity_arrays.append(spectrum["intensity array"])
-            self.precursor_mz.append(precursor_mz)
-            self.precursor_charge.append(precursor_charge)
-            self.scan_id.append(self._counter)
-        else:
-            raise ValueError("Invalid precursor charge")
+            return MassSpectrum(
+                filename=str(self.path),
+                scan_id=self._counter,
+                mz=spectrum["m/z array"],
+                intensity=spectrum["intensity array"],
+                precursor_mz=precursor_mz,
+                precursor_charge=precursor_charge,
+                label=label,
+            )
+
+        raise ValueError("Invalid precursor charge")
 
 
-def _parse_scan_id(scan_str):
+def _parse_scan_id(scan_str: str | int) -> int:
     """Remove the string prefix from the scan ID.
 
     Adapted from:
diff --git a/depthcharge/data/peptide_datasets.py b/depthcharge/data/peptide_datasets.py
new file mode 100644
index 0000000..5fa476b
--- /dev/null
+++ b/depthcharge/data/peptide_datasets.py
@@ -0,0 +1,65 @@
+"""Datasets for working with peptide sequences."""
+from collections.abc import Iterable
+
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+
+from ..tokenizers import PeptideTokenizer
+
+
+class PeptideDataset(TensorDataset):
+    """A dataset for peptide sequences.
+
+    Parameters
+    ----------
+    tokenizer : PeptideTokenizer
+        A tokenizer specifying how to transform peptide sequences.
+        into tokens.
+    sequences : Iterable[str]
+        The peptide sequences in a format compatible with
+        your tokenizer. ProForma is preferred.
+    charges : torch.Tensor,
+        The charge state for each peptide.
+    *args : torch.Tensor, optional
+        Additional values to include during data loading.
+    """
+
+    def __init__(
+        self,
+        tokenizer: PeptideTokenizer,
+        sequences: Iterable[str],
+        charges: torch.Tensor,
+        *args: torch.Tensor,
+    ) -> None:
+        """Initialize a PeptideDataset."""
+        tokens = tokenizer.tokenize(sequences)
+        super().__init__(tokens, charges, *args)
+
+    @property
+    def tokens(self) -> torch.Tensor:
+        """The peptide sequence tokens."""
+        return self.tensors[0]
+
+    @property
+    def charges(self) -> torch.Tensor:
+        """The peptide charges."""
+        return self.tensors[1]
+
+    def loader(self, *args: tuple, **kwargs: dict) -> DataLoader:
+        """A PyTorch DataLoader for peptides.
+
+        Parameters
+        ----------
+        *args : tuple
+            Arguments passed initialize a torch.utils.data.DataLoader,
+            excluding ``dataset``.
+        **kwargs : dict
+            Keyword arguments passed initialize a torch.utils.data.DataLoader,
+            excluding ``dataset``.
+
+        Returns
+        -------
+        torch.utils.data.DataLoader
+            A DataLoader for the peptide.
+        """
+        return DataLoader(self, *args, **kwargs)
diff --git a/depthcharge/data/preprocessing.py b/depthcharge/data/preprocessing.py
index 065a2ed..ecf7ebd 100644
--- a/depthcharge/data/preprocessing.py
+++ b/depthcharge/data/preprocessing.py
@@ -1,26 +1,16 @@
 """Preprocessing functions for mass spectra.
 
 These functions can be used with datasets (`SpectrumDataset` and
-`AnnotatedSpectrumDataset`) or our `SpectrumDataModule`.
+`AnnotatedSpectrumDataset`).
 
 One or more preprocessing function can be applied to each mass spectrum during
-batch loading using the `preprocessing_fn` parameter of these classes. To apply
-preprocessing steps sequentially, merely pass a list of functions to this
+dataset preparation using the `preprocessing_fn` parameter of these classes. To
+apply preprocessing steps sequentially, pass a list of functions to this
 argument.
 
 We can also define custom preprocessing functions. All preprocessing functions
-must accept 4 keyword arguments:
-    - `mz_array` : torch.Tensor of shape (n_peaks,)
-        The m/z values of the peaks in the spectrum.
-    - `int_array` : torch.Tensor of shape (n_peaks,)
-        The intensity values of the peaks in the spectrum.
-    - `precursor_mz` : float
-        The precursor m/z.
-    - `precursor_charge` : int
-        The precursor charge.
-
-Preprocessing functions always return the processed `mz_array` and `int_array`.
-
+must accept a ``MassSpectrum`` as their only argument and return a
+``MassSpectrum``.
 
 ### Examples
 
@@ -30,14 +20,14 @@
 SpectrumDataset(
     ...,
     preprocessing_fn=[
-        preprocessing.remove_precursor_preak,
+        preprocessing.remove_precursor_peak,
         preprocessing.sqrt_and_norm,
     ],
 )
 ```
 
 Apply a custom function:
-```
+```Python
 def my_func(mz_array, int_array, precursor_mz, precursor_charge):
     return mz_array, torch.log(int_array)
 
@@ -48,70 +38,83 @@ def my_func(mz_array, int_array, precursor_mz, precursor_charge):
 ```
 
 """
-from typing import Tuple
-
-import torch
-
-
-def sqrt_and_norm(
-    mz_array: torch.Tensor,
-    int_array: torch.Tensor,
-    precursor_mz: float,
-    precursor_charge: int,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Square root the intensities and scale to unit norm.
-
-    Parameters
-    ----------
-    mz_array : torch.Tensor of shape (n_peaks,)
-        The m/z values of the peaks in the spectrum.
-    int_array : torch.Tensor of shape (n_peaks,)
-        The intensity values of the peaks in the spectrum.
-    precursor_mz : float
-        The precursor m/z.
-    precursor_charge : int
-        The precursor charge.
-
-    Returns
-    -------
-    mz_array : torch.Tensor of shape (n_peaks,)
-        The m/z values of the peaks in the spectrum.
-    int_array : torch.Tensor of shape (n_peaks,)
-        The intensity values of the peaks in the spectrum.
-    """
-    int_array = torch.sqrt(int_array)
-    int_array /= torch.linalg.norm(int_array)
-    return mz_array, int_array
-
-
-def remove_precursor_peak(
-    mz_array: torch.Tensor,
-    int_array: torch.Tensor,
-    precursor_mz: float,
-    tol: float = 1.5,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Remove peaks near the precursor m/z.
-
-    Parameters
-    ----------
-    mz_array : torch.Tensor of shape (n_peaks,)
-        The m/z values of the peaks in the spectrum.
-    int_array : torch.Tensor of shape (n_peaks,)
-        The intensity values of the peaks in the spectrum.
-    precursor_mz : float
-        The precursor m/z.
-    precursor_charge : int
-        The precursor charge.
-    tol : float, optional
-        The tolerance to remove the precursor peak in m/z.
-
-    Returns
-    -------
-    mz_array : torch.Tensor of shape (n_peaks,)
-        The m/z values of the peaks in the spectrum.
-    int_array : torch.Tensor of shape (n_peaks,)
-        The intensity values of the peaks in the spectrum.
-    """
-    bounds = (precursor_mz - tol, precursor_mz + tol)
-    outside = (mz_array >= bounds[0]) | (mz_array <= bounds[1])
-    return mz_array[outside], int_array[outside]
+from collections.abc import Callable
+from functools import wraps
+
+import numpy as np
+
+from ..primitives import MassSpectrum
+
+
+def scale_to_unit_norm(spectrum: MassSpectrum) -> MassSpectrum:
+    """Scale intensities to unit norm."""
+    spectrum._inner._intensity = (
+        spectrum.intensity / np.sqrt(spectrum.intensity**2).sum()
+    )
+    return spectrum
+
+
+def _spectrum_utils_fn(func: str) -> Callable:
+    """Wrap spectrum_utils.spectrum.MmsmsSpectrum preprocessing methods."""
+
+    @wraps(getattr(MassSpectrum, func))
+    def wrapper(
+        *args: tuple,
+        **kwargs: dict,
+    ) -> Callable:
+        """Wrapper for spectrum_utils MsmsSpectrum methods.
+
+        Parameters
+        ----------
+        func: Callable
+            The preprocessing function. This should exactly match a
+            MsmsSpectrum method.
+        *args: tuple
+            Arguments that are passed to the MsmsSpectrum method.
+        **kwargs : dict
+            Keyword arguments that are passed to the MsmsSpectrum method.
+
+        Returns
+        -------
+        Callable
+            A valid deptcharge preprocessing function.
+        """
+
+        @wraps(wrapper)
+        def preprocess(spec: MassSpectrum) -> MassSpectrum:
+            """The wrapped preprocessing function.
+
+            Parameters
+            ----------
+            spec : MassSpectrum
+                The mass spectrum to preprocess
+
+            Returns
+            -------
+            MassSpectrum
+                The processed mass spectrum.
+            """
+            # Call the spectrum_utils method:
+            getattr(spec, func)(*args, **kwargs)
+            return spec
+
+        return preprocess
+
+    return wrapper
+
+
+def _add_spectrum_utils_methods() -> None:
+    """Update this module with the spectrum_utils MsmsSpectrum methods."""
+    sus_methods = [
+        "filter_intensity",
+        "remove_precursor_peak",
+        "round",
+        "scale_intensity",
+        "set_mz_range",
+    ]
+
+    for method in sus_methods:
+        globals()[method] = _spectrum_utils_fn(method)
+
+
+_add_spectrum_utils_methods()
diff --git a/depthcharge/data/spectrum_datasets.py b/depthcharge/data/spectrum_datasets.py
new file mode 100644
index 0000000..05833f4
--- /dev/null
+++ b/depthcharge/data/spectrum_datasets.py
@@ -0,0 +1,711 @@
+"""Parse mass spectra into an HDF5 file format."""
+from __future__ import annotations
+
+import hashlib
+import logging
+import uuid
+from collections.abc import Callable, Iterable
+from os import PathLike
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any
+
+import dill
+import h5py
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, Dataset
+
+from .. import utils
+from ..primitives import MassSpectrum
+from ..tokenizers import PeptideTokenizer
+from . import preprocessing
+from .parsers import MgfParser, MzmlParser, MzxmlParser
+
+LOGGER = logging.getLogger(__name__)
+
+
+class SpectrumDataset(Dataset):
+    """Store and access a collection of mass spectra.
+
+    Parses one or more MS data files, adding the mass spectra to an HDF5
+    file that indexes the mass spectra from one or multiple files.
+    This allows depthcharge to access random mass spectra from many different
+    files quickly and without loading them all into memory.
+
+    Parameters
+    ----------
+    ms_data_files : PathLike or list of PathLike, optional
+        The mzML, mzXML, or MGF files to include in this collection. Files
+        can be added later using the ``.add_file()`` method.
+    ms_level : int, optional
+        The level of tandem mass spectra to use.
+    preprocessing_fn : Callable or Iterable[Callable], optional
+        The function(s) used to preprocess the mass spectra. ``None``,
+        the default, filters for the top 200 peaks above m/z 140,
+        square root transforms the intensities and scales them to unit norm.
+        See the preprocessing module for details and additional options.
+    valid_charge : Iterable[int], optional
+        Only consider spectra with the specified precursor charges. If `None`,
+        any precursor charge is accepted.
+    index_path : PathLike, optional.
+        The name and path of the HDF5 file index. If the path does
+        not contain the `.h5` or `.hdf5` extension, `.hdf5` will be added.
+        If ``None``, a file will be created in a temporary directory.
+    overwrite : bool, optional
+        Overwrite previously indexed files? If ``False`` and new files are
+        provided, they will be appended to the collection.
+
+    Attributes
+    ----------
+    ms_files : list of str
+    path : Path
+    ms_level : int
+    valid_charge : Optional[Iterable[int]]
+    annotated : bool
+    overwrite : bool
+    n_spectra : int
+    n_peaks : int
+    """
+
+    _annotated = False
+
+    def __init__(
+        self,
+        ms_data_files: PathLike | Iterable[PathLike] = None,
+        ms_level: int = 2,
+        preprocessing_fn: Callable | Iterable[Callable] | None = None,
+        valid_charge: Iterable[int] | None = None,
+        index_path: PathLike | None = None,
+        overwrite: bool = False,
+    ) -> None:
+        """Initialize a SpectrumIndex."""
+        self._tmpdir = None
+        if index_path is None:
+            # Create a random temporary file:
+            self._tmpdir = TemporaryDirectory()
+            index_path = Path(self._tmpdir.name) / f"{uuid.uuid4()}.hdf5"
+
+        index_path = Path(index_path)
+        if index_path.suffix not in [".h5", ".hdf5"]:
+            index_path = Path(index_path)
+            index_path = Path(str(index_path) + ".hdf5")
+
+        # Set attributes and check parameters:
+        self._path = index_path
+        self._ms_level = utils.check_positive_int(ms_level, "ms_level")
+        self._valid_charge = valid_charge
+        self._overwrite = bool(overwrite)
+        self._handle = None
+        self._file_offsets = np.array([0])
+        self._file_map = {}
+        self._locs = {}
+        self._offsets = None
+
+        if preprocessing_fn is not None:
+            self._preprocessing_fn = utils.listify(preprocessing_fn)
+        else:
+            self._preprocessing_fn = [
+                preprocessing.set_mz_range(min_mz=140),
+                preprocessing.filter_intensity(max_num_peaks=200),
+                preprocessing.scale_intensity(scaling="root"),
+                preprocessing.scale_to_unit_norm,
+            ]
+
+        # Create the file if it doesn't exist.
+        if not self.path.exists() or self.overwrite:
+            with h5py.File(self.path, "w") as index:
+                index.attrs["ms_level"] = self.ms_level
+                index.attrs["n_spectra"] = 0
+                index.attrs["n_peaks"] = 0
+                index.attrs["annotated"] = self.annotated
+                index.attrs["preprocessing_fn"] = _hash_obj(
+                    tuple(self.preprocessing_fn)
+                )
+        else:
+            self._validate_index()
+
+        # Now parse spectra.
+        if ms_data_files is not None:
+            ms_data_files = utils.listify(ms_data_files)
+            for ms_file in ms_data_files:
+                self.add_file(ms_file)
+
+    def _reindex(self) -> None:
+        """Update the file mappings and offsets."""
+        offsets = [0]
+        for idx in range(len(self._handle)):
+            grp = self._handle[str(idx)]
+            offsets.append(grp.attrs["n_spectra"])
+            self._file_map[grp.attrs["path"]] = idx
+
+        self._file_offsets = np.cumsum([0] + offsets)
+
+        # Build a map of 1D indices to 2D locations:
+        grp_idx = 0
+        for lin_idx in range(offsets[-1]):
+            grp_idx += lin_idx >= offsets[grp_idx + 1]
+            row_idx = lin_idx - offsets[grp_idx]
+            self._locs[lin_idx] = (grp_idx, row_idx)
+
+        self._offsets = None
+        _ = self.offsets  # Reinitialize the offsets.
+
+    def _validate_index(self) -> None:
+        """Validate that the index is appropriate for this dataset."""
+        preproc_hash = _hash_obj(tuple(self.preprocessing_fn))
+        with self:
+            try:
+                assert self._handle.attrs["ms_level"] == self.ms_level
+                assert self._handle.attrs["preprocessing_fn"] == preproc_hash
+                if self._annotated:
+                    assert self._handle.attrs["annotated"]
+            except (KeyError, AssertionError):
+                raise ValueError(
+                    f"'{self.path}' already exists, but was created with "
+                    "incompatible parameters. Use 'overwrite=True' to "
+                    "overwrite it."
+                )
+
+            self._reindex()
+
+    def _get_parser(
+        self,
+        ms_data_file: PathLike,
+    ) -> MzmlParser | MzxmlParser | MgfParser:
+        """Get the parser for the MS data file.
+
+        Parameters
+        ----------
+        ms_data_file : PathLike
+            The mass spectrometry data file to be parsed.
+
+        Returns
+        -------
+        MzmlParser, MzxmlParser, or MgfParser
+            The appropriate parser for the file.
+        """
+        kw_args = {
+            "ms_level": self.ms_level,
+            "valid_charge": self.valid_charge,
+            "preprocessing_fn": self.preprocessing_fn,
+        }
+
+        if ms_data_file.suffix.lower() == ".mzml":
+            return MzmlParser(ms_data_file, **kw_args)
+
+        if ms_data_file.suffix.lower() == ".mzxml":
+            return MzxmlParser(ms_data_file, **kw_args)
+
+        if ms_data_file.suffix.lower() == ".mgf":
+            return MgfParser(ms_data_file, **kw_args)
+
+        raise ValueError("Only mzML, mzXML, and MGF files are supported.")
+
+    def _assemble_metadata(
+        self,
+        parser: MzmlParser | MzxmlParser | MgfParser,
+    ) -> dict:
+        """Assemble the metadata.
+
+        Parameters
+        ----------
+        parser : MzmlParser, MzxmlParser, MgfParser
+            The parser to use.
+
+        Returns
+        -------
+        numpy.ndarray of shape (n_spectra,)
+            The file metadata.
+        """
+        meta_types = [
+            ("precursor_mz", np.float32),
+            ("precursor_charge", np.uint8),
+            ("offset", np.uint64),
+            ("scan_id", np.uint32),
+        ]
+
+        metadata = np.empty(parser.n_spectra, dtype=meta_types)
+        metadata["precursor_mz"] = parser.precursor_mz
+        metadata["precursor_charge"] = parser.precursor_charge
+        metadata["offset"] = parser.offset
+        metadata["scan_id"] = parser.scan_id
+        return metadata
+
+    def add_file(self, ms_data_file: PathLike) -> None:
+        """Add a mass spectrometry data file to the index.
+
+        Parameters
+        ----------
+        ms_data_file : PathLike
+            The mass spectrometry data file to add. It must be in an mzML or
+            MGF file format and use an ``.mzML``, ``.mzXML``, or ``.mgf`` file
+            extension.
+        """
+        ms_data_file = Path(ms_data_file)
+        if str(ms_data_file) in self._file_map:
+            return
+
+        # Invalidate current offsets:
+        self._offsets = None
+
+        # Read the file:
+        parser = self._get_parser(ms_data_file)
+        parser.read()
+
+        # Create the tables:
+        metadata = self._assemble_metadata(parser)
+
+        spectrum_types = [
+            ("mz_array", np.float64),
+            ("intensity_array", np.float32),
+        ]
+
+        spectra = np.zeros(parser.n_peaks, dtype=spectrum_types)
+        spectra["mz_array"] = parser.mz_arrays
+        spectra["intensity_array"] = parser.intensity_arrays
+
+        # Write the tables:
+        with h5py.File(self.path, "a") as index:
+            group_index = len(index)
+            group = index.create_group(str(group_index))
+            group.attrs["path"] = str(ms_data_file)
+            group.attrs["n_spectra"] = parser.n_spectra
+            group.attrs["n_peaks"] = parser.n_peaks
+            group.attrs["id_type"] = parser.id_type
+
+            # Update overall stats:
+            index.attrs["n_spectra"] += parser.n_spectra
+            index.attrs["n_peaks"] += parser.n_peaks
+
+            # Add the datasets:
+            group.create_dataset(
+                "metadata",
+                data=metadata,
+            )
+
+            group.create_dataset(
+                "spectra",
+                data=spectra,
+            )
+
+            try:
+                group.create_dataset(
+                    "annotations",
+                    data=parser.annotations,
+                    dtype=h5py.string_dtype(),
+                )
+            except (KeyError, AttributeError, TypeError):
+                pass
+
+            self._file_map[str(ms_data_file)] = group_index
+            end_offset = self._file_offsets[-1] + parser.n_spectra
+            self._file_offsets = np.append(self._file_offsets, [end_offset])
+
+            # Update the locations:
+            grp_idx = len(self._file_offsets) - 2
+            for row_idx in range(parser.n_spectra):
+                lin_idx = row_idx + self._file_offsets[-2]
+                self._locs[lin_idx] = (grp_idx, row_idx)
+
+    def get_spectrum(self, idx: int) -> MassSpectrum:
+        """Access a mass spectrum.
+
+        Parameters
+        ----------
+        idx : int
+            The index of the index of the mass spectrum to look-up.
+
+        Returns
+        -------
+        tuple of numpy.ndarray
+            The m/z values, intensity values, precurosr m/z, precurosr charge,
+            and the spectrum annotation.
+        """
+        group_index, row_index = self._locs[idx]
+        if self._handle is None:
+            raise RuntimeError("Use the context manager for access.")
+
+        grp = self._handle[str(group_index)]
+        metadata = grp["metadata"]
+        spectra = grp["spectra"]
+        offsets = self.offsets[str(group_index)][row_index : row_index + 2]
+
+        start_offset = offsets[0]
+        if offsets.shape[0] == 2:
+            stop_offset = offsets[1]
+        else:
+            stop_offset = spectra.shape[0]
+
+        spectrum = spectra[start_offset:stop_offset]
+        precursor = metadata[row_index]
+        return MassSpectrum(
+            filename=grp.attrs["path"],
+            scan_id=f"{grp.attrs['id_type']}={metadata[row_index]['scan_id']}",
+            mz=np.array(spectrum["mz_array"]),
+            intensity=np.array(spectrum["intensity_array"]),
+            precursor_mz=precursor["precursor_mz"],
+            precursor_charge=precursor["precursor_charge"],
+        )
+
+    def get_spectrum_id(self, idx: int) -> tuple[str, str]:
+        """Get the identifier for a mass spectrum.
+
+        Parameters
+        ----------
+        idx : int
+            The index of the mass spectrum in the SpectrumIndex.
+
+        Returns
+        -------
+        ms_data_file : str
+            The mass spectrometry data file from which the mass spectrum was
+            originally parsed.
+        identifier : str
+            The mass spectrum identifier, per PSI recommendations.
+        """
+        group_index, row_index = self._locs[idx]
+        if self._handle is None:
+            raise RuntimeError("Use the context manager for access.")
+
+        grp = self._handle[str(group_index)]
+        ms_data_file = grp.attrs["path"]
+        identifier = grp["metadata"][row_index]["scan_id"]
+        prefix = grp.attrs["id_type"]
+        return ms_data_file, f"{prefix}={identifier}"
+
+    def loader(self, *args: tuple, **kwargs: dict) -> DataLoader:
+        """A PyTorch DataLoader for the mass spectra.
+
+        Parameters
+        ----------
+        *args : tuple
+            Arguments passed initialize a torch.utils.data.DataLoader,
+            excluding ``dataset`` and ``collate_fn``.
+        **kwargs : dict
+            Keyword arguments passed initialize a torch.utils.data.DataLoader,
+            excluding ``dataset`` and ``collate_fn``.
+
+        Returns
+        -------
+        torch.utils.data.DataLoader
+            A DataLoader for the mass spectra.
+        """
+        return DataLoader(self, *args, collate_fn=self.collate_fn, **kwargs)
+
+    def __len__(self) -> int:
+        """The number of spectra in the index."""
+        return self.n_spectra
+
+    def __del__(self) -> None:
+        """Cleanup the temporary directory."""
+        if self._tmpdir is not None:
+            self._tmpdir.cleanup()
+
+    def __getitem__(self, idx: int) -> MassSpectrum:
+        """Access a mass spectrum.
+
+        Parameters
+        ----------
+        idx : int
+            The overall index of the mass spectrum to retrieve.
+
+        Returns
+        -------
+        tuple of numpy.ndarray
+            The m/z values, intensity values, precurosr m/z, precurosr charge,
+            and the annotation (if available).
+        """
+        if self._handle is None:
+            with self:
+                return self.get_spectrum(idx)
+
+        return self.get_spectrum(idx)
+
+    def __enter__(self) -> SpectrumDataset:
+        """Open the index file for reading."""
+        if self._handle is None:
+            self._handle = h5py.File(
+                self.path,
+                "r",
+                rdcc_nbytes=int(3e8),
+                rdcc_nslots=1024000,
+            )
+        return self
+
+    def __exit__(self, *args: str) -> None:
+        """Close the HDF5 file."""
+        self._handle.close()
+        self._handle = None
+
+    @property
+    def ms_files(self) -> list[str]:
+        """The files currently in the index."""
+        return list(self._file_map.keys())
+
+    @property
+    def path(self) -> Path:
+        """The path to the underyling HDF5 index file."""
+        return self._path
+
+    @property
+    def ms_level(self) -> int:
+        """The MS level of tandem mass spectra in the collection."""
+        return self._ms_level
+
+    @property
+    def preprocessing_fn(self) -> list[Callable]:
+        """The functions for preprocessing MS data."""
+        return self._preprocessing_fn
+
+    @property
+    def valid_charge(self) -> list[int]:
+        """Valid precursor charges for spectra to be included."""
+        return self._valid_charge
+
+    @property
+    def annotated(self) -> bool:
+        """Whether or not the index contains spectrum annotations."""
+        return self._annotated
+
+    @property
+    def overwrite(self) -> bool:
+        """Overwrite a previous index?."""
+        return self._overwrite
+
+    @property
+    def n_spectra(self) -> int:
+        """The total number of mass spectra in the index."""
+        if self._handle is None:
+            with self:
+                return self._handle.attrs["n_spectra"]
+
+        return self._handle.attrs["n_spectra"]
+
+    @property
+    def n_peaks(self) -> int:
+        """The total number of mass peaks in the index."""
+        if self._handle is None:
+            with self:
+                return self._handle.attrs["n_peaks"]
+
+        return self._handle.attrs["n_peaks"]
+
+    @property
+    def offsets(self) -> dict[str, np.array]:
+        """The offsets denoting where each spectrum starts."""
+        if self._offsets is not None:
+            return self._offsets
+
+        self._offsets = {
+            k: v["metadata"]["offset"] for k, v in self._handle.items()
+        }
+
+        return self._offsets
+
+    @staticmethod
+    def collate_fn(
+        batch: Iterable[MassSpectrum],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """The collate function for a SpectrumDataset.
+
+        The mass spectra must be padded so that they fit nicely as a tensor.
+        However, the padded elements are ignored during the subsequent steps.
+
+        Parameters
+        ----------
+        batch : tuple of tuple of torch.Tensor
+            A batch of data from an SpectrumDataset.
+
+        Returns
+        -------
+        spectra : torch.Tensor of shape (batch_size, n_peaks, 2)
+            The mass spectra to sequence, where ``X[:, :, 0]`` are the m/z
+            values and ``X[:, :, 1]`` are their associated intensities.
+        precursors : torch.Tensor of shape (batch_size, 2)
+            The precursor mass and charge state.
+        """
+        return _collate_fn(batch)[0:2]
+
+
+class AnnotatedSpectrumDataset(SpectrumDataset):
+    """Store and access a collection of annotated mass spectra.
+
+    This class parses one or more MGF file and converts it to our HDF5 index
+    format. This allows us to access spectra from many different files quickly
+    and without loading them all into memory.
+
+    Parameters
+    ----------
+    index_path : str
+        The name and path of the HDF5 file index. If the path does
+        not contain the `.h5` or `.hdf5` extension, `.hdf5` will be added.
+    ms_data_files : str or list of str, optional
+        The MGF to include in this collection.
+    ms_level : int, optional
+        The level of tandem mass spectra to use.
+    preprocessing_fn : Callable or Iterable[Callable], optional
+        The function(s) used to preprocess the mass spectra. ``None``,
+        the default, filters for the top 200 peaks above m/z 140,
+        square root transforms the intensities and scales them to unit norm.
+        See the preprocessing module for details and additional options.
+    valid_charge : Iterable[int], optional
+        Only consider spectra with the specified precursor charges. If `None`,
+        any precursor charge is accepted.
+    overwrite : bool
+        Overwrite previously indexed files? If ``False`` and new files are
+        provided, they will be appended to the collection.
+
+    Attributes
+    ----------
+    ms_files : list of str
+    path : Path
+    ms_level : int
+    valid_charge : Optional[Iterable[int]]
+    overwrite : bool
+    n_spectra : int
+    n_peaks : int
+    """
+
+    _annotated = True
+
+    def __init__(
+        self,
+        tokenizer: PeptideTokenizer,
+        ms_data_files: PathLike | Iterable[PathLike] = None,
+        ms_level: int = 2,
+        preprocessing_fn: Callable | Iterable[Callable] | None = None,
+        valid_charge: Iterable[int] | None = None,
+        index_path: PathLike | None = None,
+        overwrite: bool = False,
+    ) -> None:
+        """Initialize an AnnotatedSpectrumIndex."""
+        self.tokenizer = tokenizer
+        super().__init__(
+            ms_data_files=ms_data_files,
+            ms_level=ms_level,
+            preprocessing_fn=preprocessing_fn,
+            valid_charge=valid_charge,
+            index_path=index_path,
+            overwrite=overwrite,
+        )
+
+    def _get_parser(self, ms_data_file: str) -> MgfParser:
+        """Get the parser for the MS data file."""
+        if ms_data_file.suffix.lower() == ".mgf":
+            return MgfParser(
+                ms_data_file,
+                ms_level=self.ms_level,
+                annotations=True,
+                preprocessing_fn=self.preprocessing_fn,
+            )
+
+        raise ValueError("Only MGF files are currently supported.")
+
+    def get_spectrum(self, idx: int) -> MassSpectrum:
+        """Access a mass spectrum.
+
+        Parameters
+        ----------
+        idx : int
+            The index of the mass spectrum in the AnnotatedSpectrumIndex.
+
+        Returns
+        -------
+        MassSpectrum
+            The mass spectrum, labeled with its annotation.
+        """
+        spectrum = super().get_spectrum(idx)
+        group_index, row_index = self._locs[idx]
+        grp = self._handle[str(group_index)]
+        annotations = grp["annotations"]
+        spectrum.label = annotations[row_index].decode()
+        return spectrum
+
+    @property
+    def annotations(self) -> np.ndarray[str]:
+        """Retrieve all of the annotations in the index."""
+        annotations = []
+        for grp in self._handle.values():
+            try:
+                annotations.append(grp["annotations"])
+            except KeyError:
+                pass
+
+        return np.concatenate(annotations)
+
+    def collate_fn(
+        self,
+        batch: Iterable[MassSpectrum],
+    ) -> tuple[torch.Tensor, torch.Tensor, np.ndarray[str]]:
+        """The collate function for a AnnotatedSpectrumDataset.
+
+        The mass spectra must be padded so that they fit nicely as a tensor.
+        However, the padded elements are ignored during the subsequent steps.
+
+        Parameters
+        ----------
+        batch : tuple of tuple of torch.Tensor
+            A batch of data from an AnnotatedSpectrumDataset.
+
+        Returns
+        -------
+        spectra : torch.Tensor of shape (batch_size, n_peaks, 2)
+            The mass spectra to sequence, where ``X[:, :, 0]`` are the m/z
+            values and ``X[:, :, 1]`` are their associated intensities.
+        precursors : torch.Tensor of shape (batch_size, 2)
+            The precursor mass and charge state.
+        annotations : np.ndarray[str]
+            The spectrum annotations.
+        """
+        spectra, precursors, annotations = _collate_fn(batch)
+        tokens = self.tokenizer.tokenize(annotations)
+        return spectra, precursors, tokens
+
+
+def _collate_fn(
+    batch: Iterable[MassSpectrum],
+) -> tuple[torch.Tensor, torch.Tensor, list[str | None]]:
+    """The collate function for a SpectrumDataset.
+
+    The mass spectra must be padded so that they fit nicely as a tensor.
+    However, the padded elements are ignored during the subsequent steps.
+
+    Parameters
+    ----------
+    batch : tuple of tuple of torch.Tensor
+        A batch of data from an SpectrumDataset.
+
+    Returns
+    -------
+    spectra : torch.Tensor of shape (batch_size, n_peaks, 2)
+        The mass spectra to sequence, where ``X[:, :, 0]`` are the m/z
+        values and ``X[:, :, 1]`` are their associated intensities.
+    precursors : torch.Tensor of shape (batch_size, 2)
+        The precursor mass and charge state.
+    annotations : np.ndarray of shape (batch_size,)
+        The annotations, if any exist.
+    """
+    spectra = []
+    masses = []
+    charges = []
+    annotations = []
+    for spec in batch:
+        spectra.append(spec.to_tensor())
+        masses.append(spec.precursor_mass)
+        charges.append(spec.precursor_charge)
+        annotations.append(spec.label)
+
+    precursors = torch.vstack([torch.tensor(masses), torch.tensor(charges)])
+    spectra = torch.nn.utils.rnn.pad_sequence(
+        spectra,
+        batch_first=True,
+    )
+    return spectra, precursors.T.float(), annotations
+
+
+def _hash_obj(obj: Any) -> str:  # noqa: ANN401
+    """SHA1 hash for a picklable object."""
+    out = hashlib.sha1()
+    out.update(dill.dumps(obj))
+    return out.hexdigest()
diff --git a/depthcharge/encoders/__init__.py b/depthcharge/encoders/__init__.py
new file mode 100644
index 0000000..a1496dd
--- /dev/null
+++ b/depthcharge/encoders/__init__.py
@@ -0,0 +1,6 @@
+"""Avalailable encoders."""
+from .sinusoidal import (
+    FloatEncoder,
+    PeakEncoder,
+    PositionalEncoder,
+)
diff --git a/depthcharge/components/encoders.py b/depthcharge/encoders/sinusoidal.py
similarity index 57%
rename from depthcharge/components/encoders.py
rename to depthcharge/encoders/sinusoidal.py
index a578b80..219b54c 100644
--- a/depthcharge/components/encoders.py
+++ b/depthcharge/encoders/sinusoidal.py
@@ -1,9 +1,9 @@
 """Simple encoders for input into Transformers and the like."""
 import math
 
-import torch
 import einops
 import numpy as np
+import torch
 
 
 class FloatEncoder(torch.nn.Module):
@@ -11,16 +11,21 @@ class FloatEncoder(torch.nn.Module):
 
     Parameters
     ----------
-    dim_model : int
+    d_model : int
         The number of features to output.
-    min_wavelength : float
+    min_wavelength : float, optional
         The minimum wavelength to use.
-    max_wavelength : float
+    max_wavelength : float, optional
         The maximum wavelength to use.
     """
 
-    def __init__(self, dim_model, min_wavelength=0.001, max_wavelength=10000):
-        """Initialize the MassEncoder"""
+    def __init__(
+        self,
+        d_model: int,
+        min_wavelength: float = 0.001,
+        max_wavelength: float = 10000,
+    ) -> None:
+        """Initialize the MassEncoder."""
         super().__init__()
 
         # Error checking:
@@ -31,33 +36,31 @@ def __init__(self, dim_model, min_wavelength=0.001, max_wavelength=10000):
             raise ValueError("'max_wavelength' must be greater than 0.")
 
         # Get dimensions for equations:
-        d_sin = math.ceil(dim_model / 2)
-        d_cos = dim_model - d_sin
+        d_sin = math.ceil(d_model / 2)
+        d_cos = d_model - d_sin
 
         base = min_wavelength / (2 * np.pi)
         scale = max_wavelength / min_wavelength
         sin_exp = torch.arange(0, d_sin).float() / (d_sin - 1)
-        cos_exp = (torch.arange(d_sin, dim_model).float() - d_sin) / (
-            d_cos - 1
-        )
+        cos_exp = (torch.arange(d_sin, d_model).float() - d_sin) / (d_cos - 1)
         sin_term = base * (scale**sin_exp)
         cos_term = base * (scale**cos_exp)
 
         self.register_buffer("sin_term", sin_term)
         self.register_buffer("cos_term", cos_term)
 
-    def forward(self, X):
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
         """Encode m/z values.
 
         Parameters
         ----------
-        X : torch.Tensor of shape (batch_size, n_masses)
+        X : torch.Tensor of shape (batch_size, n_float)
             The masses to embed.
 
         Returns
         -------
-        torch.Tensor of shape (batch_size, n_masses, dim_model)
-            The encoded features for the mass spectra.
+        torch.Tensor of shape (batch_size, n_float, d_model)
+            The encoded features for the floating point numbers.
         """
         sin_mz = torch.sin(X[:, :, None] / self.sin_term)
         cos_mz = torch.cos(X[:, :, None] / self.cos_term)
@@ -69,63 +72,53 @@ class PeakEncoder(torch.nn.Module):
 
     Parameters
     ----------
-    dim_model : int
+    d_model : int
         The number of features to output.
-    dim_intensity : int, optional
-        The number of features to use for intensity. The remaining features
-        will be used to encode the m/z values.
-    min_wavelength : float, optional
-        The minimum wavelength to use.
-    max_wavelength : float, optional
-        The maximum wavelength to use.
-    learned_intensity_encoding : bool, optional
-        Use a learned intensity encoding as opposed to a sinusoidal encoding.
-        Note that for the sinusoidal encoding, this encoder expects values
-        between [0, 1].
+    min_mz_wavelength : float, optional
+        The minimum wavelength to use for m/z.
+    max_mz_wavelength : float, optional
+        The maximum wavelength to use for m/z.
+    min_intensity_wavelength : float, optional
+        The minimum wavelength to use for intensity. The default assumes
+        intensities between [0, 1].
+    max_intensity_wavelength : float, optional
+        The maximum wavelength to use for intensity. The default assumes
+        intensities between [0, 1].
     """
 
     def __init__(
         self,
-        dim_model,
-        dim_intensity=None,
-        min_wavelength=0.001,
-        max_wavelength=10000,
-        learned_intensity_encoding=True,
-    ):
-        """Initialize the MzEncoder"""
+        d_model: int,
+        min_mz_wavelength: float = 0.001,
+        max_mz_wavelength: float = 10000,
+        min_intensity_wavelength: float = 1e-6,
+        max_intensity_wavelength: float = 1,
+        *,
+        _legacy: bool = False,
+    ) -> None:
+        """Initialize the MzEncoder."""
         super().__init__()
-        self.dim_model = dim_model
-        self.dim_mz = dim_model
-        self.learned_intensity_encoding = learned_intensity_encoding
-        if dim_intensity is not None:
-            if dim_intensity >= dim_model:
-                raise ValueError(
-                    "'dim_intensity' must be less than 'dim_model'"
-                )
-
-            self.dim_mz -= dim_intensity
-            self.dim_intensity = dim_intensity
-        else:
-            self.dim_intensity = dim_model
+        self.d_model = d_model
+        self._legacy = _legacy
 
         self.mz_encoder = FloatEncoder(
-            dim_model=self.dim_mz,
-            min_wavelength=min_wavelength,
-            max_wavelength=max_wavelength,
+            d_model=self.d_model,
+            min_wavelength=min_mz_wavelength,
+            max_wavelength=max_mz_wavelength,
         )
 
-        if self.learned_intensity_encoding:
-            self.int_encoder = torch.nn.Linear(
-                1, self.dim_intensity, bias=False
-            )
-        else:
-            self.int_encoder = FloatEncoder(
-                dim_model=self.dim_intensity,
-                min_wavelength=1e-6,
-                max_wavelength=1,
-            )
-
-    def forward(self, X):
+        self.int_encoder = FloatEncoder(
+            d_model=self.d_model,
+            min_wavelength=min_intensity_wavelength,
+            max_wavelength=max_intensity_wavelength,
+        )
+        self.combiner = torch.nn.Linear(2 * d_model, d_model, bias=False)
+
+        # Likely removing this soon:
+        if self._legacy:
+            self.int_encoder = torch.nn.Linear(1, self.d_model, bias=False)
+
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
         """Encode m/z values and intensities.
 
         Note that we expect intensities to fall within the interval [0, 1].
@@ -141,22 +134,23 @@ def forward(self, X):
 
         Returns
         -------
-        torch.Tensor of shape (n_spectra, n_peaks, dim_model)
+        torch.Tensor of shape (n_spectra, n_peaks, d_model)
             The encoded features for the mass spectra.
         """
-        m_over_z = X[:, :, 0]
-        encoded = self.mz_encoder(m_over_z)
-
-        if self.learned_intensity_encoding:
-            int_input = X[:, :, [1]]
-        else:
-            int_input = X[:, :, 1]
-
-        intensity = self.int_encoder(int_input)
-        if self.dim_intensity == self.dim_model:
-            return encoded + intensity
+        if self._legacy:
+            encoded_mz = self.mz_encoder(X[:, :, 0])
+            encoded_int = self.int_encoder(X[:, :, [1]])
+            return encoded_mz + encoded_int
+
+        encoded = torch.cat(
+            [
+                self.mz_encoder(X[:, :, 0]),
+                self.int_encoder(X[:, :, 1]),
+            ],
+            dim=2,
+        )
 
-        return torch.cat([encoded, intensity], dim=2)
+        return self.combiner(encoded)
 
 
 class PositionalEncoder(FloatEncoder):
@@ -164,7 +158,7 @@ class PositionalEncoder(FloatEncoder):
 
     Parameters
     ----------
-    dim_model : int
+    d_model : int
         The number of features to output.
     min_wavelength : float, optional
         The shortest wavelength in the geometric progression.
@@ -172,15 +166,20 @@ class PositionalEncoder(FloatEncoder):
         The longest wavelength in the geometric progression.
     """
 
-    def __init__(self, dim_model, min_wavelength=1, max_wavelength=10000):
-        """Initialize the MzEncoder"""
+    def __init__(
+        self,
+        d_model: int,
+        min_wavelength: float = 1.0,
+        max_wavelength: float = 1e5,
+    ) -> None:
+        """Initialize the MzEncoder."""
         super().__init__(
-            dim_model=dim_model,
+            d_model=d_model,
             min_wavelength=min_wavelength,
             max_wavelength=max_wavelength,
         )
 
-    def forward(self, X):
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
         """Encode positions in a sequence.
 
         Parameters
diff --git a/depthcharge/components/feedforward.py b/depthcharge/feedforward.py
similarity index 53%
rename from depthcharge/components/feedforward.py
rename to depthcharge/feedforward.py
index 4c35135..a3814d6 100644
--- a/depthcharge/components/feedforward.py
+++ b/depthcharge/feedforward.py
@@ -1,16 +1,18 @@
 """A flexible feed-forward neural network."""
-import torch
+from collections.abc import Iterable
+
 import numpy as np
+import torch
 
 
 class FeedForward(torch.nn.Module):
-    """Create a feed forward neural net with leaky GELU activations
+    """Create a feed forward neural net with leaky GELU activations.
 
     Parameters
     ----------
-    in_dim : int
+    in_features : int
         The input dimensionality.
-    out_dim : int
+    out_features : int
         The output dimensionality.
     layers : int or tuple of int.
         If an int, layer sizes are linearly interpolated between the input and
@@ -19,24 +21,38 @@ class FeedForward(torch.nn.Module):
     dropout : float, optional
         If greater than zero, add dropout layers with the specified
         probability.
+    activation: torch.nn.Module, optional
+        The activation function to place between layers.
     append : torch.nn.Module or None, optional
         A final layer to append, such as a sigmoid or tanh.
     """
 
-    def __init__(self, in_dim, out_dim, layers, dropout=0, append=None):
-        """Initiazlize a FeedForward network"""
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        layers: int | Iterable[int],
+        dropout: float = 0,
+        activation: torch.nn.Module = torch.nn.LeakyReLU(),
+        append: torch.nn.Module | None = None,
+    ) -> None:
+        """Initiazlize a FeedForward network."""
         super().__init__()
         try:
-            sizes = np.array([in_dim] + list(layers) + [out_dim])
+            sizes = np.array(
+                [in_features] + list(layers) + [out_features],
+            )
         except TypeError:
-            sizes = np.ceil(np.linspace(in_dim, out_dim, int(layers) + 1))
+            sizes = np.ceil(
+                np.linspace(in_features, out_features, int(layers) + 1),
+            )
 
         sizes = sizes.astype(int)
         stack = []
         for idx in range(len(sizes) - 1):
             stack.append(torch.nn.Linear(sizes[idx], sizes[idx + 1]))
             if idx < len(sizes) - 2:
-                stack.append(torch.nn.LeakyReLU())
+                stack.append(activation)
                 if dropout:
                     stack.append(torch.nn.Dropout(dropout))
 
@@ -45,6 +61,17 @@ def __init__(self, in_dim, out_dim, layers, dropout=0, append=None):
 
         self.layers = torch.nn.Sequential(*stack)
 
-    def forward(self, X):
-        """The forward pass"""
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        """The forward pass.
+
+        Parameters
+        ----------
+        X : torch.Tensor of shape (..., in_features)
+            The input tensor.
+
+        Returns
+        -------
+        torch.Tensor of shape (..., out_features)
+            The output tensor.
+        """
         return self.layers(X)
diff --git a/depthcharge/masses.py b/depthcharge/masses.py
deleted file mode 100644
index 57ea262..0000000
--- a/depthcharge/masses.py
+++ /dev/null
@@ -1,99 +0,0 @@
-"""Amino acid masses and other useful mass spectrometry calculations"""
-import re
-
-
-class PeptideMass:
-    """A simple class for calculating peptide masses
-
-    Parameters
-    ----------
-    residues: Dict or str {"massivekb", "canonical"}, optional
-        The amino acid dictionary and their masses. By default this is only
-        the 20 canonical amino acids, with cysteine carbamidomethylated. If
-        "massivekb", this dictionary will include the modifications found in
-        MassIVE-KB. Additionally, a dictionary can be used to specify a custom
-        collection of amino acids and masses.
-    """
-
-    canonical = {
-        "G": 57.021463735,
-        "A": 71.037113805,
-        "S": 87.032028435,
-        "P": 97.052763875,
-        "V": 99.068413945,
-        "T": 101.047678505,
-        "C+57.021": 103.009184505 + 57.02146,
-        "L": 113.084064015,
-        "I": 113.084064015,
-        "N": 114.042927470,
-        "D": 115.026943065,
-        "Q": 128.058577540,
-        "K": 128.094963050,
-        "E": 129.042593135,
-        "M": 131.040484645,
-        "H": 137.058911875,
-        "F": 147.068413945,
-        # "U": 150.953633405,
-        "R": 156.101111050,
-        "Y": 163.063328575,
-        "W": 186.079312980,
-        # "O": 237.147726925,
-    }
-
-    # Modfications found in MassIVE-KB
-    massivekb = {
-        # N-terminal mods:
-        "+42.011": 42.010565,  # Acetylation
-        "+43.006": 43.005814,  # Carbamylation
-        "-17.027": -17.026549,  # NH3 loss
-        "+43.006-17.027": (43.006814 - 17.026549),
-        # AA mods:
-        "M+15.995": canonical["M"] + 15.994915,  # Met Oxidation
-        "N+0.984": canonical["N"] + 0.984016,  # Asn Deamidation
-        "Q+0.984": canonical["Q"] + 0.984016,  # Gln Deamidation
-    }
-
-    # Constants
-    hydrogen = 1.007825035
-    oxygen = 15.99491463
-    h2o = 2 * hydrogen + oxygen
-    proton = 1.00727646688
-
-    def __init__(self, residues="canonical"):
-        """Initialize the PeptideMass object"""
-        if residues == "canonical":
-            self.masses = self.canonical
-        elif residues == "massivekb":
-            self.masses = self.canonical
-            self.masses.update(self.massivekb)
-        else:
-            self.masses = residues
-
-    def __len__(self):
-        """Return the length of the residue dictionary"""
-        return len(self.masses)
-
-    def mass(self, seq, charge=None):
-        """Calculate a peptide's mass or m/z.
-
-        Parameters
-        ----------
-        seq : list or str
-            The peptide sequence, using tokens defined in ``self.residues``.
-        charge : int, optional
-            The charge used to compute m/z. Otherwise the neutral peptide mass
-            is calculated
-
-        Returns
-        -------
-        float
-            The computed mass or m/z.
-        """
-        if isinstance(seq, str):
-            seq = re.split(r"(?<=.)(?=[A-Z])", seq)
-
-        calc_mass = sum([self.masses[aa] for aa in seq]) + self.h2o
-        if charge is not None:
-            calc_mass = (calc_mass / charge) + self.proton
-
-        return calc_mass
diff --git a/depthcharge/primitives.py b/depthcharge/primitives.py
new file mode 100644
index 0000000..2c3ff26
--- /dev/null
+++ b/depthcharge/primitives.py
@@ -0,0 +1,423 @@
+"""Fundamental dataclasses for depthcharge."""
+from __future__ import annotations
+
+import re
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import numpy as np
+import selfies as sf
+import torch
+from numpy.typing import ArrayLike
+from PIL.PngImagePlugin import PngImageFile
+from pyteomics import proforma
+from pyteomics.proforma import GenericModification, MassModification
+from rdkit import Chem
+from rdkit.Chem import Draw
+from spectrum_utils.spectrum import MsmsSpectrum
+
+from .constants import PROTON
+
+MSKB_TO_UNIMOD = {
+    "+42.011": "[Acetyl]-",
+    "+43.006": "[Carbamyl]-",
+    "-17.027": "[Ammonia-loss]-",
+    "+43.006-17.027": "[+25.980265]-",  # Not in Unimod
+    "M+15.995": "M[Oxidation]",
+    "N+0.984": "N[Deamidated]",
+    "Q+0.984": "Q[Deamidated]",
+    "C+57.021": "C[Carbamidomethyl]",
+}
+
+
+@dataclass
+class Peptide:
+    """A peptide sequence with or without modification and/or charge.
+
+    Parameters
+    ----------
+    sequence : str
+        The bare amino acid sequence.
+    modifications : iterable of str, float, or None, optional
+        The modification at each amino acid. This should be the length of
+        ``sequence`` + 2, where index 0 and -1 are used for N- and C-terminal
+        modification respectively. Use ``None`` in the iterable to specify
+        unmodified positions. When ``modifications`` is ``None``, it is assumed
+        that no modifications are present.
+    charge : int, optional
+        The charge of the peptide.
+    """
+
+    sequence: str
+    modifications: Sequence[str | float | None] | None = None
+    charge: int | None = None
+
+    def __post_init__(self) -> None:
+        """Validate the provided parameters."""
+        if self.modifications is not None:
+            if len(self.modifications) != len(self.sequence) + 2:
+                raise ValueError(
+                    "'modifications' must be two elements longer than "
+                    "'sequences' to account for terminal modifications."
+                )
+
+        if self.modifications is None:
+            self.modifications = [None] * (len(self.sequence) + 2)
+
+        # Parse modifications into Pyteomics' format
+        parsed = [None] * (len(self.sequence) + 2)
+        for idx, mod in enumerate(self.modifications):
+            if mod is None:
+                continue
+
+            try:
+                mod = [MassModification(mod)]
+            except ValueError:
+                try:
+                    mod = [GenericModification(mod)]
+                except (AttributeError, TypeError):
+                    pass
+            except TypeError:
+                pass
+
+            parsed[idx] = mod
+
+        self.modifications = parsed
+        n_mod = self.modifications[0]
+        c_mod = self.modifications[-1]
+
+        self.proforma = proforma.to_proforma(
+            sequence=list(zip(self.sequence, self.modifications[1:-1])),
+            n_term=[n_mod] if n_mod is not None else n_mod,
+            c_term=[c_mod] if c_mod is not None else c_mod,
+            charge_state=self.charge,
+        )
+
+    def split(self) -> list[str]:
+        """Split the modified peptide for tokenization."""
+        if self.modifications is None:
+            return list(self.sequence)
+
+        out = []
+        for idx, (aa, mods) in enumerate(
+            zip(f"-{self.sequence}-", self.modifications)
+        ):
+            if mods is None:
+                if idx and (idx < len(self.modifications) - 1):
+                    out.append(aa)
+
+                continue
+
+            if len(mods) == 1:
+                try:
+                    modstr = f"[{mods[0].name}]"
+                except (AttributeError, ValueError):
+                    modstr = f"[{mods[0].mass:+0.6f}]"
+            else:
+                modstr = f"[{sum([m.mass for m in mods]):+0.6f}]"
+
+            if not idx:
+                out.append(f"{modstr}-")
+            else:
+                out.append(f"{aa}{modstr}")
+
+        return out
+
+    @classmethod
+    def from_proforma(
+        cls,
+        sequence: str,
+    ) -> Peptide:
+        """Create a Peptide from a ProForma 2.0 string.
+
+        Parameters
+        ----------
+        sequence : str
+            A ProForma 2.0-compliant string.
+
+        Returns
+        -------
+        Peptide
+            The parsed ProForma peptide.
+        """
+        pep, meta = proforma.parse(sequence)
+        try:
+            charge = meta["charge_state"].charge
+        except AttributeError:
+            charge = None
+
+        static_mods = {}
+        for mod_rule in meta["fixed_modifications"]:
+            for res in mod_rule.targets:
+                static_mods[res] = mod_rule.modification_tag
+
+        seq = [None] * len(pep)
+        mods = [None] * (len(pep) + 2)
+        mods[0] = meta["n_term"]
+        mods[-1] = meta["c_term"]
+        for idx, (res, var_mods) in enumerate(pep):
+            seq[idx] = res
+            if res in static_mods:
+                if var_mods is None:
+                    var_mods = []
+
+                var_mods.insert(0, static_mods[res])
+
+            mods[idx + 1] = var_mods
+
+        return cls("".join(seq), mods, charge)
+
+    @classmethod
+    def from_massivekb(
+        cls,
+        sequence: str,
+        charge: int | None = None,
+    ) -> Peptide:
+        """Create a Peptide from MassIVE-KB annotations.
+
+        MassIVE-KB includes N-term carbamylation, NH3-loss, acetylation,
+        as well as M oxidation, and deamidation of N and Q, in a
+        manner that does not comply with the ProForma standard.
+
+        Parameters
+        ----------
+        sequence : str
+            The peptide sequence from MassIVE-KB
+        charge : int, optional
+            The charge state of the peptide.
+
+        Returns
+        -------
+        Peptide
+            The parsed MassIVE peptide after conversion to a ProForma
+            format.
+        """
+        sequence = cls.massivekb_to_proforma(sequence, charge)
+        return cls.from_proforma(sequence)
+
+    @classmethod
+    def massivekb_to_proforma(
+        cls, sequence: str, charge: int | None = None
+    ) -> str:
+        """Convert a MassIVE-KB peptide sequence to ProForma.
+
+        MassIVE-KB includes N-term carbamylation, NH3-loss, acetylation,
+        as well as M oxidation, and deamidation of N and Q, in a
+        manner that does not comply with ProForma.
+
+        Parameters
+        ----------
+        sequence : str
+            The peptide sequence from MassIVE-KB
+        charge : int, optional
+            The charge state of the peptide.
+
+        Returns
+        -------
+        str
+            The parsed MassIVE peptide after conversion to a ProForma
+            format.
+        """
+        sequence = "".join(
+            [
+                MSKB_TO_UNIMOD.get(aa, aa)
+                for aa in re.split(r"(?<=.)(?=[A-Z])", sequence)
+            ]
+        )
+        if charge is not None:
+            sequence += f"/{charge}"
+
+        return sequence
+
+
+@dataclass
+class PeptideIons:
+    """The ions generated by a peptide.
+
+    Parameters
+    ----------
+    tokens : list[str]
+        The string tokens that comprise the peptide sequence.
+    precursor : float
+        The monoisotopic m/z of the precursor ion.
+    fragments : torch.Tensor[float]
+        The generated fragment ions originated from the peptide.
+    """
+
+    tokens: list[str]
+    precursor: float
+    fragments: torch.Tensor[float]
+
+    @property
+    def sequence(self) -> str:
+        """The peptide sequence."""
+        return "".join(self.tokens)
+
+    @property
+    def b_ions(self) -> torch.Tensor[float]:
+        """The b ion series."""
+        return self.fragments[0, :, :]
+
+    @property
+    def y_ions(self) -> torch.Tensor[float]:
+        """The y ion series."""
+        return self.fragments[1, :, :]
+
+
+@dataclass
+class Molecule:
+    """A representation of a molecule.
+
+    Parameters
+    ----------
+    smiles : str
+        A SMILES string defining the molecule.
+    charge : int, optional
+        The charge of the molecule.
+    """
+
+    smiles: str
+    charge: int | None = None
+
+    def __post_init__(self) -> None:
+        """Validate parameters."""
+        self._mol = Chem.MolFromSmiles(self.smiles)
+        if self._mol is None:
+            raise ValueError("The SMILES was invalid.")
+
+        # Canonnicalize the SMILES string:
+        self.smiles = Chem.MolToSmiles(self._mol)
+
+    def show(self, **kwargs: dict) -> PngImageFile:
+        """Show the molecule in 2D.
+
+        Parameters
+        ----------
+        **kwargs : dict
+            Keyword arguments passed to ``rdkit.Chem.Draw.MolToImage``
+        """
+        return Draw.MolToImage(self._mol, **kwargs)
+
+    def to_selfies(self) -> str:
+        """Convert SMILES to a SELFIES representaion."""
+        return sf.encoder(self.smiles)
+
+    @classmethod
+    def from_selfies(
+        cls,
+        selfies: str,
+        charge: int | None = None,
+    ) -> Molecule:
+        """Create a molecule from a SELFIES string.
+
+        Parameters
+        ----------
+        selfies : str
+            The SELFIES string defining the molecule.
+        charge : int, optional
+            The charge of the molecule.
+
+        Returns
+        -------
+        Molecule
+            The parsed Molecule.
+        """
+        return cls(sf.decoder(selfies), charge)
+
+
+class MassSpectrum(MsmsSpectrum):
+    """A mass spectrum.
+
+    Parameters
+    ----------
+    filename: str
+        The file from which the spectrum originated.
+    scan_id: str
+        The Hupo PSI standard scan identifier.
+    mz: array of shape (n_peaks,)
+        The m/z values.
+    intensity: array of shape (n_peaks, )
+        The intensity values.
+    retention_time: float, optional
+        The measured retention time.
+    ion_mobility: float, optional
+        The measured ion mobility.
+    precursor_mz: float, optional
+        The precursor ion m/z, if applicable.
+    precursor_charge: int, optional
+        The precursor charge, if applicable.
+    label: str, optional
+        A label for the mass spectrum. This is typically an
+        annotation, such as the generating peptide sequence,
+        but is distinct from spectrum_utils' annotation.
+    """
+
+    def __init__(
+        self,
+        filename: str,
+        scan_id: str,
+        mz: ArrayLike,
+        intensity: ArrayLike,
+        retention_time: float | None = None,
+        ion_mobility: float | None = None,
+        precursor_mz: float | None = None,
+        precursor_charge: int | None = None,
+        label: str | None = None,
+    ) -> None:
+        """Initialize a MassSpectrum."""
+        self.filename = filename
+        self.scan_id = scan_id
+        self.label = label
+
+        # Not currently supported by spectrum_utils:
+        self.ion_mobility = ion_mobility
+
+        # spectrum_utils requires a precursor. Will remove after
+        # I make a PR:
+        if precursor_mz is None:
+            precursor_mz = np.nan
+
+        if precursor_charge is None:
+            precursor_charge = np.nan
+
+        super().__init__(
+            identifier=str(scan_id),
+            precursor_mz=precursor_mz,
+            precursor_charge=precursor_charge,
+            mz=mz,
+            intensity=intensity,
+            retention_time=retention_time,
+        )
+
+    @property
+    def usi(self) -> str:
+        """The Universal Spectrum Identifier."""
+        index_type, scan = self.scan_id.split("=")
+        fname = re.sub(
+            ".(gz|tar|tar.gz|zip|bz2|tar.bz2)$",
+            "",
+            self.filename,
+            flags=re.IGNORECASE,
+        )
+        fname = re.sub(
+            ".(raw|mzml|mzxml|mgf|d|wiff)$",
+            "",
+            fname,
+            flags=re.IGNORECASE,
+        )
+        return ":".join([fname, index_type, scan])
+
+    @property
+    def precursor_mass(self) -> float:
+        """The monoisotopic mass."""
+        return (self.precursor_mz - PROTON) * self.precursor_charge
+
+    def to_tensor(self) -> torch.tensor:
+        """Combine the m/z and intensity arrays into a single tensor.
+
+        Returns
+        -------
+        torch.tensor of shape (n_peaks, 2)
+            The mass spectrum information.
+        """
+        return torch.tensor(np.vstack([self.mz, self.intensity]).T)
diff --git a/depthcharge/similarity.py b/depthcharge/similarity.py
deleted file mode 100644
index a73f2cd..0000000
--- a/depthcharge/similarity.py
+++ /dev/null
@@ -1,76 +0,0 @@
-"""Similarity measures between mass spectra."""
-import numba as nb
-import numpy as np
-
-
-@nb.njit
-def gsp(spectrum_x, spectrum_y, tol):
-    """Compute the dot product between two spectra.
-
-    This code was adapted from Wout:
-    https://github.com/bittremieux/GLEAMS/blob/master/gleams/feature/spectrum.py#L154
-
-
-    Parameters
-    ----------
-    spectrum_x : np.ndarray
-    spectrum_y : np.ndarray
-        The spectra to compare. Each row should be a peak, with the first
-        column indicating the m/z and the second column indicating the
-        intensities.
-    tol : float
-        The fragment m/z tolerance used to match peaks in both spectra with
-        each other.
-
-    Returns
-    -------
-    float
-        The dot product between both spectra.
-    """
-    # Find the matching peaks between both spectra.
-    peak_match_scores, peak_match_idx = [], []
-    peak_other_i = 0
-    mz_other, intensity_other = spectrum_y.T
-    for peak_i, (peak_mz, peak_intensity) in enumerate(spectrum_x):
-        # Advance while there is an excessive mass difference.
-        while (
-            peak_other_i < len(mz_other) - 1
-            and peak_mz - tol > mz_other[peak_other_i]
-        ):
-            peak_other_i += 1
-        # Match the peaks within the fragment mass window if possible.
-        peak_other_window_i = peak_other_i
-        while (
-            peak_other_window_i < len(mz_other)
-            and abs(peak_mz - (mz_other[peak_other_window_i])) <= tol
-        ):
-            peak_match_scores.append(
-                peak_intensity * intensity_other[peak_other_window_i]
-            )
-            peak_match_idx.append((peak_i, peak_other_window_i))
-            peak_other_window_i += 1
-
-    score = 0
-    if len(peak_match_scores) > 0:
-        # Use the most prominent peak matches to compute the score (sort in
-        # descending order).
-        peak_match_scores_arr = np.asarray(peak_match_scores)
-        peak_match_order = np.argsort(peak_match_scores_arr)[::-1]
-        peak_match_scores_arr = peak_match_scores_arr[peak_match_order]
-        peak_match_idx_arr = np.asarray(peak_match_idx)[peak_match_order]
-        peaks_used, peaks_used_other = set(), set()
-        for peak_match_score, peak_i, peak_other_i in zip(
-            peak_match_scores_arr,
-            peak_match_idx_arr[:, 0],
-            peak_match_idx_arr[:, 1],
-        ):
-            if (
-                peak_i not in peaks_used
-                and peak_other_i not in peaks_used_other
-            ):
-                score += peak_match_score
-                # Make sure these peaks are not used anymore.
-                peaks_used.add(peak_i)
-                peaks_used_other.add(peak_other_i)
-
-    return score
diff --git a/depthcharge/tokenizers/__init__.py b/depthcharge/tokenizers/__init__.py
new file mode 100644
index 0000000..1a76264
--- /dev/null
+++ b/depthcharge/tokenizers/__init__.py
@@ -0,0 +1,3 @@
+"""Deptcharge tokenizers."""
+from .peptides import PeptideTokenizer
+from .tokenizer import Tokenizer
diff --git a/depthcharge/tokenizers/peptides.py b/depthcharge/tokenizers/peptides.py
new file mode 100644
index 0000000..f69c5e1
--- /dev/null
+++ b/depthcharge/tokenizers/peptides.py
@@ -0,0 +1,440 @@
+"""Tokenizers for peptides."""
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable
+
+import numba as nb
+import numpy as np
+import torch
+from pyteomics.proforma import GenericModification, MassModification
+
+from .. import utils
+from ..constants import H2O, PROTON
+from ..primitives import MSKB_TO_UNIMOD, Peptide, PeptideIons
+from .tokenizer import Tokenizer
+
+
+class PeptideTokenizer(Tokenizer):
+    """A tokenizer for ProForma peptide sequences.
+
+    Parse and tokenize ProForma-compliant peptide sequences. Additionally,
+    use this class to calculate fragment and precursor ion masses.
+
+    Parameters
+    ----------
+    residues : dict[str, float], optional
+        Residues and modifications to add to the vocabulary beyond the
+        standard 20 amino acids.
+    replace_isoleucine_with_leucine : bool
+        Replace I with L residues, because they are isobaric and often
+        indistinguishable by mass spectrometry.
+    reverse : bool
+        Reverse the sequence for tokenization, C-terminus to N-terminus.
+
+    Attributes
+    ----------
+    residues : numba.typed.Dict[str, float]
+        The residues and modifications and their associated masses.
+        terminal modifcations are indicated by `-`.
+    index : SortedDict{str, int}
+        The mapping of residues and modifications to integer representations.
+    reverse_index : list[None | str]
+        The ordered residues and modifications where the list index is the
+        integer representation for a token.
+    stop_token : str
+        The stop token.
+    """
+
+    residues = nb.typed.Dict.empty(
+        nb.types.unicode_type,
+        nb.types.float64,
+    )
+    residues.update(
+        G=57.021463735,
+        A=71.037113805,
+        S=87.032028435,
+        P=97.052763875,
+        V=99.068413945,
+        T=101.047678505,
+        C=103.009184505,
+        L=113.084064015,
+        I=113.084064015,
+        N=114.042927470,
+        D=115.026943065,
+        Q=128.058577540,
+        K=128.094963050,
+        E=129.042593135,
+        M=131.040484645,
+        H=137.058911875,
+        F=147.068413945,
+        R=156.101111050,
+        Y=163.063328575,
+        W=186.079312980,
+    )
+
+    # The peptide parsing function:
+    _parse_peptide = Peptide.from_proforma
+
+    def __init__(
+        self,
+        residues: dict[str, float] | None = None,
+        replace_isoleucine_with_leucine: bool = False,
+        reverse: bool = False,
+    ) -> None:
+        """Initialize a PeptideTokenizer."""
+        self.replace_isoleucine_with_leucine = replace_isoleucine_with_leucine
+        self.reverse = reverse
+        self.residues = self.residues.copy()
+        if residues is not None:
+            self.residues.update(residues)
+
+        if self.replace_isoleucine_with_leucine:
+            del self.residues["I"]
+
+        super().__init__(list(self.residues.keys()))
+
+    def split(self, sequence: str) -> list[str]:
+        """Split a ProForma peptide sequence.
+
+        Parameters
+        ----------
+        sequence : str
+            The peptide sequence.
+
+        Returns
+        -------
+        list[str]
+            The tokens that compprise the peptide sequence.
+        """
+        pep = self._parse_peptide(sequence)
+        if self.replace_isoleucine_with_leucine:
+            pep.sequence = pep.sequence.replace("I", "L")
+
+        pep = pep.split()
+        if self.reverse:
+            pep.reverse()
+
+        return pep
+
+    def ions(  # noqa: C901
+        self,
+        sequences: Iterable[str],
+        precursor_charges: Iterable[int] | str,
+        max_fragment_charge: int | None = None,
+    ) -> tuple[torch.Tensor[float], list[torch.Tensor[float]]]:
+        """Calculate the m/z for the precursor and fragment ions.
+
+        Currently depthcharge only support b and y ions.
+
+        Parameters
+        ----------
+        sequences : Iterable[str],
+            The peptide sequences.
+        precursor_charges : Iterable[int] or None, optional
+            The charge of each precursor ion. If ``None``, the charge state
+            is expected to be found in the peptide strings.
+        max_fragment_charge : int or None, optional
+            The maximum charge for fragment ions. The default is to consider
+            up to the ``max(precursor_charge - 1, 1)``.
+
+        Returns
+        -------
+        list of PeptideIons
+            The precursor and fragment ions generated by the peptide.
+        """
+        sequences = utils.listify(sequences)
+        if max_fragment_charge is None:
+            max_fragment_charge = np.inf
+
+        if precursor_charges is None:
+            precursor_charges = [None] * len(sequences)
+        else:
+            precursor_charges = utils.listify(precursor_charges)
+
+        if len(sequences) != len(precursor_charges):
+            raise ValueError(
+                "The number of sequences and precursor charges did not match."
+            )
+
+        out = []
+        for seq, charge in zip(sequences, precursor_charges):
+            if isinstance(seq, str):
+                if self.replace_isoleucine_with_leucine:
+                    seq = seq.replace("I", "L")
+
+                try:
+                    pep = Peptide.from_proforma(seq)
+                except ValueError:
+                    pep = Peptide.from_massivekb(seq)
+
+                tokens = pep.split()
+                if charge is None:
+                    charge = max(pep.charge - 1, 1)
+            else:
+                tokens = seq
+
+            if charge is None:
+                raise ValueError(
+                    f"No charge was provided for {seq}",
+                )
+
+            try:
+                prec = _calc_precursor_mass(
+                    nb.typed.List(tokens),
+                    charge,
+                    self.residues,
+                )
+            except KeyError as err:
+                raise ValueError(
+                    f"Unrecognized token(s) in {''.join(tokens)}"
+                ) from err
+
+            frags = _calc_fragment_masses(
+                nb.typed.List(tokens),
+                min(charge, max_fragment_charge),
+                self.residues,
+            )
+
+            ions = PeptideIons(
+                tokens=tokens,
+                precursor=prec,
+                fragments=torch.tensor(frags),
+            )
+
+            out.append(ions)
+
+        return out
+
+    @classmethod
+    def from_proforma(
+        cls,
+        sequences: Iterable[str],
+        replace_isoleucine_with_leucine: bool = True,
+        reverse: bool = True,
+    ) -> PeptideTokenizer:
+        """Create a tokenizer with the observed peptide modications.
+
+        Modifications are parsed from ProForma 2.0-compliant peptide strings
+        and added to the vocabulary.
+
+        Parameters
+        ----------
+        sequences : Iterable[str]
+            The peptides from which to parse modifications.
+        replace_isoleucine_with_leucine : bool
+            Replace I with L residues, because they are isobaric and often
+            indistinguishable by mass spectrometry.
+        reverse : bool
+            Reverse the sequence for tokenization, C-terminus to N-terminus.
+
+        Returns
+        -------
+        PeptideTokenizer
+            A tokenizer for peptides with the observed modifications.
+        """
+        if isinstance(sequences, str):
+            sequences = [sequences]
+
+        # Parse modifications:
+        new_res = cls.residues.copy()
+        for peptide in sequences:
+            parsed = Peptide.from_proforma(peptide).split()
+            for token in parsed:
+                if token in new_res.keys():
+                    continue
+
+                if token == "-":
+                    continue
+
+                match = re.search(r"(.*)\[(.*)\]", token)
+                try:
+                    res, mod = match.groups()
+                    if res and res != "-":
+                        res_mass = new_res[res]
+                    else:
+                        res_mass = 0
+                except (AttributeError, KeyError) as err:
+                    raise ValueError("Unrecognized token {token}.") from err
+
+                try:
+                    mod = MassModification(mod)
+                except ValueError:
+                    mod = GenericModification(mod)
+
+                new_res[token] = res_mass + mod.mass
+
+        return cls(new_res, replace_isoleucine_with_leucine, reverse)
+
+    @staticmethod
+    def from_massivekb(
+        replace_isoleucine_with_leucine: bool = True,
+        reverse: bool = True,
+    ) -> MskbPeptideTokenizer:
+        """Create a tokenizer with the observed peptide modications.
+
+        Modifications are parsed from MassIVE-KB peptide strings
+        and added to the vocabulary.
+
+        Parameters
+        ----------
+        replace_isoleucine_with_leucine : bool
+            Replace I with L residues, because they are isobaric and often
+            indistinguishable by mass spectrometry.
+        reverse : bool
+            Reverse the sequence for tokenization, C-terminus to N-terminus.
+
+        Returns
+        -------
+        MskbPeptideTokenizer
+            A tokenizer for peptides with the observed modifications.
+        """
+        return MskbPeptideTokenizer.from_proforma(
+            [f"{mod}A" for mod in MSKB_TO_UNIMOD.values()],
+            replace_isoleucine_with_leucine,
+            reverse,
+        )
+
+
+class MskbPeptideTokenizer(PeptideTokenizer):
+    """A tokenizer for MassIVE-KB peptide sequences.
+
+    MassIVE-KB includes N-term carbamylation, NH3-loss, acetylation, as well as
+    M oxidation, and deamidation of N and Q, in a manner that does not comply
+    with the ProForma standard.
+
+    Parameters
+    ----------
+    replace_isoleucine_with_leucine : bool
+        Replace I with L residues, because they are isobaric and often
+        indistinguishable by mass spectrometry.
+    reverse : bool
+        Reverse the sequence for tokenization, C-terminus to N-terminus.
+
+    Attributes
+    ----------
+    residues : SortedDict[str, float]
+        The residues and modifications and their associated masses.
+        terminal modifcations are indicated by `-`.
+    index : SortedDicte{str, int}
+        The mapping of residues and modifications to integer representations.
+    reverse_index : list[None | str]
+        The ordered residues and modifications where the list index is the
+        integer representation for a token.
+    stop_token : str
+        The stop token.
+
+    """
+
+    _parse_peptide = Peptide.from_massivekb
+
+
+@nb.njit
+def _calc_precursor_mass(
+    tokens: list[str],
+    charge: int,
+    masses: nb.typed.Dict,
+) -> float:
+    """Calculate the precursor mass of a peptide sequence.
+
+    Parameters
+    ----------
+    tokens : list of str
+        The tokenized peptide sequence.
+    charge : int
+        The charge state to consider. Use 'None' to get the neutral mass.
+    masses : nb.typed.Dict
+        The mass dictionary to use.
+
+    Returns
+    -------
+    float
+        The precurosr monoisotopic m/z.
+    """
+    mass = sum([masses[t] for t in tokens]) + H2O
+    if charge is not None:
+        mass = _mass2mz(mass, charge)
+
+    return mass
+
+
+@nb.njit
+def _calc_fragment_masses(
+    tokens: list[str],
+    charge: int,
+    masses: nb.typed.Dict,
+) -> np.ndarray[float]:
+    """Calculate the b and y ions for a peptide sequence.
+
+    Parameters
+    ----------
+    tokens : list of str
+        The tokenized peptide sequence.
+    charge : int
+        The charge state to consider. Use 'None' to get the neutral mass.
+    n_isotopes : int
+        The number of C13 isotopes to return, starting from the
+        monoisotopic mass.
+    masses : nb.typed.Dict
+        The mass dictionary to use.
+
+    Returns
+    -------
+    np.ndarray of shape (2, len(seq) - 1, charge)
+        The m/z of the predicted b and y ions.
+    """
+    # Handle terminal mods:
+    seq = np.empty(len(tokens))
+    n_mod = False
+    c_mod = False
+    for idx, token in enumerate(tokens):
+        if not idx and token.endswith("-"):
+            n_mod = True
+
+        if idx == (len(tokens) - 1) and token.startswith("-"):
+            c_mod = True
+
+        seq[idx] = masses[token]
+
+    if n_mod:
+        seq[1] += seq[0]
+        seq = seq[1:]
+
+    if c_mod:
+        seq[-2] += seq[-1]
+        seq = seq[:-1]
+
+    # Calculate fragments:
+    max_charge = min(charge, 2)
+    n_ions = len(seq) - 1
+    ions = np.empty((2, n_ions, max_charge))
+    b_mass = 0
+    y_mass = H2O
+    for idx in range(n_ions):
+        b_mass += seq[idx]
+        y_mass += seq[-(idx + 1)]
+        for cur_charge in range(1, max_charge + 1):
+            z_idx = cur_charge - 1
+            ions[0, idx, z_idx] = _mass2mz(b_mass, cur_charge)
+            ions[1, idx, z_idx] = _mass2mz(y_mass, cur_charge)
+
+    return ions
+
+
+@nb.njit
+def _mass2mz(mass: float, charge: int) -> float:
+    """Calculate the m/z.
+
+    Parameters
+    ----------
+    mass : float
+        The neutral mass.
+    charge : int
+        The charge.
+
+    Returns
+    -------
+    float
+       The m/z
+    """
+    return (mass / charge) + PROTON
diff --git a/depthcharge/tokenizers/tokenizer.py b/depthcharge/tokenizers/tokenizer.py
new file mode 100644
index 0000000..e232cb5
--- /dev/null
+++ b/depthcharge/tokenizers/tokenizer.py
@@ -0,0 +1,134 @@
+"""A base Tokenizer class."""
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Iterable, Sequence
+
+import torch
+from sortedcontainers import SortedDict, SortedSet
+from torch import nn
+
+from .. import utils
+
+
+class Tokenizer(ABC):
+    """An abstract base class for Depthcharge tokenizers.
+
+    Parameters
+    ----------
+    tokens : Sequence[str]
+        The tokens to consider.
+    stop_token : str, optional
+        The stop token to use.
+    """
+
+    def __init__(self, tokens: Sequence[str], stop_token: str = "$") -> None:
+        """Initialize a tokenizer."""
+        self.stop_token = stop_token
+
+        tokens = SortedSet(tokens)
+        if self.stop_token in tokens:
+            raise ValueError(
+                f"Stop token {stop_token} already exists in tokens.",
+            )
+
+        tokens.add(self.stop_token)
+        self.index = SortedDict({k: i + 1 for i, k in enumerate(tokens)})
+        self.reverse_index = [None] + list(tokens)
+        self.stop_int = self.index[self.stop_token]
+
+    def __len__(self) -> int:
+        """The number of tokens."""
+        return len(self.index)
+
+    @abstractmethod
+    def split(self, sequence: str) -> list[str]:
+        """Split a sequence into the constituent string tokens."""
+
+    def tokenize(
+        self,
+        sequences: Iterable[str],
+        to_strings: bool = False,
+        add_stop: bool = False,
+    ) -> torch.Tensor | list[list[str]]:
+        """Tokenize the input sequences.
+
+        Parameters
+        ----------
+        sequences : Iterable[str]
+            The sequences to tokenize.
+        to_strings : bool, optional
+            Return each as a list of token strings rather than a
+            tensor. This is useful for debugging.
+        add_stop : bool, optional
+            Append the stop token tothe end of the sequence.
+
+        Returns
+        -------
+        torch.Tensor of shape (n_sequences, max_length) or list[list[str]]
+            Either a tensor containing the integer values for each
+            token, padded with 0's, or the list of tokens comprising
+            each sequence.
+        """
+        try:
+            out = []
+            for seq in utils.listify(sequences):
+                tokens = self.split(seq)
+                if add_stop and tokens[-1] != self.stop_token:
+                    tokens.append(self.stop_token)
+
+                if to_strings:
+                    out.append(tokens)
+                    continue
+
+                out.append(torch.tensor([self.index[t] for t in tokens]))
+
+            if to_strings:
+                return out
+
+            if isinstance(sequences, str):
+                return out[0]
+
+            return nn.utils.rnn.pad_sequence(out, batch_first=True)
+        except KeyError as err:
+            raise ValueError("Unrecognized token") from err
+
+    def detokenize(
+        self,
+        tokens: torch.Tensor,
+        join: bool = True,
+        trim_stop_token: bool = True,
+    ) -> list[str] | list[list[str]]:
+        """Retreive sequences from tokens.
+
+        Parameters
+        ----------
+        tokens : torch.Tensor of shape (n_sequences, max_length)
+            The zero-padded tensor of integerized tokens to decode.
+        join : bool, optional
+            Join tokens into strings?
+        trim_stop_token : bool, optional
+            Remove the stop token from the end of a sequence.
+
+        Returns
+        -------
+        list[str] or list[list[str]]
+            The decoded sequences each as a string or list or strings.
+        """
+        decoded = []
+        for row in tokens:
+            seq = [
+                self.reverse_index[i]
+                for i in row
+                if self.reverse_index[i] is not None
+            ]
+
+            if trim_stop_token and seq[-1] == self.stop_token:
+                seq.pop(-1)
+
+            if join:
+                seq = "".join(seq)
+
+            decoded.append(seq)
+
+        return decoded
diff --git a/depthcharge/transformers/__init__.py b/depthcharge/transformers/__init__.py
new file mode 100644
index 0000000..04b3b7d
--- /dev/null
+++ b/depthcharge/transformers/__init__.py
@@ -0,0 +1,6 @@
+"""Transformer models."""
+from .peptides import (
+    PeptideTransformerDecoder,
+    PeptideTransformerEncoder,
+)
+from .spectra import SpectrumTransformerEncoder
diff --git a/depthcharge/transformers/peptides.py b/depthcharge/transformers/peptides.py
new file mode 100644
index 0000000..95b47d9
--- /dev/null
+++ b/depthcharge/transformers/peptides.py
@@ -0,0 +1,289 @@
+"""Transformer models for peptides."""
+import torch
+
+from ..encoders import FloatEncoder, PositionalEncoder
+from ..tokenizers import PeptideTokenizer
+
+
+class _PeptideTransformer(torch.nn.Module):
+    """A transformer base class for peptide sequences.
+
+    Parameters
+    ----------
+    n_tokens : int or PeptideTokenizer
+        The number of tokens used to tokenize peptide sequences.
+    d_model : int
+        The latent dimensionality to represent the amino acids in a peptide
+        sequence.
+    positional_encoder : PositionalEncoder or bool
+        The positional encodings to use for the amino acid sequence. If
+        ``True``, the default positional encoder is used. ``False`` disables
+        positional encodings, typically only for ablation tests.
+    max_charge : int
+        The maximum precursor charge to embed.
+    """
+
+    def __init__(
+        self,
+        n_tokens: int | PeptideTokenizer,
+        d_model: int,
+        positional_encoder: PositionalEncoder | bool,
+        max_charge: int,
+    ) -> None:
+        super().__init__()
+        try:
+            n_tokens = len(n_tokens)
+        except TypeError:
+            pass
+
+        if callable(positional_encoder):
+            self.positional_encoder = positional_encoder
+        elif positional_encoder:
+            self.positional_encoder = PositionalEncoder(d_model)
+        else:
+            self.positional_encoder = torch.nn.Identity()
+
+        self.charge_encoder = torch.nn.Embedding(max_charge, d_model)
+        self.aa_encoder = torch.nn.Embedding(
+            n_tokens + 1,
+            d_model,
+            padding_idx=0,
+        )
+
+    @property
+    def device(self) -> torch.device:
+        """The current device for the model."""
+        return next(self.parameters()).device
+
+
+class PeptideTransformerEncoder(_PeptideTransformer):
+    """A transformer encoder for peptide sequences.
+
+    Parameters
+    ----------
+    n_tokens : int or PeptideTokenizer
+        The number of tokens used to tokenize peptide sequences.
+    d_model : int
+        The latent dimensionality to represent the amino acids in a peptide
+        sequence.
+    nhead : int, optional
+        The number of attention heads in each layer. ``d_model`` must be
+        divisible by ``nhead``.
+    dim_feedforward : int, optional
+        The dimensionality of the fully connected layers in the Transformer
+        layers of the model.
+    n_layers : int, optional
+        The number of Transformer layers.
+    dropout : float, optional
+        The dropout probability for all layers.
+    positional_encoder : PositionalEncoder or bool, optional
+        The positional encodings to use for the amino acid sequence. If
+        ``True``, the default positional encoder is used. ``False`` disables
+        positional encodings, typically only for ablation tests.
+    max_charge : int, optional
+        The maximum charge state for peptide sequences.
+    """
+
+    def __init__(
+        self,
+        n_tokens: int | PeptideTokenizer,
+        d_model: int = 128,
+        nhead: int = 8,
+        dim_feedforward: int = 1024,
+        n_layers: int = 1,
+        dropout: float = 0,
+        positional_encoder: PositionalEncoder | bool = True,
+        max_charge: int = 5,
+    ) -> None:
+        """Initialize a PeptideEncoder."""
+        super().__init__(
+            n_tokens=n_tokens,
+            d_model=d_model,
+            positional_encoder=positional_encoder,
+            max_charge=max_charge,
+        )
+
+        # The Transformer layers:
+        layer = torch.nn.TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=dim_feedforward,
+            batch_first=True,
+            dropout=dropout,
+        )
+
+        self.transformer_encoder = torch.nn.TransformerEncoder(
+            layer,
+            num_layers=n_layers,
+        )
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        charges: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Predict the next amino acid for a collection of sequences.
+
+        Parameters
+        ----------
+        tokens : torch.Tensor of size (batch_size, peptide_length)
+            The integer tokens describing each peptide sequence, padded
+            to the maximum peptide length in the batch with 0s.
+        charges : torch.Tensor of size (batch_size,)
+            The charge state of each peptide.
+
+        Returns
+        -------
+        latent : torch.Tensor of shape (n_sequences, len_sequence, d_model)
+            The latent representations for the spectrum and each of its
+            peaks.
+        mem_mask : torch.Tensor
+            The memory mask specifying which elements were padding in X.
+        """
+        # Encode everything:
+        encoded = self.aa_encoder(tokens)
+        charges = self.charge_encoder(charges - 1)[:, None]
+        encoded = torch.cat([charges, encoded], dim=1)
+
+        # Create mask
+        mask = ~encoded.sum(dim=2).bool()
+
+        # Add positional encodings
+        encoded = self.positional_encoder(encoded)
+
+        # Run through the model:
+        latent = self.transformer_encoder(encoded, src_key_padding_mask=mask)
+        return latent, mask
+
+
+class PeptideTransformerDecoder(_PeptideTransformer):
+    """A transformer decoder for peptide sequences.
+
+    Parameters
+    ----------
+    n_tokens : int or PeptideTokenizer
+        The number of tokens used to tokenize peptide sequences.
+    d_model : int, optional
+        The latent dimensionality to represent peaks in the mass spectrum.
+    nhead : int, optional
+        The number of attention heads in each layer. ``d_model`` must be
+        divisible by ``nhead``.
+    dim_feedforward : int, optional
+        The dimensionality of the fully connected layers in the Transformer
+        layers of the model.
+    n_layers : int, optional
+        The number of Transformer layers.
+    dropout : float, optional
+        The dropout probability for all layers.
+    positional_encoder : PositionalEncoder or bool, optional
+        The positional encodings to use for the amino acid sequence. If
+        ``True``, the default positional encoder is used. ``False`` disables
+        positional encodings, typically only for ablation tests.
+    max_charge : int, optional
+        The maximum charge state for peptide sequences.
+    """
+
+    def __init__(
+        self,
+        n_tokens: int | PeptideTokenizer,
+        d_model: int = 128,
+        nhead: int = 8,
+        dim_feedforward: int = 1024,
+        n_layers: int = 1,
+        dropout: float = 0,
+        positional_encoder: PositionalEncoder | bool = True,
+        max_charge: int = 5,
+    ) -> None:
+        """Initialize a PeptideDecoder."""
+        super().__init__(
+            n_tokens=n_tokens,
+            d_model=d_model,
+            positional_encoder=positional_encoder,
+            max_charge=max_charge,
+        )
+
+        # Additional model components
+        self.mass_encoder = FloatEncoder(d_model)
+        layer = torch.nn.TransformerDecoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=dim_feedforward,
+            batch_first=True,
+            dropout=dropout,
+        )
+
+        self.transformer_decoder = torch.nn.TransformerDecoder(
+            layer,
+            num_layers=n_layers,
+        )
+
+        self.final = torch.nn.Linear(
+            d_model,
+            self.aa_encoder.num_embeddings - 1,
+        )
+
+    def forward(
+        self,
+        tokens: torch.Tensor | None,
+        precursors: torch.Tensor,
+        memory: torch.Tensor,
+        memory_key_padding_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Predict the next amino acid for a collection of sequences.
+
+        Parameters
+        ----------
+        tokens : list of str, torch.Tensor, or None
+            The partial peptide sequences for which to predict the next
+            amino acid. Optionally, these may be the token indices instead
+            of a string.
+        precursors : torch.Tensor of size (batch_size, 2)
+            The precursor mass (axis 0) and charge (axis 1).
+        memory : torch.Tensor of shape (batch_size, n_peaks, d_model)
+            The representations from a ``TransformerEncoder``, such as a
+            ``SpectrumEncoder``.
+        memory_key_padding_mask : torch.Tensor of shape (batch_size, n_peaks)
+            The mask that indicates which elements of ``memory`` are padding.
+
+        Returns
+        -------
+        scores : torch.Tensor of size (batch_size, len_sequence, n_amino_acids)
+            The raw output for the final linear layer. These can be Softmax
+            transformed to yield the probability of each amino acid for the
+            prediction.
+
+        """
+        # Prepare sequences
+        if tokens is None:
+            tokens = torch.tensor([[]]).to(self.device)
+
+        # Encode everything:
+        tokens = self.aa_encoder(tokens)
+        masses = self.mass_encoder(precursors[:, None, 0])
+        charges = self.charge_encoder(precursors[:, 1].int() - 1)
+        precursors = masses + charges[:, None, :]
+
+        # Feed through model:
+        tgt = torch.cat([precursors, tokens], dim=1)
+        tgt_key_padding_mask = tgt.sum(axis=2) == 0
+        tgt = self.positional_encoder(tgt)
+        tgt_mask = generate_tgt_mask(tgt.shape[1]).type_as(precursors)
+        preds = self.transformer_decoder(
+            tgt=tgt,
+            memory=memory,
+            tgt_mask=tgt_mask,
+            tgt_key_padding_mask=tgt_key_padding_mask,
+            memory_key_padding_mask=memory_key_padding_mask.to(self.device),
+        )
+        return self.final(preds)
+
+
+def generate_tgt_mask(sz: int) -> torch.Tensor:
+    """Generate a square mask for the sequence.
+
+    Parameters
+    ----------
+    sz : int
+        The length of the target sequence.
+    """
+    return ~torch.triu(torch.ones(sz, sz, dtype=torch.bool)).transpose(0, 1)
diff --git a/depthcharge/transformers/spectra.py b/depthcharge/transformers/spectra.py
new file mode 100644
index 0000000..d2ea493
--- /dev/null
+++ b/depthcharge/transformers/spectra.py
@@ -0,0 +1,105 @@
+"""Tranformer models to handle mass spectra."""
+from collections.abc import Callable
+
+import torch
+
+from ..encoders import PeakEncoder
+
+
+class SpectrumTransformerEncoder(torch.nn.Module):
+    """A Transformer encoder for input mass spectra.
+
+    Parameters
+    ----------
+    d_model : int, optional
+        The latent dimensionality to represent peaks in the mass spectrum.
+    nhead : int, optional
+        The number of attention heads in each layer. ``d_model`` must be
+        divisible by ``nhead``.
+    dim_feedforward : int, optional
+        The dimensionality of the fully connected layers in the Transformer
+        layers of the model.
+    n_layers : int, optional
+        The number of Transformer layers.
+    dropout : float, optional
+        The dropout probability for all layers.
+    peak_encoder : PeakEncoder or bool, optional
+        Sinusoidal encodings m/z and intensityvalues of each peak.
+    """
+
+    def __init__(
+        self,
+        d_model: int = 128,
+        nhead: int = 8,
+        dim_feedforward: int = 1024,
+        n_layers: int = 1,
+        dropout: float = 0,
+        peak_encoder: PeakEncoder | Callable | bool = True,
+    ) -> None:
+        """Initialize a SpectrumEncoder."""
+        super().__init__()
+
+        self.latent_spectrum = torch.nn.Parameter(torch.randn(1, 1, d_model))
+
+        if callable(peak_encoder):
+            self.peak_encoder = peak_encoder
+        elif peak_encoder:
+            self.peak_encoder = PeakEncoder(d_model)
+        else:
+            self.peak_encoder = torch.nn.Linear(2, d_model)
+
+        # The Transformer layers:
+        layer = torch.nn.TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=dim_feedforward,
+            batch_first=True,
+            dropout=dropout,
+        )
+
+        self.transformer_encoder = torch.nn.TransformerEncoder(
+            layer,
+            num_layers=n_layers,
+        )
+
+    def forward(
+        self,
+        spectra: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Embed a mass spectrum.
+
+        Parameters
+        ----------
+        spectra : torch.Tensor of shape (n_spectra, n_peaks, 2)
+            The spectra to embed. Axis 0 represents a mass spectrum, axis 1
+            contains the peaks in the mass spectrum, and axis 2 is essentially
+            a 2-tuple specifying the m/z-intensity pair for each peak. These
+            should be zero-padded, such that all of the spectra in the batch
+            are the same length.
+
+        Returns
+        -------
+        latent : torch.Tensor of shape (n_spectra, n_peaks + 1, d_model)
+            The latent representations for the spectrum and each of its
+            peaks.
+        mem_mask : torch.Tensor
+            The memory mask specifying which elements were padding in X.
+        """
+        zeros = ~spectra.sum(dim=2).bool()
+        mask = [
+            torch.tensor([[False]] * spectra.shape[0]).type_as(zeros),
+            zeros,
+        ]
+        mask = torch.cat(mask, dim=1)
+        peaks = self.peak_encoder(spectra)
+
+        # Add the spectrum representation to each input:
+        latent_spectra = self.latent_spectrum.expand(peaks.shape[0], -1, -1)
+
+        peaks = torch.cat([latent_spectra, peaks], dim=1)
+        return self.transformer_encoder(peaks, src_key_padding_mask=mask), mask
+
+    @property
+    def device(self) -> torch.device:
+        """The current device for the model."""
+        return next(self.parameters()).device
diff --git a/depthcharge/utils.py b/depthcharge/utils.py
index 2d171fb..a0f4511 100644
--- a/depthcharge/utils.py
+++ b/depthcharge/utils.py
@@ -1,47 +1,8 @@
-"""Common utility functions"""
-import pandas as pd
+"""Common utility functions."""
+from typing import Any
 
-try:
-    from tensorboard.backend.event_processing.event_accumulator import (
-        EventAccumulator,
-    )
-except ImportError:
-    EventAccumulator = None
 
-
-def read_tensorboard_scalars(path):
-    """Read scalars from Tensorboard logs.
-
-    Parameters
-    ----------
-    path : str
-        The path of the scalar log file.
-
-    Returns
-    -------
-    pandas.DataFrame
-        A dataframe containing the scalar values.
-    """
-    if EventAccumulator is None:
-        raise ImportError(
-            "Install the 'tensorboard' optional dependency to "
-            "use this function. https://www.tensorflow.org/tensorboard"
-        )
-
-    event = EventAccumulator(path)
-    event.Reload()
-    data = []
-    for tag in event.Tags()["scalars"]:
-        tag_df = pd.DataFrame(
-            event.Scalars(tag), columns=["wall_time", "step", "value"]
-        )
-        tag_df["tag"] = tag
-        data.append(tag_df)
-
-    return pd.concat(data)
-
-
-def listify(obj):
+def listify(obj: Any) -> list[Any]:  # noqa: ANN401
     """Turn an object into a list, but don't split strings."""
     try:
         assert not isinstance(obj, str)
@@ -52,8 +13,7 @@ def listify(obj):
     return list(obj)
 
 
-# For Parameter Checking ------------------------------------------------------
-def check_int(integer, name):
+def check_int(integer: int, name: str) -> int:
     """Verify that an object is an integer, or coercible to one.
 
     Parameters
@@ -62,6 +22,11 @@ def check_int(integer, name):
         The integer to check.
     name : str
         The name to print in the error message if it fails.
+
+    Returns
+    -------
+    int
+        The coerced integer.
     """
     if isinstance(integer, int):
         return integer
@@ -74,7 +39,7 @@ def check_int(integer, name):
     return coerced
 
 
-def check_positive_int(integer, name):
+def check_positive_int(integer: int, name: str) -> int:
     """Verify that an object is an integer and positive.
 
     Parameters
@@ -83,6 +48,11 @@ def check_positive_int(integer, name):
         The integer to check.
     name : str
         The name to print in the error message if it fails.
+
+    Returns
+    -------
+    int
+        The coerced integer.
     """
     try:
         integer = check_int(integer, name)
diff --git a/docs/api/datasets.md b/docs/api/datasets.md
new file mode 100644
index 0000000..a358582
--- /dev/null
+++ b/docs/api/datasets.md
@@ -0,0 +1,5 @@
+# Datasets (`depthcharge.data`)
+
+::: depthcharge.data.SpectrumDataset
+::: depthcharge.data.AnnotatedSpectrumDataset
+::: depthcharge.data.PeptideDataset
diff --git a/docs/api/encoders.md b/docs/api/encoders.md
new file mode 100644
index 0000000..c15720f
--- /dev/null
+++ b/docs/api/encoders.md
@@ -0,0 +1,5 @@
+# Encoders (`depthcharge.encoders`)
+
+::: depthcharge.encoders.PositionalEncoder
+::: depthcharge.encoders.FloatEncoder
+::: depthcharge.encoders.PeakEncoder
diff --git a/docs/api/index.md b/docs/api/index.md
new file mode 100644
index 0000000..5cf3097
--- /dev/null
+++ b/docs/api/index.md
@@ -0,0 +1,41 @@
+# The depthcharge Python API
+
+The depthcharge package provides utilities and classes to parse, store, and use mass spectra, peptides, and small molecules in Transformer models.
+Although depthcharge is primarly focused Transformers, many of these classes can be used as building blocks for models of your own choice.
+
+
+## Tokenizers
+
+Tokenizers split string inputs into tokens, such as peptide sequences or SMILES strings, into the pieces that we want to model with our neural networks.
+For example, peptide sequences need to be broken into constituent amino acids, with or without modifications (the tokens).
+In addition to specifying how tokens are created, the tokenizer classes also aid in mass calculations that are useful for modeling mass spectrometry data.
+
+Tokenizers live in the [tokenizers submodule](tokenizers).
+
+## Datasets
+
+Datasets are subclasses of PyTorch Datasets, specifically designed to store and retrieve mass spectrometry data, including the mass spectra themselves and analyte representations.
+
+Datasets live in the [datasets submodule](datsets).
+
+## Encoders
+
+Often the most useful representation for an input is not the raw parsed format.
+In depthcharge, we provide a collection of sinusoidal encoders to represent general floating point numbers, positions in a sequence, and peaks in a mass spectrum.
+
+Encoders live in the [encoders submodule](encoders).
+
+## Transformers
+
+The Transformer architecture is a powerful tool for modeling sequence data, such as the amino acids of a peptide or even the peaks in a mass spectrum.
+In depthcharge, we provide Transformer models specifically tailored for modeling different pieces of mass spectrometry and related data.
+These classes are compound classes: they often combine the appropriate encoders with a traditional Transformer model to directly model the data of interest.
+
+Transformers live in the [transformers submodule](transformers).
+
+## Primitives
+
+Although most users will not interact with our primitive classes directly, they are the building blocks for representing mass spectrometry data types.
+These classes are not primitives in the traditional sense of the word; rather, they define how to parse and store data such as peptides, mass spectra, and small molecules.
+
+[Primitives](primitives) are available as top-level classes in Depthcharge.
diff --git a/docs/api/primitives.md b/docs/api/primitives.md
new file mode 100644
index 0000000..d603e5c
--- /dev/null
+++ b/docs/api/primitives.md
@@ -0,0 +1,6 @@
+# Primitives
+
+::: depthcharge.MassSpectrum
+::: depthcharge.Peptide
+::: depthcharge.PeptideIons
+::: depthcharge.Molecule
diff --git a/docs/api/tokenizers.md b/docs/api/tokenizers.md
new file mode 100644
index 0000000..92778b5
--- /dev/null
+++ b/docs/api/tokenizers.md
@@ -0,0 +1,4 @@
+# Tokenizers (`depthcharge.tokenizers`)
+
+::: depthcharge.tokenizers.PeptideTokenizer
+::: depthcharge.tokenizers.Tokenizer
diff --git a/docs/api/transformers.md b/docs/api/transformers.md
new file mode 100644
index 0000000..67c23c3
--- /dev/null
+++ b/docs/api/transformers.md
@@ -0,0 +1,5 @@
+# Transformers (`depthcharge.transformers`)
+
+::: depthcharge.transformers.SpectrumTransformerEncoder
+::: depthcharge.transformers.PeptideTransformerEncoder
+::: depthcharge.transformers.PeptideTransformerDecoder
diff --git a/docs/getting-started/installation.md b/docs/getting-started/installation.md
new file mode 100644
index 0000000..b613a68
--- /dev/null
+++ b/docs/getting-started/installation.md
@@ -0,0 +1,31 @@
+# Installating Depthcharge
+
+Depthcharge can currently be installed with pip, and will eventually be made available through [Bioconda](https://bioconda.github.io).
+
+## Requirements
+
+### Python
+Depthcharge is a Python package and requires Python >= 3.10.
+To check what version of Python you have installed, you can run:
+
+``` sh
+$ python --version
+```
+
+If you do not have Python installed, we recommend Miniconda, which also comes with the conda package manager.
+See the [Miniconda documentation](https://docs.conda.io/en/latest/miniconda.html) for details.
+
+### PyTorch (*optional*)
+Depthcharge creates PyTorch modules which you can use to create deep learning models for your application.
+PyTorch can be installed automatically during Depthcharge installation; however, to get the most out of Depthcharge and PyTorch, you'll likely want to install PyTorch manually to account for any GPUs and their drivers that you may have installed.
+To install PyTorch, follow the installation instructions found in the [PyTorch documentation](https://pytorch.org/get-started/locally/).
+
+
+## Install depthcharge
+
+Depthcharge can be install with `pip`, directly from PyPI:
+```sh
+$ pip install depthcharge-ms
+```
+
+We have plans to also add it to Bioconda!
diff --git a/docs/index.md b/docs/index.md
new file mode 100644
index 0000000..7715ac1
--- /dev/null
+++ b/docs/index.md
@@ -0,0 +1,21 @@
+![depthcharge logo](./static/logo-dark.png#gh-dark-mode-only)
+
+Depthcharge is a deep learning toolkit for building Transformer models to analyze mass spectrometry data.
+
+## About
+
+Many deep learning tools have been developed for the analysis of mass spectra or mass spectrometry analytes, like peptides and small molecules.
+However, each one has had to reinvent the wheel.
+
+Depthcharge aims to provide a flexible, but opinionated, framework for rapidly prototyping deep learning models for mass spectrometry data.
+Think of depthcharge as a set of building blocks to get you started on a new deep learning project focused around mass spectrometry data.
+Depthcharge delivers these building blocks in the form of PyTorch modules, which can be readily used to assemble customized deep learning models for your task.
+
+
+<style>
+  .md-typeset h1 {
+    visibility: hidden;
+    font-size: 2px;
+    margin: 0px;
+  }
+</style>
diff --git a/docs/stylesheets/extra.css b/docs/stylesheets/extra.css
index b55513c..b4fb1cd 100644
--- a/docs/stylesheets/extra.css
+++ b/docs/stylesheets/extra.css
@@ -1,12 +1,52 @@
-:root {
-  --md-primary-fg-color:             #404040;
+ [data-md-color-scheme="depthcharge"]{
+  --md-default-fg-color:             #EEEEEE;
+  --md-default-fg-color--light:      #EEEEEE;
+  --md-primary-fg-color:             #101010;
   --md-primary-fg-color--light:      #01BCA3;
-  --md-primary-fg-color--dark:       #404040;
+  --md-primary-fg-color--dark:       #303030;
   --md-accent-fg-color:              #01BCA3;
   --md-accent-fg-color--transparent: #01BCA3;
+  --md-accent-fg-color--transparent: #01BCA3;
+  --md-typeset-a-color:              #01BCA3;
+  --md-default-bg-color:             #1d1f21;
+  --md-primary-bg-color--light:      #1d1f21;
+  --md-primary-bg-color--dark:       #1d1f21;
+  --md-footer-bg-color--dark:        #101010;
+  --md-typeset-color:                #EEEEEE;
+  --md-typeset-table-color:          #01BCA3;
 
-}
+  --md-code-fg-color: #eeeeee;
+  --md-code-bg-color: #101010;
+  --md-code-hl-color: #ffff0080;
+  --md-code-hl-number-color: #D19A66;
+  --md-code-hl-special-color: #db1457;
+  --md-code-hl-function-color: #61AFEF;
+  --md-code-hl-constant-color: #6e59d9;
+  --md-code-hl-keyword-color: #C678DD;
+  --md-code-hl-string-color: #98C379;
+  --md-code-hl-name-color: #E06C75;
+  --md-code-hl-punctuation-color: #ABB2BF;
+  --md-code-hl-operator-color: #56B6C2;
+  --md-code-hl-comment-color: #7F848E;
+  --md-code-hl-generic-color: var(--md-default-fg-color--light);
+  --md-code-hl-variable-color: var(--md-default-fg-color--light);
+ }
 
 p a {
     text-decoration: underline;
 }
+
+.md-typeset .doc code {
+  background-color: var(--md-default-bg-color);
+  color: var(--md-accent-fg-color);
+}
+
+.doc .highlight .n {
+  color: var(--md-accent-fg-color);
+}
+
+.doc .highlight .o,
+.doc .highlight .p,
+.doc .highlight .kc {
+  color: var(--md-default-fg-color)
+}
diff --git a/mkdocs.yml b/mkdocs.yml
index c8b2ff5..19c853f 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -1,18 +1,30 @@
 site_name: depthcharge
 repo_url: https://github.com/wfondrie/depthcharge
 nav:
-  - Home: "README.md"
+  - Getting Started:
+      - Introduction: "index.md"
+      - Installation: "getting-started/installation.md"
+  - API:
+      - Overview: "api/index.md"
+      - Tokenizers: "api/tokenizers.md"
+      - Datasets: "api/datasets.md"
+      - Encoders: "api/encoders.md"
+      - Transformers: "api/transformers.md"
+      - Primitives: "api/primitives.md"
 
 theme:
   name: material
   features:
     - navigation.tabs
-    - navigation.sections
-    - toc.integrate
+    #- navigation.sections
+    #- toc.integrate
+    - content.code.copy
   logo: static/icon.svg
   favicon: static/icon.svg
   font:
-    text: Roboto
+    text: Fira Code
+  palette:
+    scheme: depthcharge
 
 extra_css:
   - stylesheets/extra.css
@@ -22,14 +34,18 @@ plugins:
       default_handler: python
       handlers:
         python:
-          selection:
+          options:
             docstring_style: numpy
-          rendering:
             show_root_heading: true
             heading_level: 3
             docstring_section_style: spacy
             show_source: false
             merge_init_into_class: true
+            docstring_options:
+              ignore_init_summary: true
+              trim_doctest_flags: true
+            options:
+              members: true
 
 markdown_extensions:
   - attr_list
@@ -37,6 +53,8 @@ markdown_extensions:
   - footnotes
   - pymdownx.highlight:
       anchor_linenums: true
+      line_spans: __span
+      pygments_lang_class: true
   - pymdownx.inlinehilite
   - pymdownx.snippets
   - pymdownx.superfences
diff --git a/pyproject.toml b/pyproject.toml
index 11f14dc..14d10d8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,16 +1,85 @@
 [build-system]
-requires = ["setuptools>=42", "wheel", "setuptools_scm[toml]>=3.4"]
+requires = ["setuptools>=61", "wheel", "setuptools_scm[toml]>=3.4"]
 build-backend = "setuptools.build_meta"
 
+[project]
+name = "depthcharge-ms"
+authors = [{name = "William E. Fondrie", email = "fondriew@gmail.com"}]
+description = "A mass spectrometry toolkit for deep learning with Transformer models."
+license = {text = "Apache 2.0"}
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: Apache Software License",
+    "Operating System :: OS Independent",
+    "Topic :: Scientific/Engineering :: Bio-Informatics",
+]
+requires-python = ">=3.10"
+dependencies = [
+    "torch>=1.11.0",
+    "pyteomics>=4.4.2",
+    "numpy>=1.18.1",
+    "numba>=0.48.0",
+    "lxml>=4.9.1",
+    "h5py>=3.7.0",
+    "einops>=0.4.1",
+    "tqdm>=4.65.0",
+    "lark>=1.1.4",
+    "selfies>=2.1.1",
+    "sortedcontainers>=2.4.0",
+    "dill>=0.3.6",
+    "rdkit>=2023.03.1",
+    "pillow>=9.4.0",
+    "spectrum-utils>=0.4.1",
+]
+dynamic = ["version"]
+
+[project.readme]
+file = "REAMDE.md"
+content-type = "text/markdown"
+
+[project.urls]
+Homepage = "https://github.com/wfondrie/depthcharge"
+Documentation = "https://wfondrie.github.io/depthcharge"
+
+[project.optional-dependencies]
+docs = [
+     "mkdocs",
+     "mkdocs-material",
+     "mkdocstrings[python]>=0.18",
+]
+dev = [
+    "pre-commit>=2.7.1"
+]
+
+[tool.setuptools]
+include-package-data = false
+
+[tool.setuptools.packages]
+find = {namespaces = false}
+
 [tool.setuptools_scm]
 
 [tool.ruff]
+select = ["E", "F", "W", "C", "I", "D", "UP", "N", "ANN", "T20"]
+
+# ANN101 Missing type annotation for `self` in method
+# D213 Multi-line docstring summary should start at the second lin
+# D203 1 blank line required before class docstring
+# D100 Missing docstring in public module
+# ANN102 Missing type annotation for `cls` in classmethod
+# D401 First line of docstring should be in imperative mood
+ignore = ["D213", "ANN101", "D203", "D100", "ANN102", "D401"]
+fix = true
+
 [tool.ruff.per-file-ignores]
-"__init__.py" = ["F401"]
+"*tests/*.py" = ["ANN", "N806"]
+"__init__.py" = ["F401", "D104"]
+"depthcharge/encoders/sinusoidal.py" = ["N803"]
+"depthcharge/feedforward.py" = ["N803"]
 
 [tool.black]
 line-length = 79
-target-version = ['py37']
+target-version = ['py310']
 include = '\.pyi?$'
 exclude = '''
 
diff --git a/setup.cfg b/setup.cfg
deleted file mode 100644
index 0b7b247..0000000
--- a/setup.cfg
+++ /dev/null
@@ -1,50 +0,0 @@
-[metadata]
-name = depthcharge-ms
-author = William E. Fondrie
-author_email = fondriew@gmail.com
-description =
-    A deep learning toolkit for proteomics
-long_description = file: README.md
-long_description_content_type = text/markdown
-url = https://github.com/wfondrie/depthcharge
-project_urls =
-    Documentation = https://depthcharge.readthedocs.io
-    Bug Tracker = https://github.com/wfondrie/depthcharge/issues
-    Discussion Board = https://github.com/wfondrie/depthcharge/discussions
-license = Apache 2.0
-classifiers =
-    Programming Language :: Python :: 3
-    License :: OSI Approved :: Apache Software License
-    Operating System :: OS Independent
-    Topic :: Scientific/Engineering :: Bio-Informatics
-
-[options]
-packages = find:
-python_requires = >=3.8
-install_requires =
-    torch>=1.11.0
-    pyteomics>=4.4.2
-    numpy>=1.18.1
-    pandas>=1.0.3
-    numba>=0.48.0
-    lxml>=4.9.1
-    h5py>=3.7.0
-    einops>=0.4.1
-    tqdm>=4.65.0
-
-[options.extras_require]
-docs =
-    mkdocs
-    mkdocs-material
-    mkdocstrings
-    mkdocstrings-python
-dev =
-    pre-commit>=2.7.1
-    black>=19.10b0
-tensorboard =
-    tesnsorboard>=2.11.0
-
-
-[options.entry_points]
-console_scripts =
-    depthcharge = depthcharge.depthcharge:main
diff --git a/setup.py b/setup.py
deleted file mode 100644
index 1bcd1f2..0000000
--- a/setup.py
+++ /dev/null
@@ -1,4 +0,0 @@
-"""Setup the depthcharge package."""
-import setuptools
-
-setuptools.setup()
diff --git a/static/logo-dark.png b/static/logo-dark.png
index 07c7262..36622be 100644
Binary files a/static/logo-dark.png and b/static/logo-dark.png differ
diff --git a/static/logo-light.png b/static/logo-light.png
index 5dd44cb..37f3c5d 100644
Binary files a/static/logo-light.png and b/static/logo-light.png differ


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 05f370f..b3c6755 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,8 +1,8 @@
-"""Pytest fixtures"""
+"""Pytest fixtures."""
 from pathlib import Path
 
-import pytest
 import numpy as np
+import pytest
 from pyteomics.mass import calculate_mass
 
 DATA_DIR = Path(__file__).parent / ".." / "data"
@@ -36,7 +36,7 @@ def mgf_medium(tmp_path):
 @pytest.fixture
 def mgf_small(tmp_path):
     """An MGF file with 2 annotated spectra."""
-    peptides = ["LESLIEK", "PEPTIDEK"]
+    peptides = ["LESLIEK", "EDITHR"]
     mgf_file = tmp_path / "small.mgf"
     return _create_mgf(peptides, mgf_file)
 
@@ -107,6 +107,8 @@ def _create_mgf(peptides, mgf_file, add_problems=False, random_state=42):
 
     Returns
     -------
+    PathLike
+        The MGF file.
     """
     rng = np.random.default_rng(random_state)
     peptides = list(peptides)
@@ -132,7 +134,7 @@ def _create_mgf(peptides, mgf_file, add_problems=False, random_state=42):
 
 
 def _random_peptides(n_peptides, random_state=42):
-    """Create random peptides
+    """Create random peptides.
 
     Parameters
     ----------
diff --git a/tests/unit_tests/test_components/test_transformers.py b/tests/unit_tests/test_components/test_transformers.py
deleted file mode 100644
index 5ae1f6a..0000000
--- a/tests/unit_tests/test_components/test_transformers.py
+++ /dev/null
@@ -1,72 +0,0 @@
-"""Test the transformer components"""
-import torch
-from depthcharge.masses import PeptideMass
-from depthcharge.components.transformers import (
-    SpectrumEncoder,
-    PeptideEncoder,
-    PeptideDecoder,
-)
-
-
-def test_spectrum_encoder():
-    """Test that a spectrum encoder will run."""
-    spectra = torch.tensor(
-        [
-            [[100.1, 0.1], [200.2, 0.2], [300.3, 0.3]],
-            [[400.4, 0.4], [500, 0.5], [0, 0]],
-        ]
-    )
-
-    model = SpectrumEncoder(8, 1, 12, dim_intensity=None)
-    emb, mask = model(spectra)
-    assert emb.shape == (2, 4, 8)
-    assert mask.sum() == 1
-
-    model = SpectrumEncoder(8, 1, 12, dim_intensity=4)
-    emb, mask = model(spectra)
-    assert emb.shape == (2, 4, 8)
-    assert mask.sum() == 1
-
-
-def test_peptide_encoder():
-    """Test that a peptide encoder will run"""
-    peptides = ["LESLIEK", "PEPTIDER", "EDITPEPR"]
-    charges = torch.tensor([2, 3, 3])
-
-    model = PeptideEncoder(8, 1, 12, max_charge=3)
-    emb, mask = model(peptides, charges)
-
-    vocab_size = len(PeptideMass("canonical").masses.keys()) + 1
-    assert model.vocab_size == vocab_size
-    assert emb.shape == (3, 10, 8)
-    assert mask.sum() == 1
-
-    sums = emb.sum(dim=1)
-    assert (
-        sums[
-            1,
-            :,
-        ]
-        != sums[2, :]
-    ).all()
-
-
-def test_peptide_decoder():
-    """Test that a peptide decoder will run"""
-    spectra = torch.tensor(
-        [
-            [[100.1, 0.1], [200.2, 0.2], [300.3, 0.3]],
-            [[400.4, 0.4], [500, 0.5], [0, 0]],
-        ]
-    )
-
-    peptides = ["LESLIEK", "PEPTIDER"]
-    precursors = torch.tensor([[100.0, 2], [200.0, 3]])
-
-    encoder = SpectrumEncoder(8, 1, 12)
-    memory, mem_mask = encoder(spectra)
-
-    decoder = PeptideDecoder(8, 1, 12, max_charge=3)
-    scores, tokens = decoder(peptides, precursors, memory, mem_mask)
-    assert scores.shape == (2, 10, decoder.vocab_size + 1)
-    assert tokens.shape == (2, 9)
diff --git a/tests/unit_tests/test_data/test_datasets.py b/tests/unit_tests/test_data/test_datasets.py
index 8dcd377..33984ac 100644
--- a/tests/unit_tests/test_data/test_datasets.py
+++ b/tests/unit_tests/test_data/test_datasets.py
@@ -1,12 +1,23 @@
-"""Test the datasets"""
+"""Test the datasets."""
+import functools
 import shutil
 
+import pytest
+import torch
+
 from depthcharge.data import (
-    AnnotatedSpectrumIndex,
-    SpectrumIndex,
     AnnotatedSpectrumDataset,
+    PeptideDataset,
     SpectrumDataset,
 )
+from depthcharge.primitives import MassSpectrum
+from depthcharge.tokenizers import PeptideTokenizer
+
+
+@pytest.fixture(scope="module")
+def tokenizer():
+    """Use a tokenizer for every test."""
+    return PeptideTokenizer()
 
 
 def test_spectrum_id(mgf_small, tmp_path):
@@ -14,14 +25,274 @@ def test_spectrum_id(mgf_small, tmp_path):
     mgf_small2 = tmp_path / "mgf_small2.mgf"
     shutil.copy(mgf_small, mgf_small2)
 
-    index = SpectrumIndex(tmp_path / "index.hdf5", [mgf_small, mgf_small2])
-    dataset = SpectrumDataset(index)
-    assert dataset.get_spectrum_id(0) == (str(mgf_small), "index=0")
-    assert dataset.get_spectrum_id(3) == (str(mgf_small2), "index=1")
+    dataset = SpectrumDataset([mgf_small, mgf_small2])
+    with dataset:
+        assert dataset.get_spectrum_id(0) == (str(mgf_small), "index=0")
+        assert dataset.get_spectrum_id(3) == (str(mgf_small2), "index=1")
+
+    dataset = AnnotatedSpectrumDataset(tokenizer, [mgf_small, mgf_small2])
+    with dataset:
+        assert dataset.get_spectrum_id(0) == (str(mgf_small), "index=0")
+        assert dataset.get_spectrum_id(3) == (str(mgf_small2), "index=1")
+
+
+def test_indexing(mgf_small, tmp_path):
+    """Test retrieving spectra."""
+    mgf_small2 = tmp_path / "mgf_small2.mgf"
+    shutil.copy(mgf_small, mgf_small2)
+
+    dataset = SpectrumDataset([mgf_small, mgf_small2])
+    spec = dataset[0]
+    assert isinstance(spec, MassSpectrum)
+    assert spec.label is None
 
-    index = AnnotatedSpectrumIndex(
-        tmp_path / "index2.hdf5", [mgf_small, mgf_small2]
+    dataset = AnnotatedSpectrumDataset(tokenizer, [mgf_small, mgf_small2])
+    spec = dataset[0]
+    assert isinstance(spec, MassSpectrum)
+    assert spec.label == "LESLIEK"
+    assert spec.to_tensor().shape[1] == 2
+    assert dataset[3].label == "EDITHR"
+
+
+def test_valid_charge(mgf_medium, tmp_path):
+    """Test that the valid_charge argument works."""
+    mkindex = functools.partial(
+        SpectrumDataset,
+        index_path=tmp_path / "index.hdf5",
+        ms_data_files=mgf_medium,
+        overwrite=True,
     )
-    dataset = AnnotatedSpectrumDataset(index)
-    assert dataset.get_spectrum_id(0) == (str(mgf_small), "index=0")
-    assert dataset.get_spectrum_id(3) == (str(mgf_small2), "index=1")
+
+    index = mkindex()
+    assert index.n_spectra == 100
+
+    index = mkindex(valid_charge=[0, 2, 3, 4])
+    assert index.n_spectra == 100
+
+    index = mkindex(valid_charge=[2, 3, 4])
+    assert index.n_spectra == 99
+
+    index = mkindex(valid_charge=[2, 3])
+    assert index.n_spectra == 98
+
+
+def test_mgf(mgf_small, tmp_path):
+    """Test the mgf index."""
+    mgf_small2 = tmp_path / "mgf_small2.mgf"
+    shutil.copy(mgf_small, mgf_small2)
+
+    index = SpectrumDataset([mgf_small, mgf_small2])
+    assert index.ms_files == [str(mgf_small), str(mgf_small2)]
+    assert index.ms_level == 2
+    assert not index.annotated
+    assert not index.overwrite
+    assert index.n_spectra == 4
+    assert index.n_peaks == 66
+
+    with index:
+        assert index.get_spectrum_id(0) == (str(mgf_small), "index=0")
+        assert index.get_spectrum_id(3) == (str(mgf_small2), "index=1")
+
+    index = AnnotatedSpectrumDataset(tokenizer, [mgf_small, mgf_small2])
+    assert index.ms_files == [str(mgf_small), str(mgf_small2)]
+    assert index.ms_level == 2
+    assert index.annotated
+    assert not index.overwrite
+    assert index.n_spectra == 4
+    assert index.n_peaks == 66
+
+    with index:
+        assert index.get_spectrum_id(0) == (str(mgf_small), "index=0")
+        assert index.get_spectrum_id(3) == (str(mgf_small2), "index=1")
+
+
+def test_mzml_index(real_mzml, tmp_path):
+    """Test an mzML index."""
+    real_mzml2 = tmp_path / "real_mzml2.mzML"
+    shutil.copy(real_mzml, real_mzml2)
+
+    index = SpectrumDataset([real_mzml, real_mzml2], 2, [])
+    assert index.ms_files == [str(real_mzml), str(real_mzml2)]
+    assert index.ms_level == 2
+    assert not index.annotated
+    assert not index.overwrite
+    assert index.n_spectra == 8
+    assert index.n_peaks == 726
+
+    with index:
+        assert index.get_spectrum_id(0) == (str(real_mzml), "scan=501")
+        assert index.get_spectrum_id(7) == (str(real_mzml2), "scan=510")
+
+    # MS3
+    index = SpectrumDataset([real_mzml, real_mzml2], 3, [])
+    assert index.ms_files == [str(real_mzml), str(real_mzml2)]
+    assert index.ms_level == 3
+    assert index.n_spectra == 6
+    assert index.n_peaks == 194
+
+    with index:
+        assert index.get_spectrum_id(0) == (str(real_mzml), "scan=502")
+        assert index.get_spectrum_id(5) == (str(real_mzml2), "scan=508")
+
+    # MS1
+    index = SpectrumDataset([real_mzml, real_mzml2], 1, [])
+    assert index.ms_files == [str(real_mzml), str(real_mzml2)]
+    assert index.ms_level == 1
+    assert index.n_spectra == 8
+    assert index.n_peaks == 4316
+
+    with index:
+        assert index.get_spectrum_id(0) == (str(real_mzml), "scan=500")
+        assert index.get_spectrum_id(5) == (str(real_mzml2), "scan=503")
+
+
+def test_mzxml_index(real_mzxml, tmp_path):
+    """Test an mzXML index."""
+    real_mzxml2 = tmp_path / "real_mzxml2.mzXML"
+    shutil.copy(real_mzxml, real_mzxml2)
+
+    index = SpectrumDataset([real_mzxml, real_mzxml2], 2, [])
+    assert index.ms_files == [str(real_mzxml), str(real_mzxml2)]
+    assert index.ms_level == 2
+    assert not index.annotated
+    assert not index.overwrite
+    assert index.n_spectra == 8
+    assert index.n_peaks == 726
+
+    with index:
+        assert index.get_spectrum_id(0) == (str(real_mzxml), "scan=501")
+        assert index.get_spectrum_id(7) == (str(real_mzxml2), "scan=510")
+
+    # MS3
+    index = SpectrumDataset([real_mzxml, real_mzxml2], 3, [])
+    assert index.ms_files == [str(real_mzxml), str(real_mzxml2)]
+    assert index.ms_level == 3
+    assert index.n_spectra == 6
+    assert index.n_peaks == 194
+
+    with index:
+        assert index.get_spectrum_id(0) == (str(real_mzxml), "scan=502")
+        assert index.get_spectrum_id(5) == (str(real_mzxml2), "scan=508")
+
+    # MS1
+    index = SpectrumDataset([real_mzxml, real_mzxml2], 1, [])
+    assert index.ms_files == [str(real_mzxml), str(real_mzxml2)]
+    assert index.ms_level == 1
+    assert index.n_spectra == 8
+    assert index.n_peaks == 4316
+
+    with index:
+        assert index.get_spectrum_id(0) == (str(real_mzxml), "scan=500")
+        assert index.get_spectrum_id(5) == (str(real_mzxml2), "scan=503")
+
+
+def test_spectrum_index_reuse(mgf_small, tmp_path):
+    """Reuse a previously created (annotated) spectrum index."""
+    plain_index = tmp_path / "plain.hdf5"
+    ann_index = tmp_path / "ann.hdf5"
+
+    index = SpectrumDataset(mgf_small, index_path=plain_index)
+    index2 = SpectrumDataset(index_path=plain_index)
+    assert index.ms_level == index2.ms_level
+    assert index.annotated == index2.annotated
+    assert not index2.annotated
+    assert index.n_peaks == index2.n_peaks
+    assert index.n_spectra == index2.n_spectra
+
+    index3 = AnnotatedSpectrumDataset(
+        tokenizer, mgf_small, index_path=ann_index
+    )
+    index4 = AnnotatedSpectrumDataset(tokenizer, index_path=ann_index)
+    assert index3.ms_level == index4.ms_level
+    assert index3.annotated == index4.annotated
+    assert index4.annotated
+    assert index3.n_peaks == index4.n_peaks
+    assert index3.n_spectra == index4.n_spectra
+
+    # An annotated spectrum dataset may be loaded as a spectrum dataset,
+    # but not vice versa:
+    SpectrumDataset(index_path=ann_index)
+    with pytest.raises(ValueError):
+        AnnotatedSpectrumDataset(tokenizer, index_path=plain_index)
+
+    # Verify we invalidate correctly:
+    with pytest.raises(ValueError):
+        SpectrumDataset(index_path=plain_index, ms_level=1)
+
+    with pytest.raises(ValueError):
+        SpectrumDataset(index_path=plain_index, preprocessing_fn=[])
+
+    with pytest.raises(ValueError):
+        AnnotatedSpectrumDataset(tokenizer, index_path=ann_index, ms_level=3)
+
+    with pytest.raises(ValueError):
+        AnnotatedSpectrumDataset(
+            tokenizer, index_path=plain_index, preprocessing_fn=[]
+        )
+
+
+def test_spectrum_indexing_bug(tmp_path, mgf_small):
+    """Test that we've fixed reindexing upon reload."""
+    dset1 = AnnotatedSpectrumDataset(
+        PeptideTokenizer(), mgf_small, 2, index_path=tmp_path / "test.hdf5"
+    )
+
+    dset2 = AnnotatedSpectrumDataset(
+        PeptideTokenizer(), mgf_small, 2, index_path=tmp_path / "test.hdf5"
+    )
+
+    assert dset1._locs == dset2._locs
+
+
+def test_preprocessing_fn(mgf_small):
+    """Test preprocessing functions."""
+    dset = SpectrumDataset(mgf_small)
+    loader = dset.loader(batch_size=1, num_workers=0)
+
+    spec, *_ = next(iter(loader))
+    assert (spec[:, :, 1] < 1).all()
+
+    dset = SpectrumDataset(mgf_small, preprocessing_fn=[])
+    loader = dset.loader(batch_size=1, num_workers=0)
+
+    spec, *_ = next(iter(loader))
+    assert (spec[:, :, 1] == 1).all()
+
+    def my_func(spec):
+        """A simple test function."""
+        spec.intensity[:] = 2.0
+        return spec
+
+    dset = SpectrumDataset(mgf_small, preprocessing_fn=my_func)
+    loader = dset.loader(batch_size=1, num_workers=0)
+    spec, *_ = next(iter(loader))
+    assert (spec[:, :, 1] == 2).all()
+
+
+def test_peptide_dataset():
+    """Test the peptide dataset."""
+    tokenizer = PeptideTokenizer()
+    seqs = ["LESLIEK", "EDITHR"]
+    charges = torch.tensor([2, 3])
+    dset = PeptideDataset(tokenizer, seqs, charges)
+    torch.testing.assert_close(dset[0][0], tokenizer.tokenize("LESLIEK"))
+    torch.testing.assert_close(dset[1][0][:6], tokenizer.tokenize("EDITHR"))
+    assert dset[0][1].item() == 2
+    assert dset[1][1].item() == 3
+    assert len(dset) == 2
+
+    seqs = ["LESLIEK", "EDITHR"]
+    charges = torch.tensor([2, 3])
+    target = torch.tensor([1.1, 2.2])
+    other = torch.tensor([[1, 1], [2, 2]])
+    dset = PeptideDataset(tokenizer, seqs, charges, target, other)
+    torch.testing.assert_close(dset[0][0], tokenizer.tokenize("LESLIEK"))
+    torch.testing.assert_close(dset[1][0][:6], tokenizer.tokenize("EDITHR"))
+    assert dset[0][1].item() == 2
+    assert dset[1][1].item() == 3
+    torch.testing.assert_close(dset[0][2], torch.tensor(1.1))
+    torch.testing.assert_close(dset[1][3], other[1, :])
+    assert len(dset) == 2
+
+    torch.testing.assert_close(dset.tokens, tokenizer.tokenize(seqs))
+    torch.testing.assert_close(dset.charges, charges)
diff --git a/tests/unit_tests/test_data/test_index.py b/tests/unit_tests/test_data/test_index.py
deleted file mode 100644
index 618eea7..0000000
--- a/tests/unit_tests/test_data/test_index.py
+++ /dev/null
@@ -1,143 +0,0 @@
-"""Test that the MS data file indices work correctly"""
-import shutil
-import functools
-
-from depthcharge.data import AnnotatedSpectrumIndex, SpectrumIndex
-
-
-def test_valid_charge(mgf_medium, tmp_path):
-    """Test that the valid_charge argument works"""
-    mkindex = functools.partial(
-        SpectrumIndex,
-        index_path=tmp_path / "index.hdf5",
-        ms_data_files=mgf_medium,
-        overwrite=True,
-    )
-
-    index = mkindex()
-    assert index.n_spectra == 100
-
-    index = mkindex(valid_charge=[0, 2, 3, 4])
-    assert index.n_spectra == 100
-
-    index = mkindex(valid_charge=[2, 3, 4])
-    assert index.n_spectra == 99
-
-    index = mkindex(valid_charge=[2, 3])
-    assert index.n_spectra == 98
-
-
-def test_mgf_index(mgf_small, tmp_path):
-    """Test the mgf index."""
-    mgf_small2 = tmp_path / "mgf_small2.mgf"
-    shutil.copy(mgf_small, mgf_small2)
-
-    index = SpectrumIndex(tmp_path / "index.hdf5", [mgf_small, mgf_small2])
-    assert index.ms_files == [str(mgf_small), str(mgf_small2)]
-    assert index.ms_level == 2
-    assert not index.annotated
-    assert not index.overwrite
-    assert index.n_spectra == 4
-    assert index.n_peaks == 92
-
-    with index:
-        assert index.get_spectrum_id(0) == (str(mgf_small), "index=0")
-        assert index.get_spectrum_id(3) == (str(mgf_small2), "index=1")
-
-    index = AnnotatedSpectrumIndex(
-        tmp_path / "index2.hdf5", [mgf_small, mgf_small2]
-    )
-    assert index.ms_files == [str(mgf_small), str(mgf_small2)]
-    assert index.ms_level == 2
-    assert index.annotated
-    assert not index.overwrite
-    assert index.n_spectra == 4
-    assert index.n_peaks == 92
-
-    with index:
-        assert index.get_spectrum_id(0) == (str(mgf_small), "index=0")
-        assert index.get_spectrum_id(3) == (str(mgf_small2), "index=1")
-
-
-def test_mzml_index(real_mzml, tmp_path):
-    """Test an mzML index"""
-    real_mzml2 = tmp_path / "real_mzml2.mzML"
-    shutil.copy(real_mzml, real_mzml2)
-
-    index = SpectrumIndex(tmp_path / "index.hdf5", [real_mzml, real_mzml2])
-    assert index.ms_files == [str(real_mzml), str(real_mzml2)]
-    assert index.ms_level == 2
-    assert not index.annotated
-    assert not index.overwrite
-    assert index.n_spectra == 8
-    assert index.n_peaks == 726
-
-    with index:
-        assert index.get_spectrum_id(0) == (str(real_mzml), "scan=501")
-        assert index.get_spectrum_id(7) == (str(real_mzml2), "scan=510")
-
-    # MS3
-    index = SpectrumIndex(tmp_path / "index2.hdf5", [real_mzml, real_mzml2], 3)
-    assert index.ms_files == [str(real_mzml), str(real_mzml2)]
-    assert index.ms_level == 3
-    assert index.n_spectra == 6
-    assert index.n_peaks == 194
-
-    with index:
-        assert index.get_spectrum_id(0) == (str(real_mzml), "scan=502")
-        assert index.get_spectrum_id(5) == (str(real_mzml2), "scan=508")
-
-    # MS1
-    index = SpectrumIndex(tmp_path / "index3.hdf5", [real_mzml, real_mzml2], 1)
-    assert index.ms_files == [str(real_mzml), str(real_mzml2)]
-    assert index.ms_level == 1
-    assert index.n_spectra == 8
-    assert index.n_peaks == 4316
-
-    with index:
-        assert index.get_spectrum_id(0) == (str(real_mzml), "scan=500")
-        assert index.get_spectrum_id(5) == (str(real_mzml2), "scan=503")
-
-
-def test_mzxml_index(real_mzxml, tmp_path):
-    """Test an mzXML index"""
-    real_mzxml2 = tmp_path / "real_mzxml2.mzXML"
-    shutil.copy(real_mzxml, real_mzxml2)
-
-    index = SpectrumIndex(tmp_path / "index.hdf5", [real_mzxml, real_mzxml2])
-    assert index.ms_files == [str(real_mzxml), str(real_mzxml2)]
-    assert index.ms_level == 2
-    assert not index.annotated
-    assert not index.overwrite
-    assert index.n_spectra == 8
-    assert index.n_peaks == 726
-
-    with index:
-        assert index.get_spectrum_id(0) == (str(real_mzxml), "scan=501")
-        assert index.get_spectrum_id(7) == (str(real_mzxml2), "scan=510")
-
-    # MS3
-    index = SpectrumIndex(
-        tmp_path / "index2.hdf5", [real_mzxml, real_mzxml2], 3
-    )
-    assert index.ms_files == [str(real_mzxml), str(real_mzxml2)]
-    assert index.ms_level == 3
-    assert index.n_spectra == 6
-    assert index.n_peaks == 194
-
-    with index:
-        assert index.get_spectrum_id(0) == (str(real_mzxml), "scan=502")
-        assert index.get_spectrum_id(5) == (str(real_mzxml2), "scan=508")
-
-    # MS1
-    index = SpectrumIndex(
-        tmp_path / "index3.hdf5", [real_mzxml, real_mzxml2], 1
-    )
-    assert index.ms_files == [str(real_mzxml), str(real_mzxml2)]
-    assert index.ms_level == 1
-    assert index.n_spectra == 8
-    assert index.n_peaks == 4316
-
-    with index:
-        assert index.get_spectrum_id(0) == (str(real_mzxml), "scan=500")
-        assert index.get_spectrum_id(5) == (str(real_mzxml2), "scan=503")
diff --git a/tests/unit_tests/test_data/test_loaders.py b/tests/unit_tests/test_data/test_loaders.py
index a45795e..b06cded 100644
--- a/tests/unit_tests/test_data/test_loaders.py
+++ b/tests/unit_tests/test_data/test_loaders.py
@@ -1,123 +1,68 @@
-"""Test PyTorch DataLoaders"""
+"""Test PyTorch DataLoaders."""
 import torch
+
 from depthcharge.data import (
-    SpectrumIndex,
-    AnnotatedSpectrumIndex,
-    SpectrumDataset,
     AnnotatedSpectrumDataset,
+    PeptideDataset,
+    SpectrumDataset,
 )
+from depthcharge.tokenizers import PeptideTokenizer
 
 
-def test_spectrum_index_init(mgf_small, tmp_path):
-    """Test initialization of with a SpectrumIndex"""
-    index = SpectrumIndex(tmp_path / "index.hdf5", mgf_small)
-    dset = SpectrumDataset(index)
-    loader = torch.utils.data.DataLoader(
-        dset,
-        batch_size=1,
-        collate_fn=dset.collate_fn,
-        num_workers=0,
-    )
+def test_spectrum_loader(mgf_small):
+    """Test initialization of with a SpectrumIndex."""
+    dset = SpectrumDataset(mgf_small, 2)
+    loader = dset.loader(batch_size=1, num_workers=0)
 
     batch = next(iter(loader))
     assert len(batch) == 2
-    assert batch[0].shape[0] == 1
-    assert batch[0].shape[2] == 2
+    assert batch[0].shape == (1, 13, 2)
     assert batch[1].shape == (1, 2)
 
-    dset = SpectrumDataset(index, n_peaks=3)
-    loader = torch.utils.data.DataLoader(
-        dset,
-        batch_size=1,
-        collate_fn=dset.collate_fn,
-        num_workers=0,
-    )
+    dset = SpectrumDataset(mgf_small, 2, [])
+    loader = dset.loader(batch_size=1, num_workers=0)
 
     batch = next(iter(loader))
     assert len(batch) == 2
-    assert batch[0].shape[0] == 1
-    assert batch[0].shape[1] == 3
-    assert batch[0].shape[2] == 2
+    assert batch[0].shape == (1, 14, 2)
     assert batch[1].shape == (1, 2)
 
 
-def test_ann_spectrum_index_init(mgf_small, tmp_path):
-    """Test initialization of with a SpectrumIndex"""
-    index = AnnotatedSpectrumIndex(tmp_path / "index.hdf5", mgf_small)
-    dset = AnnotatedSpectrumDataset(index)
-    loader = torch.utils.data.DataLoader(
-        dset,
-        batch_size=1,
-        num_workers=0,
-        collate_fn=dset.collate_fn,
-    )
+def test_ann_spectrum_loader(mgf_small):
+    """Test initialization of with a SpectrumIndex."""
+    tokenizer = PeptideTokenizer()
+    dset = AnnotatedSpectrumDataset(tokenizer, mgf_small)
+    loader = dset.loader(batch_size=1, num_workers=0)
 
     batch = next(iter(loader))
     assert len(batch) == 3
-    assert batch[0].shape[0] == 1
-    assert batch[0].shape[2] == 2
+    assert batch[0].shape == (1, 13, 2)
     assert batch[1].shape == (1, 2)
-    assert batch[2].shape == (1,)
-
-
-def test_spectrum_index_reuse(mgf_small, tmp_path):
-    """Reuse a previously created (annotated) spectrum index."""
-    index = SpectrumIndex(tmp_path / "index.hdf5", mgf_small)
-    index2 = SpectrumIndex(tmp_path / "index.hdf5")
-    assert index.ms_level == index2.ms_level
-    assert index.annotated == index2.annotated
-    assert not index2.annotated
-    assert index.n_peaks == index2.n_peaks
-    assert index.n_spectra == index2.n_spectra
-
-    index = AnnotatedSpectrumIndex(
-        tmp_path / "annotated_index.hdf5", mgf_small
-    )
-    index2 = AnnotatedSpectrumIndex(tmp_path / "annotated_index.hdf5")
-    assert index.ms_level == index2.ms_level
-    assert index.annotated == index2.annotated
-    assert index2.annotated
-    assert index.n_peaks == index2.n_peaks
-    assert index.n_spectra == index2.n_spectra
-
+    assert batch[2].shape == (1, 7)
 
-def test_preprocessing_fn(mgf_small, tmp_path):
-    """Test preprocessing functions."""
-    index = SpectrumIndex(tmp_path / "index.hdf5", mgf_small)
-    dset = SpectrumDataset(index)
-    loader = torch.utils.data.DataLoader(
-        dset,
-        batch_size=1,
-        collate_fn=dset.collate_fn,
-        num_workers=0,
-    )
 
-    spec, *_ = next(iter(loader))
-    assert (spec[:, :, 1] < 1).all()
+def test_peptide_loader():
+    """Test our peptid data loader."""
+    seqs = ["LESLIE", "EDITH", "PEPTIDE"]
+    charges = torch.tensor([5, 3, 1])
+    tokenizer = PeptideTokenizer()
+    dset = PeptideDataset(tokenizer, seqs, charges)
+    loader = dset.loader(batch_size=2, num_workers=0)
 
-    dset = SpectrumDataset(index, preprocessing_fn=[])
-    loader = torch.utils.data.DataLoader(
-        dset,
-        batch_size=1,
-        collate_fn=dset.collate_fn,
-        num_workers=0,
+    batch = next(iter(loader))
+    assert len(batch) == 2
+    torch.testing.assert_close(
+        batch[0],
+        tokenizer.tokenize(seqs)[:2, :],
     )
+    torch.testing.assert_close(batch[1], torch.tensor(charges[:2], dtype=int))
 
-    spec, *_ = next(iter(loader))
-    assert (spec[:, :, 1] == 1).all()
+    args = (torch.tensor([1, 2, 3]), torch.tensor([[1, 1], [2, 2], [3, 3]]))
+    dset = PeptideDataset(tokenizer, seqs, charges, *args)
+    loader = dset.loader(batch_size=2, num_workers=0)
 
-    def my_func(mz_array, int_array, precursor_mz, precursor_charge):
-        """A simple test function"""
-        int_array[:] = 2.0
-        return mz_array, int_array
-
-    dset = SpectrumDataset(index, preprocessing_fn=my_func)
-    loader = torch.utils.data.DataLoader(
-        dset,
-        batch_size=1,
-        collate_fn=dset.collate_fn,
-        num_workers=0,
-    )
-
-    spec, *_ = next(iter(loader))
-    assert (spec[:, :, 1] == 2).all()
+    batch = next(iter(loader))
+    assert len(batch) == 4
+    torch.testing.assert_close(batch[1], torch.tensor(charges[:2], dtype=int))
+    torch.testing.assert_close(batch[2], torch.tensor([1, 2]))
+    torch.testing.assert_close(batch[3], args[1][:2, :])
diff --git a/tests/unit_tests/test_data/test_parsers.py b/tests/unit_tests/test_data/test_parsers.py
index 840fbda..19f715f 100644
--- a/tests/unit_tests/test_data/test_parsers.py
+++ b/tests/unit_tests/test_data/test_parsers.py
@@ -1,64 +1,57 @@
 """Test that parsers work."""
 
 import numpy as np
-from depthcharge.data.parsers import MzmlParser, MgfParser, MzxmlParser
+
+from depthcharge.data.parsers import MgfParser, MzmlParser, MzxmlParser
 
 SMALL_MGF_MZS = [
     114.09134044,
-    831.48221068,
+    147.11280416,
     243.13393353,
-    718.3981467,
+    276.15539725,
     330.16596194,
-    589.35555361,
+    389.23946123,
     443.25002591,
     502.32352521,
     556.33408989,
-    389.23946123,
+    589.35555361,
     685.37668298,
-    276.15539725,
+    718.3981467,
     813.47164599,
-    147.11280416,
-    98.06004032,
-    928.46220351,
-    49.53365839,
-    464.73473999,
-    227.1026334,
-    831.40943966,
-    114.05495494,
-    416.20835806,
-    324.15539725,
-    702.36684657,
-    162.58133686,
-    351.68706152,
-    425.20307572,
-    605.31408272,
-    213.10517609,
-    303.16067959,
-    538.2871397,
-    504.26640425,
-    269.64720808,
-    252.63684036,
-    653.31408272,
-    391.18234028,
-    327.16067959,
-    196.09480837,
-    782.35667581,
-    276.15539725,
-    391.68197614,
-    138.58133686,
-    910.45163882,
-    147.11280416,
-    455.72945765,
-    74.06004032,
+    831.48221068,
+    65.52857301,
+    88.06311432,
+    123.04204452,
+    130.04986955,
+    156.59257025,
+    175.11895217,
+    179.58407651,
+    207.11640948,
+    230.10791575,
+    245.07681258,
+    263.65844147,
+    298.63737167,
+    312.17786403,
+    321.17191298,
+    358.16087656,
+    376.68792719,
+    385.69320953,
+    413.2255425,
+    459.20855502,
+    526.30960648,
+    596.26746688,
+    641.3365495,
+    752.36857791,
+    770.37914259,
 ]
 
 
 def test_mgf_and_base(mgf_small):
-    """MGF file with a missing charge"""
+    """MGF file with a missing charge."""
     parser = MgfParser(mgf_small).read()
     np.testing.assert_allclose(
         parser.precursor_mz,
-        np.array([416.24474357, 310.15891881]),
+        np.array([416.24474357, 257.464565]),
     )
     np.testing.assert_equal(
         parser.precursor_charge,
@@ -73,7 +66,9 @@ def test_mgf_and_base(mgf_small):
         np.array(SMALL_MGF_MZS),
     )
 
-    np.testing.assert_allclose(parser.intensity_arrays, np.ones(46, dtype=int))
+    np.testing.assert_allclose(
+        parser.intensity_arrays, np.ones(len(SMALL_MGF_MZS), dtype=int)
+    )
 
     parser = MgfParser(mgf_small, valid_charge=[2]).read()
     assert parser.precursor_charge.shape == (1,)
@@ -90,26 +85,26 @@ def test_mgf_and_base(mgf_small):
 
 
 def test_mzml(real_mzml):
-    """A simple mzML test"""
-    parser = MzmlParser(real_mzml).read()
+    """A simple mzML test."""
+    parser = MzmlParser(real_mzml, 2).read()
     prev_len = len(parser.mz_arrays)
     assert len(parser.intensity_arrays) == prev_len
     assert len(parser.precursor_charge) == 4
 
-    parser = MzmlParser(real_mzml, valid_charge=[3]).read()
+    parser = MzmlParser(real_mzml, 2, valid_charge=[3]).read()
     assert len(parser.precursor_charge) == 3
     assert len(parser.intensity_arrays) < prev_len
     assert len(parser.mz_arrays) < prev_len
 
 
 def test_mzxml(real_mzxml):
-    """A simple mzML test"""
-    parser = MzxmlParser(real_mzxml).read()
+    """A simple mzML test."""
+    parser = MzxmlParser(real_mzxml, 2).read()
     prev_len = len(parser.mz_arrays)
     assert len(parser.intensity_arrays) == prev_len
     assert len(parser.precursor_charge) == 4
 
-    parser = MzxmlParser(real_mzxml, valid_charge=[3]).read()
+    parser = MzxmlParser(real_mzxml, 2, valid_charge=[3]).read()
     assert len(parser.precursor_charge) == 3
     assert len(parser.intensity_arrays) < prev_len
     assert len(parser.mz_arrays) < prev_len
diff --git a/tests/unit_tests/test_components/test_encoders.py b/tests/unit_tests/test_encoders/test_sinusoidal.py
similarity index 70%
rename from tests/unit_tests/test_components/test_encoders.py
rename to tests/unit_tests/test_encoders/test_sinusoidal.py
index ce3e85d..45377eb 100644
--- a/tests/unit_tests/test_components/test_encoders.py
+++ b/tests/unit_tests/test_encoders/test_sinusoidal.py
@@ -1,17 +1,17 @@
-"""Test the encoders"""
+"""Test the encoders."""
+import numpy as np
 import pytest
 import torch
-import numpy as np
 
-from depthcharge.components.encoders import (
-    PositionalEncoder,
+from depthcharge.encoders import (
     FloatEncoder,
     PeakEncoder,
+    PositionalEncoder,
 )
 
 
 def test_positional_encoder():
-    """Test the positional encoder"""
+    """Test the positional encoder."""
     enc = PositionalEncoder(8, 1, 8)
     X = torch.zeros(1, 9, 8)
     Y = enc(X)
@@ -42,15 +42,13 @@ def test_positional_encoder():
 
 
 def test_float_encoder():
-    """Test the float encodings"""
+    """Test the float encodings."""
     enc = FloatEncoder(8, 0.1, 10)
     X = torch.tensor([[0, 0.1, 10, 0.256]])
     Y = enc(X)
     period = torch.cat([torch.zeros(4), torch.ones(4)], axis=0)
     torch.testing.assert_close(Y[0, 0, :], period)
     torch.testing.assert_close(Y[0, 1, (0, 4)], torch.tensor([0.0, 1]))
-
-    print(Y[0, 2])
     torch.testing.assert_close(Y[0, 2, (3, 7)], torch.tensor([0.0, 1]))
 
     # Check for things in-between the expected period:
@@ -67,9 +65,9 @@ def test_invalid_float_encoder():
         FloatEncoder(8, 10, 0)
 
 
-def test_default_peak_encoder():
-    """Test the peak encodings"""
-    enc = PeakEncoder(8)
+def test_legacy_peak_encoder():
+    """Test the peak encodings."""
+    enc = PeakEncoder(8, _legacy=True)
 
     # We want to test that the addition between the peak and
     # m/z encoding is happening:
@@ -83,30 +81,10 @@ def test_default_peak_encoder():
 
 
 def test_both_sinusoid():
-    """Test that both encoders are sinusoidal"""
-    enc = PeakEncoder(8, learned_intensity_encoding=False)
-
+    """Test that both encoders are sinusoidal."""
+    enc = PeakEncoder(8)
     X = torch.tensor([[[0.0, 0], [0, 1]]])
     Y = enc(X)
 
     assert Y.shape == (1, 2, 8)
     assert isinstance(enc.int_encoder, FloatEncoder)
-
-
-def test_split_dims():
-    """Test that dimensions are split correctly"""
-    enc = PeakEncoder(8, 4)
-
-    # We want to test that the addition between the peak and
-    # m/z encoding is happening:
-    enc.int_encoder.weight = torch.nn.Parameter(torch.ones(4, 1) * 5)
-
-    X = torch.tensor([[[0.0, 0], [0, 1]]])
-    Y = enc(X)
-
-    assert Y.shape == (1, 2, 8)
-    torch.testing.assert_close(Y[0, 1, 4:], torch.ones(4) * 5)
-    assert (Y[0, 1, :4] <= 1).all()
-
-    with pytest.raises(ValueError):
-        PeakEncoder(8, 9)
diff --git a/tests/unit_tests/test_feedforward.py b/tests/unit_tests/test_feedforward.py
new file mode 100644
index 0000000..51fd9d5
--- /dev/null
+++ b/tests/unit_tests/test_feedforward.py
@@ -0,0 +1,56 @@
+"""Test the feedforward model."""
+import torch
+
+from depthcharge.feedforward import FeedForward
+
+
+def test_init():
+    """Test initialization."""
+    X = torch.randn(20, 3, 10)
+
+    model = FeedForward(10, 1, 2)
+    assert isinstance(model.layers[0], torch.nn.Linear)
+    assert isinstance(model.layers[1], torch.nn.LeakyReLU)
+    assert isinstance(model.layers[2], torch.nn.Linear)
+    assert model.layers[0].in_features == 10
+    assert model.layers[0].out_features == 6
+    assert model.layers[2].in_features == 6
+    assert model.layers[2].out_features == 1
+    assert len(model.layers) == 3
+    Y = model(X)
+    assert Y.shape == (20, 3, 1)
+
+    model = FeedForward(10, 1, 2, append=torch.nn.GELU())
+    assert len(model.layers) == 4
+    assert isinstance(model.layers[-1], torch.nn.GELU)
+    Y = model(X)
+    assert Y.shape == (20, 3, 1)
+
+    model = FeedForward(10, 1, 2, dropout=0.1)
+    assert len(model.layers) == 4
+    assert isinstance(model.layers[2], torch.nn.Dropout)
+    assert model.layers[2].p == 0.1
+    Y = model(X)
+    assert Y.shape == (20, 3, 1)
+
+    model = FeedForward(10, 1, (50, 3))
+    assert isinstance(model.layers[0], torch.nn.Linear)
+    assert isinstance(model.layers[1], torch.nn.LeakyReLU)
+    assert isinstance(model.layers[2], torch.nn.Linear)
+    assert isinstance(model.layers[3], torch.nn.LeakyReLU)
+    assert isinstance(model.layers[4], torch.nn.Linear)
+    assert model.layers[0].in_features == 10
+    assert model.layers[0].out_features == 50
+    assert model.layers[2].in_features == 50
+    assert model.layers[2].out_features == 3
+    assert model.layers[4].in_features == 3
+    assert model.layers[4].out_features == 1
+    assert len(model.layers) == 5
+    Y = model(X)
+    assert Y.shape == (20, 3, 1)
+
+    model = FeedForward(10, 1, 2, activation=torch.nn.GELU())
+    assert len(model.layers) == 3
+    assert isinstance(model.layers[1], torch.nn.GELU)
+    Y = model(X)
+    assert Y.shape == (20, 3, 1)
diff --git a/tests/unit_tests/test_masses.py b/tests/unit_tests/test_masses.py
deleted file mode 100644
index 5e9e465..0000000
--- a/tests/unit_tests/test_masses.py
+++ /dev/null
@@ -1,69 +0,0 @@
-"""Test that peptide masses are calculated correctly."""
-import math
-from functools import partial
-
-from pyteomics import mass
-from depthcharge.masses import PeptideMass
-
-
-def test_string():
-    """Test peptide strings"""
-    pep = PeptideMass(residues="massivekb")
-    aa_mass = dict(mass.std_aa_mass)
-    aa_mass["a"] = 42.010565
-    aa_mass["o"] = 15.994915
-
-    pymass = partial(
-        mass.fast_mass,
-        ion_type="M",
-        aa_mass=aa_mass,
-    )
-
-    isclose = partial(math.isclose, rel_tol=1e-6)
-
-    seq = "LESLIEK"
-    assert isclose(mass.calculate_mass(seq), mass.fast_mass(seq))
-
-    assert isclose(pep.mass(seq), pymass(seq))
-    assert isclose(pep.mass(seq, 1), pymass(seq, charge=1))
-    assert isclose(pep.mass(seq, 2), pymass(seq, charge=2))
-    assert isclose(pep.mass(seq, 3), pymass(seq, charge=3))
-
-    seq = "+42.011LESLIEM+15.995K"
-    seq2 = "aLESLIEMoK"
-    assert isclose(pep.mass(seq), pymass(seq2))
-    assert isclose(pep.mass(seq, 1), pymass(seq2, charge=1))
-    assert isclose(pep.mass(seq, 2), pymass(seq2, charge=2))
-    assert isclose(pep.mass(seq, 3), pymass(seq2, charge=3))
-
-
-def test_list():
-    """Test peptides as a list."""
-    pep = PeptideMass(residues="massivekb")
-    aa_mass = dict(mass.std_aa_mass)
-    aa_mass["a"] = 42.010565
-    aa_mass["o"] = 15.994915
-
-    pymass = partial(
-        mass.fast_mass,
-        ion_type="M",
-        aa_mass=aa_mass,
-    )
-
-    isclose = partial(math.isclose, rel_tol=1e-6)
-
-    seq = "LESLIEK"
-    lseq = list(seq)
-    assert isclose(mass.calculate_mass(lseq), mass.fast_mass(seq))
-
-    assert isclose(pep.mass(lseq), pymass(seq))
-    assert isclose(pep.mass(lseq, 1), pymass(seq, charge=1))
-    assert isclose(pep.mass(lseq, 2), pymass(seq, charge=2))
-    assert isclose(pep.mass(lseq, 3), pymass(seq, charge=3))
-
-    seq = "aLESLIEMoK"
-    lseq = ["+42.011", "L", "E", "S", "L", "I", "E", "M+15.995", "K"]
-    assert isclose(pep.mass(lseq), pymass(seq))
-    assert isclose(pep.mass(lseq, 1), pymass(seq, charge=1))
-    assert isclose(pep.mass(lseq, 2), pymass(seq, charge=2))
-    assert isclose(pep.mass(lseq, 3), pymass(seq, charge=3))
diff --git a/tests/unit_tests/test_primitives.py b/tests/unit_tests/test_primitives.py
new file mode 100644
index 0000000..13b579b
--- /dev/null
+++ b/tests/unit_tests/test_primitives.py
@@ -0,0 +1,126 @@
+"""Test that our fundamental dataclasses work."""
+import numpy as np
+import pytest
+import torch
+from pyteomics.proforma import ProFormaError
+
+from depthcharge.primitives import (
+    MassSpectrum,
+    Molecule,
+    Peptide,
+    PeptideIons,
+)
+
+
+def test_peptide_init():
+    """Test our peptide dataclass."""
+    peptide = "LESLIEK"
+    assert Peptide(peptide).split() == list(peptide)
+    assert Peptide(peptide, None, None).charge is None
+    assert Peptide(peptide, None, 2).charge == 2
+
+    mods = [-10.0, None, None, 71.0, None, None, "Phospho", None, +1]
+    expected = [
+        "[-10.000000]-",
+        "L",
+        "E",
+        "S[+71.000000]",
+        "L",
+        "I",
+        "E[Phospho]",
+        "K",
+        "-[+1.000000]",
+    ]
+
+    parsed = Peptide(peptide, mods)
+    assert parsed.split() == expected
+
+    mods = [None, None, None, 71.0, None, None, "Phospho", None, None]
+    expected = [
+        "L",
+        "E",
+        "S[+71.000000]",
+        "L",
+        "I",
+        "E[Phospho]",
+        "K",
+    ]
+    parsed = Peptide(peptide, mods)
+    assert parsed.split() == expected
+
+
+def test_peptide_from_proforma():
+    """Test proforma parsing."""
+    parsed = Peptide.from_proforma("LESLIEK/2")
+    assert parsed.split() == list("LESLIEK")
+    assert parsed.charge == 2
+
+    parsed = Peptide.from_proforma("<[Phospho]@C,T>CAT")
+    assert parsed.split() == ["C[Phospho]", "A", "T[Phospho]"]
+
+    parsed = Peptide.from_proforma("ED[Phospho]IT[Obs:+1]H")
+    assert parsed.split() == ["E", "D[Phospho]", "I", "T[+1.000000]", "H"]
+
+    # If multiple mods exist on a single AA, the mass should be summed.
+    parsed = Peptide.from_proforma("<[Phospho]@C,T>C[+1]AT")
+    assert parsed.split() == ["C[+80.966331]", "A", "T[Phospho]"]
+
+    reparsed = Peptide.from_proforma(parsed.proforma)
+    assert parsed.split() == reparsed.split()
+
+
+def test_peptide_from_massivekb():
+    """Test MassIVE-KB parsing."""
+    parsed = Peptide.from_massivekb("LESLIEK")
+    assert parsed.split() == list("LESLIEK")
+
+    parsed = Peptide.from_massivekb("+42.011EDITH", 2)
+    assert parsed.charge == 2
+    assert parsed.split() == ["[Acetyl]-", "E", "D", "I", "T", "H"]
+
+    with pytest.raises(ProFormaError):
+        Peptide.from_massivekb("LES+79LIEK")
+
+
+def test_peptide_ions():
+    """Test peptide ions."""
+    ions = PeptideIons(["A", "B"], 42.0, torch.tensor([[[0]], [[1]]]))
+    assert ions.sequence == "AB"
+    assert ions.b_ions == torch.tensor([[0]])
+    assert ions.y_ions == torch.tensor([[1]])
+
+
+def test_molecule_init():
+    """Test that molecule initialize correctly."""
+    etoh1 = Molecule("CCO")
+    etoh2 = Molecule("C(O)C")
+    assert etoh1.smiles == etoh2.smiles
+    assert etoh1.to_selfies() == etoh2.to_selfies()
+
+    shown = etoh1.show()
+    assert shown is not None
+
+    etoh1_selfie = etoh1.to_selfies()
+    assert Molecule.from_selfies(etoh1_selfie) == etoh1
+
+    with pytest.raises(ValueError):
+        Molecule("not a smiles string")
+
+
+def test_mass_spectrum():
+    """Test that the mass spectrum is created correctly."""
+    mzs = np.arange(10)
+    intensities = np.arange(10) + 100
+
+    spec = MassSpectrum("test", "scan=1", mzs, intensities)
+    assert spec.usi == "test:scan:1"
+
+    spec = MassSpectrum(
+        "test",
+        "scan=1",
+        mzs,
+        intensities,
+        precursor_mz=10.0,
+        precursor_charge=3,
+    )
+    assert spec.to_tensor().shape == (10, 2)
diff --git a/tests/unit_tests/test_tokenizers/test_peptides.py b/tests/unit_tests/test_tokenizers/test_peptides.py
new file mode 100644
index 0000000..94a891c
--- /dev/null
+++ b/tests/unit_tests/test_tokenizers/test_peptides.py
@@ -0,0 +1,179 @@
+"""Test peptide tokenizers."""
+import math
+from functools import partial
+
+import pytest
+import torch
+from pyteomics import mass
+
+from depthcharge.tokenizers.peptides import PeptideTokenizer
+
+# Calculated using Pyteomics:
+# These are [b_ions, y_ions]
+LESLIEK_PLUS_ONE = [
+    [
+        114.09134044390001,
+        243.13393353187,
+        330.16596193614004,
+        443.25002591327006,
+        556.3340898904,
+        685.37668297837,
+    ],
+    [
+        147.11280416447,
+        276.15539725243997,
+        389.23946122957,
+        502.3235252067,
+        589.3555536109699,
+        718.39814669894,
+    ],
+]
+
+LESLIEK_PLUS_TWO = [
+    [
+        57.54930845533501,
+        122.07060499932,
+        165.58661920145502,
+        222.12865119002004,
+        278.670683178585,
+        343.19197972257,
+    ],
+    [
+        74.06004031561999,
+        138.581336859605,
+        195.12336884817,
+        251.66540083673502,
+        295.18141503886994,
+        359.702711582855,
+    ],
+]
+
+
+def test_proforma_init():
+    """Test initialization."""
+    seqs = ["[+5.]-LES[Phospho]LIE[-10.0]K"]
+    expected_tokens = [
+        ("S[Phospho]", 87.032028435 + 79.966331),
+        ("E[-10.000000]", 129.042593135 - 10.0),
+        ("[+5.000000]-", 5.0),
+    ]
+
+    proforma = PeptideTokenizer.from_proforma(
+        sequences=seqs,
+        replace_isoleucine_with_leucine=True,
+        reverse=False,
+    )
+
+    for key, val in expected_tokens:
+        assert proforma.residues[key] == val
+
+    tokens = proforma.tokenize(seqs, to_strings=True)[0]
+    expected = [
+        "[+5.000000]-",
+        "L",
+        "E",
+        "S[Phospho]",
+        "L",
+        "L",
+        "E[-10.000000]",
+        "K",
+    ]
+
+    assert tokens == expected
+
+    proforma = PeptideTokenizer.from_proforma(
+        sequences=seqs,
+        replace_isoleucine_with_leucine=False,
+        reverse=False,
+    )
+    tokens = proforma.tokenize(["LESLIEK"], to_strings=True)[0]
+    assert tokens == list("LESLIEK")
+    tokens = proforma.tokenize(["LESLIEK"])
+    orig = proforma.detokenize(tokens)
+    assert orig == ["LESLIEK"]
+
+    proforma = PeptideTokenizer.from_proforma(
+        sequences=seqs,
+        replace_isoleucine_with_leucine=False,
+        reverse=True,
+    )
+    tokens = proforma.tokenize(["LESLIEK"], to_strings=True)[0]
+    assert tokens == list("KEILSEL")
+    tokens = proforma.tokenize(["LESLIEK"])
+    orig = proforma.detokenize(tokens)
+    assert orig == ["KEILSEL"]
+
+    tokens = proforma.tokenize("LESLIEK", True, True)[0]
+    assert "".join(tokens) == "KEILSEL$"
+
+
+def test_mskb_init():
+    """Test that the MassIVE-KB dataset works."""
+    seqs = ["+42.011EDITH"]
+    mskb = PeptideTokenizer.from_massivekb(False, False)
+    tokens = mskb.tokenize(seqs, to_strings=True)[0]
+    assert tokens == ["[Acetyl]-", "E", "D", "I", "T", "H"]
+
+
+def test_precursor_ions():
+    """Test calculation of precurosr m/z."""
+    tokenizer = PeptideTokenizer()
+
+    aa_mass = dict(mass.std_aa_mass)
+    aa_mass["a"] = 42.010565
+    aa_mass["o"] = 15.994915
+    pymass = partial(mass.fast_mass, ion_type="M", aa_mass=aa_mass)
+    close = partial(math.isclose, rel_tol=1e-6)
+
+    seq = "LESLIEK"
+    assert close(tokenizer.ions(seq, 1)[0].precursor, pymass(seq, charge=1))
+    assert close(tokenizer.ions(seq, 2)[0].precursor, pymass(seq, charge=2))
+    assert close(tokenizer.ions(seq, 3)[0].precursor, pymass(seq, charge=3))
+
+    seq = "[Acetyl]-LESLIM[Oxidation]K"
+    with pytest.raises(ValueError):
+        tokenizer.ions(seq, 1)
+
+    tokenizer = PeptideTokenizer.from_proforma([seq])
+    seq2 = "aLESLIMoK"
+    assert close(tokenizer.ions(seq, 1)[0].precursor, pymass(seq2, charge=1))
+    assert close(tokenizer.ions(seq, 2)[0].precursor, pymass(seq2, charge=2))
+    assert close(tokenizer.ions(seq, 3)[0].precursor, pymass(seq2, charge=3))
+
+
+def test_fragment_ions():
+    """Test ion calculations."""
+    tokenizer = PeptideTokenizer()
+    ions = tokenizer.ions(["LESLIEK"], [1])[0]
+    expected = torch.tensor(LESLIEK_PLUS_ONE)[:, :, None]
+    torch.testing.assert_close(ions.fragments, expected, check_dtype=False)
+
+    ions = tokenizer.ions(["LESLIEK"], [2])[0]
+    expected = torch.cat(
+        [
+            torch.tensor(LESLIEK_PLUS_ONE)[:, :, None],
+            torch.tensor(LESLIEK_PLUS_TWO)[:, :, None],
+        ],
+        dim=2,
+    )
+    torch.testing.assert_close(ions.fragments, expected, check_dtype=False)
+
+    ions = tokenizer.ions(["LESLIEK/1"], None)[0]
+    expected = torch.tensor(LESLIEK_PLUS_ONE)[:, :, None]
+    torch.testing.assert_close(ions.fragments, expected, check_dtype=False)
+
+    ions = tokenizer.ions(["LESLIEK/3"], None)[0]
+    expected = torch.cat(
+        [
+            torch.tensor(LESLIEK_PLUS_ONE)[:, :, None],
+            torch.tensor(LESLIEK_PLUS_TWO)[:, :, None],
+        ],
+        dim=2,
+    )
+    torch.testing.assert_close(ions.fragments, expected, check_dtype=False)
+
+    tokenizer = PeptideTokenizer.from_proforma(["[+10]-LESLIEK"])
+    ions = tokenizer.ions(["[+10.000000]-LESLIEK"], 1)[0]
+    expected = torch.tensor(LESLIEK_PLUS_ONE)[:, :, None]
+    expected[0, :, :] += 10
+    torch.testing.assert_close(ions.fragments, expected, check_dtype=False)
diff --git a/tests/unit_tests/test_transformers/test_peptide_transformers.py b/tests/unit_tests/test_transformers/test_peptide_transformers.py
new file mode 100644
index 0000000..b345b6a
--- /dev/null
+++ b/tests/unit_tests/test_transformers/test_peptide_transformers.py
@@ -0,0 +1,49 @@
+"""Test the peptide transformers."""
+import torch
+
+from depthcharge.tokenizers import PeptideTokenizer
+from depthcharge.transformers import (
+    PeptideTransformerDecoder,
+    PeptideTransformerEncoder,
+    SpectrumTransformerEncoder,
+)
+
+
+def test_peptide_encoder():
+    """Test that a peptide encoder will run."""
+    tokenizer = PeptideTokenizer()
+    peptides = tokenizer.tokenize(["LESLIEK", "PEPTIDER", "EDITHYKK"])
+    charges = torch.tensor([2, 3, 3])
+
+    model = PeptideTransformerEncoder(tokenizer, 8, 1, 12, max_charge=3)
+    emb, mask = model(peptides, charges)
+
+    # Axis 1 should be 1 longer than the longest peptide.
+    assert emb.shape == (3, 9, 8)
+    assert mask.sum() == 1
+
+    res = emb.sum(dim=1)
+    assert (res[1, :] != res[2, :]).all()
+
+
+def test_peptide_decoder():
+    """Test that a peptide decoder will run."""
+    tokenizer = PeptideTokenizer()
+    n_tokens = len(tokenizer)
+
+    spectra = torch.tensor(
+        [
+            [[100.1, 0.1], [200.2, 0.2], [300.3, 0.3]],
+            [[400.4, 0.4], [500, 0.5], [0, 0]],
+        ]
+    )
+
+    peptides = tokenizer.tokenize(["LESLIEK", "PEPTIDER"])
+    precursors = torch.tensor([[100.0, 2], [200.0, 3]])
+
+    encoder = SpectrumTransformerEncoder(8, 1, 12)
+    memory, mem_mask = encoder(spectra)
+
+    decoder = PeptideTransformerDecoder(n_tokens, 8, 1, 12, max_charge=3)
+    scores = decoder(peptides, precursors, memory, mem_mask)
+    assert scores.shape == (2, 9, len(tokenizer))
diff --git a/tests/unit_tests/test_transformers/test_spectrum_transformers.py b/tests/unit_tests/test_transformers/test_spectrum_transformers.py
new file mode 100644
index 0000000..843b1d6
--- /dev/null
+++ b/tests/unit_tests/test_transformers/test_spectrum_transformers.py
@@ -0,0 +1,30 @@
+"""Test the spectrum transformers."""
+import torch
+
+from depthcharge.encoders import PeakEncoder
+from depthcharge.transformers import SpectrumTransformerEncoder
+
+
+def test_spectrum_encoder():
+    """Test that a spectrum encoder will run."""
+    spectra = torch.tensor(
+        [
+            [[100.1, 0.1], [200.2, 0.2], [300.3, 0.3]],
+            [[400.4, 0.4], [500, 0.5], [0, 0]],
+        ]
+    )
+
+    model = SpectrumTransformerEncoder(8, 1, 12)
+    emb, mask = model(spectra)
+    assert emb.shape == (2, 4, 8)
+    assert mask.sum() == 1
+
+    model = SpectrumTransformerEncoder(8, 1, 12, peak_encoder=PeakEncoder(8))
+    emb, mask = model(spectra)
+    assert emb.shape == (2, 4, 8)
+    assert mask.sum() == 1
+
+    model = SpectrumTransformerEncoder(8, 1, 12, peak_encoder=False)
+    emb, mask = model(spectra)
+    assert emb.shape == (2, 4, 8)
+    assert mask.sum() == 1
diff --git a/tests/unit_tests/test_utils.py b/tests/unit_tests/test_utils.py
deleted file mode 100644
index dd0a05d..0000000
--- a/tests/unit_tests/test_utils.py
+++ /dev/null
@@ -1,14 +0,0 @@
-"""Test that the utility function work."""
-import pytest
-
-
-def test_no_tensorboard(monkeypatch):
-    """Test that import works without tensorboard"""
-    import sys
-
-    sys.modules["tensorboard"] = None
-
-    import depthcharge as dc
-
-    with pytest.raises(ImportError):
-        dc.utils.read_tensorboard_scalars(".")
diff --git a/tests/unit_tests/test_version.py b/tests/unit_tests/test_version.py
index 89afbf7..7a7f932 100644
--- a/tests/unit_tests/test_version.py
+++ b/tests/unit_tests/test_version.py
@@ -1,7 +1,8 @@
-"""Test getting the version"""
-import depthcharge
+"""Test getting the version."""
 from importlib.metadata import PackageNotFoundError
 
+import depthcharge
+
 
 def test_version():
     """Just make sure its something."""
@@ -12,7 +13,10 @@ def test_version():
 
 
 def test_bad_version(monkeypatch):
+    """Test for a bad version."""
+
     def bad_version(*args, **kwargs):
+        """A bad version."""
         raise PackageNotFoundError
 
     monkeypatch.setattr("depthcharge.version.version", bad_version)