wehub-resource-sync
265 lines
7.7 KiB
Python
265 lines
7.7 KiB
Python
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()
|