data/fairseq/examples/speech_to_speech/benchmarking/data_utils.py

from fairseq import tasks
import numpy as np
import logging
import random
from fairseq import options
import torch
import os
import soundfile as sf

from fairseq.data.audio.audio_utils import (
    get_waveform,
    parse_path,
)

logging.basicConfig()
logging.root.setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

random.seed(1)
np.random.seed(1)
random_number_generator = np.random.RandomState(30)


def generate_random_data_sample(T, B=1, D=80):
    """Generate random data sample given the T, B, D values"""
    net_input = {
        "src_tokens": torch.tensor(random_number_generator.randn(B, T, D)).float(),
        "src_lengths": torch.tensor([T]),
    }
    return {"net_input": net_input}


def generate_random_dataset(T_range_min, T_range_max, B=1, D=80, dataset_size=100):
    """Generate random dataset with T values within a given range, B, D"""
    T_values = [random.randint(T_range_min, T_range_max) for i in range(dataset_size)]
    dataset = []
    for t in T_values:
        dataset.append(generate_random_data_sample(t, B, D))
    return dataset, sum(T_values) / dataset_size


def load_dataset_npy(file_name, dataset_size=None):
    """Load dataset from a .npy file."""
    data = np.load(file_name, allow_pickle=True)
    if dataset_size:
        data = data[:dataset_size]
    return data


def load_dataset_raw_to_waveforms(
    file_name,
    dataset_size=None,
    need_waveform=True,
    sample_rate=16000,
    read_using_soundfile=False,
):
    """Load raw dataset from w2v tsv file. Optionally get waveforms"""
    data = []
    with open(file_name, "r") as fp:
        lines = fp.readlines()
        data = [
            os.path.join(lines[0].strip(), line.strip().split("\t")[0])
            for line in lines[1:]
        ]

    if dataset_size:
        data = data[:dataset_size]

    if not need_waveform:
        return data

    features = []
    if read_using_soundfile:
        for _i, d in enumerate(data):
            wav = sf.read(d)[0]
            if wav.ndim == 2:
                wav = wav.mean(-1)
            features.append(torch.from_numpy(wav).float().view(1, -1))
    else:
        for i, d in enumerate(data):
            _path, slice_ptr = parse_path(d)
            if len(slice_ptr) == 0:
                feat = get_waveform(
                    _path, always_2d=True, output_sample_rate=sample_rate
                )[0]
                features.append(
                    {
                        "id": i,
                        "net_input": {
                            "src_tokens": torch.tensor(feat),
                            "src_lengths": torch.tensor([feat.shape[1]]),
                        },
                    }
                )
            else:
                raise Exception("Currently unsupported data format")
    return features


def load_dataset_task(
    args,
    batch_size=1,
    limit_size=None,
    ref_dataset=None,
):
    """Loads dataset based on args by creating a task"""
    if not args.data or not args.subset or not args.task:
        raise Exception(
            "Please provide necessary arguments to load the dataset - data, subset and task"
        )
    task = tasks.setup_task(args)

    task.load_dataset(args.subset)
    if not limit_size:
        limit_size = len(task.dataset(args.subset))

    iter = task.get_batch_iterator(
        dataset=task.dataset(args.subset), max_sentences=batch_size
    ).next_epoch_itr(shuffle=False)
    dataset = []
    for i, sample in enumerate(iter):
        sample = {
            "id": task.datasets[args.subset].ids[sample["id"].item()],
            "net_input": {
                "src_tokens": sample["net_input"]["src_tokens"],
                "src_lengths": sample["net_input"]["src_lengths"],
            },
        }
        dataset.append(sample)
        if i == limit_size - 1:
            break

    if ref_dataset:
        try:
            ids = get_ids_from_dataset(ref_dataset)
        except Exception as e:
            raise Exception(f"{e} - Cannot extract ids from reference dataset")

        filtered_dataset = []
        for sample in dataset:
            if (
                sample["id"] in ids
                or sample["id"][5:] in ids
                or f"dev_{sample['id']}" in ids
            ):
                filtered_dataset.append(sample)
        dataset = filtered_dataset

    max_len, min_len, avg_len = get_dataset_stats(dataset)
    print(
        f"{args.subset} dataset stats : num_samples={len(dataset)} max_len = {max_len} min_len = {min_len} avg_len = {avg_len}"
    )

    return dataset


def randomly_sample_subset(dataset, size=500):
    """Randomly sample subset from a dataset"""
    random_indices = [random.randint(0, len(dataset) - 1) for i in range(size)]
    return [dataset[i] for i in random_indices]


def get_short_data_subset(dataset, size=500):
    """Get a subset of desired size by sorting based on src_lengths"""
    return sort_dataset(dataset)[:size]


def get_long_data_subset(dataset, size=500):
    """Get a subset of desired size by sorting based on src_lengths descending"""
    return sort_dataset(dataset, reverse=True)[:size]


def sort_dataset(dataset, reverse=False):
    return sorted(
        dataset, key=lambda x: x["net_input"]["src_lengths"].item(), reverse=reverse
    )


def save_dataset_npy(dataset, file_name):
    """Save a dataset as .npy file"""
    np.save(file_name, dataset)


def get_dataset_stats(dataset):
    """Get stats about dataset based on src_lengths of samples"""
    max_len = 0
    min_len = 100000
    avg_len = 0
    for d in dataset:
        max_len = max(max_len, d["net_input"]["src_lengths"].item())
        min_len = min(min_len, d["net_input"]["src_lengths"].item())
        avg_len += d["net_input"]["src_lengths"].item()

    return max_len, min_len, avg_len / len(dataset)


def make_parser():
    """
    Additional args:
        1. Provide the dataset dir path using --data.
        2. Loading the dataset doesn't require config, provide --config-yaml to apply additional feature transforms
    """
    parser = options.get_speech_generation_parser()
    parser.add_argument(
        "--subset",
        default=None,
        type=str,
        required=True,
        help="Subset to use for dataset generation",
    )
    parser.add_argument(
        "--dataset-save-dir",
        default=None,
        type=str,
        required=False,
        help="Dir path in which the datasets are to be saved",
    )
    parser.add_argument(
        "--ref-dataset",
        default=None,
        type=str,
        required=False,
        help="If provided, the ids in the reference dataset will be used to filter the new dataset generated.",
    )
    parser.add_argument("--dataset-save-token", default="", type=str, required=False)

    options.add_generation_args(parser)
    return parser


def get_ids_from_dataset(dataset):
    return {sample["id"]: 1 for sample in dataset}


def cli_main():
    parser = make_parser()
    args = options.parse_args_and_arch(parser)
    dataset = load_dataset_task(args)

    random_dataset = randomly_sample_subset(dataset)
    short_dataset = get_short_data_subset(dataset)
    long_dataset = get_long_data_subset(dataset)

    if args.dataset_save_token:
        args.dataset_save_token = f"_{args.dataset_save_token}_"

    if args.dataset_save_dir:
        save_dataset_npy(
            random_dataset,
            f"{args.dataset_save_dir}/random_dataset{args.dataset_save_token}w_ids.npy",
        )
        save_dataset_npy(
            short_dataset,
            f"{args.dataset_save_dir}/short_dataset{args.dataset_save_token}w_ids.npy",
        )
        save_dataset_npy(
            long_dataset,
            f"{args.dataset_save_dir}/long_dataset{args.dataset_save_token}w_ids.npy",
        )


if __name__ == "__main__":
    cli_main()