voicebox/src/utils/data.py

import math
from copy import deepcopy
import torch
import torch.nn as nn
import torch.nn.functional as F
import random
import hashlib

from torch.utils.data import Dataset, DataLoader, TensorDataset
from torch.nn.utils.rnn import pad_sequence

from typing import Union, List, Sized, Iterable

from tqdm import tqdm

from src.utils.filesystem import ensure_dir
from src.utils.device import DataParallelWrapper
from src.pipelines.pipeline import Pipeline
from src.models.speaker.speaker import SpeakerVerificationModel
from src.models.speech.speech import SpeechRecognitionModel
from src.constants import *

################################################################################
# Data-loading utilities
################################################################################


class DatasetWrapper(Dataset):
    """
    Most data utilities here involve re-assigning or computing targets to
    train or evaluate adversarial attacks. This class wraps an existing
    dataset to overwrite its stored inputs and targets as necessary.
    """
    def __init__(self, dataset, inputs, targets):

        super().__init__()

        self.dataset = dataset
        self.inputs = inputs
        self.targets = targets

        ref_batch = next(iter(dataset))

        if isinstance(ref_batch, tuple):
            self.format = 'tuple'
        else:
            self.format = 'dict'

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):

        if self.format == 'tuple':
            x, y, *features = self.dataset[idx]
            return self.inputs[idx], self.targets[idx], *features

        else:
            batch = self.dataset[idx]
            batch['x'] = self.inputs[idx]
            batch['y'] = self.targets[idx]
            return batch


def pad_batch_power_2(batch):
    """
    Given a batch of tensors, pad to nearest power of 2 to maximum length. Used
    as a `collate_fn` argument to Pytorch `DataLoader` objects.

    Parameters:
    -----------
    batch

    Returns:
    --------
    """

    # get tensors
    (x, y) = zip(*batch)

    n_batch = len(x)

    if n_batch < 1:
        return torch.Tensor([]), None

    if n_batch == 1:
        return x[0:1], y

    if type(y[0]) != str:
        y = torch.stack(y, dim=0)

    # compute maximum length
    dtype, device = x[0].dtype, x[0].device
    lengths = [x_i.shape[-1] for x_i in x]
    max_len = max(lengths)
    next_pow_2 = 2**(max_len - 1).bit_length()

    # pad inputs
    shape = next(iter(x)).shape[1:-1]
    batch_padded = torch.zeros(
        (n_batch, *shape, next_pow_2),
        dtype=dtype,
        device=device
    )

    for i in range(n_batch):
        batch_padded[i, ..., :lengths[i]] = x[i]

    return batch_padded, y


def text_to_tensor(
        text: Union[str, List[str]],
        labels: list,
        return_lengths: bool = True,
        max_length: int = None,
        padding_value: int = -1):
    """
    Convert one or more string transcripts to padded tensor form (character
    indices), and optionally return sequence lengths.

    Parameters:
    -----------
    text (str):            a string or list of string transcripts

    labels (list):         list of characters, ordered by index

    return_lengths (bool): if True, return sequence lengths

    max_length (int):      if given, trim/pad all sequences to length

    padding_value (int):   value with which to perform length padding

    Returns:
    --------
    sequences (Tensor): tensor containing padded index sequences

    lengths (Tensor):   tensor containing sequence lengths
    """

    if isinstance(text, str):
        text = [text]

    # convert from characters to token indices
    char_to_idx = {labels[i].upper(): i for i in range(len(labels))}

    lengths = []
    tensors = []

    for t in text:

        lengths.append(len(t))
        token_indices = [char_to_idx[c] for c in t.upper()]
        tensors.append(
            torch.as_tensor(token_indices, dtype=torch.long)
        )

    # pad and return
    tensors = pad_sequence(
        tensors,
        batch_first=True,
        padding_value=padding_value
    )  # (n_batch, max_seq_len)

    if max_length is not None:
        if tensors.shape[-1] > max_length:
            tensors = tensors[..., :max_length]
        elif tensors.shape[-1] < max_length:
            tensors = F.pad(
                tensors,
                (0, max_length - tensors.shape[-1]),
                value=padding_value)

    lengths = torch.as_tensor(lengths, dtype=torch.long)

    if return_lengths:
        return tensors, lengths
    else:
        return tensors


def padded_transcript_length(
        transcript: torch.Tensor,
        padding_value: int = -1):
    """
    Given one or more transcripts in index sequence format, determine lengths
    by looking for padding value.

    Parameters:
    -----------
    transcript (Tensor): tensor containing one or more index sequences

    padding_value (int): value used to pad sequence tensors

    Returns:
    --------
    lengths (Tensor): tensor containing un-padded length of each sequence
    """

    # find first occurence of padding value in each transcript tensor
    mask = transcript == padding_value
    mask_max_values, mask_max_indices = torch.max(mask, dim=-1)

    # if the max-mask is zero, there is no padding in the tensor
    mask_max_indices[mask_max_values == 0] = transcript.shape[-1]

    return mask_max_indices.long()


def move_to_device_recursive(d: dict, device: Union[str, torch.device]):
    """Move all tensors in a dictionary object to given device"""
    for k, v in d.items():
        if isinstance(v, dict):
            d[k] = move_to_device_recursive(v, device)
        elif isinstance(v, torch.Tensor):
            d[k] = v.to(device)
        elif isinstance(v, tuple) and len(v) > 0 and isinstance(
                next(iter(v)), torch.Tensor):
            v_new = tuple()
            for v_i in v:
                v_new += (v_i.to(device),)
            d[k] = v_new
        elif isinstance(v, list) and len(v) > 0 and isinstance(
                next(iter(v)), torch.Tensor):
            v_new = []
            for v_i in v:
                v_new += [v_i.to(device)]
            d[k] = v_new
    return d


def dataset_to_device(data: Dataset, device: Union[str, torch.device]):
    """Move datasets directly to given device. May cause memory issues!"""

    data.__dict__ = move_to_device_recursive(data.__dict__, device)


@torch.no_grad()
def create_embedding_dataset(data_train: Dataset,
                             data_test: Dataset,
                             pipeline: Pipeline,
                             select_train: str = 'random',
                             select_test: str = 'random',
                             targeted: bool = True,
                             target_class: int = None,
                             num_per_class_train: int = None,
                             num_per_class_test: int = None,
                             num_embeddings_train: int = 5,
                             exclude_class: Union[int, List] = None,
                             exclude_success: bool = False,
                             use_cache: bool = False,
                             shuffle: bool = False,
                             **kwargs
                             ):
    """
    Given training and test datasets holding audio and labels, compute
    embeddings according to (potentially reassigned) labels. It is assumed
    that the provided train and test targets both index the same set of classes,
    i.e. that label 0 in the train set refers to the same class as label 0 in
    the test set. This function considers three possible cases:

      1. `targeted` == True, `target_class` != None

         In this case, a single target class from the training set is assigned
         to all training and test instances. Embeddings corresponding to the
         target class are computed and assigned based on the `select` parameter.

      2.  `targeted` == True, `target_class` == None

         In this case, targets are randomly reassigned within both the training
         and test sets. Embeddings corresponding to the targets are computed and
         assigned based on the `select` parameter.

      3.  `targeted` == False

         In this case, targets remain unchanged and the embedding of each
         instance is computed directly.

    Parameters
    ----------
    data_train (Dataset):       a Dataset object holding audio and labels

    data_test (Dataset):        a Dataset object holding audio and labels

    pipeline (Pipeline):        a Pipeline object; must wrap a
                                SpeakerVerificationModel object

    select_train (str):         method of selecting target embeddings for train
                                set; must be one of `random`, `single`, 'same',
                                `centroid`, or 'median'

    select_test (str):          method of selecting target embeddings for test
                                set; must be one of `random`, `single`, same',
                                `centroid`, or 'median'

    targeted (bool):            if False, target classes will not be reassigned

    target_class (int):         if given, reassign all targets to this class

    num_per_class_train (int):  if given, perform stratified sampling with this
                                number of instances drawn per class

    num_per_class_test (int):   if given, perform stratified sampling with this
                                number of instances drawn per class

    num_embeddings_train (int): if given and attack is targeted, train on only
                                this many distinct embeddings of the target
                                speaker, and evaluate on other embeddings of the
                                target speaker

    exclude_class (int, list):  if given, exclude all instances from this class
                                or in this list of classes

    exclude_success (bool):     if True, drop or replace instances for which the
                                initial prediction achieves the desired
                                adversarial outcome (i.e. matches the target in
                                the case of a targeted attack, and evades the
                                target in the case of an untargeted attack)

    use_cache (bool):           if True, save to and load from disk using hash-
                                based lookup

    shuffle (bool):             if True, shuffle data and targets; may result in
                                mismatch if dataset contains other features (e.g.
                                pitch, periodicity)

    Returns
    -------
    train (dict): dictionary containing Dataset with audio and embedding
                  targets, example audio of each target class, original targets
                  indices, and reassigned target indices

    test (dict):  dictionary containing Dataset with audio and embedding
                  targets, example audio of each target class, original targets
                  indices, and reassigned target indices
    """

    if not num_per_class_train:
        num_per_class_train = 0
    if not num_per_class_test:
        num_per_class_test = 0

    if targeted and target_class is not None and num_per_class_train:
        assert num_per_class_train > num_embeddings_train, \
            f'For targets drawn from training set, number of embeddings ' \
            f'reserved for training ({num_embeddings_train}) must be less ' \
            f'than number of embeddings computed per class ' \
            f'({num_per_class_train})'

    if exclude_success:
        raise NotImplementedError(f'Target correction not yet implemented; '
                                  f'use `NullAttack` to measure trivial '
                                  f'success rates for now')

    # ensure pipeline is capable of producing embeddings
    assert isinstance(pipeline.model, SpeakerVerificationModel)

    # match devices
    ref_batch_train = next(iter(data_train))
    ref_batch_test = next(iter(data_test))
    if isinstance(ref_batch_train, tuple):
        example_input, *_ = ref_batch_train
    elif isinstance(ref_batch_train, dict):
        example_input = ref_batch_train['x']
    else:
        raise ValueError(f'Dataset must provide batches in tuple or dictionary'
                         f' format')
    orig_device = example_input.device

    # check that model produces embeddings with valid shape
    try:
        embedding_shape = list(
            pipeline.model(example_input.to(pipeline.device)).shape
        )
        assert len(embedding_shape) == 3 and embedding_shape[0] == 1
    except AssertionError:
        raise RuntimeError(f'Speaker verification model must produce '
                           f'embeddings of shape '
                           f'(n_batch, n_segments, embedding_dim)')

    assert isinstance(data_train, Sized) and isinstance(data_test, Sized), \
        f"Datasets must have length attribute accessible via `len()`"

    # check embedding selection method
    assert select_train in ['random', 'single', 'centroid', 'same', 'median'], \
        f"invalid value for `select_train` {select_train}"
    assert select_test in ['random', 'single', 'centroid', 'same', 'median'], \
        f"invalid value for `select_test` {select_test}"

    assert not targeted or select_train != 'same', \
        f'`same` embedding selection only valid for untargeted mode'

    # check for optional `batch_size` argument; otherwise, use batch size of 1
    batch_size = kwargs.get('batch_size', 1)

    # creating embedding datasets is time-consuming; to avoid repeated
    # computation, we can store the generated dataset under a hash
    hash_str = str(pipeline.model)
    hash_str += str(data_train.__class__.__name__)
    hash_str += str(data_test.__class__.__name__)
    hash_str += select_train + select_test
    hash_str += str(targeted) + str(target_class)
    hash_str += str(num_per_class_train) + str(num_per_class_test)
    hash_str += str(exclude_class) + str(exclude_success)

    # obtain hash and convert to filename
    dataset_hash = hashlib.md5(hash_str.encode()).digest()
    dataset_file = str(dataset_hash).replace("\'", "")[1:].replace("\\", ".")
    dataset_file += ".pt"

    # check whether a cached embedding dataset with matching hash exists
    embeddings_cache_dir = Path(CACHE_DIR) / 'embeddings'
    ensure_dir(embeddings_cache_dir)
    cached_datasets = embeddings_cache_dir.glob('*.pt')

    # if dataset is already cached, load and return
    if use_cache and dataset_file in [d.name for d in cached_datasets]:

        dataset = torch.load(embeddings_cache_dir / dataset_file)

        # check for valid dataset structure
        try:
            assert isinstance(dataset, dict)
            assert 'train' in dataset and 'test' in dataset

            return move_to_device_recursive(
                dataset['train'], orig_device
            ), move_to_device_recursive(dataset['test'], orig_device)

        except AssertionError:
            raise RuntimeWarning(f'Invalid dataset structure; will re-compute '
                                 f'and overwrite existing dataset '
                                 f'{dataset_file}')

    # shuffle data
    rand_idx_train = torch.randperm(
        len(data_train)) if shuffle else torch.arange(len(data_train))
    rand_idx_test = torch.randperm(
        len(data_test)) if shuffle else torch.arange(len(data_test))

    # separate data and labels
    if isinstance(ref_batch_train, tuple):
        inputs_train, labels_train, *_ = data_train[:]
    else:
        inputs_train, labels_train = data_train[:]['x'], data_train[:]['y']
    inputs_train_shuffled = inputs_train[rand_idx_train]
    labels_train_shuffled = labels_train[rand_idx_train]

    if isinstance(ref_batch_test, tuple):
        inputs_test, labels_test, *_ = data_test[:]
    else:
        inputs_test, labels_test = data_test[:]['x'], data_test[:]['y']
    inputs_test_shuffled = inputs_test[rand_idx_test]
    labels_test_shuffled = labels_test[rand_idx_test]

    # if target is given, check that it is present in training data
    if target_class is not None:
        assert target_class in labels_train, \
            f'Target class {target_class} is not present in training data'

    # determine train and test labels
    unique_labels_train = [l.item() for l in torch.unique(labels_train)]
    unique_labels_test = [l.item() for l in torch.unique(labels_test)]

    # filter excluded classes (if given) from train and test sets
    if isinstance(exclude_class, List):
        unique_labels_train = [
            l for l in unique_labels_train if l not in exclude_class]
        unique_labels_test = [
            l for l in unique_labels_test if l not in exclude_class]

    elif exclude_class is not None:
        unique_labels_train = [
            l for l in unique_labels_train if not l == exclude_class]
        unique_labels_test = [
            l for l in unique_labels_test if not l == exclude_class]

    # prepare to store one example audio input per label (speaker)
    audio_train = {}
    audio_test = {}

    # prepare to store training and test embeddings by label
    embeddings_train = {}
    embeddings_test = {}

    def compute_embeddings_by_label(
            unique_labels: list,
            inputs: torch.Tensor,
            labels: torch.Tensor,
            saved_audio: dict,
            saved_embeddings: dict,
            num_per_class):
        """
        Compute an embedding for every instance in the given dataset and store
        by label in a dictionary; store one audio example per label in a
        dictionary.
        """

        # compute embeddings over training set and sort by label
        for label in tqdm(
                unique_labels,
                total=len(unique_labels),
                desc="Computing embeddings for dataset"):

            # select training instances of class, allowing for a limit on the
            # number of embeddings stored per class
            x_l = inputs[labels == label]
            n_l = num_per_class if num_per_class else len(x_l)

            # store one audio example per training label
            saved_audio[label] = x_l[0:1]

            # store embeddings per training label
            n_batches = math.ceil(n_l / batch_size)
            saved_embeddings[label] = []
            for i in range(n_batches):
                saved_embeddings[label].append(
                    pipeline.model(
                        x_l[i*batch_size:(i+1)*batch_size].to(pipeline.device)
                    ).to('cpu')  # store intermediate results on CPU
                )

            saved_embeddings[label] = torch.cat(
                saved_embeddings[label], dim=0)[:n_l]

    # compute embeddings over training and test datasets and store by label
    compute_embeddings_by_label(
        unique_labels_train,
        inputs_train_shuffled,
        labels_train_shuffled,
        audio_train,
        embeddings_train,
        num_per_class_train
    )
    compute_embeddings_by_label(
        unique_labels_test,
        inputs_test_shuffled,
        labels_test_shuffled,
        audio_test,
        embeddings_test,
        num_per_class_test
    )

    # filter datasets to remove excluded and target labels
    if targeted and target_class is not None:
        unique_labels_train = [
            l for l in unique_labels_train if not l == target_class]
        unique_labels_test = [
            l for l in unique_labels_test if not l == target_class]

    def reassign_labels(
            unique_labels: list,
            inputs: torch.Tensor,
            labels_orig: torch.Tensor,
            num_per_class):
        """
        Reassign targets, as detailed in documentation above.
        """

        labels_new = torch.full(labels_orig.shape, -1, dtype=labels_orig.dtype)

        # reassign label-by-label
        for i, label in enumerate(
                tqdm(
                    unique_labels,
                    total=len(unique_labels),
                    desc="Reassigning labels for dataset")):

            # select all training instances with label
            idx_l = labels_orig == label
            x_l = inputs[idx_l]

            # store original targets
            y_orig_l = torch.full((len(x_l), ), label)

            # use a placeholder to allow for deletion of rows; overwrite with
            # valid labels and delete rows where -1 remains
            y_new_l = torch.full((len(x_l), ), -1)

            # limit number of instances per class if specified
            n_l = num_per_class if num_per_class else len(x_l)

            # targeted attacks require that the given targets be reassigned
            if targeted:

                # if target class is provided, reassign targets to given class
                if target_class is not None:
                    y_new_l[:n_l] = target_class

                # if no target class is given, randomly reassign targets; ensure
                # that no target is unchanged and new targets are evenly
                # distributed
                else:
                    remaining_labels = [
                        l for l in unique_labels if l != label]

                    for j in range(n_l):
                        y_new_l[j] = random.choice(remaining_labels)

            # otherwise, classes remain unchanged
            else:
                y_new_l[:n_l] = y_orig_l[:n_l]

            # update data and labels, deleting rows corresponding to
            # extraneous inputs (according to `num_per_class`)
            labels_new[idx_l] = y_new_l

        keep_idx = labels_new != -1
        inputs = inputs[keep_idx]
        labels_orig = labels_orig[keep_idx]
        labels_new = labels_new[keep_idx]

        return inputs, labels_orig, labels_new, keep_idx

    # reassign training and test labels if necessary (see documentation
    # above); remove instances of target class and those not required by
    # `num_per_class`, if given
    (
        inputs_train_shuffled,
        labels_train_shuffled,
        labels_train_reassigned,
        select_idx_train
    ) = reassign_labels(
        unique_labels_train,
        inputs_train_shuffled,
        labels_train_shuffled,
        num_per_class_train
    )
    (
        inputs_test_shuffled,
        labels_test_shuffled,
        labels_test_reassigned,
        select_idx_test
    ) = reassign_labels(
        unique_labels_test,
        inputs_test_shuffled,
        labels_test_shuffled,
        num_per_class_test
    )

    # prepare to store target embeddings corresponding to reassigned labels
    embedding_targets_train = torch.empty(
        (len(labels_train_reassigned), *embedding_shape[1:]))
    embedding_targets_test = torch.empty(
        (len(labels_test_reassigned), *embedding_shape[1:]))

    def assign_embeddings(
            labels_new: torch.Tensor,
            embeddings_by_label: dict,
            embedding_targets: torch.Tensor,
            select: str,
            is_train: bool = True
    ):

        # iterate over dataset and associate embedding targets with
        # reassigned labels
        labels_to_assign = [l.item() for l in torch.unique(labels_new)]

        for label in labels_to_assign:

            # find indices for which embeddings of given label are to
            # be assigned
            idx_l = labels_new == label
            n_l = int(torch.sum(idx_l * 1).item())

            if n_l == 0:
                continue

            # obtain all embeddings corresponding to given label
            embeddings_l = embeddings_by_label[label]

            # if untargeted, assign ground-truth embeddings for each instance
            if select == 'same':
                y_emb_l = embeddings_l

            else:

                # separate train and test embeddings of given speaker
                if num_embeddings_train:

                    # for targeted attacks, allow training/testing on separate
                    # small subsets of a speaker's utterances
                    if targeted:
                        assert num_embeddings_train <= len(embeddings_l), \
                            f"`num_embeddings_train` {num_embeddings_train} " \
                            f"is greater than the number of utterances for " \
                            f"speaker {label}"

                        if is_train:
                            embeddings_l = embeddings_l[:num_embeddings_train]
                        else:
                            embeddings_l = embeddings_l[num_embeddings_train:]

                # using `select` parameter, assign embeddings
                y_emb_l = []

                for i in range(n_l):

                    if select == 'single':  # use single embedding
                        y_emb_l.append(embeddings_l[0:1])

                    elif select == 'random':  # use random embeddings
                        emb_idx = random.randint(0, len(embeddings_l) - 1)
                        y_emb_l.append(embeddings_l[emb_idx:emb_idx+1])

                    elif select == 'centroid':  # average over embeddings

                        _, n_segments, embedding_dim = embedding_shape

                        # duplicate over all segments
                        centroid = embeddings_l.mean(dim=(0, 1)).reshape(
                            (1, 1, embedding_dim)
                        ).repeat(1, n_segments, 1)

                        y_emb_l.append(centroid)

                    elif select == 'median':  # median over embeddings

                        _, n_segments, embedding_dim = embedding_shape

                        # duplicate over all segments
                        median = embeddings_l.reshape(
                            n_l*n_segments, -1
                        ).median(dim=0)[0].reshape(
                            (1, 1, embedding_dim)
                        ).repeat(1, n_segments, 1)

                        y_emb_l.append(median)

                    else:
                        raise ValueError(f'Invalid embedding selection method '
                                         f'{select}')

                y_emb_l = torch.cat(y_emb_l, dim=0)

            embedding_targets[idx_l] = y_emb_l

    # with labels finalized, assign embedding targets
    assign_embeddings(
        labels_train_reassigned,
        embeddings_train,
        embedding_targets_train,
        select_train,
        True
    )
    assign_embeddings(
        labels_test_reassigned,
        embeddings_train if targeted and target_class is not None else embeddings_test,
        embedding_targets_test,
        select_test,
        False
    )

    # account for shuffling
    final_idx_train = rand_idx_train[select_idx_train]
    final_idx_test = rand_idx_test[select_idx_test]

    from src.data.dataset import VoiceBoxDataset

    if isinstance(data_train, VoiceBoxDataset):
        data_train_final = data_train.overwrite_dataset(
            inputs_train_shuffled,
            embedding_targets_train,
            final_idx_train
        )
    else:
        data_train_final = DatasetWrapper(
            data_train,
            inputs_train_shuffled,
            embedding_targets_train)

    if isinstance(data_test, VoiceBoxDataset):
        data_test_final = data_test.overwrite_dataset(
            inputs_test_shuffled,
            embedding_targets_test,
            final_idx_test
        )
    else:
        data_test_final = DatasetWrapper(
            data_test,
            inputs_test_shuffled,
            embedding_targets_test)

    # store data and embeddings, audio examples, original targets, and
    # reassigned targets
    train = {
        'dataset': data_train_final,
        'id_to_audio': audio_train,
        'true_id': labels_train_shuffled,
        'target_id': labels_train_reassigned
    }
    test = {
        'dataset': data_test_final,
        'id_to_audio': audio_test,
        'true_id': labels_test_shuffled,
        'target_id': labels_test_reassigned
    }

    if use_cache:
        dataset = {
            'train': train,
            'test': test
        }
        torch.save(dataset, embeddings_cache_dir / dataset_file)

    # restore device and return
    return move_to_device_recursive(
        train, orig_device
    ), move_to_device_recursive(
        test, orig_device
    )


@torch.no_grad()
def create_transcription_dataset(data_train: Dataset,
                                 data_test: Dataset,
                                 pipeline: Pipeline,
                                 targeted: bool = True,
                                 target_transcription: str = None,
                                 output_format: str = 'transcript',
                                 shuffle: bool = False,
                                 **kwargs):
    """
    Given training and test datasets holding audio assign string transcriptions
    for performing speech recognition attacks.

      1. `targeted` == True, `target_transcription` != None

         In this case, a single transcription target is assigned to all instances.

      2.  `targeted` == True, `target_transcription` == None

         In this case, ground-truth transcriptions are randomly reassigned as
         targets within both the training and test sets.

      3.  `targeted` == False

         In this case, ground-truth transcriptions are used as targets.

    Parameters
    ----------
    data_train (Dataset):       a Dataset object holding audio

    data_test (Dataset):        a Dataset object holding audio

    targeted (bool):            if False, target classes will not be reassigned

    target_transcription (str): if given, reassign all targets to the given
                                transcription string

    shuffle (bool):             if True, shuffle data and targets; may result in
                                mismatch if dataset contains other features (e.g.
                                pitch, periodicity)

    Returns
    -------
    train (dict):

    test (dict):
    """

    # check output format (string transcripts or frame-wise token probabilities)
    assert output_format in ['transcript', 'emission'], \
        f'Invalid output format; must be one of `transcript` or `emission`'

    # check for valid model type
    assert isinstance(pipeline.model, SpeechRecognitionModel)

    assert isinstance(data_train, Sized) and isinstance(data_test, Sized), \
        f"Datasets must have length attribute accessible via `len()`"

    # match devices
    ref_batch_train = next(iter(data_train))
    ref_batch_test = next(iter(data_test))
    if isinstance(ref_batch_train, tuple):
        example_input, *_ = ref_batch_train
    elif isinstance(ref_batch_train, dict):
        example_input = ref_batch_train['x']
    else:
        raise ValueError(f'Dataset must provide batches in tuple or dictionary'
                         f' format')

    orig_device = example_input.device

    # check for optional `batch_size` argument; otherwise, use batch size of 1
    batch_size = kwargs.get('batch_size', 1)

    # shuffle data

    rand_idx_train = torch.randperm(
        len(data_train)) if shuffle else torch.arange(len(data_train))
    rand_idx_test = torch.randperm(
        len(data_test)) if shuffle else torch.arange(len(data_test))

    if isinstance(ref_batch_train, tuple):
        inputs_train, *_ = data_train[:]
    else:
        inputs_train = data_train[:]['x']
    inputs_train_shuffled = inputs_train[rand_idx_train]

    if isinstance(ref_batch_test, tuple):
        inputs_test, *_ = data_test[:]
    else:
        inputs_test = data_test[:]['x']
    inputs_test_shuffled = inputs_test[rand_idx_test]

    # if targeted and target transcription provided, simply assign and return
    if targeted and target_transcription is not None:

        assert output_format == 'transcript', \
            f"Target transcript provided; cannot use emission targets"

        # check that target transcript contains character set compatible with
        # pipeline, and does not contain 'blank' character
        valid_characters = deepcopy(pipeline.model.get_labels())
        try:
            del valid_characters[pipeline.model.get_blank_idx()]
        except (IndexError, TypeError):
            pass
        assert all([c in valid_characters for c in target_transcription]), \
            f'Target transcription contains invalid characters'

        single_target = text_to_tensor(
            target_transcription,
            pipeline.model.get_labels(),
            return_lengths=False
        )

        targets_train = single_target.repeat(len(inputs_train_shuffled), 1)
        targets_test = single_target.repeat(len(inputs_test_shuffled), 1)

    # otherwise, compute transcriptions using given pipeline
    else:

        def transcribe(dataset: torch.Tensor):

            results = []

            n_batches = math.ceil(len(dataset) / batch_size)
            for batch_idx in tqdm(
                    range(n_batches),
                    total=n_batches,
                    desc="Computing transcriptions for dataset"):

                x = dataset[
                    batch_idx*batch_size:(batch_idx+1)*batch_size
                    ].to(pipeline.device)

                if output_format == 'transcript':
                    results.extend(pipeline.model.transcribe(x))
                elif output_format == 'emission':
                    results.extend(
                        torch.split(
                            pipeline.model(x).to(orig_device), 1, dim=0))

            if output_format == 'emission':

                # pad to max emission length
                results = pad_sequence(results, batch_first=True).squeeze(1)

            elif output_format == 'transcript':
                results = text_to_tensor(
                    results,
                    pipeline.model.get_labels(),
                    return_lengths=False)

            return results

        targets_train = transcribe(inputs_train_shuffled)
        targets_test = transcribe(inputs_test_shuffled)

        # if targeted, permute transcriptions such that no input retains its
        # original transcription
        if targeted:

            # use derangements with a fixed iteration budget; expected number
            # of iterations required to shuffle with no fixed points is e (~3)
            def derange(x: torch.Tensor):

                max_iter = 10
                orig_shape = x.shape
                x = x.reshape(x.shape[0], -1)

                for i in range(max_iter):

                    rand_idx = torch.randperm(len(x))
                    equal = torch.sum(
                        1.0 * (x == x[rand_idx]),
                        dim=-1
                    ) >= x.shape[-1]

                    if not equal.sum().item():
                        break

                return x[rand_idx].reshape(orig_shape)

            targets_train = derange(targets_train)
            targets_test = derange(targets_test)

    # compute transcript lengths
    if output_format == 'transcript':
        lengths_train = padded_transcript_length(targets_train)
        lengths_test = padded_transcript_length(targets_test)
    elif output_format == 'emission':
        lengths_train = torch.full(
            size=(len(inputs_train_shuffled),),
            fill_value=targets_train.shape[1],
            dtype=torch.long
        )
        lengths_test = torch.full(
            size=(len(inputs_test_shuffled),),
            fill_value=targets_test.shape[1],
            dtype=torch.long
        )
    else:
        raise ValueError(f'Invalid value for `output_format`')

    from src.data.dataset import VoiceBoxDataset

    if isinstance(data_train, VoiceBoxDataset):
        data_train_final = data_train.overwrite_dataset(
            inputs_train_shuffled,
            targets_train,
            rand_idx_train)
    else:
        data_train_final = DatasetWrapper(
            data_train,
            inputs_train_shuffled,
            targets_train)

    if isinstance(data_test, VoiceBoxDataset):
        data_test_final = data_test.overwrite_dataset(
            inputs_test_shuffled,
            targets_test,
            rand_idx_test
        )
    else:
        data_test_final = DatasetWrapper(
            data_test,
            inputs_test_shuffled,
            targets_test)

    train = {
        'dataset': data_train_final,
        'targets': targets_train,
        'target_lengths': lengths_train
    }

    test = {
        'dataset': data_test_final,
        'targets': targets_test,
        'target_lengths': lengths_test
    }

    return move_to_device_recursive(
        train, orig_device
    ), move_to_device_recursive(
        test, orig_device
    )