storytold/storyteller-ml
2
0
Fork 0
mirror of https://github.com/storytold/storyteller-ml.git synced 2026-10-16 00:13:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Introduce code from Rust monorepo. This is a mess

Browse Source
This commit is contained in:
Brandon Thomas
2022-02-28 05:43:54 -05:00
parent 921fb49f9a
commit ba326bfcfd
1058 changed files with 192192 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tts/NeMo/scripts/asr_language_modeling/neural_rescorer/create_tarred_transformer_lm_dataset.py
+318
View File
@@ -0,0 +1,318 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Execution :
python create_tarred_tokenized_text_lm_dataset.py \
--text_path=<comma seperated text filepaths> \
--data_root=<path to output directory> \
--tokenizer_name="bert-base-cased" \
--tokenizer_vocab_file=<path to vocab file for tokenizer> \
--num_shards=64 \
--chunk_size=8192 \
--chunk_write_buffer=512 \
--lower_case \
--log
"""
import argparse
import glob
import json
import logging
import os
import tarfile
import joblib
import numpy as np
from tqdm import tqdm
from nemo.collections.nlp.modules.common.tokenizer_utils import get_tokenizer
parser = argparse.ArgumentParser(description='Tarred Tokenized dataset for text language modelling')
# Data path arguments
parser.add_argument('--text_path', required=True, default=None, type=str, help='Text paths, seperated by commas')
parser.add_argument('--data_root', required=True, default=None, type=str, help='Output directory')
# General arguments
parser.add_argument(
'--chunk_write_buffer',
default=128,
type=int,
help='Number of chunks of `chunk_size` to buffer for parallel tokenization and serial write to disk',
)
parser.add_argument('--lower_case', action='store_true', help='Whether to lower case the corpus')
parser.add_argument('--log', action='store_true', help='Whether to print logs to terminal')
# Tokenizer arguments
parser.add_argument('--tokenizer_name', required=False, default=None, type=str, help='Tokenizer name for resolution')
parser.add_argument(
'--tokenizer_model', required=False, default=None, type=str, help='Path to tokenizer model for sentencepiece'
)
parser.add_argument('--tokenizer_vocab_file', required=False, type=str, default=None, help='Path to a vocab file')
parser.add_argument(
'--tokenizer_special_tokens', default=None, type=str, nargs='+', help='List of special tokens for the tokenizer'
)
# Tarred dataset arguments
parser.add_argument('--num_shards', default=1, type=int, help='Number of shards for the tarfile')
parser.add_argument('--chunk_size', default=8192, type=int, help='Number of rows of data concatenated into a vector')
parser.set_defaults(log=False, lower_case=False)
args = parser.parse_args()
def __build_dataset_from_text(texts: str, lower_case: bool, chunk_size: int):
if ',' in texts:
texts = texts.split(',')
else:
texts = [texts]
num_lines = 0
text_dataset = []
for text in texts:
with open(text, 'r') as in_reader:
reader = tqdm(iter(lambda: in_reader.readline(), ''), desc="Read 0 lines", unit=' lines')
for i, line in enumerate(reader):
# Clean text line
line = line.replace("\n", "").strip()
if lower_case:
line = line.lower()
if line:
text_dataset.append(line)
num_lines += 1
if num_lines % 100000 == 0:
reader.set_description(f"Read {num_lines} lines")
if num_lines % chunk_size == 0:
yield text_dataset, num_lines
# Empty cache
text_dataset = []
logging.info(f"Finished extracting manifest : {text}")
logging.info("Finished extracting all manifests ! Number of sentences : {}".format(num_lines))
if len(text_dataset) != 0:
yield text_dataset, num_lines
def __tokenize_str(texts, tokenizer):
tokenized_text = []
for text in texts:
tok_text = tokenizer.text_to_ids(text)
tokenized_text.extend(tok_text)
return tokenized_text
def __tokenize_text(
text_paths, tokenizer, tokenized_cachedir, lower_case: bool = False, chunk_size=8192, write_buffer: int = -1
):
if write_buffer < 1:
write_buffer = max(os.cpu_count() - write_buffer, 1)
logging.info(f"Using write chunk buffer of size {write_buffer}")
if not os.path.exists(tokenized_cachedir):
os.makedirs(tokenized_cachedir)
# global parameters
global_chunk_idx = 0
chunk_paths = []
chunk_lens = []
# buffer parameters
data_cache = []
chunk_idx = 0
text_generator = iter(__build_dataset_from_text(text_paths, lower_case=lower_case, chunk_size=chunk_size))
global_num_lines = 0
last_batch = False
with joblib.Parallel(n_jobs=-2, verbose=10) as parallel:
while True:
try:
data, num_lines = next(text_generator)
data_cache.append(data)
global_num_lines = num_lines
except StopIteration:
last_batch = True
# Update counters
chunk_idx += 1
if (chunk_idx == write_buffer) or last_batch:
# write the chunks into disk after parallel tokenization
tokenized_data_list = parallel(
joblib.delayed(__tokenize_str)(chunk, tokenizer) for chunk in data_cache
)
# Sequential write cache
for chunk in tokenized_data_list:
fp = os.path.join(tokenized_cachedir, f"chunk_{global_chunk_idx}.npy")
chunk = np.asarray(chunk, dtype=np.int64)
np.save(fp, chunk, allow_pickle=False)
chunk_paths.append(fp)
chunk_lens.append(len(chunk))
global_chunk_idx += 1
logging.info(f"Wrote a total of {global_chunk_idx} chunks to file...")
# reset buffers
data_cache.clear()
del data_cache
data_cache = []
chunk_idx = 0
if last_batch:
logging.info("Finished tokenizing last chunk")
break
logging.info(
f"Chunking {global_num_lines} rows into {global_num_lines // chunk_size} tasks (each chunk contains {chunk_size} elements)"
)
return chunk_paths, chunk_lens
def __create_chunk(data_root, chunk_path, shard_id, compute_metrics=False):
"""Creates a tarball containing the tokenized text chunks.
"""
tar = tarfile.open(os.path.join(data_root, f'text_{shard_id}.tar'), mode='a')
# We squash the filename since we do not preserve directory structure of tokenized text in the tarball.
base, ext = os.path.splitext(chunk_path)
base = base.replace(os.pathsep, '_')
# Need the following replacement as long as WebDataset splits on first period
base = base.replace('.', '_')
squashed_filename = f'{base}{ext}'
tar.add(chunk_path, arcname=squashed_filename)
tar.close()
if compute_metrics:
data = np.load(chunk_path, allow_pickle=False)
chunk_len = len(data)
return (chunk_len,)
else:
return None
def __write_tarred_tokenized_text_dataset(data_root, num_shards, chunk_paths, chunk_lens):
num_chunks = len(chunk_paths)
if chunk_lens is not None:
num_text = sum(chunk_lens)
shard_counts = {chunk_id: chunk_len for chunk_id, chunk_len in enumerate(chunk_lens)}
compute_metrics = False
else:
num_text = 0
shard_counts = {}
compute_metrics = True
for chunk_id, chunk_path in enumerate(tqdm(chunk_paths, desc='Writing chunk ', unit=' chunks')):
shard_id = chunk_id % num_shards
metrics = __create_chunk(data_root, chunk_path, shard_id, compute_metrics=compute_metrics)
if metrics is not None:
num_text += metrics[0]
shard_counts[chunk_id] = metrics[0]
# write metadata
metadata_path = os.path.join(data_root, 'metadata.json')
with open(metadata_path, 'w') as f:
metadata = {'num_chunks': num_chunks, 'num_text': num_text, 'shard_count': shard_counts}
json.dump(metadata, f, indent=4)
logging.info("Metadata writen..")
def main():
text_path = args.text_path
data_root = args.data_root
if args.log:
logging.basicConfig(level=logging.INFO)
tokenized_cachedir = os.path.join(data_root, '_tokenized_dataset_cachedir')
if os.path.exists(tokenized_cachedir):
logging.warning(
f'Tokenized cache directory {tokenized_cachedir} already potentially contains files.'
f'In such a case, please be aware that the tarfiles will be **appended** instead of overridden!'
)
if not os.path.exists(data_root):
os.makedirs(data_root)
chunk_paths = None
chunk_lens = None
if os.path.exists(tokenized_cachedir):
paths = glob.glob(os.path.join(tokenized_cachedir, "*.npy"))
if len(paths) > 0:
logging.info("Cached tokenized numpy files found, skipping re-tokenization of dataset")
chunk_paths = paths
chunk_lens = None
if chunk_paths is None:
if args.tokenizer_name is None:
raise ValueError("`tokenizer_name` name is required when tokenizing the dataset for the first time.")
if args.tokenizer_vocab_file is None:
raise ValueError("`tokenizer_vocab_file` is required when constructing the tokenized dataset")
tokenizer = get_tokenizer(
tokenizer_name=args.tokenizer_name,
tokenizer_model=args.tokenizer_model,
vocab_file=args.tokenizer_vocab_file,
special_tokens=args.tokenizer_special_tokens,
)
logging.info("Built tokenizer")
# tokenize text data into sub-words
chunk_paths, chunk_lens = __tokenize_text(
text_paths=text_path,
tokenizer=tokenizer,
tokenized_cachedir=tokenized_cachedir,
lower_case=args.lower_case,
chunk_size=args.chunk_size,
write_buffer=args.chunk_write_buffer,
)
logging.info(f"Tokenized dataset into sub-words and serialized cache at {tokenized_cachedir}")
# Write tarred dataset
__write_tarred_tokenized_text_dataset(
data_root, num_shards=args.num_shards, chunk_paths=chunk_paths, chunk_lens=chunk_lens
)
logging.info('Done preparing tokenized dataset!')
if __name__ == "__main__":
main()
tts/NeMo/scripts/asr_language_modeling/neural_rescorer/eval_neural_rescorer.py
+288
View File
@@ -0,0 +1,288 @@
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
This script would evaluate a neural language model (Transformer) trained with
`examples/nlp/language_modeling/transformer_lm.py' as a rescorer for ASR systems.
Given a trained TransformerLMModel `.nemo` file, this script can be used to re-score the beams obtained from a beam
search decoder of an ASR model.
USAGE:
1. Obtain `.tsv` file with beams and their corresponding scores. Scores can be from a regular beam search decoder or
in fusion with an N-gram LM scores. For a given beam size `beam_size` and a number of examples
for evaluation `num_eval_examples`, it should contain (`beam_size` x `num_eval_examples`) lines of
form `beam_candidate_text \t score`. This file can be generated by `scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram.py`.
2. Rescore the candidates:
python eval_neural_rescorer.py
--lm_model=[path to .nemo file of the LM]
--beams_file=[path to beams .tsv file]
--beam_size=[size of the beams]
--eval_manifest=[path to eval manifest .json file]
--batch_size=[batch size used for inference on the LM model]
--alpha=[the value for the parameter rescorer_alpha]
--beta=[the value for the parameter rescorer_beta]
You may find more info on how to use this script at:
https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/asr_language_modeling.html
"""
import contextlib
import json
from argparse import ArgumentParser
import editdistance
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import tqdm
from nemo.collections.nlp.models.language_modeling import TransformerLMModel
from nemo.utils import logging
class BeamScoresDataset(torch.utils.data.Dataset):
"""
Dataset to read the score file containing the beams and their score
Args:
data_path: path to the beams file
tokenizer: tokenizer of the LM model
manifest_path: manifest `.json` file which contains the ground truths transcripts
beam_size: the number of beams per sample
max_seq_length: the maximum length of sequences
"""
def __init__(self, data_path, tokenizer, manifest_path, beam_size=128, max_seq_length=256):
self.data = pd.read_csv(data_path, delimiter="\t", header=None)
self.tokenizer = tokenizer
self.ground_truths = []
with open(manifest_path, 'r') as f_orig:
for line in f_orig:
item = json.loads(line)
self.ground_truths.append(item['text'])
self.beam_size = beam_size
self.max_seq_length = max_seq_length
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
text = str(self.data[0][idx])
tokens = [self.tokenizer.bos_id] + self.tokenizer.text_to_ids(text) + [self.tokenizer.eos_id]
input_ids = [self.tokenizer.pad_id] * self.max_seq_length
input_ids[: len(tokens)] = tokens
input_ids = np.array(input_ids)
input_mask = (input_ids != self.tokenizer.pad_id).astype(np.float32)
acoustic_score = self.data[1][idx]
dist = editdistance.eval(text.split(), self.ground_truths[idx // self.beam_size].split())
ref_len = len(self.ground_truths[idx // self.beam_size].split())
len_in_chars = len(str(self.data[0][idx]))
return input_ids, input_mask, acoustic_score, dist, ref_len, len_in_chars, idx
def linear_search_wer(
dists, scores1, scores2, total_len, coef_range=[0, 10], coef_steps=10000, param_name='parameter'
):
"""
performs linear search to find the best coefficient when two set of scores are getting linearly fused.
Args:
dists: Tesnor of the distances between the ground truth and the candidates with shape of [number of samples, beam size]
scores1: Tensor of the first set of scores with shape of [number of samples, beam size]
scores2: Tensor of the second set of scores with shape of [number of samples, beam size]
total_len: The total length of all samples
coef_range: the search range for the coefficient
coef_steps: the number of steps that the search range would get divided into
param_name: the name of the parameter to be used in the figure
Output:
(best coefficient found, best WER achived)
"""
scale = scores1.mean().abs().item() / scores2.mean().abs().item()
left = coef_range[0] * scale
right = coef_range[1] * scale
coefs = np.linspace(left, right, coef_steps)
best_wer = 10000
best_coef = left
wers = []
for coef in coefs:
scores = scores1 + coef * scores2
wer = compute_wer(dists, scores, total_len)
wers.append(wer)
if wer < best_wer:
best_wer = wer
best_coef = coef
plt.plot(coefs, wers)
plt.title(f'WER% after rescoring with different values of {param_name}')
plt.ylabel('WER%')
plt.xlabel(param_name)
plt.show()
return best_coef, best_wer
def compute_wer(dists, scores, total_len):
"""
Sorts the candidates based on the scores and calculates the WER with the new top candidates.
Args:
dists: Tensor of the distances between the ground truth and the candidates with shape of [number of samples, beam size]
scores: Tensor of the scores for candidates with shape of [number of samples, beam size]
total_len: The total length of all samples
Output:
WER with the new scores
"""
indices = scores.max(dim=1, keepdim=True)[1]
wer = dists.gather(dim=1, index=indices).sum() / total_len
wer = wer.item()
return wer
def main():
parser = ArgumentParser()
parser.add_argument("--lm_model_file", type=str, required=True, help="path to LM model .nemo file")
parser.add_argument("--beams_file", type=str, required=True, help="path to beams .tsv file")
parser.add_argument(
"--eval_manifest", type=str, required=True, help="path to the evaluation `.json` manifest file"
)
parser.add_argument("--beam_size", type=int, required=True, help="number of beams per candidate")
parser.add_argument("--batch_size", type=int, default=256, help="inference batch size")
parser.add_argument("--alpha", type=float, default=None, help="parameter alpha of the fusion")
parser.add_argument("--beta", type=float, default=None, help="parameter beta of the fusion")
parser.add_argument(
"--scores_output_file", default=None, type=str, help="The optional path to store the rescored beams"
)
parser.add_argument(
"--device", default="cuda", type=str, help="The device to load the model onto to calculate the scores"
)
parser.add_argument(
"--use_amp", action="store_true", help="Whether to use AMP if available to calculate the scores"
)
args = parser.parse_args()
device = args.device
if device.startswith("cuda") and not torch.cuda.is_available():
logging.info(f"cuda is not available! switched to cpu.")
device = "cpu"
if args.lm_model_file.endswith(".nemo"):
logging.info("Attempting to initialize from .nemo file")
model = TransformerLMModel.restore_from(
restore_path=args.lm_model_file, map_location=torch.device(device)
).eval()
else:
raise NotImplementedError(f"Only supports .nemo files, but got: {args.model}")
max_seq_length = model.encoder._embedding.position_embedding.pos_enc.shape[0]
dataset = BeamScoresDataset(args.beams_file, model.tokenizer, args.eval_manifest, args.beam_size, max_seq_length)
data_loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=args.batch_size)
if args.use_amp:
if torch.cuda.is_available() and hasattr(torch.cuda, 'amp') and hasattr(torch.cuda.amp, 'autocast'):
logging.info("AMP is enabled!\n")
autocast = torch.cuda.amp.autocast
else:
@contextlib.contextmanager
def autocast():
yield
logging.info(f"Rescoring with beam_size: {args.beam_size}")
logging.info("Calculating the scores...")
with autocast():
with torch.no_grad():
am_scores, lm_scores, dists, ref_lens, lens_in_chars = [], [], [], [], []
for batch in tqdm.tqdm(data_loader):
input_ids, input_mask, acoustic_score, dist, ref_len, len_in_chars, idx = batch
max_len_in_batch = input_mask.sum(dim=0).argmin().item()
input_ids, input_mask = input_ids[:, :max_len_in_batch], input_mask[:, :max_len_in_batch]
if torch.cuda.is_available():
input_ids, input_mask = input_ids.to(device), input_mask.to(device)
dist, acoustic_score, len_in_chars = (
dist.to(device),
acoustic_score.to(device),
len_in_chars.to(device),
)
log_probs = model.forward(input_ids[:, :-1], input_mask[:, :-1])
target_log_probs = log_probs.gather(2, input_ids[:, 1:].unsqueeze(2)).squeeze(2)
neural_lm_score = torch.sum(target_log_probs * input_mask[:, 1:], dim=-1)
am_scores.append(acoustic_score)
lm_scores.append(neural_lm_score)
dists.append(dist)
ref_lens.append(ref_len)
lens_in_chars.append(len_in_chars)
am_scores = torch.cat(am_scores).view(-1, args.beam_size)
lm_scores = torch.cat(lm_scores).view(-1, args.beam_size)
dists = torch.cat(dists).view(-1, args.beam_size)
ref_lens = torch.cat(ref_lens).view(-1, args.beam_size)
lens_in_chars = torch.cat(lens_in_chars).view(-1, args.beam_size).to(am_scores.dtype)
total_len = ref_lens[:, 0].sum()
model_wer = dists[:, 0].sum() / total_len
ideal_wer = dists.min(dim=1)[0].sum() / total_len
if args.alpha is None:
logging.info("Linear search for alpha...")
coef1, _ = linear_search_wer(
dists=dists, scores1=am_scores, scores2=lm_scores, total_len=total_len, param_name='alpha'
)
coef1 = np.round(coef1, 3)
logging.info(f"alpha={coef1} achieved the best WER.")
logging.info(f"------------------------------------------------")
else:
coef1 = args.alpha
scores = am_scores + coef1 * lm_scores
if args.beta is None:
logging.info("Linear search for beta...")
coef2, _ = linear_search_wer(
dists=dists, scores1=scores, scores2=lens_in_chars, total_len=total_len, param_name='beta'
)
coef2 = np.round(coef2, 3)
logging.info(f"beta={coef2} achieved the best WER.")
logging.info(f"------------------------------------------------")
else:
coef2 = args.beta
new_scores = am_scores + coef1 * lm_scores + coef2 * lens_in_chars
rescored_wer = compute_wer(dists, new_scores, total_len)
logging.info(f"Input beams WER: {np.round(model_wer.item() * 100, 2)}%")
logging.info(f"------------------------------------------------")
logging.info(f" +LM rescoring WER: {np.round(rescored_wer * 100, 2)}%")
logging.info(f" with alpha={coef1}, beta={coef2}")
logging.info(f"------------------------------------------------")
logging.info(f"Best possible WER: {np.round(ideal_wer.item() * 100, 2)}%")
logging.info(f"------------------------------------------------")
new_scores_flatten = new_scores.flatten()
if args.scores_output_file is not None:
logging.info(f'Saving the candidates with their new scores at `{args.scores_output_file}`...')
with open(args.scores_output_file, "w") as fout:
for sample_id in range(len(dataset)):
fout.write(f"{dataset.data[0][sample_id]}\t{new_scores_flatten[sample_id]}\n")
if __name__ == '__main__':
main()
tts/NeMo/scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram.py
+345
View File
@@ -0,0 +1,345 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# This script would evaluate an N-gram language model trained with KenLM library (https://github.com/kpu/kenlm) in
# fusion with beam search decoders on top of a trained ASR model. NeMo's beam search decoders are capable of using the
# KenLM's N-gram models to find the best candidates. This script supports both character level and BPE level
# encodings and models which is detected automatically from the type of the model.
# You may train the LM model with 'scripts/ngram_lm/train_kenlm.py'.
#
# USAGE: python eval_beamsearch_ngram.py --nemo_model_file <path to the .nemo file of the model> \
# --input_manifest <path to the evaluation JSON manifest file \
# --kenlm_model_file <path to the binary KenLM model> \
# --beam_width <list of the beam widths> \
# --beam_alpha <list of the beam alphas> \
# --beam_width <list of the beam betas> \
# --preds_output_folder <optional folder to store the predictions> \
# --decoding_mode beam_search_ngram
# ...
#
# You may find more info on how to use this script at:
# https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/asr_language_modeling.html
# Please check train_kenlm.py to find out why we need TOKEN_OFFSET for BPE-based models
TOKEN_OFFSET = 100
import argparse
import contextlib
import json
import os
import pickle
import editdistance
import kenlm_utils
import numpy as np
import torch
from sklearn.model_selection import ParameterGrid
from tqdm.auto import tqdm
import nemo
import nemo.collections.asr as nemo_asr
from nemo.utils import logging
def beam_search_eval(
all_probs,
target_transcripts,
vocab,
ids_to_text_func=None,
preds_output_file=None,
lm_path=None,
beam_alpha=1.0,
beam_beta=0.0,
beam_width=128,
beam_batch_size=128,
progress_bar=True,
):
# creating the beam search decoder
beam_search_lm = nemo_asr.modules.BeamSearchDecoderWithLM(
vocab=vocab,
beam_width=beam_width,
alpha=beam_alpha,
beta=beam_beta,
lm_path=lm_path,
num_cpus=max(os.cpu_count(), 1),
input_tensor=False,
)
wer_dist_first = cer_dist_first = 0
wer_dist_best = cer_dist_best = 0
words_count = 0
chars_count = 0
sample_idx = 0
if preds_output_file:
out_file = open(preds_output_file, 'w')
if progress_bar:
it = tqdm(
range(int(np.ceil(len(all_probs) / beam_batch_size))),
desc=f"Beam search decoding with width={beam_width}, alpha={beam_alpha}, beta={beam_beta}",
ncols=120,
)
else:
it = range(int(np.ceil(len(all_probs) / beam_batch_size)))
for batch_idx in it:
# disabling type checking
with nemo.core.typecheck.disable_checks():
probs_batch = all_probs[batch_idx * beam_batch_size : (batch_idx + 1) * beam_batch_size]
beams_batch = beam_search_lm.forward(log_probs=probs_batch, log_probs_length=None,)
for beams_idx, beams in enumerate(beams_batch):
target = target_transcripts[sample_idx + beams_idx]
target_split_w = target.split()
target_split_c = list(target)
words_count += len(target_split_w)
chars_count += len(target_split_c)
wer_dist_min = cer_dist_min = 10000
for candidate_idx, candidate in enumerate(beams):
if ids_to_text_func is not None:
# For BPE encodings, need to shift by TOKEN_OFFSET to retrieve the original sub-word ids
pred_text = ids_to_text_func([ord(c) - TOKEN_OFFSET for c in candidate[1]])
else:
pred_text = candidate[1]
pred_split_w = pred_text.split()
wer_dist = editdistance.eval(target_split_w, pred_split_w)
pred_split_c = list(pred_text)
cer_dist = editdistance.eval(target_split_c, pred_split_c)
wer_dist_min = min(wer_dist_min, wer_dist)
cer_dist_min = min(cer_dist_min, cer_dist)
if candidate_idx == 0:
# first candidate
wer_dist_first += wer_dist
cer_dist_first += cer_dist
score = candidate[0]
if preds_output_file:
out_file.write('{}\t{}\n'.format(pred_text, score))
wer_dist_best += wer_dist_min
cer_dist_best += cer_dist_min
sample_idx += len(probs_batch)
if preds_output_file:
out_file.close()
logging.info(f"Stored the predictions of beam search decoding at '{preds_output_file}'.")
if lm_path:
logging.info(
'WER/CER with beam search decoding and N-gram model = {:.2%}/{:.2%}'.format(
wer_dist_first / words_count, cer_dist_first / chars_count
)
)
else:
logging.info(
'WER/CER with beam search decoding = {:.2%}/{:.2%}'.format(
wer_dist_first / words_count, cer_dist_first / chars_count
)
)
logging.info(
'Best WER/CER in candidates = {:.2%}/{:.2%}'.format(wer_dist_best / words_count, cer_dist_best / chars_count)
)
logging.info(f"=================================================================================")
def main():
parser = argparse.ArgumentParser(
description='Evaluate an ASR model with beam search decoding and n-gram KenLM language model.'
)
parser.add_argument(
"--nemo_model_file", required=True, type=str, help="The path of the '.nemo' file of the ASR model"
)
parser.add_argument(
"--kenlm_model_file", required=False, default=None, type=str, help="The path of the KenLM binary model file"
)
parser.add_argument("--input_manifest", required=True, type=str, help="The manifest file of the evaluation set")
parser.add_argument(
"--preds_output_folder", default=None, type=str, help="The optional folder where the predictions are stored"
)
parser.add_argument(
"--probs_cache_file", default=None, type=str, help="The cache file for storing the outputs of the model"
)
parser.add_argument(
"--acoustic_batch_size", default=16, type=int, help="The batch size to calculate log probabilities"
)
parser.add_argument(
"--device", default="cuda", type=str, help="The device to load the model onto to calculate log probabilities"
)
parser.add_argument(
"--use_amp", action="store_true", help="Whether to use AMP if available to calculate log probabilities"
)
parser.add_argument(
"--decoding_mode",
choices=["greedy", "beamsearch", "beamsearch_ngram"],
default="beamsearch_ngram",
type=str,
help="The decoding scheme to be used for evaluation.",
)
parser.add_argument(
"--beam_width",
required=True,
type=int,
nargs="+",
help="The width or list of the widths for the beam search decoding",
)
parser.add_argument(
"--beam_alpha",
required=True,
type=float,
nargs="+",
help="The alpha parameter or list of the alphas for the beam search decoding",
)
parser.add_argument(
"--beam_beta",
required=True,
type=float,
nargs="+",
help="The beta parameter or list of the betas for the beam search decoding",
)
parser.add_argument(
"--beam_batch_size", default=128, type=int, help="The batch size to be used for beam search decoding"
)
args = parser.parse_args()
asr_model = nemo_asr.models.EncDecCTCModelBPE.restore_from(
args.nemo_model_file, map_location=torch.device(args.device)
)
target_transcripts = []
with open(args.input_manifest, 'r') as manifest_file:
audio_file_paths = []
for line in tqdm(manifest_file, desc=f"Reading Manifest {args.input_manifest} ...", ncols=120):
data = json.loads(line)
target_transcripts.append(data['text'])
audio_file_paths.append(data['audio_filepath'])
if args.probs_cache_file and os.path.exists(args.probs_cache_file):
logging.info(f"Found a pickle file of probabilities at '{args.probs_cache_file}'.")
logging.info(f"Loading the cached pickle file of probabilities from '{args.probs_cache_file}' ...")
with open(args.probs_cache_file, 'rb') as probs_file:
all_probs = pickle.load(probs_file)
if len(all_probs) != len(audio_file_paths):
raise ValueError(
f"The number of samples in the probabilities file '{args.probs_cache_file}' does not "
f"match the manifest file. You may need to delete the probabilities cached file."
)
else:
if args.use_amp:
if torch.cuda.is_available() and hasattr(torch.cuda, 'amp') and hasattr(torch.cuda.amp, 'autocast'):
logging.info("AMP is enabled!\n")
autocast = torch.cuda.amp.autocast
else:
@contextlib.contextmanager
def autocast():
yield
with autocast():
with torch.no_grad():
all_logits = asr_model.transcribe(audio_file_paths, batch_size=args.acoustic_batch_size, logprobs=True)
all_probs = [kenlm_utils.softmax(logits) for logits in all_logits]
if args.probs_cache_file:
logging.info(f"Writing pickle files of probabilities at '{args.probs_cache_file}'...")
with open(args.probs_cache_file, 'wb') as f_dump:
pickle.dump(all_probs, f_dump)
wer_dist_greedy = 0
cer_dist_greedy = 0
words_count = 0
chars_count = 0
for batch_idx, probs in enumerate(all_probs):
preds = np.argmax(probs, axis=1)
preds_tensor = torch.tensor(preds, device='cpu').unsqueeze(0)
pred_text = asr_model._wer.ctc_decoder_predictions_tensor(preds_tensor)[0]
pred_split_w = pred_text.split()
target_split_w = target_transcripts[batch_idx].split()
pred_split_c = list(pred_text)
target_split_c = list(target_transcripts[batch_idx])
wer_dist = editdistance.eval(target_split_w, pred_split_w)
cer_dist = editdistance.eval(target_split_c, pred_split_c)
wer_dist_greedy += wer_dist
cer_dist_greedy += cer_dist
words_count += len(target_split_w)
chars_count += len(target_split_c)
logging.info('Greedy WER/CER = {:.2%}/{:.2%}'.format(wer_dist_greedy / words_count, cer_dist_greedy / chars_count))
encoding_level = kenlm_utils.SUPPORTED_MODELS.get(type(asr_model).__name__, None)
if not encoding_level:
logging.warning(
f"Model type '{type(asr_model).__name__}' may not be supported. Would try to train a char-level LM."
)
encoding_level = 'char'
vocab = asr_model.decoder.vocabulary
ids_to_text_func = None
if encoding_level == "subword":
vocab = [chr(idx + TOKEN_OFFSET) for idx in range(len(vocab))]
ids_to_text_func = asr_model.tokenizer.ids_to_text
# delete the model to free the memory
del asr_model
if args.decoding_mode == "beamsearch_ngram":
if not os.path.exists(args.kenlm_model_file):
raise FileNotFoundError(f"Could not find the KenLM model file '{args.kenlm_model_file}'.")
lm_path = args.kenlm_model_file
else:
lm_path = None
# 'greedy' decoding_mode would skip the beam search decoding
if args.decoding_mode in ["beamsearch_ngram", "beamsearch"]:
params = {'beam_width': args.beam_width, 'beam_alpha': args.beam_alpha, 'beam_beta': args.beam_beta}
hp_grid = ParameterGrid(params)
hp_grid = list(hp_grid)
logging.info(f"==============================Starting the beam search decoding===============================")
logging.info(f"Grid search size: {len(hp_grid)}")
logging.info(f"It may take some time...")
logging.info(f"==============================================================================================")
if args.preds_output_folder and not os.path.exists(args.preds_output_folder):
os.mkdir(args.preds_output_folder)
for hp in hp_grid:
if args.preds_output_folder:
preds_output_file = os.path.join(
args.preds_output_folder,
f"preds_out_width{hp['beam_width']}_alpha{hp['beam_alpha']}_beta{hp['beam_beta']}.tsv",
)
else:
preds_output_file = None
beam_search_eval(
all_probs=all_probs,
target_transcripts=target_transcripts,
vocab=vocab,
ids_to_text_func=ids_to_text_func,
preds_output_file=preds_output_file,
lm_path=lm_path,
beam_width=hp["beam_width"],
beam_alpha=hp["beam_alpha"],
beam_beta=hp["beam_beta"],
beam_batch_size=args.beam_batch_size,
progress_bar=True,
)
if __name__ == '__main__':
main()
tts/NeMo/scripts/asr_language_modeling/ngram_lm/install_beamsearch_decoders.sh
+6
View File
@@ -0,0 +1,6 @@
git clone https://github.com/NVIDIA/OpenSeq2Seq -b ctc-decoders
mv OpenSeq2Seq/decoders .
rm -rf OpenSeq2Seq
cd decoders
./setup.sh
cd ..
tts/NeMo/scripts/asr_language_modeling/ngram_lm/kenlm_utils.py
+112
View File
@@ -0,0 +1,112 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""
Utility methods to be used for training N-gram LM with KenLM in train_kenlm.py
"""
import json
import os
import numpy as np
from joblib import Parallel, delayed
from tqdm.auto import tqdm
# List of the supported models to be used with N-gram LM and beam search decoding
SUPPORTED_MODELS = {
'EncDecCTCModelBPE': 'subword',
'EncDecCTCModel': 'char',
}
def softmax(x):
e = np.exp(x - np.max(x))
return e / e.sum(axis=-1).reshape([x.shape[0], 1])
def read_train_file(path, lowercase: bool = False):
lines_read = 0
text_dataset = []
with open(path, 'r') as f:
reader = tqdm(iter(lambda: f.readline(), ''), desc="Read 0 lines", unit=' lines')
for i, line in enumerate(reader):
if path.endswith('.json'):
line = json.loads(line)['text']
line = line.replace("\n", "").strip()
if lowercase:
line = line.lower()
if line:
text_dataset.append(line)
lines_read += 1
if lines_read % 100000 == 0:
reader.set_description(f"Read {lines_read} lines")
return text_dataset
def tokenize_str(texts, tokenizer, offset):
tokenized_text = []
for text in texts:
tok_text = tokenizer.text_to_ids(text)
tok_text = [chr(token + offset) for token in tok_text]
tokenized_text.append(tok_text)
return tokenized_text
def tokenize_text(data, tokenizer, path, chunk_size=8192, buffer_size=32, token_offset=100):
dataset_len = len(data)
print(
f"Chunking {dataset_len} rows into {dataset_len / float(chunk_size):0.4f} tasks (each chunk contains {chunk_size} elements)"
)
current_step = 0
if os.path.exists(path):
print(f"Deleting previous file : {path}")
os.remove(path)
with Parallel(n_jobs=-2, verbose=10) as parallel:
while True:
start = current_step * chunk_size
end = min((current_step + buffer_size) * chunk_size, dataset_len)
tokenized_data = parallel(
delayed(tokenize_str)(data[start : start + chunk_size], tokenizer, token_offset)
for start in range(start, end, chunk_size)
)
# Write dataset
write_dataset(tokenized_data, path)
current_step += len(tokenized_data)
print(f"Finished writing {len(tokenized_data)} chunks to {path}. Current chunk index = {current_step}")
del tokenized_data
if end >= dataset_len:
break
def write_dataset(chunks, path):
basedir = os.path.dirname(path)
if not os.path.exists(basedir):
os.makedirs(basedir, exist_ok=True)
with open(path, 'a+', encoding='utf-8') as f:
for chunk_idx in tqdm(range(len(chunks)), desc='Chunk ', total=len(chunks), unit=' chunks'):
for text in chunks[chunk_idx]:
line = ' '.join(text)
f.write(f"{line}\n")
tts/NeMo/scripts/asr_language_modeling/ngram_lm/train_kenlm.py
+154
View File
@@ -0,0 +1,154 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# This script would train an N-gram language model with KenLM library (https://github.com/kpu/kenlm) which can be used
# with the beam search decoders on top of the ASR models. This script supports both character level and BPE level
# encodings and models which is detected automatically from the type of the model.
# After the N-gram model is trained, and stored in the binary format, you may use
# 'scripts/ngram_lm/eval_beamsearch_ngram.py' to evaluate it on an ASR model.
#
# You need to install the KenLM library and also the beam search decoders to use this feature. Please refer
# to 'scripts/ngram_lm/install_beamsearch_decoders.sh' on how to install them.
#
# USAGE: python train_kenlm.py --nemo_model_file <path to the .nemo file of the model> \
# --train_file <path to the training text or JSON manifest file \
# --kenlm_bin_path <path to the bin folder of KenLM library> \
# --kenlm_model_file <path to store the binary KenLM model> \
# --ngram_length <order of N-gram model>
#
# After training is done, the binary LM model is stored at the path specified by '--kenlm_model_file'.
# You may find more info on how to use this script at:
# https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/asr_language_modeling.html
import argparse
import logging
import os
import subprocess
import sys
import kenlm_utils
import torch
import nemo.collections.asr as nemo_asr
from nemo.utils import logging
"""
NeMo's beam search decoders only support char-level encodings. In order to make it work with BPE-level encodings, we
use a trick to encode the sub-word tokens of the training data as unicode characters and train a char-level KenLM.
TOKEN_OFFSET is the offset in the unicode table to be used to encode the BPE sub-words. This encoding scheme reduces
the required memory significantly, and the LM and its binary blob format require less storage space.
"""
TOKEN_OFFSET = 100
CHUNK_SIZE = 8192
CHUNK_BUFFER_SIZE = 512
def main():
parser = argparse.ArgumentParser(
description='Train an N-gram language model with KenLM to be used with beam search decoder of ASR models.'
)
parser.add_argument(
"--train_file",
required=True,
type=str,
help="Path to the training file, it can be a text file or JSON manifest",
)
parser.add_argument(
"--nemo_model_file", required=True, type=str, help="The path of the '.nemo' file of the ASR model"
)
parser.add_argument(
"--kenlm_model_file", required=True, type=str, help="The path to store the KenLM binary model file"
)
parser.add_argument("--ngram_length", required=True, type=int, help="The order of N-gram LM")
parser.add_argument("--kenlm_bin_path", required=True, type=str, help="The path to the bin folder of KenLM")
parser.add_argument(
"--do_lowercase", action='store_true', help="Whether to apply lower case conversion on the training text"
)
args = parser.parse_args()
""" TOKENIZER SETUP """
logging.info(f"Loading nemo model '{args.nemo_model_file}' ...")
model = nemo_asr.models.ASRModel.restore_from(args.nemo_model_file, map_location=torch.device('cpu'))
encoding_level = kenlm_utils.SUPPORTED_MODELS.get(type(model).__name__, None)
if not encoding_level:
logging.warning(
f"Model type '{type(model).__name__}' may not be supported. Would try to train a char-level LM."
)
encoding_level = 'char'
""" DATASET SETUP """
logging.info(f"Encoding the train file '{args.train_file}' ...")
dataset = kenlm_utils.read_train_file(args.train_file, lowercase=args.do_lowercase)
encoded_train_file = f"{args.kenlm_model_file}.tmp.txt"
if encoding_level == "subword":
kenlm_utils.tokenize_text(
dataset,
model.tokenizer,
path=encoded_train_file,
chunk_size=CHUNK_SIZE,
buffer_size=CHUNK_BUFFER_SIZE,
token_offset=TOKEN_OFFSET,
)
# --discount_fallback is needed for training KenLM for BPE-based models
discount_arg = "--discount_fallback"
else:
with open(encoded_train_file, 'w', encoding='utf-8') as f:
for line in dataset:
f.write(f"{line}\n")
discount_arg = ""
del model
arpa_file = f"{args.kenlm_model_file}.tmp.arpa"
""" LMPLZ ARGUMENT SETUP """
kenlm_args = [
os.path.join(args.kenlm_bin_path, 'lmplz'),
"-o",
f"{args.ngram_length}",
"--text",
encoded_train_file,
"--arpa",
arpa_file,
discount_arg,
]
ret = subprocess.run(kenlm_args, capture_output=False, text=True, stdout=sys.stdout, stderr=sys.stderr)
if ret.returncode != 0:
raise RuntimeError("Training KenLM was not successful!")
""" BINARY BUILD """
logging.info(f"Running binary_build command \n\n{' '.join(kenlm_args)}\n\n")
kenlm_args = [
os.path.join(args.kenlm_bin_path, "build_binary"),
"trie",
arpa_file,
args.kenlm_model_file,
]
ret = subprocess.run(kenlm_args, capture_output=False, text=True, stdout=sys.stdout, stderr=sys.stderr)
if ret.returncode != 0:
raise RuntimeError("Training KenLM was not successful!")
os.remove(encoded_train_file)
logging.info(f"Deleted the temporary encoded training file '{encoded_train_file}'.")
os.remove(arpa_file)
logging.info(f"Deleted the arpa file '{arpa_file}'.")
if __name__ == '__main__':
main()
tts/NeMo/scripts/average_model_checkpoints.py
+171
View File
@@ -0,0 +1,171 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Copyright 2017 Johns Hopkins University (Shinji Watanabe)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
# Changes to script
Change the script to import the NeMo model class you would like to load a checkpoint for,
then update the model constructor to use this model class. This can be found by the line:
<<< Change model class here ! >>>
By default, this script imports and creates the `EncDecCTCModelBPE` class but it can be
changed to any NeMo Model.
# Run the script
## Saving a .nemo model file (loaded with ModelPT.restore_from(...))
HYDRA_FULL_ERROR=1 python average_model_checkpoints.py \
--config-path="<path to config directory>" \
--config-name="<config name>" \
name=<name of the averaged checkpoint> \
+checkpoint_dir=<OPTIONAL: directory of checkpoint> \
+checkpoint_paths=\"[/path/to/ptl_1.ckpt,/path/to/ptl_2.ckpt,/path/to/ptl_3.ckpt,...]\"
## Saving an averaged pytorch checkpoint (loaded with torch.load(...))
HYDRA_FULL_ERROR=1 python average_model_checkpoints.py \
--config-path="<path to config directory>" \
--config-name="<config name>" \
name=<name of the averaged checkpoint> \
+checkpoint_dir=<OPTIONAL: directory of checkpoint> \
+checkpoint_paths=\"[/path/to/ptl_1.ckpt,/path/to/ptl_2.ckpt,/path/to/ptl_3.ckpt,...]\" \
+save_ckpt_only=true
"""
import os
import pytorch_lightning as pl
import torch
from omegaconf import OmegaConf, open_dict
# Change this import to the model you would like to average
from nemo.collections.asr.models import EncDecCTCModelBPE
from nemo.core.config import hydra_runner
from nemo.utils import logging
def process_config(cfg: OmegaConf):
if 'name' not in cfg or cfg.name is None:
raise ValueError("`cfg.name` must be provided to save a model checkpoint")
if 'checkpoint_paths' not in cfg or cfg.checkpoint_paths is None:
raise ValueError(
"`cfg.checkpoint_paths` must be provided as a list of one or more str paths to "
"pytorch lightning checkpoints"
)
save_ckpt_only = False
with open_dict(cfg):
name_prefix = cfg.name
checkpoint_paths = cfg.pop('checkpoint_paths')
if 'checkpoint_dir' in cfg:
checkpoint_dir = cfg.pop('checkpoint_dir')
else:
checkpoint_dir = None
if 'save_ckpt_only' in cfg:
save_ckpt_only = cfg.pop('save_ckpt_only')
if type(checkpoint_paths) not in (list, tuple):
checkpoint_paths = str(checkpoint_paths).replace("[", "").replace("]", "")
checkpoint_paths = checkpoint_paths.split(",")
checkpoint_paths = [ckpt_path.strip() for ckpt_path in checkpoint_paths]
if checkpoint_dir is not None:
checkpoint_paths = [os.path.join(checkpoint_dir, path) for path in checkpoint_paths]
return name_prefix, checkpoint_paths, save_ckpt_only
@hydra_runner(config_path=None, config_name=None)
def main(cfg):
name_prefix, checkpoint_paths, save_ckpt_only = process_config(cfg)
if not save_ckpt_only:
trainer = pl.Trainer(**cfg.trainer)
# <<< Change model class here ! >>>
# Model architecture which will contain the averaged checkpoints
# Change the model constructor to the one you would like (if needed)
model = EncDecCTCModelBPE(cfg=cfg.model, trainer=trainer)
""" < Checkpoint Averaging Logic > """
# load state dicts
n = len(checkpoint_paths)
avg_state = None
logging.info(f"Averaging {n} checkpoints ...")
for ix, path in enumerate(checkpoint_paths):
checkpoint = torch.load(path, map_location='cpu')
if 'state_dict' in checkpoint:
checkpoint = checkpoint['state_dict']
if ix == 0:
# Initial state
avg_state = checkpoint
logging.info(f"Initialized average state dict with checkpoint : {path}")
else:
# Accumulated state
for k in avg_state:
avg_state[k] = avg_state[k] + checkpoint[k]
logging.info(f"Updated average state dict with state from checkpoint : {path}")
for k in avg_state:
if str(avg_state[k].dtype).startswith("torch.int"):
# For int type, not averaged, but only accumulated.
# e.g. BatchNorm.num_batches_tracked
pass
else:
avg_state[k] = avg_state[k] / n
# Save model
if save_ckpt_only:
ckpt_name = name_prefix + '-averaged.ckpt'
torch.save(avg_state, ckpt_name)
logging.info(f"Averaged pytorch checkpoint saved as : {ckpt_name}")
else:
# Set model state
logging.info("Loading averaged state dict in provided model")
model.load_state_dict(avg_state, strict=True)
ckpt_name = name_prefix + '-averaged.nemo'
model.save_to(ckpt_name)
logging.info(f"Averaged model saved as : {ckpt_name}")
if __name__ == '__main__':
main()
tts/NeMo/scripts/dataset_processing/get_commonvoice_data.py
+178
View File
@@ -0,0 +1,178 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Copyright (c) 2020, SeanNaren. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# To convert mp3 files to wav using sox, you must have installed sox with mp3 support
# For example sudo apt-get install libsox-fmt-mp3
import argparse
import csv
import json
import logging
import multiprocessing
import os
import subprocess
import sys
import tarfile
from multiprocessing.pool import ThreadPool
from pathlib import Path
from typing import List
import sox
from sox import Transformer
from tqdm import tqdm
parser = argparse.ArgumentParser(description='Downloads and processes Mozilla Common Voice dataset.')
parser.add_argument("--data_root", default='CommonVoice_dataset/', type=str, help="Directory to store the dataset.")
parser.add_argument('--manifest_dir', default='./', type=str, help='Output directory for manifests')
parser.add_argument("--num_workers", default=multiprocessing.cpu_count(), type=int, help="Workers to process dataset.")
parser.add_argument('--sample_rate', default=16000, type=int, help='Sample rate')
parser.add_argument('--n_channels', default=1, type=int, help='Number of channels for output wav files')
parser.add_argument(
'--files_to_process',
nargs='+',
default=['test.tsv', 'dev.tsv', 'train.tsv'],
type=str,
help='list of *.csv file names to process',
)
parser.add_argument(
'--version',
default='cv-corpus-5.1-2020-06-22',
type=str,
help='Version of the dataset (obtainable via https://commonvoice.mozilla.org/en/datasets',
)
parser.add_argument(
'--language',
default='en',
type=str,
help='Which language to download.(default english,'
'check https://commonvoice.mozilla.org/en/datasets for more language codes',
)
args = parser.parse_args()
COMMON_VOICE_URL = (
f"https://voice-prod-bundler-ee1969a6ce8178826482b88e843c335139bd3fb4.s3.amazonaws.com/"
"{}/{}.tar.gz".format(args.version, args.language)
)
def create_manifest(data: List[tuple], output_name: str, manifest_path: str):
output_file = Path(manifest_path) / output_name
output_file.parent.mkdir(exist_ok=True, parents=True)
with output_file.open(mode='w') as f:
for wav_path, duration, text in tqdm(data, total=len(data)):
f.write(
json.dumps({'audio_filepath': os.path.abspath(wav_path), "duration": duration, 'text': text}) + '\n'
)
def process_files(csv_file, data_root, num_workers):
""" Read *.csv file description, convert mp3 to wav, process text.
Save results to data_root.
Args:
csv_file: str, path to *.csv file with data description, usually start from 'cv-'
data_root: str, path to dir to save results; wav/ dir will be created
"""
wav_dir = os.path.join(data_root, 'wav/')
os.makedirs(wav_dir, exist_ok=True)
audio_clips_path = os.path.dirname(csv_file) + '/clips/'
def process(x):
file_path, text = x
file_name = os.path.splitext(os.path.basename(file_path))[0]
text = text.lower().strip()
audio_path = os.path.join(audio_clips_path, file_path)
output_wav_path = os.path.join(wav_dir, file_name + '.wav')
tfm = Transformer()
tfm.rate(samplerate=args.sample_rate)
tfm.channels(n_channels=args.n_channels)
tfm.build(input_filepath=audio_path, output_filepath=output_wav_path)
duration = sox.file_info.duration(output_wav_path)
return output_wav_path, duration, text
logging.info('Converting mp3 to wav for {}.'.format(csv_file))
with open(csv_file) as csvfile:
reader = csv.DictReader(csvfile, delimiter='\t')
next(reader, None) # skip the headers
data = [(row['path'], row['sentence']) for row in reader]
with ThreadPool(num_workers) as pool:
data = list(tqdm(pool.imap(process, data), total=len(data)))
return data
def main():
data_root = args.data_root
os.makedirs(data_root, exist_ok=True)
target_unpacked_dir = os.path.join(data_root, "CV_unpacked")
if os.path.exists(target_unpacked_dir):
logging.info('Find existing folder {}'.format(target_unpacked_dir))
else:
logging.info("Could not find Common Voice, Downloading corpus...")
commands = [
'wget',
'--user-agent',
'"Mozilla/5.0 (Windows NT 10.0; WOW64) '
'AppleWebKit/537.36 (KHTML, like Gecko) Chrome/51.0.2704.103 Safari/537.36"',
'-P',
data_root,
f'{COMMON_VOICE_URL}',
]
commands = " ".join(commands)
subprocess.run(commands, shell=True, stderr=sys.stderr, stdout=sys.stdout, capture_output=False)
filename = f"{args.language}.tar.gz"
target_file = os.path.join(data_root, os.path.basename(filename))
os.makedirs(target_unpacked_dir, exist_ok=True)
logging.info("Unpacking corpus to {} ...".format(target_unpacked_dir))
tar = tarfile.open(target_file)
tar.extractall(target_unpacked_dir)
tar.close()
folder_path = os.path.join(target_unpacked_dir, args.version + f'/{args.language}/')
for csv_file in args.files_to_process:
data = process_files(
csv_file=os.path.join(folder_path, csv_file),
data_root=os.path.join(data_root, os.path.splitext(csv_file)[0]),
num_workers=args.num_workers,
)
logging.info('Creating manifests...')
create_manifest(
data=data,
output_name=f'commonvoice_{os.path.splitext(csv_file)[0]}_manifest.json',
manifest_path=args.manifest_dir,
)
if __name__ == "__main__":
main()
tts/NeMo/scripts/dataset_processing/get_hi-mia_data.py
+185
View File
@@ -0,0 +1,185 @@
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# USAGE: python get_aishell_data.py --data_root=<where to put data>
import argparse
import json
import logging as _logging
import os
import tarfile
import urllib.request
from glob import glob
import librosa as l
from sklearn.model_selection import StratifiedShuffleSplit
from tqdm import tqdm
parser = argparse.ArgumentParser(description='HI-MIA Data download')
parser.add_argument("--data_root", required=True, default=None, type=str)
parser.add_argument("--log_level", default=20, type=int)
args = parser.parse_args()
logging = _logging.getLogger(__name__)
logging.addHandler(_logging.StreamHandler())
logging.setLevel(args.log_level)
URL = {
'dev': "http://www.openslr.org/resources/85/dev.tar.gz",
'test': "http://www.openslr.org/resources/85/test.tar.gz",
'train': "http://www.openslr.org/resources/85/train.tar.gz",
}
def __maybe_download_file(destination: str, source: str):
"""
Downloads source to destination if it doesn't exist.
If exists, skips download
Args:
destination: local filepath
source: url of resource
Returns:
"""
source = URL[source]
if not os.path.exists(destination) and not os.path.exists(os.path.splitext(destination)[0]):
logging.info("{0} does not exist. Downloading ...".format(destination))
urllib.request.urlretrieve(source, filename=destination + '.tmp')
os.rename(destination + '.tmp', destination)
logging.info("Downloaded {0}.".format(destination))
elif os.path.exists(destination):
logging.info("Destination {0} exists. Skipping.".format(destination))
elif os.path.exists(os.path.splitext(destination)[0]):
logging.warning(
"Assuming extracted folder %s contains the extracted files from %s. Will not download.",
os.path.basename(destination),
destination,
)
return destination
def __extract_all_files(filepath: str, data_root: str, data_dir: str):
if not os.path.exists(data_dir):
extract_file(filepath, data_root)
audio_dir = os.path.join(data_dir, 'wav')
for subfolder, _, filelist in os.walk(audio_dir):
for ftar in filelist:
extract_file(os.path.join(subfolder, ftar), subfolder)
else:
logging.info('Skipping extracting. Data already there %s' % data_dir)
def extract_file(filepath: str, data_dir: str):
try:
tar = tarfile.open(filepath)
tar.extractall(data_dir)
tar.close()
except Exception:
logging.info('Not extracting. Maybe already there?')
def __remove_tarred_files(filepath: str, data_dir: str):
if os.path.exists(data_dir) and os.path.isfile(filepath):
logging.info('Deleting %s' % filepath)
os.remove(filepath)
def write_file(name, lines, idx):
with open(name, 'w') as fout:
for i in idx:
dic = lines[i]
json.dump(dic, fout)
fout.write('\n')
logging.info("wrote %s", name)
def __process_data(data_folder: str, data_set: str):
"""
To generate manifest
Args:
data_folder: source with wav files
Returns:
"""
fullpath = os.path.abspath(data_folder)
scp = glob(fullpath + '/**/*.wav', recursive=True)
out = os.path.join(fullpath, data_set + '_all.json')
utt2spk = os.path.join(fullpath, 'utt2spk')
utt2spk_file = open(utt2spk, 'w')
id = -2 # speaker id
if os.path.exists(out):
logging.warning(
"%s already exists and is assumed to be processed. If not, please delete %s and rerun this script",
out,
out,
)
return
speakers = []
lines = []
with open(out, 'w') as outfile:
for line in tqdm(scp):
line = line.strip()
y, sr = l.load(line, sr=None)
if sr != 16000:
y, sr = l.load(line, sr=16000)
l.output.write_wav(line, y, sr)
dur = l.get_duration(y=y, sr=sr)
if data_set == 'test':
speaker = line.split('/')[-1].split('.')[0].split('_')[0]
else:
speaker = line.split('/')[id]
speaker = list(speaker)
speaker = ''.join(speaker)
speakers.append(speaker)
meta = {"audio_filepath": line, "duration": float(dur), "label": speaker}
lines.append(meta)
json.dump(meta, outfile)
outfile.write("\n")
utt2spk_file.write(line.split('/')[-1] + "\t" + speaker + "\n")
utt2spk_file.close()
if data_set != 'test':
sss = StratifiedShuffleSplit(n_splits=1, test_size=0.1, random_state=42)
for train_idx, test_idx in sss.split(speakers, speakers):
print(len(train_idx))
out = os.path.join(fullpath, 'train.json')
write_file(out, lines, train_idx)
out = os.path.join(fullpath, 'dev.json')
write_file(out, lines, test_idx)
def main():
data_root = args.data_root
for data_set in URL.keys():
# data_set = 'data_aishell'
logging.info("\n\nWorking on: {0}".format(data_set))
file_path = os.path.join(data_root, data_set + ".tgz")
logging.info("Getting {0}".format(data_set))
__maybe_download_file(file_path, data_set)
logging.info("Extracting {0}".format(data_set))
data_folder = os.path.join(data_root, data_set)
__extract_all_files(file_path, data_root, data_folder)
__remove_tarred_files(file_path, data_folder)
logging.info("Processing {0}".format(data_set))
__process_data(data_folder, data_set)
logging.info('Done!')
if __name__ == "__main__":
main()
tts/NeMo/scripts/dataset_processing/get_librispeech_data.py
+171
View File
@@ -0,0 +1,171 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# USAGE: python get_librispeech_data.py --data_root=<where to put data>
# --data_set=<datasets_to_download> --num_workers=<number of parallel workers>
# where <datasets_to_download> can be: dev_clean, dev_other, test_clean,
# test_other, train_clean_100, train_clean_360, train_other_500 or ALL
# You can also put more than one data_set comma-separated:
# --data_set=dev_clean,train_clean_100
import argparse
import fnmatch
import functools
import json
import logging
import multiprocessing
import os
import subprocess
import tarfile
import urllib.request
from sox import Transformer
from tqdm import tqdm
parser = argparse.ArgumentParser(description='LibriSpeech Data download')
parser.add_argument("--data_root", required=True, default=None, type=str)
parser.add_argument("--data_sets", default="dev_clean", type=str)
parser.add_argument("--num_workers", default=4, type=int)
args = parser.parse_args()
URLS = {
'TRAIN_CLEAN_100': ("http://www.openslr.org/resources/12/train-clean-100.tar.gz"),
'TRAIN_CLEAN_360': ("http://www.openslr.org/resources/12/train-clean-360.tar.gz"),
'TRAIN_OTHER_500': ("http://www.openslr.org/resources/12/train-other-500.tar.gz"),
'DEV_CLEAN': "http://www.openslr.org/resources/12/dev-clean.tar.gz",
'DEV_OTHER': "http://www.openslr.org/resources/12/dev-other.tar.gz",
'TEST_CLEAN': "http://www.openslr.org/resources/12/test-clean.tar.gz",
'TEST_OTHER': "http://www.openslr.org/resources/12/test-other.tar.gz",
}
def __maybe_download_file(destination: str, source: str):
"""
Downloads source to destination if it doesn't exist.
If exists, skips download
Args:
destination: local filepath
source: url of resource
Returns:
"""
source = URLS[source]
if not os.path.exists(destination):
logging.info("{0} does not exist. Downloading ...".format(destination))
urllib.request.urlretrieve(source, filename=destination + '.tmp')
os.rename(destination + '.tmp', destination)
logging.info("Downloaded {0}.".format(destination))
else:
logging.info("Destination {0} exists. Skipping.".format(destination))
return destination
def __extract_file(filepath: str, data_dir: str):
try:
tar = tarfile.open(filepath)
tar.extractall(data_dir)
tar.close()
except Exception:
logging.info('Not extracting. Maybe already there?')
def __process_transcript(file_path: str, dst_folder: str):
"""
Converts flac files to wav from a given transcript, capturing the metadata.
Args:
file_path: path to a source transcript with flac sources
dst_folder: path where wav files will be stored
Returns:
a list of metadata entries for processed files.
"""
entries = []
root = os.path.dirname(file_path)
with open(file_path, encoding="utf-8") as fin:
for line in fin:
id, text = line[: line.index(" ")], line[line.index(" ") + 1 :]
transcript_text = text.lower().strip()
# Convert FLAC file to WAV
flac_file = os.path.join(root, id + ".flac")
wav_file = os.path.join(dst_folder, id + ".wav")
if not os.path.exists(wav_file):
Transformer().build(flac_file, wav_file)
# check duration
duration = subprocess.check_output("soxi -D {0}".format(wav_file), shell=True)
entry = {}
entry['audio_filepath'] = os.path.abspath(wav_file)
entry['duration'] = float(duration)
entry['text'] = transcript_text
entries.append(entry)
return entries
def __process_data(data_folder: str, dst_folder: str, manifest_file: str, num_workers: int):
"""
Converts flac to wav and build manifests's json
Args:
data_folder: source with flac files
dst_folder: where wav files will be stored
manifest_file: where to store manifest
num_workers: number of parallel workers processing files
Returns:
"""
if not os.path.exists(dst_folder):
os.makedirs(dst_folder)
files = []
entries = []
for root, dirnames, filenames in os.walk(data_folder):
for filename in fnmatch.filter(filenames, '*.trans.txt'):
files.append(os.path.join(root, filename))
with multiprocessing.Pool(num_workers) as p:
processing_func = functools.partial(__process_transcript, dst_folder=dst_folder)
results = p.imap(processing_func, files)
for result in tqdm(results, total=len(files)):
entries.extend(result)
with open(manifest_file, 'w') as fout:
for m in entries:
fout.write(json.dumps(m) + '\n')
def main():
data_root = args.data_root
data_sets = args.data_sets
num_workers = args.num_workers
if data_sets == "ALL":
data_sets = "dev_clean,dev_other,train_clean_100,train_clean_360,train_other_500,test_clean,test_other"
for data_set in data_sets.split(','):
logging.info("\n\nWorking on: {0}".format(data_set))
filepath = os.path.join(data_root, data_set + ".tar.gz")
logging.info("Getting {0}".format(data_set))
__maybe_download_file(filepath, data_set.upper())
logging.info("Extracting {0}".format(data_set))
__extract_file(filepath, data_root)
logging.info("Processing {0}".format(data_set))
__process_data(
os.path.join(os.path.join(data_root, "LibriSpeech"), data_set.replace("_", "-"),),
os.path.join(os.path.join(data_root, "LibriSpeech"), data_set.replace("_", "-"),) + "-processed",
os.path.join(data_root, data_set + ".json"),
num_workers=num_workers,
)
logging.info('Done!')
if __name__ == "__main__":
main()
tts/NeMo/scripts/dataset_processing/get_openslr_rir_data.py
+140
View File
@@ -0,0 +1,140 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# USAGE: python get_openslr_rir_data.py --data_root=<where to put data>
# Data is downloaded from OpenSLR's "Room Impulse Response and Noise Database"
# RIRs in multichannel files are separated into single channel files and
# a json file that can be used as in input to NeMo is created
import argparse
import glob
import json
import logging
import os
import subprocess
import urllib.request
from shutil import copy, move
from zipfile import ZipFile
parser = argparse.ArgumentParser(description='OpenSLR RIR Data download and process')
parser.add_argument("--data_root", required=True, default=None, type=str)
args = parser.parse_args()
URLS = {
'SLR28': ("http://www.openslr.org/resources/28/rirs_noises.zip"),
}
def __maybe_download_file(destination: str, source: str):
"""
Downloads source to destination if it doesn't exist.
If exists, skips download
Args:
destination: local filepath
source: url of resource
Returns:
"""
source = URLS[source]
if not os.path.exists(destination):
logging.info("{0} does not exist. Downloading ...".format(destination))
urllib.request.urlretrieve(source, filename=destination + '.tmp')
os.rename(destination + '.tmp', destination)
logging.info("Downloaded {0}.".format(destination))
else:
logging.info("Destination {0} exists. Skipping.".format(destination))
return destination
def __extract_file(filepath: str, data_dir: str):
try:
with ZipFile(filepath, 'r') as zipObj:
zipObj.extractall(data_dir)
except Exception:
logging.info('Not extracting. Maybe already there?')
def __process_data(data_folder: str, dst_folder: str, manifest_file: str):
"""
Converts flac to wav and build manifests's json
Args:
data_folder: source with flac files
dst_folder: where wav files will be stored
manifest_file: where to store manifest
Returns:
"""
if not os.path.exists(dst_folder):
os.makedirs(dst_folder)
real_rir_list = os.path.join(data_folder, "RIRS_NOISES", "real_rirs_isotropic_noises", "rir_list")
rirfiles = []
with open(real_rir_list, "r") as rir_f:
for line in rir_f:
rirfiles.append(os.path.join(data_folder, line.rstrip().split(" ")[4]))
real_rir_folder = os.path.join(dst_folder, "real_rirs")
if not os.path.exists(real_rir_folder):
os.makedirs(real_rir_folder)
# split multi-channel rir files to single channel
for rir_f in rirfiles:
n_chans = int(subprocess.check_output("soxi -c {0}".format(rir_f), shell=True))
if n_chans == 1:
copy(rir_f, real_rir_folder)
else:
for chan in range(1, n_chans + 1):
chan_file_name = os.path.join(
real_rir_folder, os.path.splitext(os.path.basename(rir_f))[0] + "-" + str(chan) + ".wav"
)
_ = subprocess.check_output(f"sox {rir_f} {chan_file_name} remix {chan}", shell=True)
# move simulated rirs to processed
if not os.path.exists(os.path.join(dst_folder, "simulated_rirs")):
move(os.path.join(data_folder, "RIRS_NOISES", "simulated_rirs"), dst_folder)
os.chdir(dst_folder)
all_rirs = glob.glob("**/*.wav", recursive=True)
with open(manifest_file, 'w') as man_f:
entry = {}
for rir in all_rirs:
rir_file = os.path.join(dst_folder, rir)
duration = subprocess.check_output("soxi -D {0}".format(rir_file), shell=True)
entry['audio_filepath'] = rir_file
entry['duration'] = float(duration)
entry['offset'] = 0
entry['text'] = '_'
man_f.write(json.dumps(entry) + "\n")
print('Done!')
def main():
data_root = os.path.abspath(args.data_root)
data_set = "slr28"
logging.getLogger().setLevel(logging.INFO)
logging.info("\n\nWorking on: {0}".format(data_set))
filepath = os.path.join(data_root, data_set + ".zip")
logging.info("Getting {0}".format(data_set))
__maybe_download_file(filepath, data_set.upper())
logging.info("Extracting {0}".format(data_set))
__extract_file(filepath, data_root)
logging.info("Processing {0}".format(data_set))
__process_data(
data_root,
os.path.join(os.path.join(data_root, "processed")),
os.path.join(os.path.join(data_root, "processed", "rir.json")),
)
logging.info('Done!')
if __name__ == "__main__":
main()
tts/NeMo/scripts/dataset_processing/ljspeech/README
+14
View File
@@ -0,0 +1,14 @@
These scripts are used for extracting features from LJSpeech.
You will need to run the following scripts (note that these will call the other
scripts as necessary):
- `create_manifests_and_textfiles.py`: Creates the manifest and text files that
are used by the Montreal Forced Aligner (MFA) library.
- `extract_ljspeech_phonemes_and_durs.sh`: Extracts phonemes and alignments via
MFA and calculates phoneme durations.
Needs to be run with Conda active since it creates a Conda environment.
Also note that you may need to install OpenBlas manually (e.g. sudo apt-get
install libopenblas-dev), sometimes MFA can't find it if installed via Conda.
(This script calls `create_token2idx_dict.py` and `calculate_durs.py`.)
- `extract_ljspeech_energy_pitch.py`
tts/NeMo/scripts/dataset_processing/ljspeech/calculate_durs.py
+126
View File
@@ -0,0 +1,126 @@
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Calculates durations for LJSpeech based on MFA TextGrid alignments.
"""
import argparse
import glob
import json
import os
import pickle
from math import ceil
import numpy as np
import tgt
import torch
from tqdm import tqdm
parser = argparse.ArgumentParser(description="Calculates phoneme durations for LJSpeech from TextGrids.")
parser.add_argument('--ljspeech_dir', required=True, default=None, type=str)
parser.add_argument(
'--mappings',
required=False,
default=None,
type=str,
help='JSON file of mappings created with create_token2idx_dict.py',
)
parser.add_argument('--sr', required=False, default=22050, type=int)
parser.add_argument('--hop_length', required=False, default=256, type=int)
args = parser.parse_args()
def calculate_durations(textgrid, phone2idx):
tokens = []
durs = []
frames_per_second = args.sr / args.hop_length
tg = tgt.read_textgrid(textgrid, include_empty_intervals=True)
data_tier = tg.get_tier_by_name("phones")
# Get total frames
total_frames = ceil((data_tier.end_time - data_tier.start_time) * frames_per_second)
# Find start and end frames of each token
se_in_frames = np.array([(frames_per_second * d.start_time, frames_per_second * d.end_time) for d in data_tier])
se_in_frames = np.round(se_in_frames)
durs = (se_in_frames[:, 1] - se_in_frames[:, 0]).astype(int)
blank_set = ('sil', 'sp', 'spn', '', '<unk>')
blank_token = " "
# merge repeated blank tokens
tokens, durations = [], []
for i in range(len(data_tier)):
x = data_tier[i].text
if x == 'spn':
return None, None, None
x = blank_token if x in blank_set else x
if len(tokens) and tokens[-1] == blank_token and x == blank_token:
durations[-1] += durs[i]
else:
tokens.append(x)
durations.append(durs[i])
tokens_enc = [phone2idx[token] for token in tokens]
tokens_enc, durations = torch.LongTensor(tokens_enc), torch.LongTensor(durations)
# Add rounding error to final token
durations[-1] += total_frames - durations.sum()
return tokens, tokens_enc, durations
def main():
textgrid_list = glob.glob(os.path.join(args.ljspeech_dir, 'alignments/wavs/*.TextGrid'))
# Create target_dir if necessary
target_dir = os.path.join(args.ljspeech_dir, 'phoneme_durations/')
print(f"Calculating phoneme durations, files will be in: {target_dir}")
if not os.path.exists(target_dir):
print(f"Creating {target_dir}")
os.mkdir(target_dir)
# Read phoneme to idx mappings
phone2idx = None
if args.mappings:
with open(args.mappings, 'r') as f:
mappings = json.load(f)
phone2idx = mappings['phone2idx']
oov_samples = []
# Iterate through all TextGrid files
for textgrid in tqdm(textgrid_list):
basename = os.path.splitext(os.path.basename(textgrid))[0][5:] # Chop off 'wavs_' prefix
phones_mfa, tokens_mfa, durs = calculate_durations(textgrid, phone2idx)
if phones_mfa is None:
oov_samples.append(basename)
continue
# Save to file
target_path = os.path.join(target_dir, f'{basename}.pt')
torch.save({'text_encoded': tokens_mfa, 'token_duration': durs}, target_path)
print(f"Getting rid of {len(oov_samples)} samples with OOV words.")
oov_target = os.path.join(args.ljspeech_dir, 'wavs_to_ignore.pkl')
with open(oov_target, 'wb') as f:
pickle.dump(oov_samples, f)
print(f"List of OOV samples written to: {oov_target}")
if __name__ == '__main__':
main()
tts/NeMo/scripts/dataset_processing/ljspeech/create_manifests_and_textfiles.py
+71
View File
@@ -0,0 +1,71 @@
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Creates the manifest and .txt transcript files for an LJSpeech split.
The manifest will be used for training, and the .txt files are for the MFA library to find.
"""
import argparse
import json
import os
import sox
import wget
from nemo.collections.common.parts.preprocessing import parsers
parser = argparse.ArgumentParser()
parser.add_argument('--ljspeech_base', required=True, default=None, type=str)
args = parser.parse_args()
def main():
filelist_base = 'https://raw.githubusercontent.com/NVIDIA/tacotron2/master/filelists/'
filelists = ['train', 'val', 'test']
# NeMo parser for text normalization
text_parser = parsers.make_parser(name='en')
for split in filelists:
# Download file list if necessary
filelist_path = os.path.join(args.ljspeech_base, f"ljs_audio_text_{split}_filelist.txt")
if not os.path.exists(filelist_path):
wget.download(f"{filelist_base}/ljs_audio_text_{split}_filelist.txt", out=args.ljspeech_base)
manifest_target = os.path.join(args.ljspeech_base, f"ljspeech_{split}.json")
with open(manifest_target, 'w') as f_out:
with open(filelist_path, 'r') as filelist:
print(f"\nCreating {manifest_target}...")
for line in filelist:
basename = line[6:16]
text = text_parser._normalize(line[21:].strip())
# Make sure corresponding wavfile exists and write .txt transcript
wav_path = os.path.join(args.ljspeech_base, 'wavs/', basename + '.wav')
assert os.path.exists(wav_path)
txt_path = os.path.join(args.ljspeech_base, 'wavs/', basename + '.txt')
with open(txt_path, 'w') as f_txt:
f_txt.write(text)
# Write manifest entry
entry = {
'audio_filepath': wav_path,
'duration': sox.file_info.duration(wav_path),
'text': text,
}
f_out.write(json.dumps(entry) + '\n')
if __name__ == '__main__':
main()
tts/NeMo/scripts/dataset_processing/ljspeech/create_token2idx_dict.py
+68
View File
@@ -0,0 +1,68 @@
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Creates a dictionary from token to index based on dictionary .txt given.
"""
import argparse
import json
import os
parser = argparse.ArgumentParser()
parser.add_argument('--dictionary', required=True, default=None, type=str)
parser.add_argument('--dict_out', required=True, default=None, type=str)
args = parser.parse_args()
def main():
if not os.path.exists(args.dictionary):
raise FileNotFoundError(f"Could not find dictionary file {args.dictionary}")
phonemes = set()
word2phones = {}
with open(args.dictionary, 'r') as f:
for line in f:
line = line.split()
word = line[0]
tokens = line[1:]
word2phones[word] = tokens
phonemes.update(tokens)
# Small list of additional punctuation
word2phones[','] = [' ']
word2phones[';'] = [' ']
word2phones['.'] = [' ']
word2phones['!'] = [' ']
word2phones['?'] = [' ']
word2phones['"'] = [' ']
word2phones['-'] = [' ']
phone2idx = {k: i for i, k in enumerate(phonemes)}
phone2idx[' '] = len(phone2idx)
phone2idx['sil'] = phone2idx[' '] # Silence
phone2idx['sp'] = phone2idx[' '] # Space
phone2idx['spn'] = phone2idx[' '] # OOV/unk
dicts = {
'phone2idx': phone2idx,
'word2phones': word2phones,
}
with open(args.dict_out, 'w') as f:
json.dump(dicts, f)
print(f"Total number of phone indices: {len(phone2idx)}")
if __name__ == '__main__':
main()
tts/NeMo/scripts/dataset_processing/ljspeech/extract_ljspeech_energy_pitch.py
+88
View File
@@ -0,0 +1,88 @@
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Extracts energy and pitch data from LJSpeech wav files and writes them to a directory as LJxxx-xxxx.npy files.
Assuming that the wavs are located at `<LJSpeech_base_path>/wavs`, this script will write the
energy files to `<LJSpeech_base_path>/energies/` and the pitch files to `<LJSpeech_base_path>/pitches`,
creating the relevant directories if necessary.
USAGE: python extract_ljspeech_energy_pitch.py --ljspeech_dir=<LJSpeech_base_path>
"""
import argparse
from pathlib import Path
import librosa
import numpy as np
tqdm = None
try:
from tqdm import tqdm
except ModuleNotFoundError:
pass
parser = argparse.ArgumentParser(
description="Extracts energies (L2-norm of STFT frame amplitudes) and pitches (F0) from LJSpeech data."
)
parser.add_argument("--ljspeech_dir", required=True, default=None, type=str)
args = parser.parse_args()
def main():
wavfile_dir = Path(args.ljspeech_dir) / "wavs"
wavfile_list = list(wavfile_dir.glob('*.wav'))
target_dir = Path(args.ljspeech_dir)
# Create target dir <LJSpeech_base_dir>/energies and <LJSpeech_base_dir>/pitches if necessary
if not Path(target_dir / "energies").exists():
print(f"Creating target directory: {target_dir/'energies'}")
Path(target_dir / "energies").mkdir()
if not Path(target_dir / "pitches").exists():
print(f"Creating target directory: {target_dir/'pitches'}")
Path(target_dir / "pitches").mkdir()
if tqdm is not None:
wavfile_list = tqdm(wavfile_list)
for count, file_ in enumerate(wavfile_list):
basename = Path(file_).stem
pitch_path = target_dir / "pitches" / f"{basename}.npy"
energy_path = target_dir / "energies" / f"{basename}.npy"
if pitch_path.exists() and energy_path.exists():
continue
audio, sr = librosa.load(file_, sr=22050)
# Calculate f0
# Please note that fmin and fmax are good approximates for the speaker in LJSpeech and may not generalize to
# other speakers
f0, _, _ = librosa.pyin(audio, fmin=80, fmax=800, frame_length=1024, sr=sr, fill_na=0.0)
# Save to new file
np.save(pitch_path, f0)
# Calculate energy
stft_amplitude = np.abs(librosa.stft(audio, n_fft=1024, hop_length=256, win_length=1024))
energy = np.linalg.norm(stft_amplitude, axis=0) # axis=0 since librosa.stft -> (freq bins, frames)
# Save to new file
np.save(energy_path, energy)
assert energy.shape == f0.shape
if tqdm is None and count % 1000 == 0:
print(f"Finished processing {count} wav files...")
print(f"Finished energy extraction for a total of {len(wavfile_list)} wav files.")
if __name__ == '__main__':
main()
tts/NeMo/scripts/dataset_processing/ljspeech/extract_ljspeech_phonemes_and_durs.sh
Executable
+150
View File
@@ -0,0 +1,150 @@
#!/bin/bash
# Extracts the phonemes and alignments for the LJSpeech dataset via the
# Montreal Forced Aligner (MFA) library, and computes durations per phoneme.
# Assumes you have downloaded and expanded the LJSpeech dataset, and that they
# are located at the directory specified.
#
# This script will create:
# - <LJSPEECH_BASE>/mappings.json: Contains word->phone and phone->idx mappings
# - <LJSPEECH_BASE>/wavs_to_ignore.pkl: A pickled list of wavs to ignore
# because of OOV words in their transcripts
# - <LJSPEECH_BASE>/phoneme_durations/LJ*.npz: Numpy files for each wavfile
# with fields 'tokens' (token indices) and 'durs' (durations per token)
#
# It will also create intermediate files:
# - <LJSPEECH_BASE>/alignments/wavs/LJ*.TextGrid: MFA output
#
# Note: This script will create a conda environment to set up MFA, so please
# make sure that conda is active before running this script.
#
# Note: You may need to install OpenBlas yourself if the MFA script can't find
# it. (e.g. sudo apt-get install libopenblas-dev)
#
# Example Usage:
# ./extract_ljspeech_phonemes_and_durs.sh /data/LJSpeech-1.1
ENV_NAME='aligner'
SAMPLE_RATE=22050
WINDOW_STRIDE=256
CMUDICT_URL='https://raw.githubusercontent.com/cmusphinx/cmudict/master/cmudict.dict'
SPLITS_BASE_URL='https://raw.githubusercontent.com/NVIDIA/tacotron2/master/filelists/'
# Usage info
show_help() {
cat << EOF
Usage: $(basename "$0") [-h] \
[--skip_env_setup] \
[--g2p_dict=<G2P_DICT_PATH>] \
<LJSPEECH_BASE>
Extracts phonemes and their respective durations for the LJSpeech dataset using
the Montreal Forced Aligner (MFA).
-h Help message
--skip_env_setup (Optional) Skips setting up the MFA conda environment
"aligner". Use only if you already have this set up.
--g2p_dict (Optional) Path to the grapheme to phoneme dictionary
text file, if already generated. If set, skips the G2P
step.
EOF
}
SKIP_ENV_SETUP=false
G2P_DICT=''
while :; do
case $1 in
-h|-\?|--help)
show_help
exit
;;
--g2p_dict=?*)
G2P_DICT=${1#*=}
;;
--skip_env_setup)
SKIP_ENV_SETUP=true
;;
*)
break
esac
shift
done
if [[ -z $1 ]]; then
echo "Must specify a LJSpeech base directory."
exit 1
fi
LJSPEECH_BASE=$1
# Check for conda
read -r -d '' CONDA_MESSAGE << EOM
This script requires either Anaconda or Miniconda to be active, as it installs the Montreal Forced Aligner via a conda environment.
See their documentation (https://montreal-forced-aligner.readthedocs.io/en/latest/installation.html) for more details.
EOM
if [ -z "$(which conda)" ]; then
echo $CONDA_MESSAGE
exit
fi
CONDA_PREFIX=$(conda info --base)
source $CONDA_PREFIX/etc/profile.d/conda.sh
# Set up env for MFA (env name "aligner") and install
if $SKIP_ENV_SETUP; then
echo "Skipping environment setup. Assuming env name "aligner" exists."
else
echo "Setting up conda environment for MFA (env name \"aligner\")..."
conda create -n $ENV_NAME -c conda-forge openblas python=3.8 openfst pynini ngram baumwelch
conda activate $ENV_NAME
pip install montreal-forced-aligner tgt torch
mfa thirdparty download
echo "Conda environment \"$ENV_NAME\" set up."
fi
if ! conda activate $ENV_NAME; then
echo "Could not activate environment, see Conda output above."
exit
fi
# Download CMU word-to-phoneme dict and clean out comments so they're not mistaken for tokens
if [ -z $G2P_DICT ]; then
if [ ! -f $LJSPEECH_BASE/cmudict.dict ]; then
echo "Downloading CMU dict."
wget -P $LJSPEECH_BASE $CMUDICT_URL
fi
if [ ! -f $LJSPEECH_BASE/uncommented_cmudict.dict ]; then
echo "Creating uncommented version of CMUdict."
sed 's/\ \#.*//' $LJSPEECH_BASE/cmudict.dict > $LJSPEECH_BASE/uncommented_cmudict.dict
fi
G2P_DICT=$LJSPEECH_BASE/uncommented_cmudict.dict
fi
# Run alignment
echo "Starting MFA with dictionary at: $G2P_DICT"
read -r -d '' MFA_ERROR_MSG << EOM
Could not run MFA. If it could not find OpenBlas, you may need to install it manually.
(e.g. sudo apt-get install libopenblas-dev)
EOM
if ! mfa download acoustic english; then
echo $MFA_ERROR_MSG
fi
mfa align --clean $LJSPEECH_BASE $G2P_DICT english $LJSPEECH_BASE/alignments
# Create JSON mappings from word to phonemes and phonemes to indices
echo "Creating word->phone and phone->idx mappings at $LJSPEECH_BASE/mappings.json..."
python create_token2idx_dict.py \
--dictionary=$G2P_DICT \
--dict_out=$LJSPEECH_BASE/mappings.json
# Calculate phoneme durations
echo "Calculating phoneme durations..."
python calculate_durs.py \
--ljspeech_dir=$LJSPEECH_BASE \
--mappings=$LJSPEECH_BASE/mappings.json \
--sr=$SAMPLE_RATE \
--hop_length=$WINDOW_STRIDE
echo "Phoneme durations and tokens written."
tts/NeMo/scripts/dataset_processing/process_an4_data.py
+88
View File
@@ -0,0 +1,88 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import glob
import json
import logging
import os
import subprocess
import tarfile
import librosa
import wget
parser = argparse.ArgumentParser(description="AN4 dataset download and processing")
parser.add_argument("--data_root", required=True, default=None, type=str)
args = parser.parse_args()
def build_manifest(data_root, transcripts_path, manifest_path, wav_path):
with open(transcripts_path, 'r') as fin:
with open(manifest_path, 'w') as fout:
for line in fin:
# Lines look like this:
# <s> transcript </s> (fileID)
transcript = line[: line.find('(') - 1].lower()
transcript = transcript.replace('<s>', '').replace('</s>', '')
transcript = transcript.strip()
file_id = line[line.find('(') + 1 : -2] # e.g. "cen4-fash-b"
audio_path = os.path.join(
data_root, wav_path, file_id[file_id.find('-') + 1 : file_id.rfind('-')], file_id + '.wav',
)
duration = librosa.core.get_duration(filename=audio_path)
# Write the metadata to the manifest
metadata = {
"audio_filepath": audio_path,
"duration": duration,
"text": transcript,
}
json.dump(metadata, fout)
fout.write('\n')
def main():
data_root = os.path.abspath(args.data_root)
# Convert from .sph to .wav
logging.info("Converting audio files to .wav...")
sph_list = glob.glob(os.path.join(data_root, 'an4/**/*.sph'), recursive=True)
for sph_path in sph_list:
wav_path = sph_path[:-4] + '.wav'
cmd = ['sox', sph_path, wav_path]
subprocess.run(cmd)
logging.info("Finished conversion.")
# Build manifests
logging.info("Building training manifest...")
train_transcripts = os.path.join(data_root, 'an4/etc/an4_train.transcription')
train_manifest = os.path.join(data_root, 'an4/train_manifest.json')
train_wavs = os.path.join(data_root, 'an4/wav/an4_clstk')
build_manifest(data_root, train_transcripts, train_manifest, train_wavs)
logging.info("Training manifests created.")
logging.info("Building test manifest...")
test_transcripts = os.path.join(data_root, 'an4/etc/an4_test.transcription')
test_manifest = os.path.join(data_root, 'an4/test_manifest.json')
test_wavs = os.path.join(data_root, 'an4/wav/an4test_clstk')
build_manifest(data_root, test_transcripts, test_manifest, test_wavs)
logging.info("Test manifest created.")
logging.info("Done with AN4 processing!")
if __name__ == '__main__':
main()
tts/NeMo/scripts/dataset_processing/process_degli_data.py
+70
View File
@@ -0,0 +1,70 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import sys
from nemo.collections.tts.data.datalayers import preprocess_linear_specs_dataset
sys.path.insert(0, '../')
def main():
parser = argparse.ArgumentParser(
description='Create dataset fitted for training and validating deep griffin iteration from wavefiles'
)
parser.add_argument(
"-v",
"--valid_filelist",
help="Filelist for validation set, with all validation audio files listed",
required=True,
default=None,
type=str,
)
parser.add_argument(
"-t",
"--train_filelist",
help="Filelist for train set, with all train audio files listed",
required=True,
default=None,
type=str,
)
parser.add_argument(
"-n",
"--n_fft",
help="Value for the n_fft parameter, and the filter length for the STFT",
default=512,
type=int,
)
parser.add_argument("--hop_length", help="STFT parameter", default=256, type=int)
parser.add_argument(
"-d", "--destination", help="Destination to save the preprocessed data set to", default="/tmp", type=str
)
parser.add_argument(
"-s",
"--num_snr",
help="Number of distinctive noisy samples to generate for each clear sample at the file list",
default=1,
type=int,
)
args = parser.parse_args()
preprocess_linear_specs_dataset(
args.valid_filelist, args.train_filelist, args.n_fft, args.hop_length, args.num_snr, args.destination
)
if __name__ == "__main__":
main()
tts/NeMo/scripts/dataset_processing/process_hub5_data.py
+256
View File
@@ -0,0 +1,256 @@
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# This script is heavily derived from the Patter HUB5 processing script written
# by Ryan Leary
import argparse
import glob
import json
import os
import re
import subprocess
import sys
from collections import namedtuple
from math import ceil, floor
from operator import attrgetter
import numpy as np
import scipy.io.wavfile as wavfile
from tqdm import tqdm
parser = argparse.ArgumentParser(description="Prepare HUB5 data for training/eval")
parser.add_argument(
"--data_root", default=None, type=str, required=True, help="The path to the root LDC HUB5 dataset directory.",
)
parser.add_argument(
"--dest_root",
default=None,
type=str,
required=True,
help="Path to the destination root directory for processed files.",
)
# Optional arguments
parser.add_argument(
"--min_slice_duration", default=10.0, type=float, help="Minimum audio slice duration after processing.",
)
args = parser.parse_args()
StmUtterance = namedtuple(
'StmUtterance', ['filename', 'channel', 'speaker_id', 'begin', 'end', 'label', 'transcript',],
)
STM_LINE_FMT = re.compile(r"^(\w+)\s+(\w+)\s+(\w+)\s+([0-9.]+)\s+([0-9.]+)\s+(<.*>)?\s+(.+)$")
# Transcription errors and their fixes
TRANSCRIPT_BUGS = {"en_4622-B-12079-12187": "KIND OF WEIRD BUT"}
def get_utt_id(segment):
"""
Gives utterance IDs in a form like: en_4156-a-36558-37113
"""
return "{}-{}-{}-{}".format(segment.filename, segment.channel, int(segment.begin * 100), int(segment.end * 100),)
def convert_utterances(sph_path, wav_path):
"""
Converts a sphere audio file to wav.
"""
cmd = ["sph2pipe", "-f", "wav", "-p", sph_path, wav_path]
subprocess.run(cmd)
def create_wavs(data_root, dest_root):
"""
Converts the English sph files to wav using sph2pipe.
"""
sph_root = os.path.join(data_root, "hub5e_00", "english")
sph_list = glob.glob(os.path.join(sph_root, "*.sph"))
# Iterate over each sphere file and conver to wav
for sph_path in tqdm(sph_list, desc="Converting to wav", unit="file"):
sph_name, _ = os.path.splitext(os.path.basename(sph_path))
wav_path = os.path.join(dest_root, 'full_audio_wav', sph_name + ".wav")
cmd = ["sph2pipe", "-f", "wav", "-p", sph_path, wav_path]
subprocess.run(cmd)
def process_transcripts(dataset_root):
"""
Reads in transcripts for each audio segment and processes them.
"""
stm_path = os.path.join(dataset_root, "2000_hub5_eng_eval_tr", "reference", "hub5e00.english.000405.stm",)
results = []
chars = set()
with open(stm_path, "r") as fh:
for line in fh:
# lines with ';;' are comments
if line.startswith(";;"):
continue
if "IGNORE_TIME_SEGMENT_" in line:
continue
line = line.replace("<B_ASIDE>", "").replace("<E_ASIDE>", "")
line = line.replace("(%HESITATION)", "UH")
line = line.replace("-", "")
line = line.replace("(%UH)", "UH")
line = line.replace("(%AH)", "UH")
line = line.replace("(", "").replace(")", "")
line = line.lower()
m = STM_LINE_FMT.search(line.strip())
utt = StmUtterance(*m.groups())
# Convert begin/end times to float
utt = utt._replace(begin=float(utt.begin))
utt = utt._replace(end=float(utt.end))
# Check for utterance in dict of transcript mistakes
transcript_update = TRANSCRIPT_BUGS.get(get_utt_id(utt))
if transcript_update is not None:
utt = utt._replace(transcript=transcript_update)
results.append(utt)
chars.update(list(utt.transcript))
return results, chars
def write_one_segment(dest_root, speaker_id, count, audio, sr, duration, transcript):
"""
Writes out one segment of audio, and writes its corresponding transcript
in the manifest.
Args:
dest_root: the path to the output directory root
speaker_id: ID of the speaker, used in file naming
count: number of segments from this speaker so far
audio: the segment's audio data
sr: sample rate of the audio
duration: duration of the audio
transcript: the corresponding transcript
"""
audio_path = os.path.join(dest_root, "audio", f"{speaker_id}_{count:03}.wav")
manifest_path = os.path.join(dest_root, "manifest_hub5.json")
# Write audio
wavfile.write(audio_path, sr, audio)
# Write transcript
transcript = {
"audio_filepath": audio_path,
"duration": duration,
"text": transcript,
}
with open(manifest_path, 'a') as f:
json.dump(transcript, f)
f.write('\n')
def segment_audio(info_list, dest_root, min_slice_duration):
"""
Combines audio into >= min_slice_duration segments of the same speaker,
and writes the combined transcripts into a manifest.
Args:
info_list: list of StmUtterance objects with transcript information.
dest_root: path to output destination
min_slice_duration: min number of seconds per output audio slice
"""
info_list = sorted(info_list, key=attrgetter('speaker_id', 'begin'))
prev_id = None # For checking audio concatenation
id_count = 0
sample_rate, audio_data = None, None
transcript_buffer = ''
audio_buffer = []
buffer_duration = 0.0
# Iterate through utterances to build segments
for info in info_list:
if info.speaker_id != prev_id:
# Scrap the remainder in the buffers and start next segment
prev_id = info.speaker_id
id_count = 0
sample_rate, audio_data = wavfile.read(os.path.join(dest_root, 'full_audio_wav', info.filename + '.wav'))
transcript_buffer = ''
audio_buffer = []
buffer_duration = 0.0
# Append utterance info to buffers
transcript_buffer += info.transcript
channel = 0 if info.channel.lower() == 'a' else 1
audio_buffer.append(
audio_data[floor(info.begin * sample_rate) : ceil(info.end * sample_rate), channel,]
)
buffer_duration += info.end - info.begin
if buffer_duration < min_slice_duration:
transcript_buffer += ' '
else:
# Write out segment and transcript
id_count += 1
write_one_segment(
dest_root,
info.speaker_id,
id_count,
np.concatenate(audio_buffer, axis=0),
sample_rate,
buffer_duration,
transcript_buffer,
)
transcript_buffer = ''
audio_buffer = []
buffer_duration = 0.0
def main():
data_root = args.data_root
dest_root = args.dest_root
min_slice_duration = args.min_slice_duration
if not os.path.exists(os.path.join(dest_root, 'full_audio_wav')):
os.makedirs(os.path.join(dest_root, 'full_audio_wav'))
if not os.path.exists(os.path.join(dest_root, 'audio')):
os.makedirs(os.path.join(dest_root, 'audio'))
# Create/wipe manifest contents
open(os.path.join(dest_root, "manifest_hub5.json"), 'w').close()
# Convert full audio files from .sph to .wav
create_wavs(data_root, dest_root)
# Get each audio transcript from transcript file
info_list, chars = process_transcripts(data_root)
print("Writing out vocab file", file=sys.stderr)
with open(os.path.join(dest_root, "vocab.txt"), 'w') as fh:
for x in sorted(list(chars)):
fh.write(x + "\n")
# Segment the audio data
print("Segmenting audio and writing manifest")
segment_audio(info_list, dest_root, min_slice_duration)
if __name__ == '__main__':
main()
tts/NeMo/scripts/dataset_processing/process_speech_commands_data.py
+402
View File
@@ -0,0 +1,402 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Usage:
python process_speech_commands_data.py \
--data_root=<absolute path to where the data should be stored> \
--data_version=<either 1 or 2, indicating version of the dataset> \
--class_split=<either "all" or "sub", indicates whether all 30/35 classes should be used, or the 10+2 split should be used> \
--rebalance \
--log
"""
import argparse
import glob
import json
import logging
import os
import re
import tarfile
import urllib.request
import librosa
import numpy as np
URL_v1 = "http://download.tensorflow.org/data/speech_commands_v0.01.tar.gz"
URL_v2 = "http://download.tensorflow.org/data/speech_commands_v0.02.tar.gz"
def __maybe_download_file(destination: str, source: str):
"""
Downloads source to destination if it doesn't exist.
If exists, skips download
Args:
destination: local filepath
source: url of resource
Returns:
"""
if not os.path.exists(destination):
logging.info(f"{destination} does not exist. Downloading ...")
urllib.request.urlretrieve(source, filename=destination + '.tmp')
os.rename(destination + '.tmp', destination)
logging.info(f"Downloaded {destination}.")
else:
logging.info(f"Destination {destination} exists. Skipping.")
return destination
def __extract_all_files(filepath: str, data_root: str, data_dir: str):
if not os.path.exists(data_dir):
extract_file(filepath, data_dir)
else:
logging.info(f'Skipping extracting. Data already there {data_dir}')
def extract_file(filepath: str, data_dir: str):
try:
tar = tarfile.open(filepath)
tar.extractall(data_dir)
tar.close()
except Exception:
logging.info('Not extracting. Maybe already there?')
def __process_data(
data_folder: str, dst_folder: str, rebalance: bool = False, class_split: str = "all", skip_duration: bool = False
):
"""
To generate manifest
Args:
data_folder: source with wav files and validation / test lists
dst_folder: where manifest files will be stored
rebalance: rebalance the classes to have same number of samples.
class_split: whether to use all classes as distinct labels, or to use
10 classes subset and rest of the classes as noise or background.
skip_duration: Bool whether to skip duration computation. Use this only for
colab notebooks where knowing duration is not necessary for demonstration.
Returns:
"""
if not os.path.exists(dst_folder):
os.makedirs(dst_folder)
# Used for 10 classes + silence + unknown class setup
class_subset = ["yes", "no", "up", "down", "left", "right", "on", "off", "stop", "go"]
pattern = re.compile(r"(.+\/)?(\w+)\/([^_]+)_.+wav")
all_files = glob.glob(os.path.join(data_folder, '*/*wav'))
with open(os.path.join(data_folder, 'validation_list.txt'), 'r') as fin:
validation_files = fin.readlines()
valset = set()
for entry in validation_files:
r = re.match(pattern, entry)
if r:
valset.add(r.group(3))
with open(os.path.join(data_folder, 'testing_list.txt'), 'r') as fin:
testing_files = fin.readlines()
testset = set()
for entry in testing_files:
r = re.match(pattern, entry)
if r:
testset.add(r.group(3))
logging.info("Validation and test set lists extracted")
label_count = {}
label_filepaths = {}
unknown_val_filepaths = []
unknown_test_filepaths = []
train, val, test = [], [], []
for entry in all_files:
r = re.match(pattern, entry)
if r:
label, uid = r.group(2), r.group(3)
if label == '_background_noise_' or label == 'silence':
continue
if class_split == "sub":
if label not in class_subset:
label = "unknown" # replace label
if uid in valset:
unknown_val_filepaths.append((label, entry))
elif uid in testset:
unknown_test_filepaths.append((label, entry))
sample = (label, entry)
if uid not in valset and uid not in testset:
if label in label_count:
label_count[label] += 1
else:
label_count[label] = 1
if label in label_filepaths:
label_filepaths[label] += [sample]
else:
label_filepaths[label] = [sample]
if label == 'unknown':
continue
if uid in valset:
val.append(sample)
elif uid in testset:
test.append(sample)
else:
train.append(sample)
logging.info("Prepared filepaths for dataset")
# Add silence and unknown class label samples
if class_split == "sub":
logging.info("Perforiming 10+2 class subsplit")
silence_path = os.path.join(data_folder, "silence")
if not os.path.exists(silence_path):
os.mkdir(silence_path)
silence_stride = 1000 # 0.0625 second stride
sampling_rate = 16000
folder = os.path.join(data_folder, "_background_noise_")
silence_files = []
rng = np.random.RandomState(0)
for file in os.listdir(folder):
if ".wav" in file:
load_path = os.path.join(folder, file)
y, sr = librosa.load(load_path, sr=sampling_rate)
for i in range(0, len(y) - sampling_rate, silence_stride):
file_path = "silence/{}_{}.wav".format(file[:-4], i)
y_slice = y[i : i + sampling_rate]
magnitude = rng.uniform(0.0, 1.0)
y_slice *= magnitude
out_file_path = os.path.join(data_folder, file_path)
librosa.output.write_wav(out_file_path, y_slice, sr)
silence_files.append(('silence', out_file_path))
rng = np.random.RandomState(0)
rng.shuffle(silence_files)
logging.info(f"Constructed silence set of {len(silence_files)}")
# Create the splits
rng = np.random.RandomState(0)
silence_split = 0.1
unknown_split = 0.1
# train split
num_total_samples = sum([label_count[cls] for cls in class_subset])
num_silence_samples = int(np.ceil(silence_split * num_total_samples))
# initialize sample
label_count['silence'] = 0
for silence_id in range(num_silence_samples):
label_count['silence'] += 1
if 'silence' in label_filepaths:
label_filepaths['silence'] += [silence_files[silence_id]]
else:
label_filepaths['silence'] = [silence_files[silence_id]]
train.extend(label_filepaths['silence'])
# Update train unknown set
unknown_train_samples = label_filepaths['unknown']
rng.shuffle(unknown_train_samples)
unknown_size = int(np.ceil(unknown_split * num_total_samples))
label_count['unknown'] = unknown_size
label_filepaths['unknown'] = unknown_train_samples[:unknown_size]
train.extend(label_filepaths['unknown'])
logging.info("Train set prepared")
# val set silence
num_val_samples = len(val)
num_silence_samples = int(np.ceil(silence_split * num_val_samples))
val_idx = label_count['silence'] + 1
for silence_id in range(num_silence_samples):
val.append(silence_files[val_idx + silence_id])
# Update val unknown set
rng.shuffle(unknown_val_filepaths)
unknown_size = int(np.ceil(unknown_split * num_val_samples))
val.extend(unknown_val_filepaths[:unknown_size])
logging.info("Validation set prepared")
# test set silence
num_test_samples = len(test)
num_silence_samples = int(np.ceil(silence_split * num_test_samples))
test_idx = val_idx + num_silence_samples + 1
for silence_id in range(num_silence_samples):
test.append(silence_files[test_idx + silence_id])
# Update test unknown set
rng.shuffle(unknown_test_filepaths)
unknown_size = int(np.ceil(unknown_split * num_test_samples))
test.extend(unknown_test_filepaths[:unknown_size])
logging.info("Test set prepared")
max_command = None
max_count = -1
for command, count in label_count.items():
if command == 'unknown':
continue
if count > max_count:
max_count = count
max_command = command
if rebalance:
logging.info(f"Command with maximum number of samples = {max_command} with {max_count} samples")
logging.info(f"Rebalancing dataset by duplicating classes with less than {max_count} samples...")
for command, samples in label_filepaths.items():
filepaths = [sample[1] for sample in samples]
rng = np.random.RandomState(0)
filepaths = np.asarray(filepaths)
num_samples = len(filepaths)
if num_samples < max_count:
difference = max_count - num_samples
duplication_ids = rng.choice(num_samples, difference, replace=True)
filepaths = np.append(filepaths, filepaths[duplication_ids], axis=0)
logging.info(f"Extended class label {command} from {num_samples} samples to {len(filepaths)} samples")
label_filepaths[command] = [(command, filepath) for filepath in filepaths]
del train
train = []
for label, samples in label_filepaths.items():
train.extend(samples)
manifests = [
('train_manifest.json', train),
('validation_manifest.json', val),
('test_manifest.json', test),
]
for manifest_filename, dataset in manifests:
with open(os.path.join(dst_folder, manifest_filename), 'w') as fout:
num_files = len(dataset)
pct_file = num_files // 100
file_count = 0
logging.info(f"Preparing manifest : {manifest_filename} with #{num_files} files")
for label, audio_path in dataset:
if not skip_duration:
duration = librosa.core.get_duration(filename=audio_path)
else:
duration = 0.0
# Write the metadata to the manifest
metadata = {
"audio_filepath": audio_path,
"duration": duration,
"command": label,
}
json.dump(metadata, fout)
fout.write('\n')
fout.flush()
file_count += 1
if file_count % pct_file == 0:
if not skip_duration:
logging.info(f"Finished serializing {file_count} / {num_files} into {manifest_filename}")
logging.info(f"Finished construction of manifest : {manifest_filename}")
if skip_duration:
logging.info(
f"\n<<NOTE>> Duration computation was skipped for demonstration purposes on Colaboratory.\n"
f"In order to replicate paper results and properly perform data augmentation, \n"
f"please recompute the manifest file without the `--skip_duration` flag !\n"
)
def main():
parser = argparse.ArgumentParser(description='Google Speech Command Data download')
parser.add_argument("--data_root", required=True, default=None, type=str)
parser.add_argument('--data_version', required=True, default=1, type=int, choices=[1, 2])
parser.add_argument('--class_split', required=False, default='all', type=str, choices=['all', 'sub'])
parser.add_argument('--rebalance', required=False, action='store_true')
parser.add_argument('--skip_duration', required=False, action='store_true')
parser.add_argument('--log', required=False, action='store_true')
parser.set_defaults(log=False, rebalance=False, skip_duration=False)
args = parser.parse_args()
if args.log:
logging.basicConfig(level=logging.DEBUG)
data_root = args.data_root
data_set = "google_speech_recognition_v{0}".format(args.data_version)
data_folder = os.path.join(data_root, data_set)
logging.info(f"Working on: {data_set}")
if args.data_version == 1:
URL = URL_v1
else:
URL = URL_v2
# Download and extract
if not os.path.exists(data_folder):
file_path = os.path.join(data_root, data_set + ".tar.bz2")
logging.info(f"Getting {data_set}")
__maybe_download_file(file_path, URL)
logging.info(f"Extracting {data_set}")
__extract_all_files(file_path, data_root, data_folder)
logging.info(f"Processing {data_set}")
__process_data(
data_folder,
data_folder,
rebalance=args.rebalance,
class_split=args.class_split,
skip_duration=args.skip_duration,
)
logging.info('Done!')
if __name__ == "__main__":
main()
tts/NeMo/scripts/dataset_processing/process_vad_data.py
+486
View File
@@ -0,0 +1,486 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Usage:
python process_vad_data.py \
--out_dir=<output path to where the generated manifest should be stored> \
--speech_data_root=<path where the speech data are stored> \
--background_data_root=<path where the background data are stored> \
--rebalance_method=<'under' or 'over' or 'fixed'> \
--log
(Optional --demo (for demonstration in tutorial). If you want to use your own background noise data, make sure to delete --demo)
"""
import argparse
import glob
import json
import logging
import os
import tarfile
import urllib.request
import librosa
import numpy as np
import soundfile as sf
from sklearn.model_selection import train_test_split
sr = 16000
# google speech command v2
URL = "http://download.tensorflow.org/data/speech_commands_v0.02.tar.gz"
def __maybe_download_file(destination: str, source: str):
"""
Downloads source to destination if it doesn't exist.
If exists, skips download
Args:
destination: local filepath
source: url of resource
Returns:
"""
if not os.path.exists(destination):
logging.info(f"{destination} does not exist. Downloading ...")
urllib.request.urlretrieve(source, filename=destination + '.tmp')
os.rename(destination + '.tmp', destination)
logging.info(f"Downloaded {destination}.")
else:
logging.info(f"Destination {destination} exists. Skipping.")
return destination
def extract_file(filepath: str, data_dir: str):
try:
tar = tarfile.open(filepath)
tar.extractall(data_dir)
tar.close()
except Exception:
logging.info('Not extracting. Maybe already there?')
def __extract_all_files(filepath: str, data_root: str, data_dir: str):
if not os.path.exists(data_dir):
extract_file(filepath, data_dir)
else:
logging.info(f'Skipping extracting. Data already there {data_dir}')
def split_train_val_test(data_dir, file_type, test_size=0.1, val_size=0.1, demo=False):
X = []
if file_type == "speech":
for o in os.listdir(data_dir):
if os.path.isdir(os.path.join(data_dir, o)) and o.split("/")[-1] != "_background_noise_":
X.extend(glob.glob(os.path.join(data_dir, o) + '/*.wav'))
if demo:
logging.info(
f"For Demonstration, we use {int(len(X)/100)}/{len(X)} speech data. Make sure to remove --demo flag when you actually train your model!"
)
X = np.random.choice(X, int(len(X) / 100), replace=False)
else:
for o in os.listdir(data_dir):
if os.path.isdir(os.path.join(data_dir, o)):
X.extend(glob.glob(os.path.join(data_dir, o) + '/*.wav'))
else: # for using "_background_noise_" from google speech commands as background data
if o.endswith(".wav"):
X.append(os.path.join(data_dir, o))
X_train, X_test = train_test_split(X, test_size=test_size, random_state=1)
val_size_tmp = val_size / (1 - test_size)
X_train, X_val = train_test_split(X_train, test_size=val_size_tmp, random_state=1)
with open(os.path.join(data_dir, file_type + "_training_list.txt"), "w") as outfile:
outfile.write("\n".join(X_train))
with open(os.path.join(data_dir, file_type + "_testing_list.txt"), "w") as outfile:
outfile.write("\n".join(X_test))
with open(os.path.join(data_dir, file_type + "_validation_list.txt"), "w") as outfile:
outfile.write("\n".join(X_val))
logging.info(f'Overall: {len(X)}, Train: {len(X_train)}, Validatoin: {len(X_val)}, Test: {len(X_test)}')
logging.info(f"Finish spliting train, val and test for {file_type}. Write to files!")
def process_google_speech_train(data_dir):
X = []
for o in os.listdir(data_dir):
if os.path.isdir(os.path.join(data_dir, o)) and o.split("/")[-1] != "_background_noise_":
X.extend(glob.glob(os.path.join(data_dir, o) + '/*.wav'))
short_files = [i.split(data_dir)[1] for i in files]
with open(os.path.join(data_dir, 'testing_list.txt'), 'r') as allfile:
testing_list = allfile.read().splitlines()
with open(os.path.join(data_dir, 'validation_list.txt'), 'r') as allfile:
validation_list = allfile.read().splitlines()
exist_set = set(testing_list).copy()
exist_set.update(set(validation_list))
training_list = [i for i in short_files if i not in exist_set]
with open(os.path.join(data_dir, "training_list.txt"), "w") as outfile:
outfile.write("\n".join(training_list))
logging.info(
f'Overall: {len(files)}, Train: {len(training_list)}, Validatoin: {len(validation_list)}, Test: {len(testing_list)}'
)
def write_manifest(
out_dir,
files,
prefix,
manifest_name,
start=0.0,
end=None,
duration_stride=1.0,
duration_max=None,
duration_limit=100.0,
filter_long=False,
):
"""
Given a list of files, segment each file and write them to manifest with restrictions.
Args:
out_dir: directory of generated manifest
files: list of files to be processed
prefix: label of samples
manifest_name: name of generated manifest
start: beginning of audio of generating segment
end: end of audio of generating segment
duration_stride: stride for segmenting audio samples
duration_max: duration for each segment
duration_limit: duration threshold for filtering out long audio samples
filter_long: boolean to determine whether to filter out long audio samples
Returns:
"""
seg_num = 0
skip_num = 0
if duration_max is None:
duration_max = 1e9
if not os.path.exists(out_dir):
logging.info(f'Outdir {out_dir} does not exist. Creat directory.')
os.mkdir(out_dir)
output_path = os.path.join(out_dir, manifest_name + '.json')
with open(output_path, 'w') as fout:
for file in files:
label = prefix
try:
x, _sr = librosa.load(file, sr=sr)
duration = librosa.get_duration(x, sr=sr)
except Exception:
continue
if filter_long and duration > duration_limit:
skip_num += 1
continue
offsets = []
durations = []
if duration > duration_max:
current_offset = start
while current_offset < duration:
if end is not None and current_offset > end:
break
difference = duration - current_offset
if difference < duration_max:
break
offsets.append(current_offset)
durations.append(duration_max)
current_offset += duration_stride
else:
# Duration is not long enough! Skip
skip_num += 1
for duration, offset in zip(durations, offsets):
metadata = {
'audio_filepath': file,
'duration': duration,
'label': label,
'text': '_', # for compatibility with ASRAudioText
'offset': offset,
}
json.dump(metadata, fout)
fout.write('\n')
fout.flush()
seg_num += 1
return skip_num, seg_num, output_path
def load_list_write_manifest(
data_dir,
out_dir,
filename,
prefix,
start,
end,
duration_stride=1.0,
duration_max=1.0,
duration_limit=100.0,
filter_long=True,
):
filename = prefix + '_' + filename
file_path = os.path.join(data_dir, filename)
with open(file_path, 'r') as allfile:
files = allfile.read().splitlines()
manifest_name = filename.split('_list.txt')[0] + '_manifest'
skip_num, seg_num, output_path = write_manifest(
out_dir,
files,
prefix,
manifest_name,
start,
end,
duration_stride,
duration_max,
duration_limit,
filter_long=True,
)
return skip_num, seg_num, output_path
def rebalance_json(data_dir, data_json, num, prefix):
data = []
seg = 0
for line in open(data_json, 'r'):
data.append(json.loads(line))
filename = data_json.split('/')[-1]
fout_path = os.path.join(data_dir, prefix + "_" + filename)
if len(data) >= num:
selected_sample = np.random.choice(data, num, replace=False)
else:
selected_sample = np.random.choice(data, num, replace=True)
with open(fout_path, 'a') as fout:
for i in selected_sample:
seg += 1
json.dump(i, fout)
fout.write('\n')
fout.flush()
logging.info(f'Get {seg}/{num} to {fout_path} from {data_json}')
return fout_path
def generate_variety_noise(data_dir, filename, prefix):
curr_dir = data_dir.split("_background_noise_")[0]
silence_path = os.path.join(curr_dir, "_background_noise_more")
if not os.path.exists(silence_path):
os.mkdir(silence_path)
silence_stride = 1000 # stride = 1/16 seconds
sampling_rate = 16000
silence_files = []
rng = np.random.RandomState(0)
filename = prefix + '_' + filename
file_path = os.path.join(data_dir, filename)
with open(file_path, 'r') as allfile:
files = allfile.read().splitlines()
for file in files:
y, sr = librosa.load(file, sr=sampling_rate)
for i in range(
0, len(y) - sampling_rate, silence_stride * 100
): # stride * 100 to generate less samples for demo
file_name = "{}_{}.wav".format(file.split("/")[-1], i)
y_slice = y[i : i + sampling_rate]
magnitude = rng.uniform(0.0, 1.0)
y_slice *= magnitude
out_file_path = os.path.join(silence_path, file_name)
sf.write(out_file_path, y_slice, sr)
silence_files.append(out_file_path)
new_list_file = os.path.join(silence_path, filename)
with open(new_list_file, "w") as outfile:
outfile.write("\n".join(silence_files))
logging.info(f"Generate {len(out_file_path)} background files for {file_path}. => {new_list_file} !")
return len(silence_files)
def main():
parser = argparse.ArgumentParser(description='Speech and backgound data download and preprocess')
parser.add_argument("--out_dir", required=False, default='./manifest/', type=str)
parser.add_argument("--speech_data_root", required=True, default=None, type=str)
parser.add_argument("--background_data_root", required=True, default=None, type=str)
parser.add_argument('--test_size', required=False, default=0.1, type=float)
parser.add_argument('--val_size', required=False, default=0.1, type=float)
parser.add_argument('--seg_len', required=False, default=0.63, type=float)
parser.add_argument('--log', required=False, action='store_true')
parser.add_argument('--rebalance_method', required=False, default=None, type=str)
parser.add_argument('--demo', required=False, action='store_true')
parser.set_defaults(log=False, generate=False)
args = parser.parse_args()
if not args.rebalance_method:
rebalance = False
else:
if args.rebalance_method != 'over' and args.rebalance_method != 'under' and args.rebalance_method != 'fixed':
raise NameError("Please select a valid sampling method: over/under/fixed.")
else:
rebalance = True
if args.log:
logging.basicConfig(level=logging.DEBUG)
# Download speech data
speech_data_root = args.speech_data_root
data_set = "google_speech_recognition_v2"
speech_data_folder = os.path.join(speech_data_root, data_set)
background_data_folder = args.background_data_root
logging.info(f"Working on: {data_set}")
# Download and extract speech data
if not os.path.exists(speech_data_folder):
file_path = os.path.join(speech_data_root, data_set + ".tar.bz2")
logging.info(f"Getting {data_set}")
__maybe_download_file(file_path, URL)
logging.info(f"Extracting {data_set}")
__extract_all_files(file_path, speech_data_root, speech_data_folder)
logging.info(f"Split speech data!")
# dataset provide testing.txt and validation.txt feel free to split data using that with process_google_speech_train
split_train_val_test(speech_data_folder, "speech", args.test_size, args.val_size, args.demo)
logging.info(f"Split background data!")
split_train_val_test(background_data_folder, "background", args.test_size, args.val_size)
out_dir = args.out_dir
# Process Speech manifest
logging.info(f"=== Write speech data to manifest!")
skip_num_val, speech_seg_num_val, speech_val = load_list_write_manifest(
speech_data_folder, out_dir, 'validation_list.txt', 'speech', 0.2, 0.8, args.seg_len, args.seg_len
)
skip_num_test, speech_seg_num_test, speech_test = load_list_write_manifest(
speech_data_folder, out_dir, 'testing_list.txt', 'speech', 0.2, 0.8, 0.01, args.seg_len
)
skip_num_train, speech_seg_num_train, speech_train = load_list_write_manifest(
speech_data_folder, out_dir, 'training_list.txt', 'speech', 0.2, 0.8, args.seg_len, args.seg_len
)
logging.info(f'Val: Skip {skip_num_val} samples. Get {speech_seg_num_val} segments! => {speech_val} ')
logging.info(f'Test: Skip {skip_num_test} samples. Get {speech_seg_num_test} segments! => {speech_test}')
logging.info(f'Train: Skip {skip_num_train} samples. Get {speech_seg_num_train} segments!=> {speech_train}')
# Process background manifest
# if we select to generate more background noise data
if args.demo:
logging.info("Start generating more background noise data")
generate_variety_noise(background_data_folder, 'validation_list.txt', 'background')
generate_variety_noise(background_data_folder, 'training_list.txt', 'background')
generate_variety_noise(background_data_folder, 'testing_list.txt', 'background')
background_data_folder = os.path.join(
background_data_folder.split("_background_noise_")[0], "_background_noise_more"
)
logging.info(f"=== Write background data to manifest!")
skip_num_val, background_seg_num_val, background_val = load_list_write_manifest(
background_data_folder, out_dir, 'validation_list.txt', 'background', 0, None, 0.15, args.seg_len
)
skip_num_test, background_seg_num_test, background_test = load_list_write_manifest(
background_data_folder, out_dir, 'testing_list.txt', 'background', 0, None, 0.01, args.seg_len
)
skip_num_train, background_seg_num_train, background_train = load_list_write_manifest(
background_data_folder, out_dir, 'training_list.txt', 'background', 0, None, 0.15, args.seg_len
)
logging.info(f'Val: Skip {skip_num_val} samples. Get {background_seg_num_val} segments! => {background_val}')
logging.info(f'Test: Skip {skip_num_test} samples. Get {background_seg_num_test} segments! => {background_test}')
logging.info(
f'Train: Skip {skip_num_train} samples. Get {background_seg_num_train} segments! => {background_train}'
)
min_val, max_val = min(speech_seg_num_val, background_seg_num_val), max(speech_seg_num_val, background_seg_num_val)
min_test, max_test = (
min(speech_seg_num_test, background_seg_num_test),
max(speech_seg_num_test, background_seg_num_test),
)
min_train, max_train = (
min(speech_seg_num_train, background_seg_num_train),
max(speech_seg_num_train, background_seg_num_train),
)
logging.info('Finish generating manifest!')
if rebalance:
# Random Oversampling: Randomly duplicate examples in the minority class.
# Random Undersampling: Randomly delete examples in the majority class.
if args.rebalance_method == 'under':
logging.info(f"Rebalancing number of samples in classes using {args.rebalance_method} sampling.")
logging.info(f'Val: {min_val} Test: {min_test} Train: {min_train}!')
rebalance_json(out_dir, background_val, min_val, 'balanced')
rebalance_json(out_dir, background_test, min_test, 'balanced')
rebalance_json(out_dir, background_train, min_train, 'balanced')
rebalance_json(out_dir, speech_val, min_val, 'balanced')
rebalance_json(out_dir, speech_test, min_test, 'balanced')
rebalance_json(out_dir, speech_train, min_train, 'balanced')
if args.rebalance_method == 'over':
logging.info(f"Rebalancing number of samples in classes using {args.rebalance_method} sampling.")
logging.info(f'Val: {max_val} Test: {max_test} Train: {max_train}!')
rebalance_json(out_dir, background_val, max_val, 'balanced')
rebalance_json(out_dir, background_test, max_test, 'balanced')
rebalance_json(out_dir, background_train, max_train, 'balanced')
rebalance_json(out_dir, speech_val, max_val, 'balanced')
rebalance_json(out_dir, speech_test, max_test, 'balanced')
rebalance_json(out_dir, speech_train, max_train, 'balanced')
if args.rebalance_method == 'fixed':
fixed_test, fixed_val, fixed_train = 200, 100, 500
logging.info(f"Rebalancing number of samples in classes using {args.rebalance_method} sampling.")
logging.info(f'Val: {fixed_val} Test: {fixed_test} Train: {fixed_train}!')
rebalance_json(out_dir, background_val, fixed_val, 'balanced')
rebalance_json(out_dir, background_test, fixed_test, 'balanced')
rebalance_json(out_dir, background_train, fixed_train, 'balanced')
rebalance_json(out_dir, speech_val, fixed_val, 'balanced')
rebalance_json(out_dir, speech_test, fixed_test, 'balanced')
rebalance_json(out_dir, speech_train, fixed_train, 'balanced')
else:
logging.info("Don't rebalance number of samples in classes.")
if __name__ == '__main__':
main()
tts/NeMo/scripts/export/convasr_to_enemo.py
+76
View File
@@ -0,0 +1,76 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import os
import tarfile
from shutil import copyfile
from nemo.collections.asr.models import EncDecClassificationModel, EncDecCTCModel, EncDecSpeakerLabelModel
from nemo.utils import logging
def get_parser():
parser = argparse.ArgumentParser(description="Convert .nemo file to onnx file (encoder and decoder fused)")
parser.add_argument(
"--nemo_file", default=None, type=str, required=True, help="Path to .nemo file",
)
parser.add_argument(
"--onnx_file", default=None, type=str, required=True, help="Path to the onnx output",
)
parser.add_argument(
"--enemo_file", default=None, type=str, required=True, help="Path to the enemo output, consiting of onnx file",
)
parser.add_argument(
"--model_type",
default="asr",
type=str,
choices=["asr", "speech_label", "speaker"],
help="Type of decoder used by the model.",
)
return parser
def main(
nemo_file, enemo_file, onnx_file, model_type="asr",
):
if model_type == "asr":
logging.info("Preparing ASR model")
model = EncDecCTCModel.restore_from(nemo_file)
elif model_type == "speech_label":
logging.info("Preparing Speech Label Classification model")
model = EncDecClassificationModel.restore_from(nemo_file)
elif model_type == "speaker":
logging.info("Preparing Speaker Recognition model")
model = EncDecSpeakerLabelModel.restore_from(nemo_file)
else:
raise NameError("Available model names are asr, speech_label and speaker")
logging.info("Writing onnx file")
model.export(onnx_file, onnx_opset_version=12)
logging.info("succesfully ported onnx file")
with tarfile.open(nemo_file, "r") as archive:
archive.extract("./model_config.yaml")
with tarfile.open(enemo_file, "w") as enemo_archive:
enemo_archive.add("./model_config.yaml")
copyfile(onnx_file, "model_graph.onnx")
enemo_archive.add("model_graph.onnx")
os.remove("model_graph.onnx") # cleanup extra file
if __name__ == "__main__":
args = get_parser().parse_args()
main(
args.nemo_file, args.enemo_file, args.onnx_file, model_type=args.model_type,
)
tts/NeMo/scripts/export/convasr_to_onnx.py
+66
View File
@@ -0,0 +1,66 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
from nemo.collections.asr.models import EncDecClassificationModel, EncDecCTCModel, EncDecSpeakerLabelModel
from nemo.utils import logging
def get_parser():
parser = argparse.ArgumentParser(description="Convert .nemo file to encoder decoder onnx files")
parser.add_argument(
"--nemo_file", default=None, type=str, required=True, help="Path to .nemo file",
)
parser.add_argument(
"--onnx_encoder", default=None, type=str, required=True, help="Path to the onnx encoder output.",
)
parser.add_argument(
"--onnx_decoder", default=None, type=str, required=True, help="Path to the onnx decoder output.",
)
parser.add_argument(
"--model_type",
default='asr',
type=str,
choices=['asr', 'speech_label', 'speaker'],
help="Type of decoder used by the model.",
)
return parser
def main(
nemo_file, onnx_encoder, onnx_decoder, model_type='asr',
):
if model_type == 'asr':
logging.info("Preparing encoder decoder for ASR model")
model = EncDecCTCModel.restore_from(nemo_file)
elif model_type == 'speech_label':
logging.info("Preparing encoder decoder for Speech Label Classification model")
model = EncDecClassificationModel.restore_from(nemo_file)
elif model_type == 'speaker':
logging.info("Preparing encoder decoder for Speaker Recognition model")
model = EncDecSpeakerLabelModel.restore_from(nemo_file)
else:
raise NameError("Available model names are asr, speech_label and speaker ")
logging.info("Writing onnx encoder and decoder onnx files")
model.encoder.export(onnx_encoder)
model.decoder.export(onnx_decoder, onnx_opset_version=10)
logging.info("succesfully ported onnx files")
if __name__ == "__main__":
args = get_parser().parse_args()
main(
args.nemo_file, args.onnx_encoder, args.onnx_decoder, model_type=args.model_type,
)
tts/NeMo/scripts/export/convasr_to_single_onnx.py
+62
View File
@@ -0,0 +1,62 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
from nemo.collections.asr.models import EncDecClassificationModel, EncDecCTCModel, EncDecSpeakerLabelModel
from nemo.utils import logging
def get_parser():
parser = argparse.ArgumentParser(description="Convert .nemo file to onnx file (encoder and decoder fused)")
parser.add_argument(
"--nemo_file", default=None, type=str, required=True, help="Path to .nemo file",
)
parser.add_argument(
"--onnx_file", default=None, type=str, required=True, help="Path to the onnx output",
)
parser.add_argument(
"--model_type",
default='asr',
type=str,
choices=['asr', 'speech_label', 'speaker'],
help="Type of decoder used by the model.",
)
return parser
def main(
nemo_file, onnx_file, model_type='asr',
):
if model_type == 'asr':
logging.info("Preparing ASR model")
model = EncDecCTCModel.restore_from(nemo_file)
elif model_type == 'speech_label':
logging.info("Preparing Speech Label Classification model")
model = EncDecClassificationModel.restore_from(nemo_file)
elif model_type == 'speaker':
logging.info("Preparing Speaker Recognition model")
model = EncDecSpeakerLabelModel.restore_from(nemo_file)
else:
raise NameError("Available model names are asr, speech_label and speaker")
logging.info("Writing onnx file")
model.export(onnx_file, onnx_opset_version=12)
logging.info("succesfully ported onnx file")
if __name__ == "__main__":
args = get_parser().parse_args()
main(
args.nemo_file, args.onnx_file, model_type=args.model_type,
)
tts/NeMo/scripts/freesound_download_resample/download_resample_freesound.sh
+115
View File
@@ -0,0 +1,115 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#!/bin/bash
# This is bash script actually run the downloading and resampling script.
# See instructions in freesound_download.py
# Change this arguments if you want
page_size=100 # Number of sounds per page
max_samples=200 # Maximum number of sound samples
min_filesize=0 # Minimum filesize allowed (in MB)
max_filesize=100 # Maximum filesize allowed (in MB)
if [[ $# -ne 3 ]]; then
echo "Require number of all files | data directory | resample data directory as arguments to the script"
exit 2
fi
NUM_ALL_FILES=$1
DATADIR=$2
RESAMPLE_DATADIR=$3
if [ ! -d "$DATADIR" ]; then
echo "Creating dir $DATADIR"
mkdir -p "$DATADIR"
fi
if [ ! -d "$RESAMPLE_DATADIR" ]; then
echo "Creating dir $RESAMPLE_DATADIR"
mkdir -p "$RESAMPLE_DATADIR"
fi
# we just need background categories for constructing dataset, feel free to include other (speech) categories for testing and training your VAD model
# background
categories=(
"Air brake"
"Static"
"Acoustic environment"
"Distortion"
"Tape hiss"
"Hubbub"
"Vibration"
"Cacophony"
"Throbbing"
"Reverberation"
"Inside, public space"
"Inside, small room"
"Echo"
"Outside, rural"
"Outside, natural"
"Outside, urban"
"Outside, manmade"
"Car"
"Bus"
"Traffic noise"
"Roadway noise"
"Truck"
"Emergency vehicle"
"Motorcycle"
"Aircraft engine"
"Aircraft"
"Helicopter"
"Bicycle"
"Skateboard"
"Subway, metro, underground"
"Railroad car"
"Train wagon"
"Train"
"Sailboat"
"Rowboat"
"Ship"
)
WAV_FILECOUNT="$(find $DATADIR -name '*.wav' -type f | wc -l)"
FLAC_FILECOUNT="$(find $DATADIR -name '*.flac' -type f | wc -l)"
FILECOUNT="$((WAV_FILECOUNT + FLAC_FILECOUNT))"
echo "File count: " $FILECOUNT
while((FILECOUNT <= NUM_ALL_FILES))
do
for category in "${categories[@]}"
do
python freesound_download.py --data_dir "${DATADIR}" --category "${category}" --page_size "${page_size}" --max_samples "${max_samples}" --min_filesize "${min_filesize}" --max_filesize "${max_filesize}"
ret=$?
if [ $ret -ne 0 ]; then
exit 1
fi
done
WAV_FILECOUNT="$(find $DATADIR -name '*.wav' -type f | wc -l)"
FLAC_FILECOUNT="$(find $DATADIR -name '*.flac' -type f | wc -l)"
FILECOUNT="$((WAV_FILECOUNT + FLAC_FILECOUNT))"
echo "Current file count is: " $FILECOUNT
done
# RESAMPLE
echo "Got enough data. Start resample!"
python freesound_resample.py --data_dir="${DATADIR}" --resampled_dir="${RESAMPLE_DATADIR}"
echo "Done resample data!"
tts/NeMo/scripts/freesound_download_resample/freesound_download.py
+585
View File
@@ -0,0 +1,585 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import os
import pickle
import time
try:
import librosa
import requests
import requests_oauthlib
from joblib import Parallel, delayed
from oauthlib.oauth2 import TokenExpiredError
except (ModuleNotFoundError, ImportError) as e:
raise e
try:
import freesound
except ModuleNotFoundError as e:
raise ModuleNotFoundError(
"freesound is not installed. Execute `pip install --no-cache-dir git+https://github.com/MTG/freesound-python.git` in terminal"
)
"""
Instructions
1. We will need some requirements including freesound, requests, requests_oauthlib, joblib, librosa and sox. If they are not installed, please run `pip install -r freesound_requirements.txt`
2. Create an API key for freesound.org at https://freesound.org/help/developers/
3. Create a python file called `freesound_private_apikey.py` and add lined `api_key = <your Freesound api key>` and `client_id = <your Freesound client id>`
4. Authorize by run `python freesound_download.py --authorize` and visit website, and paste response code
5. Feel free to change any arguments in download_resample_freesound.sh such as max_samples and max_filesize
6. Run `bash download_resample_freesound.sh <numbers of files you want> <download data directory> <resampled data directory>`
"""
# Import the API Key
try:
from freesound_private_apikey import api_key, client_id
print("API Key found !")
except ImportError:
raise ImportError(
"Create a python file called `freesound_private_apikey.py` and add lined `api_key = <your Freesound api key>` and `client_id = <your Freesound client id>`"
)
auth_url = 'https://freesound.org/apiv2/oauth2/authorize/'
redirect_url = 'https://freesound.org/home/app_permissions/permission_granted/'
token_url = 'https://freesound.org/apiv2/oauth2/access_token/'
scope = ["read", "write"]
BACKGROUND_CLASSES = [
"Air brake",
"Static",
"Acoustic environment",
"Distortion",
"Tape hiss",
"Hubbub",
"Vibration",
"Cacophony",
"Throbbing",
"Reverberation",
"Inside, public space",
"Inside, small room",
"Echo",
"Outside, rural",
"Outside, natural",
"Outside, urban",
"Outside, manmade",
"Car",
"Bus",
"Traffic noise",
"Roadway noise",
"Truck",
"Emergency vehicle",
"Motorcycle",
"Aircraft engine",
"Aircraft",
"Helicopter",
"Bicycle",
"Skateboard",
"Subway, metro, underground",
"Railroad car",
"Train wagon",
"Train",
"Sailboat",
"Rowboat",
"Ship",
]
SPEECH_CLASSES = [
"Male speech",
"Female speech",
"Speech synthesizer",
"Babbling",
"Conversation",
"Child speech",
"Narration",
"Laughter",
"Yawn",
"Whispering",
"Whimper",
"Baby cry",
"Sigh",
"Groan",
"Humming",
"Male singing",
"Female singing",
"Child singing",
"Children shouting",
]
def initialize_oauth():
# If token already exists, then just load it
if os.path.exists('_token.pkl'):
token = unpickle_object('_token')
oauth = requests_oauthlib.OAuth2Session(client_id, redirect_uri=redirect_url, scope=scope, token=token)
else:
# Construct a new token after OAuth2 flow
# Initialize a OAuth2 session
oauth = requests_oauthlib.OAuth2Session(client_id, redirect_uri=redirect_url, scope=scope)
authorization_url, state = oauth.authorization_url(auth_url)
print(f"Visit below website and paste access token below : \n\n{authorization_url}\n")
authorization_response = input("Paste authorization response code here :\n")
token = oauth.fetch_token(
token_url,
authorization_response=authorization_response,
code=authorization_response,
client_secret=api_key,
)
# Save the token generated
pickle_object(token, '_token')
return oauth, token
def instantiate_session():
# Reconstruct session in process, and force singular execution thread to reduce session
# connections to server
token = unpickle_object('_token')
session = requests_oauthlib.OAuth2Session(client_id, redirect_uri=redirect_url, scope=scope, token=token)
adapter = requests.adapters.HTTPAdapter(pool_connections=1, pool_maxsize=1)
session.mount('http://', adapter)
return session
def refresh_token(session):
print("Refreshing tokens...")
# Token expired, perform token refresh
extras = {'client_id': client_id, 'client_secret': api_key}
token = session.refresh_token(token_url, **extras)
print("Token refresh performed...")
# Save the refreshed token
pickle_object(token, '_token')
return session
def pickle_object(token, name):
with open(name + '.pkl', 'wb') as f:
pickle.dump(token, f)
def unpickle_object(name):
fp = name + '.pkl'
if os.path.exists(fp):
with open(fp, 'rb') as f:
token = pickle.load(f)
return token
else:
raise FileNotFoundError('Token not found!')
def is_resource_limited(e: freesound.FreesoundException):
"""
Test if the reason for a freesound exception was either rate limit
or daily limit.
If it was for either reason, sleep for an appropriate delay and return
to try again.
Args:
e: Freesound Exception object
Returns:
A boolean which describes whether the error was due to some
api limit issue, or if it was some other reason.
If false is returned, then the user should carefully check the cause
and log it.
"""
detail = e.detail['detail']
if '2000' in detail:
# This is the request limit, hold off for 1 hour and try again
print(f"Hit daily limit, sleeping for 20 minutes.")
time.sleep(60 * 20)
return True
elif '60' in detail:
# This is the request limit per minute, hold off for 1 minute and try again
print(f"Hit rate limit, sleeping for 1 minute.")
time.sleep(60)
return True
else:
return False
def prepare_client(client: freesound.FreesoundClient, token) -> freesound.FreesoundClient:
# Initialize the client with token auth
client.set_token(token['access_token'], auth_type='oauth')
print("Client ready !")
return client
def get_text_query_with_resource_limit_checks(client, query: str, filters: list, fields: str, page_size: int):
"""
Performs a text query, checks for rate / api limits, and retries.
Args:
client: FreesoundAPI client
query: query string (either exact or inexact)
filters: list of string filters
fields: String of values to recover
page_size: samples per page returned
Returns:
"""
pages = None
attempts = 20
while pages is None:
try:
pages = client.text_search(query=query, filter=" ".join(filters), fields=fields, page_size=str(page_size),)
except freesound.FreesoundException as e:
# Most probably a rate limit or a request limit
# Check if that was the case, and wait appropriate ammount of time
# for retry
was_resource_limited = is_resource_limited(e)
# If result of test False, it means that failure was due to some other reason.
# Log it, then break loop
if not was_resource_limited:
print(e.with_traceback(None))
break
attempts -= 1
# Attempt to refresh tokens if it fails multiple times
if attempts % 5 == 0 and attempts > 0:
session = instantiate_session()
refresh_token(session)
session.close()
token = unpickle_object('_token')
client = prepare_client(client, token)
if attempts <= 0:
print(f"Failed to query pages for '{query}' after 10 attempts, skipping query")
break
if pages is None:
print(f"Query attempts remaining = {attempts}")
return client, pages
def get_resource_with_auto_refresh(session, download_url):
"""
Attempts download of audio with a token refresh if necessary.
"""
try:
result = session.get(download_url)
except TokenExpiredError as e:
session = refresh_token(session)
result = session.get(download_url)
except Exception as e:
result = None
print(f"Skipping file {download_url} due to exception below\n\n")
print(e)
return result.content
def download_song(basepath, id, name, download_url):
# Cleanup name
name = name.replace(":", "-")
name = name.replace(",", "-")
name = name.replace("/", "-")
name = name.replace("\\", "-")
name = name.replace(".", "-")
name = name.replace(" ", "")
# Correct last `.` for filetype
name = name[:-4] + '.wav'
# Add file id to filename
name = f"id_{id} " + name
fp = os.path.join(basepath, name)
# Check if file, if exists already, can be loaded by librosa
# If it cannot be loaded, possibly corrupted file.
# Delete and then re-download
if os.path.exists(fp):
try:
_ = librosa.load(fp)
except Exception:
# File is currupted, delete and re-download.
os.remove(fp)
print(f"Pre-existing file {fp} was corrupt and was deleted, will be re-downloaded.")
if not os.path.exists(fp):
print("Downloading file :", name)
session = instantiate_session()
data = None
attempts = 10
try:
while data is None:
try:
# Get the sound data
data = get_resource_with_auto_refresh(session, download_url)
except freesound.FreesoundException as e:
# Most probably a rate limit or a request limit
# Check if that was the case, and wait appropriate amount of time
# for retry
was_resource_limited = is_resource_limited(e)
# If result of test False, it means that failure was due to some other reason.
# Log it, then break loop
if not was_resource_limited:
print(e)
break
attempts -= 1
if attempts <= 0:
print(f"Failed to download file {fp} after 10 attempts, skipping file")
break
if data is None:
print(f"Download attempts remaining = {attempts}")
finally:
session.close()
# Write the data to file
if data is not None:
print("Downloaded file :", name)
with open(fp, 'wb') as f:
f.write(data)
# If file size is less than 89, then this probably is a text format and not an actual audio file.
if os.path.getsize(fp) > 89:
print(f"File written : {fp}")
else:
os.remove(fp)
print(f"File corrupted and has been deleted: {fp}")
else:
print(f"File [{fp}] corrupted or faced some issue when downloading, skipped.")
# Sleep to avoid hitting rate limits
time.sleep(5)
else:
print(f"File [{fp}] already exists in dataset, skipping re-download.")
def get_songs_by_category(
client: freesound.FreesoundClient,
category: str,
data_dir: str,
max_num_samples=100,
page_size=100,
min_filesize_in_mb=0,
max_filesize_in_mb=10,
n_jobs=None,
):
"""
Download songs of a category with restrictions
Args:
client: FreesoundAPI client
category: category to be downloaded
data_dir: directory of downloaded songs
max_num_samples: maximum number of samples of this category
page_size: samples per page returned
min_filesize_in_mb: minimum filesize of the song in MB
max_filesize_in_mb: maximum filesize of the song in MB
n_jobs: number of jobs for parallel processing
Returns:
"""
# quote string to force exact match
query = f'"{category}"'
print(f"Query : {query}")
page_size = min(page_size, 150)
max_filesize = int(max_filesize_in_mb * (2 ** 20))
if min_filesize_in_mb == 0:
min_filesize_in_mb = 1
else:
min_filesize_in_mb = int(min_filesize_in_mb * (2 ** 20))
if max_num_samples < 0:
max_num_samples = int(1e6)
filters = [
'type:(wav OR flac)',
'license:("Attribution" OR "Creative Commons 0")',
f'filesize:[{min_filesize_in_mb} TO {max_filesize}]',
]
fields = "id,name,download,license"
client, pages = get_text_query_with_resource_limit_checks(
client, query=query, filters=filters, fields=fields, page_size=page_size
)
if pages is None:
print(f"Number of attempts exceeded limit, skipping query {query}")
return
num_pages = pages.count
# Check if returned empty result; if so, fallback to inexact category search
if num_pages == 0:
print(f"Found 0 samples of results for query '{query}'")
print(f"Trying less restricted query : {category}")
client, pages = get_text_query_with_resource_limit_checks(
client, query=category, filters=filters, fields=fields, page_size=page_size
)
if pages is None:
print(f"Number of attempts exceeded limit, skipping query {query}")
return
num_pages = pages.count
print(f"Found {num_pages} samples of results for query '{query}'")
category = category.replace(' ', '_')
basepath = os.path.join(data_dir, category)
if not os.path.exists(basepath):
os.makedirs(basepath)
sounds = []
sample_count = 0
# Retrieve sound license information
with open(os.path.join(basepath, 'licenses.txt'), 'w') as f:
f.write("ID,LICENSE\n")
f.flush()
while True:
for sound in pages:
if sample_count >= max_num_samples:
print(
f"Collected {sample_count} samples, which is >= max number of samples requested "
f"{max_num_samples}. Stopping for this category : {category}"
)
break
sounds.append(sound)
sample_count += 1
f.write(f"{sound.id},{sound.license}\n")
f.flush()
if sample_count >= max_num_samples:
break
try:
pages = pages.next_page()
except ValueError:
break
if n_jobs is None:
n_jobs = max(1, len(sounds))
# Parallel download all songs
with Parallel(n_jobs=n_jobs, verbose=10) as parallel:
_ = parallel(delayed(download_song)(basepath, sound.id, sound.name, sound.download) for sound in sounds)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description="Freesound download script")
parser.add_argument(
'--authorize', action='store_true', dest='auth', help='Flag to only perform OAuth2 authorization step'
)
parser.add_argument('-c', '--category', default='', type=str, help='Category required to download')
parser.add_argument('-d', '--data_dir', default='', type=str, help='Destination folder to store data')
parser.add_argument('--page_size', default=100, type=int, help='Number of sounds per page')
parser.add_argument('--max_samples', default=100, type=int, help='Maximum number of sound samples')
parser.add_argument('--min_filesize', default=0, type=int, help='Maximum filesize allowed (in MB)')
parser.add_argument('--max_filesize', default=20, type=int, help='Maximum filesize allowed (in MB)')
parser.set_defaults(auth=False)
args = parser.parse_args()
if args.auth:
""" Initialize oauth token to be used by all """
oauth, token = initialize_oauth()
oauth.close()
print("Authentication suceeded ! Token stored in `_token.pkl`")
exit(0)
if not os.path.exists('_token.pkl'):
raise FileNotFoundError(
"Please authorize the application first using " "`python freesound_download.py --authorize`"
)
if args.data_dir == '':
raise ValueError("Data dir must be passed as an argument using `--data_dir`")
data_dir = args.data_dir
page_size = args.page_size
max_num_samples = args.max_samples
min_filesize_in_mb = args.min_filesize
max_filesize_in_mb = args.max_filesize
# Initialize and authenticate client
token = unpickle_object('_token')
freesound_client = freesound.FreesoundClient()
client = prepare_client(freesound_client, token)
category = args.category
if category == '':
raise ValueError("Cannot pass empty string as it will select all of FreeSound data !")
print(f"Downloading category : {category}")
get_songs_by_category(
client,
category,
data_dir=data_dir,
max_num_samples=max_num_samples,
page_size=page_size,
min_filesize_in_mb=min_filesize_in_mb,
max_filesize_in_mb=max_filesize_in_mb,
n_jobs=30,
)
tts/NeMo/scripts/freesound_download_resample/freesound_requirements.txt
+6
View File
@@ -0,0 +1,6 @@
git+git://github.com/MTG/freesound-python.git
requests
requests_oauthlib
joblib
librosa
sox
tts/NeMo/scripts/freesound_download_resample/freesound_resample.py
+119
View File
@@ -0,0 +1,119 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import glob
import os
import time
import librosa
import sox
from joblib import Parallel, delayed
def resample_file(resampled_dir, filepath, ext, sample_rate):
"""
Resample an audio file to 16kHZ and transform to monochannel
Remove incompatible files.
Args:
resampled_dir: Directory of transformed files.
filepath: Filepath of Audio
ext: File type e.g. "wav", "flac"
Returns:
"""
head, filename = os.path.split(filepath)
_, clsname = os.path.split(head)
filename, _ = os.path.splitext(filename)
new_dir = os.path.join(resampled_dir, clsname)
if not os.path.exists(new_dir):
os.makedirs(new_dir)
new_path = os.path.join(new_dir, filename + f'.{ext}')
# check if the resampled data exists.
if os.path.exists(new_path):
print(f"Resampled file {filepath} exists. Skip it.")
return None
transform = sox.Transformer()
transform.set_output_format(file_type='wav')
transform.convert(samplerate=sample_rate, n_channels=1)
try:
transform.build(filepath, new_path)
print(f"Finished converting file {filepath}.")
return None
except sox.core.SoxError as e:
try:
# Check if the file is readable
librosa.load(filepath)
# if it is, force input format and try again
transform.set_input_format(file_type=ext)
transform.build(filepath, new_path)
return None
except Exception:
return filepath
def main():
start = time.time()
parser = argparse.ArgumentParser(description='Freesound data resample')
parser.add_argument("--data_dir", required=True, default=None, type=str)
parser.add_argument('--resampled_dir', required=True, default=None, type=str)
parser.add_argument('--sample_rate', default=16000, type=int)
args = parser.parse_args()
data_dir = args.data_dir
resampled_dir = args.resampled_dir
sample_rate = args.sample_rate
wav_files = sorted(glob.glob(os.path.join(data_dir, '*/*.wav')))
flac_files = sorted(glob.glob(os.path.join(data_dir, '*/*.flac')))
with Parallel(n_jobs=-1, verbose=10) as parallel:
wav_files_failed = parallel(
delayed(resample_file)(resampled_dir, filepath, ext='wav', sample_rate=sample_rate)
for filepath in wav_files
)
flac_files_failed = parallel(
delayed(resample_file)(resampled_dir, filepath, ext='flac', sample_rate=sample_rate)
for filepath in flac_files
)
with open('dataset_conversion_logs.txt', 'w') as f:
for file in wav_files_failed:
if file is not None:
f.write(f"{file}\n")
for file in flac_files_failed:
if file is not None:
f.write(f"{file}\n")
end = time.time()
print(f'Resample data in {data_dir} and save to {resampled_dir} takes {end-start} seconds.')
if __name__ == '__main__':
main()
tts/NeMo/scripts/nemo_legacy_import/asr_checkpoint_port.py
+72
View File
@@ -0,0 +1,72 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""This script converts old Jasper/QuartzNet models from NeMo 0.11.* to NeMo v1.0.0*
"""
import argparse
import torch
from omegaconf import DictConfig
from ruamel.yaml import YAML
import nemo.collections.asr as nemo_asr
from nemo.utils import logging
def get_parser():
parser = argparse.ArgumentParser(description="Converts old Jasper/QuartzNet models to NeMo v1.0beta")
parser.add_argument("--config_path", default=None, required=True, help="Path to model config (NeMo v1.0beta)")
parser.add_argument("--encoder_ckpt", default=None, required=True, help="Encoder checkpoint path")
parser.add_argument("--decoder_ckpt", default=None, required=True, help="Decoder checkpoint path")
parser.add_argument("--output_path", default=None, required=True, help="Output checkpoint path (should be .nemo)")
parser.add_argument(
"--model_type",
default='asr',
type=str,
choices=['asr', 'speech_label', 'speaker'],
help="Type of decoder used by the model.",
)
return parser
def main(config_path, encoder_ckpt, decoder_ckpt, output_path, model_type):
yaml = YAML(typ='safe')
with open(config_path) as f:
params = yaml.load(f)
model = None
if model_type == 'asr':
logging.info("Creating ASR NeMo 1.0 model")
model = nemo_asr.models.EncDecCTCModel(cfg=DictConfig(params['model']))
elif model_type == 'speech_label':
logging.info("Creating speech label NeMo 1.0 model")
model = nemo_asr.models.EncDecClassificationModel(cfg=DictConfig(params['model']))
else:
logging.info("Creating Speaker Recognition NeMo 1.0 model")
model = nemo_asr.models.EncDecSpeakerLabelModel(cfg=DictConfig(params['model']))
model.encoder.load_state_dict(torch.load(encoder_ckpt))
model.decoder.load_state_dict(torch.load(decoder_ckpt))
logging.info("Succesfully ported old checkpoint")
model.save_to(output_path)
logging.info("new model saved at {}".format(output_path))
if __name__ == "__main__":
args = get_parser().parse_args()
main(args.config_path, args.encoder_ckpt, args.decoder_ckpt, args.output_path, args.model_type)
tts/NeMo/scripts/neural_machine_translation/filter_langs_nmt.py
+403
View File
@@ -0,0 +1,403 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import logging
import multiprocessing as mp
import re
import shutil
import warnings
from pathlib import Path
from time import sleep
import fasttext
from tqdm import tqdm
"""
Usage:
python filter_by_language.py --input-src train.en \
--input-tgt train.de \
--output-src train_lang_filtered.en \
--output-tgt train_lang_filtered.de \
--source-lang en \
--target-lang de \
--removed-src train_garbage.en \
--removed-tgt train_garbage.de \
--fasttext-model lid.176.bin
"""
logging.basicConfig(level=logging.INFO)
def get_args():
parser = argparse.ArgumentParser(
description="It is a script for verifying language in machine translation data sets. If the script is used on "
"a parallel corpus, it verifies both a source and a target language. If number of jobs `--num-jobs` is bigger "
"than 1 than lines in an input file (or files if parallel corpus is checked) split equally between workers. "
"If `num_jobs > 1` is used, the best performance is achieved if dataset is shuffled and lines with different "
"lengths are distributed evenly in the input file. Filtered data is stored into `output_src`[, `--output-tgt`]"
" and removed lines are put into `removed_src`[, `--removed-tgt`] files. If language cannot be detected "
"(e.g. date), the line is removed. Working time on en-de wikimatrix (6.23M pairs: 700 MB German and 625 MB "
"English) from wmt20 on machine with 20 CPU cores: less than 1 minute."
)
parser.add_argument(
"--input-src",
"-s",
help="Path to the input file which has to contain text in language `source_lang`.",
required=True,
type=Path,
)
parser.add_argument(
"--input-tgt",
"-t",
help="Path to the input file which has to contain text in language `target_lang`. If not provided, data is "
"processed as monolingual.",
type=Path,
)
parser.add_argument(
"--output-src",
"-S",
help="Path to the file where filtered `input_src` will be saved.",
required=True,
type=Path,
)
parser.add_argument(
"--output-tgt", "-T", help="Path to the output target file", type=Path,
)
parser.add_argument(
"--source-lang",
"-l",
required=True,
help="Input language. For options see https://fasttext.cc/docs/en/language-identification.html.",
)
parser.add_argument(
"--target-lang",
"-L",
help="Output language. For options see https://fasttext.cc/docs/en/language-identification.html.",
)
parser.add_argument(
"--removed-src", "-r", required=True, help="Path to file where removed source lines will be saved", type=Path,
)
parser.add_argument(
"--removed-tgt", "-R", help="Path to file where removed target lines will be saved", type=Path,
)
parser.add_argument(
"--num-jobs",
"-j",
type=int,
help="Number of jobs. By default, the number of jobs is equal to the number of CPU cores.",
)
parser.add_argument(
"--fasttext-model",
"-m",
help="Path to fasttext model. The description and download links are here "
"https://fasttext.cc/docs/en/language-identification.html",
type=Path,
)
args = parser.parse_args()
if not (
args.output_tgt is None
and args.input_tgt is None
and args.target_lang is None
and args.removed_tgt is None
or args.output_tgt is not None
and args.input_tgt is not None
and args.target_lang is not None
and args.removed_tgt is not None
):
raise ValueError(
f"Arguments `input_tgt`, `output_tgt`, `target_lang`, `removed_tgt` have to be either `None` "
f"simultaneously or not `None` simultaneously. Given "
f"input_tgt={args.input_tgt}, output_tgt={args.output_tgt}, target_lang={args.target_lang}, "
f"removed_tgt={args.removed_tgt}"
)
args.input_src = args.input_src.expanduser()
if args.input_tgt is not None:
args.input_tgt = args.input_tgt.expanduser()
args.output_src = args.output_src.expanduser()
if args.output_tgt is not None:
args.output_tgt = args.output_tgt.expanduser()
args.removed_src = args.removed_src.expanduser()
if args.removed_tgt is not None:
args.removed_tgt = args.removed_tgt.expanduser()
args.fasttext_model = args.fasttext_model.expanduser()
return args
def get_lang(line, fasttext_model):
labels, _ = fasttext_model.predict(line, k=1)
lang = labels[0].split('__')[-1]
return lang
def get_edges_in_1_file(fn, num_parts):
num_lines = 0
edges = [0]
with open(fn) as f:
i = 0
for l in f:
i += len(l.encode('utf-8'))
edges.append(i)
num_lines += 1
return [edges[int(i * num_lines / num_parts)] for i in range(num_parts)] + [edges[-1]], num_lines
def get_edges_and_num_lines(src_fn, tgt_fn, num_parts):
src_edges, src_num_lines = get_edges_in_1_file(src_fn, num_parts)
assert num_parts + 1 == len(src_edges)
src_edges = [(src_edges[i], src_edges[i + 1]) for i in range(len(src_edges) - 1)]
if tgt_fn is not None:
tgt_edges, tgt_num_lines = get_edges_in_1_file(tgt_fn, num_parts)
tgt_edges = [(tgt_edges[i], tgt_edges[i + 1]) for i in range(len(tgt_edges) - 1)]
if tgt_num_lines != src_num_lines:
raise ValueError(
f"Source {repr(src_fn)} and target {repr(tgt_fn)} files have different number of lines "
f"{src_num_lines} and {tgt_num_lines} correspondingly."
)
else:
tgt_edges = [None] * num_parts
assert len(src_edges) == num_parts
return src_edges, tgt_edges, src_num_lines
def filter_pairs(
src_edges,
tgt_edges,
input_src,
input_tgt,
filtered_dir_src,
filtered_dir_tgt,
removed_dir_src,
removed_dir_tgt,
source_lang,
target_lang,
fasttext_model,
rank,
):
global counter
fasttext_model = fasttext.load_model(str(fasttext_model))
output_src = filtered_dir_src / Path(f"rank{rank}")
output_src_removed = removed_dir_src / Path(f"rank{rank}")
output_tgt = filtered_dir_tgt / Path(f"rank{rank}")
output_tgt_removed = removed_dir_tgt / Path(f"rank{rank}")
with open(input_src) as in_src, open(input_tgt) as in_tgt, open(output_src, 'w') as out_src, open(
output_tgt, 'w'
) as out_tgt, open(output_src_removed, 'w') as out_r_src, open(output_tgt_removed, 'w') as out_r_tgt:
in_src.seek(src_edges[0])
in_tgt.seek(tgt_edges[0])
src_l, tgt_l, i = in_src.readline(), in_tgt.readline(), 0
if in_src.tell() > src_edges[1] or in_tgt.tell() > tgt_edges[1]:
return
while src_l and tgt_l:
with counter.get_lock():
counter.value += 1
src_l = src_l.strip()
tgt_l = tgt_l.strip()
src_lang = get_lang(src_l, fasttext_model)
if src_lang is not None:
tgt_lang = get_lang(tgt_l, fasttext_model)
if src_lang is None or tgt_lang is None or src_lang != source_lang or tgt_lang != target_lang:
out_r_src.write(src_l + '\n')
out_r_tgt.write(tgt_l + '\n')
else:
out_src.write(src_l + '\n')
out_tgt.write(tgt_l + '\n')
if in_src.tell() >= src_edges[1]:
if in_tgt.tell() < tgt_edges[1]:
raise ValueError(
f"Edges of target and source has to be reached simultaneously, whereas "
f"in_src.tell()={in_src.tell()}, in_tgt.tell()={in_tgt.tell()}, "
f"src_edges[1]={src_edges[1]}, tgt_edges[1]={tgt_edges[1]}."
)
break
if in_tgt.tell() >= tgt_edges[1]:
raise ValueError(
f"Edges of target and source has to be reached simultaneously, whereas "
f"in_src.tell()={in_src.tell()}, in_tgt.tell()={in_tgt.tell()}, "
f"src_edges[1]={src_edges[1]}, tgt_edges[1]={tgt_edges[1]}."
)
src_l, tgt_l, i = in_src.readline(), in_tgt.readline(), i + 1
with counter.get_lock():
counter.value += 1
def filter_singles(
src_edges, input_src, filtered_dir_src, removed_dir_src, source_lang, fasttext_model, rank,
):
logging.debug("filter singles")
global counter
fasttext_model = fasttext.load_model(str(fasttext_model))
output_src = filtered_dir_src / Path(f"rank{rank}")
output_src_removed = removed_dir_src / Path(f"rank{rank}")
with open(input_src) as in_f, open(output_src, 'w') as out_f, open(output_src_removed, 'w') as out_r_f:
in_f.seek(src_edges[0])
i, line = 0, in_f.readline()
if in_f.tell() > src_edges[1]:
return
while line:
with counter.get_lock():
counter.value += 1
line = line.strip()
in_lang = get_lang(line, fasttext_model)
if in_lang is None or in_lang != source_lang:
out_r_f.write(line + '\n')
else:
out_f.write(line + '\n')
if in_f.tell() >= src_edges[1]:
break
i, line = i + 1, in_f.readline()
with counter.get_lock():
counter.value += 1
def filter_by_lang(args):
(
src_edges,
tgt_edges,
input_src,
input_tgt,
filtered_dir_src,
filtered_dir_tgt,
removed_dir_src,
removed_dir_tgt,
source_lang,
target_lang,
fasttext_model,
rank,
) = args
logging.debug(f"filter by lang input_tgt: {input_tgt}")
if input_tgt is None:
if tgt_edges is not None:
warnings.warn("If input target is not provided `tgt_edges` argument is expected to be `None`")
filter_singles(
src_edges, input_src, filtered_dir_src, removed_dir_src, source_lang, fasttext_model, rank,
)
else:
filter_pairs(
src_edges,
tgt_edges,
input_src,
input_tgt,
filtered_dir_src,
filtered_dir_tgt,
removed_dir_src,
removed_dir_tgt,
source_lang,
target_lang,
fasttext_model,
rank,
)
def _cat_results(out_file, tmp_dir):
file_name_pattern = re.compile(r"/rank([1-9][0-9]*)|0$")
with out_file.open('w') as out_f:
for f in sorted(tmp_dir.iterdir()):
if not f.is_file():
warnings.warn(f"Unexpected not file {f}")
elif not file_name_pattern.search(str(f)):
warnings.warn(f"Unexpected file {f}")
else:
with f.open('r') as in_f:
for l in in_f:
out_f.write(l)
def cat_results(out_files, tmp_dirs):
for o_f, t_d in zip(out_files, tmp_dirs):
if o_f is None or t_d is None:
if o_f is not None or t_d is not None:
warnings.warn(
f"Output file and tmp directory are expected to be `None` simultaneously whereas tmp directory "
f"is {t_d} and output file is {o_f}."
)
else:
_cat_results(o_f, t_d)
counter = None
def init(args):
global counter
counter = args
def main():
args = get_args()
tmp_dir = Path("tmp")
i = 0
while tmp_dir.exists():
tmp_dir = Path("tmp" + str(i))
i += 1
tmp_filtered = tmp_dir / Path("filtered")
tmp_filtered_src = tmp_filtered / Path("src")
tmp_filtered_src.mkdir(parents=True, exist_ok=True)
if args.input_tgt is None:
tmp_filtered_tgt = None
else:
tmp_filtered_tgt = tmp_filtered / Path("tgt")
tmp_filtered_tgt.mkdir(parents=True, exist_ok=True)
tmp_removed = tmp_dir / Path("removed")
tmp_removed_src = tmp_removed / Path("src")
tmp_removed_src.mkdir(parents=True, exist_ok=True)
if args.input_tgt is None:
tmp_removed_tgt = None
else:
tmp_removed_tgt = tmp_removed / Path("tgt")
tmp_removed_tgt.mkdir(parents=True, exist_ok=True)
num_jobs = mp.cpu_count() if args.num_jobs is None else args.num_jobs
src_edges, tgt_edges, num_lines = get_edges_and_num_lines(args.input_src, args.input_tgt, num_jobs)
global counter
counter = mp.Value('i', 0)
t = tqdm(total=num_lines, desc="processed lines / total number of lines")
with mp.Pool(num_jobs, initializer=init, initargs=(counter,)) as pool:
async_result = pool.map_async(
filter_by_lang,
[
(
se,
te,
args.input_src,
args.input_tgt,
tmp_filtered_src,
tmp_filtered_tgt,
tmp_removed_src,
tmp_removed_tgt,
args.source_lang,
args.target_lang,
args.fasttext_model,
rank,
)
for rank, (se, te) in enumerate(zip(src_edges, tgt_edges))
],
)
while not async_result.ready():
t.update(counter.value)
with counter.get_lock():
counter.value = 0
sleep(0.1)
t.update(counter.value)
cat_results(
[args.output_src, args.output_tgt, args.removed_src, args.removed_tgt],
[tmp_filtered_src, tmp_filtered_tgt, tmp_removed_src, tmp_removed_tgt],
)
shutil.rmtree(tmp_dir)
if __name__ == "__main__":
main()
tts/NeMo/scripts/speaker_recognition/rttm_to_manifest.py
+40
View File
@@ -0,0 +1,40 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
from nemo.collections.asr.parts.utils.speaker_utils import write_rttm2manifest
from nemo.utils import logging
"""
This file converts vad outputs to manifest file for speaker diarization purposes
present in vad output directory.
every vad line consists of start_time, end_time , speech/non-speech
"""
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--paths2rttm_files", help="path to vad output rttm-like files. Could be a list or a text file", required=True
)
parser.add_argument(
"--paths2audio_files",
help="path to audio files that vad was computed. Could be a list or a text file",
required=True,
)
parser.add_argument("--manifest_file", help="output manifest file name", type=str, required=True)
args = parser.parse_args()
write_rttm2manifest(args.paths2audio_files, args.paths2rttm_files, args.manifest_file)
logging.info("wrote {} file from vad output files present in {}".format(args.manifest_file, args.paths2rttm_files))
tts/NeMo/scripts/speaker_recognition/scp_to_manifest.py
+96
View File
@@ -0,0 +1,96 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import json
import logging
import os
import librosa as l
from sklearn.model_selection import StratifiedShuffleSplit
from tqdm import tqdm
"""
This scipt converts a scp file where each line contains
<absolute path of wav file>
to a manifest json file.
Args:
--scp: scp file name
--id: index of speaker label in filename present in scp file that is separated by '/'
--out: output manifest file name
--split: True / False if you would want to split the manifest file for training purposes
you may not need this for test set. output file names is <out>_<train/dev>.json
Defaults to False
"""
def write_file(name, lines, idx):
with open(name, 'w') as fout:
for i in idx:
dic = lines[i]
json.dump(dic, fout)
fout.write('\n')
logging.info("wrote", name)
def main(scp, id, out, split=False):
if os.path.exists(out):
os.remove(out)
scp_file = open(scp, 'r').readlines()
lines = []
speakers = []
with open(out, 'w') as outfile:
for line in tqdm(scp_file):
line = line.strip()
y, sr = l.load(line, sr=None)
dur = l.get_duration(y=y, sr=sr)
speaker = line.split('/')[id]
speaker = list(speaker)
speaker = ''.join(speaker)
speakers.append(speaker)
meta = {"audio_filepath": line, "offset": 0, "duration": float(dur), "label": speaker}
lines.append(meta)
json.dump(meta, outfile)
outfile.write("\n")
path = os.path.dirname(out)
if split:
sss = StratifiedShuffleSplit(n_splits=1, test_size=0.1, random_state=42)
for train_idx, test_idx in sss.split(speakers, speakers):
logging.info(len(train_idx))
out = os.path.join(path, 'train.json')
write_file(out, lines, train_idx)
out = os.path.join(path, 'dev.json')
write_file(out, lines, test_idx)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--scp", help="scp file name", type=str, required=True)
parser.add_argument(
"--id", help="field num seperated by '/' to be considered as speaker label", type=int, required=True
)
parser.add_argument("--out", help="manifest_file name", type=str, required=True)
parser.add_argument(
"--split",
help="bool if you would want to split the manifest file for training purposes",
required=False,
action='store_true',
)
args = parser.parse_args()
main(args.scp, args.id, args.out, args.split)
tts/NeMo/scripts/speech_recognition/convert_to_tarred_audio_dataset.py
+601
View File
@@ -0,0 +1,601 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
# This script converts an existing audio dataset with a manifest to
# a tarred and sharded audio dataset that can be read by the
# TarredAudioToTextDataLayer.
# Please make sure your audio_filepath DOES NOT CONTAIN '-sub'!
# Because we will use it to handle files which have duplicate filenames but with different offsets
# (see function create_shard for details)
# Usage:
1) Creating a new tarfile dataset
python convert_to_tarred_audio_dataset.py \
--manifest_path=<path to the manifest file> \
--target_dir=<path to output directory> \
--num_shards=<number of tarfiles that will contain the audio>
--max_duration=<float representing maximum duration of audio samples> \
--min_duration=<float representing minimum duration of audio samples> \
--shuffle --shuffle_seed=1
2) Concatenating more tarfiles to a pre-existing tarred dataset
python convert_to_tarred_audio_dataset.py \
--manifest_path=<path to the tarred manifest file> \
--metadata_path=<path to the metadata.yaml (or metadata_version_{X}.yaml) file> \
--target_dir=<path to output directory where the original tarfiles are contained> \
--max_duration=<float representing maximum duration of audio samples> \
--min_duration=<float representing minimum duration of audio samples> \
--shuffle --shuffle_seed=1 \
--concat_manifest_paths \
<space seperated paths to 1 or more manifest files to concatenate into the original tarred dataset>
3) Writing an empty metadata file
python convert_to_tarred_audio_dataset.py \
--target_dir=<path to output directory> \
# any other optional argument
--num_shards=8 \
--max_duration=16.7 \
--min_duration=0.01 \
--shuffle \
--shuffle_seed=1 \
--write_metadata
"""
import argparse
import copy
import json
import os
import random
import tarfile
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, List, Optional
from joblib import Parallel, delayed
from omegaconf import DictConfig, OmegaConf, open_dict
parser = argparse.ArgumentParser(
description="Convert an existing ASR dataset to tarballs compatible with TarredAudioToTextDataLayer."
)
parser.add_argument(
"--manifest_path", default=None, type=str, required=False, help="Path to the existing dataset's manifest."
)
parser.add_argument(
'--concat_manifest_paths',
nargs='+',
default=None,
type=str,
required=False,
help="Path to the additional dataset's manifests that will be concatenated with base dataset.",
)
# Optional arguments
parser.add_argument(
"--target_dir",
default='./tarred',
type=str,
help="Target directory for resulting tarballs and manifest. Defaults to `./tarred`. Creates the path if necessary.",
)
parser.add_argument(
"--metadata_path", required=False, default=None, type=str, help="Path to metadata file for the dataset.",
)
parser.add_argument(
"--num_shards",
default=-1,
type=int,
help="Number of shards (tarballs) to create. Used for partitioning data among workers.",
)
parser.add_argument(
'--max_duration',
default=None,
type=float,
help='Maximum duration of audio clip in the dataset. By default, it is None and will not filter files.',
)
parser.add_argument(
'--min_duration',
default=None,
type=float,
help='Minimum duration of audio clip in the dataset. By default, it is None and will not filter files.',
)
parser.add_argument(
"--shuffle",
action='store_true',
help="Whether or not to randomly shuffle the samples in the manifest before tarring/sharding.",
)
parser.add_argument("--shuffle_seed", type=int, help="Random seed for use if shuffling is enabled.")
parser.add_argument(
'--write_metadata',
action='store_true',
help=(
"Flag to write a blank metadata with the current call config. "
"Note that the metadata will not contain the number of shards, "
"and it must be filled out by the user."
),
)
parser.add_argument('--workers', type=int, default=1, help='Number of worker processes')
args = parser.parse_args()
@dataclass
class ASRTarredDatasetConfig:
num_shards: int = -1
shuffle: bool = False
max_duration: Optional[float] = None
min_duration: Optional[float] = None
shuffle_seed: Optional[int] = None
@dataclass
class ASRTarredDatasetMetadata:
created_datetime: Optional[str] = None
version: int = 0
num_samples_per_shard: Optional[int] = None
is_concatenated_manifest: bool = False
dataset_config: Optional[ASRTarredDatasetConfig] = ASRTarredDatasetConfig()
history: Optional[List[Any]] = field(default_factory=lambda: [])
def __post_init__(self):
self.created_datetime = self.get_current_datetime()
def get_current_datetime(self):
return datetime.now().strftime("%m-%d-%Y %H-%M-%S")
@classmethod
def from_config(cls, config: DictConfig):
obj = cls()
obj.__dict__.update(**config)
return obj
@classmethod
def from_file(cls, filepath: str):
config = OmegaConf.load(filepath)
return ASRTarredDatasetMetadata.from_config(config=config)
class ASRTarredDatasetBuilder:
"""
Helper class that constructs a tarred dataset from scratch, or concatenates tarred datasets
together and constructs manifests for them.
"""
def __init__(self):
self.config = None
def configure(self, config: ASRTarredDatasetConfig):
"""
Sets the config generated from command line overrides.
Args:
config: ASRTarredDatasetConfig dataclass object.
"""
self.config = config # type: ASRTarredDatasetConfig
if self.config.num_shards < 0:
raise ValueError("`num_shards` must be > 0. Please fill in the metadata information correctly.")
def create_new_dataset(self, manifest_path: str, target_dir: str = "./tarred/", num_workers: int = 0):
"""
Creates a new tarred dataset from a given manifest file.
Args:
manifest_path: Path to the original ASR manifest.
target_dir: Output directory.
num_workers: Integer denoting number of parallel worker processes which will write tarfiles.
Defaults to 1 - which denotes sequential worker process.
Output:
Writes tarfiles, along with the tarred dataset compatible manifest file.
Also preserves a record of the metadata used to construct this tarred dataset.
"""
if self.config is None:
raise ValueError("Config has not been set. Please call `configure(config: ASRTarredDatasetConfig)`")
if manifest_path is None:
raise FileNotFoundError("Manifest filepath cannot be None !")
config = self.config # type: ASRTarredDatasetConfig
if not os.path.exists(target_dir):
os.makedirs(target_dir)
# Read the existing manifest
entries, filtered_entries, filtered_duration = self._read_manifest(manifest_path, config)
if len(filtered_entries) > 0:
print(f"Filtered {len(filtered_entries)} files which amounts to {filtered_duration} seconds of audio.")
if config.shuffle:
random.seed(config.shuffle_seed)
print("Shuffling...")
random.shuffle(entries)
# Create shards and updated manifest entries
print(f"Number of samples added : {len(entries)}")
print(f"Remainder: {len(entries) % config.num_shards}")
start_indices = []
end_indices = []
# Build indices
for i in range(config.num_shards):
start_idx = (len(entries) // config.num_shards) * i
end_idx = start_idx + (len(entries) // config.num_shards)
print(f"Shard {i} has entries {start_idx} ~ {end_idx}")
if i == config.num_shards - 1:
# We discard in order to have the same number of entries per shard.
print(f"Have {len(entries) - end_idx} entries left over that will be discarded.")
start_indices.append(start_idx)
end_indices.append(end_idx)
with Parallel(n_jobs=num_workers, verbose=config.num_shards) as parallel:
# Call parallel tarfile construction
new_entries_list = parallel(
delayed(self._create_shard)(entries[start_idx:end_idx], target_dir, i)
for i, (start_idx, end_idx) in enumerate(zip(start_indices, end_indices))
)
# Flatten the list of list of entries to a list of entries
new_entries = [sample for manifest in new_entries_list for sample in manifest]
del new_entries_list
print("Total number of files in manifest :", len(new_entries))
# Write manifest
new_manifest_path = os.path.join(target_dir, 'tarred_audio_manifest.json')
with open(new_manifest_path, 'w') as m2:
for entry in new_entries:
json.dump(entry, m2)
m2.write('\n')
# Write metadata (default metadata for new datasets)
new_metadata_path = os.path.join(target_dir, 'metadata.yaml')
metadata = ASRTarredDatasetMetadata()
# Update metadata
metadata.dataset_config = config
metadata.num_samples_per_shard = len(new_entries) // config.num_shards
# Write metadata
metadata_yaml = OmegaConf.structured(metadata)
OmegaConf.save(metadata_yaml, new_metadata_path, resolve=True)
def create_concatenated_dataset(
self,
base_manifest_path: str,
manifest_paths: List[str],
metadata: ASRTarredDatasetMetadata,
target_dir: str = "./tarred_concatenated/",
num_workers: int = 1,
):
"""
Creates new tarfiles in order to create a concatenated dataset, whose manifest contains the data for
both the original dataset as well as the new data submitted in manifest paths.
Args:
base_manifest_path: Path to the manifest file which contains the information for the original
tarred dataset (with flattened paths).
manifest_paths: List of one or more paths to manifest files that will be concatenated with above
base tarred dataset.
metadata: ASRTarredDatasetMetadata dataclass instance with overrides from command line.
target_dir: Output directory
Output:
Writes tarfiles which with indices mapping to a "concatenated" tarred dataset,
along with the tarred dataset compatible manifest file which includes information
about all the datasets that comprise the concatenated dataset.
Also preserves a record of the metadata used to construct this tarred dataset.
"""
if not os.path.exists(target_dir):
os.makedirs(target_dir)
if base_manifest_path is None:
raise FileNotFoundError("Base manifest filepath cannot be None !")
if manifest_paths is None or len(manifest_paths) == 0:
raise FileNotFoundError("List of additional manifest filepaths cannot be None !")
config = ASRTarredDatasetConfig(**(metadata.dataset_config))
# Read the existing manifest (no filtering here)
base_entries, _, _ = self._read_manifest(base_manifest_path, config)
print(f"Read base manifest containing {len(base_entries)} samples.")
# Precompute number of samples per shard
if metadata.num_samples_per_shard is None:
num_samples_per_shard = len(base_entries) // config.num_shards
else:
num_samples_per_shard = metadata.num_samples_per_shard
print("Number of samples per shard :", num_samples_per_shard)
# Compute min and max duration and update config (if no metadata passed)
print(f"Selected max duration : {config.max_duration}")
print(f"Selected min duration : {config.min_duration}")
entries = []
for new_manifest_idx in range(len(manifest_paths)):
new_entries, filtered_new_entries, filtered_duration = self._read_manifest(
manifest_paths[new_manifest_idx], config
)
if len(filtered_new_entries) > 0:
print(
f"Filtered {len(filtered_new_entries)} files which amounts to {filtered_duration:0.2f}"
f" seconds of audio from manifest {manifest_paths[new_manifest_idx]}."
)
entries.extend(new_entries)
if config.shuffle:
random.seed(config.shuffle_seed)
print("Shuffling...")
random.shuffle(entries)
# Drop last section of samples that cannot be added onto a chunk
drop_count = len(entries) % num_samples_per_shard
total_new_entries = len(entries)
entries = entries[:-drop_count]
print(
f"Dropping {drop_count} samples from total new samples {total_new_entries} since they cannot "
f"be added into a uniformly sized chunk."
)
# Create shards and updated manifest entries
num_added_shards = len(entries) // num_samples_per_shard
print(f"Number of samples in base dataset : {len(base_entries)}")
print(f"Number of samples in additional datasets : {len(entries)}")
print(f"Number of added shards : {num_added_shards}")
print(f"Remainder: {len(entries) % num_samples_per_shard}")
start_indices = []
end_indices = []
shard_indices = []
for i in range(num_added_shards):
start_idx = (len(entries) // num_added_shards) * i
end_idx = start_idx + (len(entries) // num_added_shards)
shard_idx = i + config.num_shards
print(f"Shard {shard_idx} has entries {start_idx + len(base_entries)} ~ {end_idx + len(base_entries)}")
start_indices.append(start_idx)
end_indices.append(end_idx)
shard_indices.append(shard_idx)
with Parallel(n_jobs=num_workers, verbose=num_added_shards) as parallel:
# Call parallel tarfile construction
new_entries_list = parallel(
delayed(self._create_shard)(entries[start_idx:end_idx], target_dir, shard_idx)
for i, (start_idx, end_idx, shard_idx) in enumerate(zip(start_indices, end_indices, shard_indices))
)
# Flatten the list of list of entries to a list of entries
new_entries = [sample for manifest in new_entries_list for sample in manifest]
del new_entries_list
# Write manifest
if metadata is None:
new_version = 1 # start with `1`, where `0` indicates the base manifest + dataset
else:
new_version = metadata.version + 1
print("Total number of files in manifest :", len(base_entries) + len(new_entries))
new_manifest_path = os.path.join(target_dir, f'tarred_audio_manifest_version_{new_version}.json')
with open(new_manifest_path, 'w') as m2:
# First write all the entries of base manifest
for entry in base_entries:
json.dump(entry, m2)
m2.write('\n')
# Finally write the new entries
for entry in new_entries:
json.dump(entry, m2)
m2.write('\n')
# Preserve historical metadata
base_metadata = metadata
# Write metadata (updated metadata for concatenated datasets)
new_metadata_path = os.path.join(target_dir, f'metadata_version_{new_version}.yaml')
metadata = ASRTarredDatasetMetadata()
# Update config
config.num_shards = config.num_shards + num_added_shards
# Update metadata
metadata.version = new_version
metadata.dataset_config = config
metadata.num_samples_per_shard = num_samples_per_shard
metadata.is_concatenated_manifest = True
metadata.created_datetime = metadata.get_current_datetime()
# Attach history
current_metadata = OmegaConf.structured(base_metadata.history)
metadata.history = current_metadata
# Write metadata
metadata_yaml = OmegaConf.structured(metadata)
OmegaConf.save(metadata_yaml, new_metadata_path, resolve=True)
def _read_manifest(self, manifest_path: str, config: ASRTarredDatasetConfig):
"""Read and filters data from the manifest"""
# Read the existing manifest
entries = []
filtered_entries = []
filtered_duration = 0.0
with open(manifest_path, 'r') as m:
for line in m:
entry = json.loads(line)
if (config.max_duration is None or entry['duration'] < config.max_duration) and (
config.min_duration is None or entry['duration'] > config.min_duration
):
entries.append(entry)
else:
filtered_entries.append(entry)
filtered_duration += entry['duration']
return entries, filtered_entries, filtered_duration
def _create_shard(self, entries, target_dir, shard_id):
"""Creates a tarball containing the audio files from `entries`.
"""
new_entries = []
tar = tarfile.open(os.path.join(target_dir, f'audio_{shard_id}.tar'), mode='w', dereference=True)
count = dict()
for entry in entries:
# We squash the filename since we do not preserve directory structure of audio files in the tarball.
base, ext = os.path.splitext(entry['audio_filepath'])
base = base.replace('/', '_')
# Need the following replacement as long as WebDataset splits on first period
base = base.replace('.', '_')
squashed_filename = f'{base}{ext}'
if squashed_filename not in count:
tar.add(entry['audio_filepath'], arcname=squashed_filename)
if 'label' in entry:
base, ext = os.path.splitext(squashed_filename)
# no suffix if it's single sample or starting sub parts, -sub1 for the second subpart -sub2 -sub3 ,etc.
if squashed_filename not in count:
to_write = squashed_filename
count[squashed_filename] = 1
else:
to_write = base + "-sub" + str(count[squashed_filename]) + ext
count[squashed_filename] += 1
new_entry = {
'audio_filepath': to_write,
'duration': entry['duration'],
'text': entry['text'],
'label': entry['label'],
'offset': entry['offset'],
'shard_id': shard_id, # Keep shard ID for recordkeeping
}
else:
count[squashed_filename] = 1
new_entry = {
'audio_filepath': squashed_filename,
'duration': entry['duration'],
'text': entry['text'],
'shard_id': shard_id, # Keep shard ID for recordkeeping
}
new_entries.append(new_entry)
tar.close()
return new_entries
@classmethod
def setup_history(cls, base_metadata: ASRTarredDatasetMetadata, history: List[Any]):
if 'history' in base_metadata.keys():
for history_val in base_metadata.history:
cls.setup_history(history_val, history)
if base_metadata is not None:
metadata_copy = copy.deepcopy(base_metadata)
with open_dict(metadata_copy):
metadata_copy.pop('history', None)
history.append(metadata_copy)
def main():
manifest_path = args.manifest_path
concat_manifest_paths = args.concat_manifest_paths
target_dir = args.target_dir
metadata_path = args.metadata_path
num_shards = args.num_shards
max_duration = args.max_duration
min_duration = args.min_duration
shuffle = args.shuffle
seed = args.shuffle_seed if args.shuffle_seed else None
write_metadata = args.write_metadata
num_workers = args.workers
builder = ASRTarredDatasetBuilder()
if write_metadata:
metadata = ASRTarredDatasetMetadata()
dataset_cfg = ASRTarredDatasetConfig(
num_shards=num_shards,
shuffle=shuffle,
max_duration=max_duration,
min_duration=min_duration,
shuffle_seed=seed,
)
metadata.dataset_config = dataset_cfg
output_path = os.path.join(target_dir, 'default_metadata.yaml')
OmegaConf.save(metadata, output_path, resolve=True)
print(f"Default metadata written to {output_path}")
exit(0)
if concat_manifest_paths is None or len(concat_manifest_paths) == 0:
print("Creating new tarred dataset ...")
# Create a tarred dataset from scratch
config = ASRTarredDatasetConfig(
num_shards=num_shards,
shuffle=shuffle,
max_duration=max_duration,
min_duration=min_duration,
shuffle_seed=seed,
)
builder.configure(config)
builder.create_new_dataset(manifest_path=manifest_path, target_dir=target_dir, num_workers=num_workers)
else:
print("Concatenating multiple tarred datasets ...")
# Implicitly update config from base details
if metadata_path is not None:
metadata = ASRTarredDatasetMetadata.from_file(metadata_path)
else:
raise ValueError("`metadata` yaml file path must be provided!")
# Preserve history
history = []
builder.setup_history(OmegaConf.structured(metadata), history)
metadata.history = history
# Add command line overrides (everything other than num_shards)
metadata.dataset_config.max_duration = max_duration
metadata.dataset_config.min_duration = min_duration
metadata.dataset_config.shuffle = shuffle
metadata.dataset_config.shuffle_seed = seed
builder.configure(metadata.dataset_config)
# Concatenate a tarred dataset onto a previous one
builder.create_concatenated_dataset(
base_manifest_path=manifest_path,
manifest_paths=concat_manifest_paths,
metadata=metadata,
target_dir=target_dir,
num_workers=num_workers,
)
if __name__ == "__main__":
main()
tts/NeMo/scripts/tokenizers/process_asr_text_tokenizer.py
+282
View File
@@ -0,0 +1,282 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# USAGE: python process_asr_text_tokenizer.py --manifest=<path to train manifest files, seperated by commas> \
# --data_root="<output directory>" \
# --vocab_size=<number of tokens in vocabulary> \
# --tokenizer=<"spe" or "wpe"> \
# --log
# where <manifest> can be: train_clean_100, train_clean_360, train_other_500
# You can also put more than one data_set comma-separated:
# --manifest="train_clean_100,train_clean_360,train_other_500"
# or
# python process_asr_text_tokenizer.py --data_file=<path to train text file> \
# --data_root="<output directory>" \
# --vocab_size=<number of tokens in vocabulary> \
# --tokenizer=<"bpe" or "wpe"> \
# --log
# where <manifest> can be: train_clean_100, train_clean_360, train_other_500
# You can also put more than one data_set comma-separated:
# --manifest="train_clean_100,train_clean_360,train_other_500"
#
# Args:
# --manifest or --data_file: If your text data lies inside of an ASR manifest file,
# then use the --manifest path. If instead the text data is inside a file with separate lines
# corresponding to different text lines, then use --data_file.
# In either case, you can add commas to concatenate different manifests or different data files.
#
# --data_root: The output directory (whose subdirectories will be created if not present) where
# the tokenizers will be placed.
#
# --vocab_size: The size of the tokenizer vocabulary. Larger vocabularies can accommodate almost entire,
# words but the decoder size of any model will grow proportionally.
#
# --tokenizer: Can be either spe or wpe . spe refers to the Google sentencepiece library tokenizer.
# wpe refers to the HuggingFace BERT Word Piece tokenizer.
#
# --no_lower_case: When this flag is passed, it will force the tokenizer to create seperate tokens for
# upper and lower case characters. By default, the script will turn all the text to lower case
# before tokenization (and if upper case characters are passed during training/inference, the
# tokenizer will emit a token equivalent to Out-Of-Vocabulary). Used primarily for the
# English language.
#
# --spe_type: The sentencepiece library has a few implementations of the tokenization technique, and
# spe_type refers to these implementations. Currently supported types are unigram, bpe, char, word.
# Defaults to bpe.
#
# --spe_character_coverage: The sentencepiece library considers how much of the original vocabulary it
# should cover in its "base set" of tokens (akin to the lower and upper case characters of the
# English language). For almost all languages with small base token sets (<1000 tokens), this
# should be kept at its default of 1.0. For languages with larger vocabularies (say Japanese,
# Mandarin, Korean etc), the suggested value is 0.9995.
#
# --spe_sample_size: If the dataset is too large, consider using a sampled dataset indicated by a
# positive integer. By default, any negative value (default = -1) will use the entire dataset.
#
# --spe_train_extremely_large_corpus: When training a sentencepiece tokenizer on very large amounts of text,
# sometimes the tokenizer will run out of memory or wont be able to process so much data on RAM.
# At some point you might receive the following error - "Input corpus too large, try with
# train_extremely_large_corpus=true". If your machine has large amounts of RAM, it might still be possible
# to build the tokenizer using the above flag. Will silently fail if it runs out of RAM.
#
# --spe_max_sentencepiece_length: Limits the maximum length that any any SentencePiece subword can be.
# Using this will change the subword tokens generated.
#
# --log: Whether the script should display log messages
import argparse
import json
import logging
import os
import tokenizers
from nemo.collections.common.tokenizers.sentencepiece_tokenizer import create_spt_model
parser = argparse.ArgumentParser(description='Create tokenizer')
group = parser.add_mutually_exclusive_group(required=True)
group.add_argument("--manifest", default=None, type=str, help='Comma separated list of manifest files')
group.add_argument("--data_file", default=None, help='data file from which to create tokenizer model')
parser.add_argument("--data_root", required=True, default=None, type=str, help='Output directory')
parser.add_argument("--vocab_size", default=1024, type=int, help='Vocabulary size')
parser.add_argument("--tokenizer", default="wpe", choices=["spe", "wpe"], help='Type of tokenization to perform')
parser.add_argument(
"--spe_type",
default="bpe",
choices=['bpe', 'unigram', 'char', 'word'],
help='Type of the SentencePiece model. Can be `bpe`, `unigram`, `char` or `word`.'
'Used only if --tokenizer == `spe`',
)
parser.add_argument(
'--spe_character_coverage',
type=float,
default=1.0,
help="Character coverage percentage for SentencePiece tokenization. For languages "
"with large vocabulary, should be close to 0.9995, otherwise kept as 1.0",
)
parser.add_argument(
'--spe_sample_size',
type=int,
default=-1,
help="Samples the dataset by `sample_size` if positive integer, otherwise uses whole dataset",
)
parser.add_argument('--spe_train_extremely_large_corpus', action='store_true', help='')
parser.add_argument(
'--spe_max_sentencepiece_length',
type=int,
default=-1,
help='Limit the maximum number of tokens in each SentencePiece subword. '
'Must be a positive integer > 0. By default places no limit on subword length.',
)
parser.add_argument('--no_lower_case', dest='lower_case', action='store_false')
parser.add_argument("--log", action='store_true')
parser.set_defaults(log=False, lower_case=True, spe_train_extremely_large_corpus=False)
args = parser.parse_args()
def __build_document_from_manifests(
data_root: str, manifests: str,
):
if ',' in manifests:
manifests = manifests.split(',')
else:
manifests = [manifests]
document_dir = os.path.join(data_root, 'text_corpus')
if not os.path.exists(document_dir):
os.makedirs(document_dir)
document_path = os.path.join(document_dir, 'document.txt')
if os.path.exists(document_path):
logging.info('Corpus already exists at path : %s', document_path)
return document_path
num_lines = 0
with open(document_path, 'w') as out_writer:
for manifest in manifests:
with open(manifest, 'r') as in_reader:
for line in in_reader:
item = json.loads(line)
text = item['text']
out_writer.write(text + '\n')
out_writer.flush()
num_lines += 1
logging.info(f"Finished extracting manifest : {manifest}")
logging.info("Finished extracting all manifests ! Number of sentences : {}".format(num_lines))
return document_path
def __process_data(
text_path: str,
dst_folder: str,
vocab_size: int,
tokenizer_type: str,
spe_type: str,
spe_character_coverage: float,
spe_train_extremely_large_corpus: bool,
spe_sample_size: int,
spe_max_sentencepiece_length: int,
lower_case: bool,
):
"""
Converts flac to wav and build manifests's json
Args:
text_path: source with text lines
dst_folder: where wav files will be stored
vocab_size: vocabular size used in encoding the text
tokenizer_type: type of tokenization to perform - wpe or spe
spe_type: type of tokenization model used for spe.
spe_character_coverage: float value between 0 and 1 (as a percentage). For languages with a vast charset,
can be < 1.0, but for all other languages, it should be set as 1.0
spe_sample_size: int, default of -1. If positive integer is used, samples the dataset
by given sample size.
spe_train_extremely_large_corpus: bool. If dataset is too large, and user has sufficient RAM,
this flag can be set to try to trained the tokenizer. Will silently fail if it runs out of RAM.
spe_max_sentencepiece_length: Limits the maximum length of the SentencePiece subword that can be constructed.
By default, no limit is placed.
lower_case: whether to tokenize with lower case character set only (for english)
Returns:
"""
if tokenizer_type == 'spe':
if spe_max_sentencepiece_length > 0:
tokenizer_dir = os.path.join(dst_folder, 'tokenizer_{}_{}_v{}_max{}').format(
tokenizer_type, spe_type, vocab_size, spe_max_sentencepiece_length
)
else:
tokenizer_dir = os.path.join(dst_folder, 'tokenizer_{}_{}_v{}').format(
tokenizer_type, spe_type, vocab_size
)
if not os.path.exists(tokenizer_dir):
os.makedirs(tokenizer_dir)
if os.path.exists(os.path.join(tokenizer_dir, 'tokenizer.model')):
logging.warning("Model file already exists, overriding old model file !")
os.remove(os.path.join(tokenizer_dir, 'tokenizer.model'))
tokenizer_path, vocab_path = create_spt_model(
data_file=text_path,
vocab_size=vocab_size,
sample_size=spe_sample_size,
do_lower_case=lower_case,
output_dir=tokenizer_dir,
tokenizer_type=spe_type,
character_coverage=spe_character_coverage,
train_extremely_large_corpus=spe_train_extremely_large_corpus,
max_sentencepiece_length=spe_max_sentencepiece_length,
)
else:
tokenizer_dir = os.path.join(dst_folder, 'tokenizer_{}_v{}').format(tokenizer_type, vocab_size)
if not os.path.exists(tokenizer_dir):
os.makedirs(tokenizer_dir)
tokenizer = tokenizers.BertWordPieceTokenizer(lowercase=lower_case)
tokenizer.train(text_path, vocab_size=vocab_size)
tokenizer.save_model(tokenizer_dir)
return tokenizer_dir
def main():
data_root = args.data_root
manifests = args.manifest
data_file = args.data_file
vocab_size = args.vocab_size
tokenizer = args.tokenizer
spe_type = args.spe_type
spe_character_coverage = args.spe_character_coverage
spe_sample_size = args.spe_sample_size
spe_train_extremely_large_corpus = args.spe_train_extremely_large_corpus
spe_max_sentencepiece_length = args.spe_max_sentencepiece_length
lower_case = args.lower_case
if not os.path.exists(data_root):
os.makedirs(data_root)
if args.log:
logging.basicConfig(level=logging.INFO)
if manifests:
text_corpus_path = __build_document_from_manifests(data_root, manifests)
else:
text_corpus_path = data_file
tokenizer_path = __process_data(
text_corpus_path,
data_root,
vocab_size,
tokenizer,
spe_type,
lower_case=lower_case,
spe_character_coverage=spe_character_coverage,
spe_sample_size=spe_sample_size,
spe_train_extremely_large_corpus=spe_train_extremely_large_corpus,
spe_max_sentencepiece_length=spe_max_sentencepiece_length,
)
print("Serialized tokenizer at location :", tokenizer_path)
logging.info('Done!')
if __name__ == "__main__":
main()
tts/NeMo/scripts/voice_activity_detection/vad_overlap_posterior.py
+259
View File
@@ -0,0 +1,259 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import glob
import os
import time
from argparse import ArgumentParser
from itertools import repeat
from multiprocessing import Pool
import numpy as np
import pandas as pd
"""
This script serves two purposes:
1) gen_overlap_seq:
Generate predictions with overlapping input segments by using the frame level prediction from NeMo/examples/asr/vad_infer.py.
Then a smoothing filter is applied to decide the label for a frame spanned by multiple segments.
2）gen_seg_table:
Converting frame level prediction to speech/no-speech segment in start and end times format.
Usage:
python vad_overlap_posterior.py --gen_overlap_seq --gen_seg_table --frame_folder=<FULL PATH OF YOU STORED FRAME LEVEL PREDICTION> --method='median' --overlap=0.875 --num_workers=20 --threshold=0.8
"""
def gen_overlap_seq(frame_filepath, per_args):
"""
Given a frame level prediction, generate predictions with overlapping input segments by using it
Args:
frame_filepath : frame prediction file to be processed.
per_args:
method : Median or mean smoothing filter.
overlap : Amounts of overlap.
seg_len : Length of window for generating the frame.
shift_len : Amount of shift of window for generating the frame.
out_dir : Output dir of generated prediction.
"""
try:
method = per_args['method']
overlap = per_args['overlap']
seg_len = per_args['seg_len']
shift_len = per_args['shift_len']
out_dir = per_args['out_dir']
frame = np.loadtxt(frame_filepath)
name = os.path.basename(frame_filepath).split(".frame")[0] + "." + method
overlap_filepath = os.path.join(out_dir, name)
shift = int(shift_len / 0.01) # number of units of shift
seg = int((seg_len / 0.01 + 1)) # number of units of each window/segment
jump_on_target = int(seg * (1 - overlap)) # jump on target generated sequence
jump_on_frame = int(jump_on_target / shift) # jump on input frame sequence
if jump_on_frame < 1:
raise ValueError(
f"Note we jump over frame sequence to generate overlapping input segments. \n \
Your input makes jump_on_fram={jump_on_frame} < 1 which is invalid because it cannot jump and will stuck.\n \
Please try different seg_len, shift_len and overlap choices. \n \
jump_on_target = int(seg * (1 - overlap)) \n \
jump_on_frame = int(jump_on_frame/shift) "
)
target_len = int(len(frame) * shift)
if method == 'mean':
preds = np.zeros(target_len)
pred_count = np.zeros(target_len)
for i, og_pred in enumerate(frame):
if i % jump_on_frame != 0:
continue
start = i * shift
end = start + seg
preds[start:end] = preds[start:end] + og_pred
pred_count[start:end] = pred_count[start:end] + 1
preds = preds / pred_count
last_non_zero_pred = preds[pred_count != 0][-1]
preds[pred_count == 0] = last_non_zero_pred
elif method == 'median':
preds = [[] for _ in range(target_len)]
for i, og_pred in enumerate(frame):
if i % jump_on_frame != 0:
continue
start = i * shift
end = start + seg
for j in range(start, end):
if j <= target_len - 1:
preds[j].append(og_pred)
preds = np.array([np.median(l) for l in preds])
nan_idx = np.isnan(preds)
last_non_nan_pred = preds[~nan_idx][-1]
preds[nan_idx] = last_non_nan_pred
else:
raise ValueError("method should be either mean or median")
round_final = np.round(preds, 4)
np.savetxt(overlap_filepath, round_final, delimiter='\n')
print(f"Finished! {overlap_filepath}!")
except Exception as e:
raise (e)
def gen_seg_table(frame_filepath, per_args):
"""
Convert frame level prediction to speech/no-speech segment in start and end times format.
And save to csv file in rttm-like format
0, 10, speech
10,12, no-speech
Args:
frame_filepath : frame prediction file to be processed.
per_args :
threshold : threshold for prediction score (from 0 to 1).
shift_len : Amount of shift of window for generating the frame.
out_dir : Output dir of generated table/csv file.
"""
threshold = per_args['threshold']
shift_len = per_args['shift_len']
out_dir = per_args['out_dir']
print(f"process {frame_filepath}")
name = frame_filepath.split("/")[-1].rsplit(".", 1)[0]
sequence = np.loadtxt(frame_filepath)
start = 0
end = 0
start_list = [0]
end_list = []
state_list = []
for i in range(len(sequence) - 1):
current_sate = "non-speech" if sequence[i] <= threshold else "speech"
next_state = "non-speech" if sequence[i + 1] <= threshold else "speech"
if next_state != current_sate:
end = i * shift_len + shift_len # shift_len for handling joint
state_list.append(current_sate)
end_list.append(end)
start = (i + 1) * shift_len
start_list.append(start)
end_list.append((i + 1) * shift_len + shift_len)
state_list.append(current_sate)
seg_table = pd.DataFrame({'start': start_list, 'end': end_list, 'vad': state_list})
save_name = name + ".txt"
save_path = os.path.join(out_dir, save_name)
seg_table.to_csv(save_path, sep='\t', index=False, header=False)
if __name__ == '__main__':
start = time.time()
parser = ArgumentParser()
parser.add_argument("--gen_overlap_seq", default=False, action='store_true')
parser.add_argument("--gen_seg_table", default=False, action='store_true')
parser.add_argument("--frame_folder", type=str, required=True)
parser.add_argument(
"--method",
type=str,
required=True,
help="Use mean/median for overlapped prediction. Use frame for gen_seg_table of frame prediction",
)
parser.add_argument("--overlap_out_dir", type=str)
parser.add_argument("--table_out_dir", type=str)
parser.add_argument("--overlap", type=float, default=0.875, help="Overlap percentatge. Default is 0.875")
parser.add_argument("--seg_len", type=float, default=0.63)
parser.add_argument("--shift_len", type=float, default=0.01)
parser.add_argument("--threshold", type=float, default=0.5)
parser.add_argument("--num_workers", type=int, default=4)
args = parser.parse_args()
if args.gen_overlap_seq:
p = Pool(processes=args.num_workers)
print("Generating predictions with overlapping input segments")
frame_filepathlist = glob.glob(args.frame_folder + "/*.frame")
if not args.overlap_out_dir:
overlap_out_dir = "./overlap_smoothing_output" + "_" + args.method + "_" + str(args.overlap)
else:
overlap_out_dir = args.overlap_out_dir
if not os.path.exists(overlap_out_dir):
print(f"Creating output dir {overlap_out_dir}")
os.mkdir(overlap_out_dir)
per_args = {
"method": args.method,
"overlap": args.overlap,
"seg_len": args.seg_len,
"shift_len": args.shift_len,
"out_dir": overlap_out_dir,
}
p.starmap(gen_overlap_seq, zip(frame_filepathlist, repeat(per_args)))
p.close()
p.join()
end = time.time()
print(f"Generate overlapped prediction takes {end-start} seconds!\n Save to {overlap_out_dir}")
if args.gen_seg_table:
p = Pool(processes=args.num_workers)
print("Converting frame level prediction to speech/no-speech segment in start and end times format.")
if args.gen_overlap_seq:
print("Use overlap prediction. Change if you want to use basic frame level prediction")
frame_filepath = overlap_out_dir
shift_len = 0.01
else:
print("Use basic frame level prediction")
frame_filepath = args.frame_folder
shift_len = args.shift_len
frame_filepathlist = glob.glob(frame_filepath + "/*." + args.method)
if not args.table_out_dir:
table_out_dir = "table_output_" + str(args.threshold)
else:
table_out_dir = args.table_out_dir
if not os.path.exists(table_out_dir):
print(f"Creating rttm-like table output dir {table_out_dir}")
os.mkdir(table_out_dir)
per_args = {
"threshold": args.threshold,
"shift_len": shift_len,
"out_dir": table_out_dir,
}
p.starmap(gen_seg_table, zip(frame_filepathlist, repeat(per_args)))
p.close()
p.join()
end = time.time()
print(f"Generate rttm-like table for {frame_filepath} takes {end-start} seconds!\n Save to {table_out_dir}")
tts/NeMo/scripts/voice_activity_detection/vad_tune_threshold.py
+58
View File
@@ -0,0 +1,58 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import numpy as np
from nemo.collections.asr.parts.utils.vad_utils import vad_tune_threshold_on_dev
from nemo.utils import logging
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--threshold_range", help="range of threshold in list 'START,END,STEP' to be tuned on", required=True
)
parser.add_argument(
"--vad_pred", help="Directory of vad predictions or a file contains the paths of them.", required=True
)
parser.add_argument(
"--groundtruth_RTTM",
help="Directory of groundtruch rttm files or a file contains the paths of them",
type=str,
required=True,
)
parser.add_argument(
"--vad_pred_method",
help="suffix of prediction file. Should be either in 'frame', 'mean' or 'median'",
required=True,
)
parser.add_argument(
"--focus_metric",
help="metrics we care most when tuning threshold. Should be either in 'DetER', 'FA', 'MISS' ",
type=str,
default='DetER',
)
args = parser.parse_args()
try:
start, stop, step = [float(i) for i in args.threshold_range.split(",")]
thresholds = np.arange(start, stop, step)
except:
raise ValueError("Theshold input is invalid! Please enter it as a 'START,STOP,STEP' ")
best_threhsold = vad_tune_threshold_on_dev(
thresholds, args.vad_pred, args.groundtruth_RTTM, args.vad_pred_method, args.focus_metric
)
logging.info(f"Best threshold selected from {thresholds} is {best_threhsold}!")
tts/NeMo/scripts/voice_activity_detection/write_long_audio_manifest.py
+173
View File
@@ -0,0 +1,173 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import logging
import multiprocessing
import os
from argparse import ArgumentParser
from itertools import repeat
import librosa
import numpy as np
"""
This script is designed for inference of frame level Voice Activity Detection (VAD)
This script serves three goals:
(1) Write audio files to manifest
(2) Split audio file for avoiding CUDA memory issue
(3) Take care of joint of seperate json line for an audio file
Usage:
python write_long_audio_manifest.py --inp_dir=<FULL PATH OF FOLDER OF AUDIO FILES> --manifest_name=<NAME OF OUTPUT JSON FILE> --split_duration=300 --time_length=0.63 --num_worker=10
"""
def write_manifest(file, args_func):
"""
Given a list of files, write them to manifest with restrictions.
Args:
files : file to be processed
label : label for audio snippet.
split_duration : Max duration of each audio clip (each line in json)
time_length : Used for taking care of joint.
Length of window for generating the frame.
Returns:
res : list of generated metadata line of json for file
"""
label = args_func['label']
split_duration = args_func['split_duration']
time_length = args_func['time_length']
res = []
try:
sr = 16000
x, _sr = librosa.load(file, sr=sr)
duration = librosa.get_duration(x, sr=sr)
left = duration
current_offset = 0
status = 'single'
while left > 0:
if left <= split_duration:
if status == 'single':
write_duration = left
current_offset = 0
else:
status = 'end'
write_duration = left + time_length
current_offset -= time_length
offset_inc = left
left = 0
else:
if status == 'start' or status == 'next':
status = 'next'
else:
status = 'start'
if status == 'start':
write_duration = split_duration
offset_inc = split_duration
else:
write_duration = split_duration + time_length
current_offset -= time_length
offset_inc = split_duration + time_length
left -= split_duration
metadata = {
'audio_filepath': file,
'duration': write_duration,
'label': label,
'text': '_',
'offset': current_offset,
}
res.append(metadata)
current_offset += offset_inc
except Exception as e:
err_file = "error.log"
with open(err_file, 'w') as fout:
fout.write(file + ":" + str(e))
return res
def main():
parser = ArgumentParser()
parser.add_argument("--inp_dir", type=str, required=True, help="(full path) folder of files to be processed")
parser.add_argument(
"--inp_list", type=str, help="(full path) a file contains NAME of files inside inp_dir to be processed"
)
parser.add_argument(
"--out_dir", type=str, default=".", help="[Optional](full path) location to store generated json file"
)
parser.add_argument("--manifest_name", type=str, required=True, help="name of generated json file")
parser.add_argument("--split_duration", type=int, required=True, help="max duration of each audio clip/line")
parser.add_argument(
"--time_length", type=float, default=0.63, help="[Optional] time length of segment, default is 0.63s"
)
parser.add_argument("--num_workers", type=int, default=4, help="[Optional] number of workers for multiprocessing")
args = parser.parse_args()
if not args.inp_list:
input_audios = []
for root, dirs, files in os.walk(args.inp_dir):
for basename in files:
if basename.endswith('.wav'):
filename = os.path.join(root, basename)
input_audios.append(filename)
else:
name_list = np.loadtxt(args.inp_list, dtype='str')
input_audios = [os.path.join(args.inp_dir, name + ".wav") for name in name_list]
print(f"Number of wav files to be processed: {len(input_audios)}")
output_path = os.path.join(args.out_dir, args.manifest_name + '.json')
print(f"Save generate manifest to {output_path}!")
if not os.path.exists(args.out_dir):
logging.info(f'Outdir {args.out_dir} does not exist. Creat directory.')
os.mkdir(args.out_dir)
if os.path.exists(output_path):
print(f"Manifest {output_path} exists! Overwriting")
os.remove(output_path)
p = multiprocessing.Pool(processes=args.num_workers)
args_func = {
'label': 'infer',
'split_duration': args.split_duration,
'time_length': args.time_length,
}
results = p.starmap(write_manifest, zip(input_audios, repeat(args_func)))
with open(output_path, 'a') as fout:
for res in results:
for r in res:
json.dump(r, fout)
fout.write('\n')
fout.flush()
p.close()
print(f"Done! Save to {output_path}")
if __name__ == '__main__':
main()