mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-17 00:23:14 +00:00
Introduce code from Rust monorepo. This is a mess
This commit is contained in:
@@ -0,0 +1,31 @@
|
||||
Dataset creation tool based on CTC-Segmentation
|
||||
-----------------------------------------------
|
||||
|
||||
This tool provides functionality to align long audio files and the corresponding transcripts into shorter fragments
|
||||
that are suitable for an Automatic Speech Recognition (ASR) model training.
|
||||
|
||||
More details could be found in [this tutorial](https://github.com/NVIDIA/NeMo/blob/main/tutorials/tools/CTC_Segmentation_Tutorial.ipynb).
|
||||
|
||||
The tool is based on the [CTC Segmentation](https://github.com/lumaku/ctc-segmentation):
|
||||
**CTC-Segmentation of Large Corpora for German End-to-end Speech Recognition**
|
||||
https://doi.org/10.1007/978-3-030-60276-5_27 or pre-print https://arxiv.org/abs/2007.09127
|
||||
|
||||
```
|
||||
@InProceedings{ctcsegmentation,
|
||||
author="K{\"u}rzinger, Ludwig
|
||||
and Winkelbauer, Dominik
|
||||
and Li, Lujun
|
||||
and Watzel, Tobias
|
||||
and Rigoll, Gerhard",
|
||||
editor="Karpov, Alexey
|
||||
and Potapova, Rodmonga",
|
||||
title="CTC-Segmentation of Large Corpora for German End-to-End Speech Recognition",
|
||||
booktitle="Speech and Computer",
|
||||
year="2020",
|
||||
publisher="Springer International Publishing",
|
||||
address="Cham",
|
||||
pages="267--278",
|
||||
abstract="Recent end-to-end Automatic Speech Recognition (ASR) systems demonstrated the ability to outperform conventional hybrid DNN/HMM ASR. Aside from architectural improvements in those systems, those models grew in terms of depth, parameters and model capacity. However, these models also require more training data to achieve comparable performance.",
|
||||
isbn="978-3-030-60276-5"
|
||||
}
|
||||
```
|
||||
@@ -0,0 +1,2 @@
|
||||
ctc_segmentation==1.1.0
|
||||
num2words
|
||||
@@ -0,0 +1,124 @@
|
||||
#!/bin/bash
|
||||
|
||||
# default values for optional arguments
|
||||
MIN_SCORE=-5
|
||||
CUT_PREFIX=0
|
||||
SCRIPTS_DIR="scripts"
|
||||
OFFSET=0
|
||||
LANGUAGE='eng' # 'eng', 'ru', 'other'
|
||||
MIN_SEGMENT_LEN=20
|
||||
MAX_SEGMENT_LEN=100
|
||||
ADDITIONAL_SPLIT_SYMBOLS=''
|
||||
AUDIO_FORMAT='.mp3'
|
||||
USE_NEMO_NORMALIZATION='True'
|
||||
|
||||
for ARG in "$@"
|
||||
do
|
||||
key=$(echo $ARG | cut -f1 -d=)
|
||||
value=$(echo $ARG | cut -f2 -d=)
|
||||
|
||||
if [[ $key == *"--"* ]]; then
|
||||
v="${key/--/}"
|
||||
declare $v="${value}"
|
||||
fi
|
||||
done
|
||||
|
||||
echo "MODEL_NAME_OR_PATH = $MODEL_NAME_OR_PATH"
|
||||
echo "DATA_DIR = $DATA_DIR"
|
||||
echo "OUTPUT_DIR = $OUTPUT_DIR"
|
||||
echo "MIN_SCORE = $MIN_SCORE"
|
||||
echo "CUT_PREFIX = $CUT_PREFIX"
|
||||
echo "SCRIPTS_DIR = $SCRIPTS_DIR"
|
||||
echo "OFFSET = $OFFSET"
|
||||
echo "LANGUAGE = $LANGUAGE"
|
||||
echo "MIN_SEGMENT_LEN = $MIN_SEGMENT_LEN"
|
||||
echo "MAX_SEGMENT_LEN = $MAX_SEGMENT_LEN"
|
||||
echo "ADDITIONAL_SPLIT_SYMBOLS = $ADDITIONAL_SPLIT_SYMBOLS"
|
||||
echo "AUDIO_FORMAT = $AUDIO_FORMAT"
|
||||
echo "USE_NEMO_NORMALIZATION = $USE_NEMO_NORMALIZATION"
|
||||
|
||||
if [[ -z $MODEL_NAME_OR_PATH ]] || [[ -z $DATA_DIR ]] || [[ -z $OUTPUT_DIR ]]; then
|
||||
echo "Usage: $(basename "$0")
|
||||
--MODEL_NAME_OR_PATH=[model_name_or_path]
|
||||
--DATA_DIR=[data_dir]
|
||||
--OUTPUT_DIR=[output_dir]
|
||||
--LANGUAGE=[language (Optional)]
|
||||
--OFFSET=[offset value (Optional)]
|
||||
--CUT_PREFIX=[cut prefix in sec (Optional)]
|
||||
--SCRIPTS_DIR=[scripts_dir_path (Optional)]
|
||||
--MIN_SEGMENT_LEN=[min number of characters of the text segment for alignment (Optional)]
|
||||
--MAX_SEGMENT_LEN=[max number of characters of the text segment for alignment (Optional)]
|
||||
--ADDITIONAL_SPLIT_SYMBOLS=[Additional symbols to use for
|
||||
sentence split if eos sentence split resulted in sequence longer than --max_length.
|
||||
Use '|' as a separator between symbols, for example: ';|:|' (Optional)]
|
||||
--AUDIO_FORMAT=[choose from ['.mp3', '.wav'], input audio files format
|
||||
--USE_NEMO_NORMALIZATION Set to 'True' to use NeMo Normalization tool to convert
|
||||
numbers from written to spoken format. By default num2words package will be used. (Optional)"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
NEMO_NORMALIZATION=""
|
||||
if [[ ${USE_NEMO_NORMALIZATION,,} == "true" ]]; then
|
||||
NEMO_NORMALIZATION="--use_nemo_normalization "
|
||||
fi
|
||||
|
||||
# STEP #1
|
||||
# Prepare text and audio data for segmentation
|
||||
python $SCRIPTS_DIR/prepare_data.py \
|
||||
--in_text=$DATA_DIR/text \
|
||||
--audio_dir=$DATA_DIR/audio \
|
||||
--audio_format=$AUDIO_FORMAT \
|
||||
--output_dir=$OUTPUT_DIR/processed/ \
|
||||
--language=$LANGUAGE \
|
||||
--cut_prefix=$CUT_PREFIX \
|
||||
--model=$MODEL_NAME_OR_PATH \
|
||||
--min_length=$MIN_SEGMENT_LEN \
|
||||
--max_length=$MAX_SEGMENT_LEN \
|
||||
--additional_split_symbols=$ADDITIONAL_SPLIT_SYMBOLS $NEMO_NORMALIZATION || exit
|
||||
|
||||
# STEP #2
|
||||
# Run CTC-segmenatation
|
||||
# one might want to perform alignment with various window sizes
|
||||
# note if the alignment with the initial window size isn't found, the window size will be double to re-attempt
|
||||
# alignment
|
||||
for WINDOW in 8000 12000
|
||||
do
|
||||
python $SCRIPTS_DIR/run_ctc_segmentation.py \
|
||||
--output_dir=$OUTPUT_DIR \
|
||||
--data=$OUTPUT_DIR/processed/ \
|
||||
--model=$MODEL_NAME_OR_PATH \
|
||||
--window_len $WINDOW || exit
|
||||
done
|
||||
|
||||
# STEP #3 (Optional)
|
||||
# Verify aligned segments only if multiple WINDOWs used in the Step #2)
|
||||
python $SCRIPTS_DIR/verify_segments.py \
|
||||
--base_dir=$OUTPUT_DIR || exit
|
||||
|
||||
# STEP #4
|
||||
# Cut the original audio files based on the alignments
|
||||
# (use --alignment=$OUTPUT_DIR/segments if only 1 WINDOW size was used in the Step #2)
|
||||
# Three manifests and corresponding clips folders will be created:
|
||||
# - high scored clips
|
||||
# - low scored clips
|
||||
# - deleted segments
|
||||
python $SCRIPTS_DIR/cut_audio.py \
|
||||
--output_dir=$OUTPUT_DIR \
|
||||
--model=$MODEL_NAME_OR_PATH \
|
||||
--alignment=$OUTPUT_DIR/verified_segments \
|
||||
--threshold=$MIN_SCORE \
|
||||
--offset=$OFFSET || exit
|
||||
|
||||
# STEP #5 (Optional)
|
||||
# If multiple audio files were segmented in the step #2, this step will aggregate manifests with high scored segments
|
||||
# for all audio files into all_manifest.json
|
||||
# Also a separate manifest with samples from across all high scored segments will be credated if --num_samples > 0
|
||||
# --num_samples samples will be taken from the beginning, end and the middle of the each audio file high score manifest
|
||||
# and will be stored at sample_manifest.json
|
||||
python $SCRIPTS_DIR/process_manifests.py \
|
||||
--output_dir=$OUTPUT_DIR \
|
||||
--manifests_dir=$OUTPUT_DIR/manifests/ \
|
||||
--num_samples 0
|
||||
|
||||
exit
|
||||
|
||||
@@ -0,0 +1,313 @@
|
||||
# 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 os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
from scipy.io import wavfile
|
||||
|
||||
from nemo.collections import asr as nemo_asr
|
||||
from nemo.collections.asr.metrics.wer import WER, word_error_rate
|
||||
|
||||
parser = argparse.ArgumentParser(description="Cut audio on the segments based on segments")
|
||||
parser.add_argument("--output_dir", default='output', type=str, help='Path to output directory')
|
||||
parser.add_argument(
|
||||
"--alignment",
|
||||
type=str,
|
||||
required=True,
|
||||
help='Path to a data directory with alignments or a single .txt file with timestamps - result of the ctc-segmentation',
|
||||
)
|
||||
parser.add_argument("--threshold", type=float, default=-5, help='Minimum score value accepted')
|
||||
parser.add_argument(
|
||||
'--model',
|
||||
type=str,
|
||||
default='QuartzNet15x5Base-En',
|
||||
help='Path to model checkpoint or pre-trained CTC-based ASR model name',
|
||||
)
|
||||
parser.add_argument('--offset', type=int, default=0, help='Offset in seconds')
|
||||
parser.add_argument("--batch_size", type=int, default=64, help='Batch size for inference')
|
||||
|
||||
|
||||
def add_transcript_to_manifest(
|
||||
manifest_original: str, manifest_updated: str, asr_model: nemo_asr.models.EncDecCTCModel, batch_size: int
|
||||
) -> None:
|
||||
"""
|
||||
Adds transcripts generated by the asr_model to the manifest_original.
|
||||
|
||||
Args:
|
||||
manifest_original: path to the manifest
|
||||
manifest_updated: path to the updated manifest with transcript included
|
||||
asr_model: CTC-based ASR model, for example, QuartzNet15x5Base-En
|
||||
batch_size: Batch size for asr_model inference
|
||||
"""
|
||||
transcripts = get_transcript(manifest_original, asr_model, batch_size)
|
||||
with open(manifest_original, 'r', encoding='utf8') as f:
|
||||
with open(manifest_updated, 'w', encoding='utf8') as f_updated:
|
||||
for i, line in enumerate(f):
|
||||
info = json.loads(line)
|
||||
info['transcript'] = transcripts[i].strip()
|
||||
info['WER'] = round(word_error_rate([info['transcript']], [info['text']]) * 100, 2)
|
||||
info['CER'] = round(word_error_rate([info['transcript']], [info['text']], use_cer=True) * 100, 2)
|
||||
json.dump(info, f_updated, ensure_ascii=False)
|
||||
f_updated.write('\n')
|
||||
|
||||
|
||||
def get_transcript(manifest_path: str, asr_model: nemo_asr.models.EncDecCTCModel, batch_size: int) -> List[str]:
|
||||
"""
|
||||
Returns transcripts for audio segments in the batch
|
||||
|
||||
Args:
|
||||
manifest_path: path to the manifest for inference
|
||||
asr_model: CTC-based ASR model, for example, QuartzNet15x5Base-En
|
||||
batch_size: batch size
|
||||
|
||||
Returns: hypotheses: transcripts for the audio segments
|
||||
"""
|
||||
# batch inference
|
||||
try:
|
||||
from torch.cuda.amp import autocast
|
||||
except ImportError:
|
||||
from contextlib import contextmanager
|
||||
|
||||
@contextmanager
|
||||
def autocast(enabled=None):
|
||||
yield
|
||||
|
||||
torch.set_grad_enabled(False)
|
||||
asr_model.setup_test_data(
|
||||
test_data_config={
|
||||
'sample_rate': 16000,
|
||||
'manifest_filepath': manifest_path,
|
||||
'labels': asr_model.decoder.vocabulary,
|
||||
'batch_size': batch_size,
|
||||
'normalize_transcripts': False,
|
||||
}
|
||||
)
|
||||
asr_model.eval()
|
||||
wer = WER(vocabulary=asr_model.decoder.vocabulary)
|
||||
hypotheses = []
|
||||
for test_batch in asr_model.test_dataloader():
|
||||
if torch.cuda.is_available():
|
||||
test_batch = [x.cuda() for x in test_batch]
|
||||
with autocast():
|
||||
log_probs, encoded_len, greedy_predictions = asr_model(
|
||||
input_signal=test_batch[0], input_signal_length=test_batch[1]
|
||||
)
|
||||
hypotheses += wer.ctc_decoder_predictions_tensor(greedy_predictions)
|
||||
del test_batch
|
||||
torch.cuda.empty_cache()
|
||||
return hypotheses
|
||||
|
||||
|
||||
def process_alignment(alignment_file: str, args):
|
||||
""" Cut original audio file into audio segments based on alignment_file
|
||||
|
||||
Args:
|
||||
alignment_file: path to the file with segmented text and corresponding time stamps.
|
||||
The first line of the file contains the path to the original audio file
|
||||
args: main script args
|
||||
"""
|
||||
if not os.path.exists(alignment_file):
|
||||
raise ValueError(f'{alignment_file} not found')
|
||||
|
||||
# read the segments, note the first line contains the path to the original audio
|
||||
segments = []
|
||||
ref_text_processed = []
|
||||
ref_text_no_preprocessing = []
|
||||
with open(alignment_file, 'r') as f:
|
||||
for line in f:
|
||||
line = line.split('|')
|
||||
# read audio file name from the first line
|
||||
if len(line) == 1:
|
||||
audio_file = line[0].strip()
|
||||
continue
|
||||
ref_text_processed.append(line[1].strip())
|
||||
ref_text_no_preprocessing.append(line[2].strip())
|
||||
line = line[0].split()
|
||||
segments.append((float(line[0]) + args.offset / 1000, float(line[1]) + args.offset / 1000, float(line[2])))
|
||||
|
||||
# cut the audio into segments
|
||||
# create manifest in /tmp directory first, then populate transcript values with the batch inference results
|
||||
# and save the final manifests at output_dir
|
||||
sampling_rate, signal = wavfile.read(audio_file)
|
||||
original_duration = len(signal) / sampling_rate
|
||||
print(f'Cutting {audio_file} based on {alignment_file}')
|
||||
print(f'Original duration: {round(original_duration)}s or ~{round(original_duration / 60)}min')
|
||||
|
||||
# create directories to store high score .wav audio fragments, low scored once, and deleted
|
||||
args.output_dir = os.path.abspath(args.output_dir)
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
fragments_dir = os.path.join(args.output_dir, "high_score_clips")
|
||||
del_fragments = os.path.join(args.output_dir, 'deleted_clips')
|
||||
low_score_segments_dir = os.path.join(args.output_dir, "low_score_clips")
|
||||
|
||||
os.makedirs(fragments_dir, exist_ok=True)
|
||||
os.makedirs(low_score_segments_dir, exist_ok=True)
|
||||
os.makedirs(del_fragments, exist_ok=True)
|
||||
|
||||
base_name = os.path.basename(alignment_file).replace('_segments.txt', '')
|
||||
high_score_manifest = f'{base_name}_high_score_manifest.json'
|
||||
low_score_manifest = f'{base_name}_low_score_manifest.json'
|
||||
del_manifest = f'{base_name}_del_manifest.json'
|
||||
manifests_dir = os.path.join(args.output_dir, 'manifests')
|
||||
os.makedirs(manifests_dir, exist_ok=True)
|
||||
tmp_dir = '/tmp'
|
||||
|
||||
low_score_dur = 0
|
||||
high_score_dur = 0
|
||||
with open(os.path.join(tmp_dir, high_score_manifest), 'w', encoding='utf8') as f:
|
||||
with open(os.path.join(tmp_dir, low_score_manifest), 'w', encoding='utf8') as low_score_f:
|
||||
for i, (st, end, score) in enumerate(segments):
|
||||
segment = signal[round(st * sampling_rate) : round(end * sampling_rate)]
|
||||
duration = len(segment) / sampling_rate
|
||||
if duration > 0:
|
||||
text_processed = ref_text_processed[i].strip()
|
||||
text_no_preprocessing = ref_text_no_preprocessing[i].strip()
|
||||
if score > args.threshold:
|
||||
high_score_dur += duration
|
||||
audio_filepath = os.path.join(fragments_dir, f'{base_name}_{i:04}.wav')
|
||||
file_to_write = f
|
||||
else:
|
||||
low_score_dur += duration
|
||||
audio_filepath = os.path.join(low_score_segments_dir, f'{base_name}_{i:04}.wav')
|
||||
file_to_write = low_score_f
|
||||
|
||||
wavfile.write(audio_filepath, sampling_rate, segment)
|
||||
|
||||
transcript = 'n/a' # asr_model.transcribe(paths2audio_files=[audio_filepath], batch_size=1)[0]
|
||||
info = {
|
||||
'audio_filepath': audio_filepath,
|
||||
'duration': duration,
|
||||
'text': text_processed,
|
||||
'text_no_preprocessing': text_no_preprocessing,
|
||||
'score': round(score, 2),
|
||||
'transcript': transcript.strip(),
|
||||
}
|
||||
json.dump(info, file_to_write, ensure_ascii=False)
|
||||
file_to_write.write('\n')
|
||||
|
||||
add_transcript_to_manifest(
|
||||
os.path.join(tmp_dir, high_score_manifest),
|
||||
os.path.join(manifests_dir, high_score_manifest),
|
||||
asr_model,
|
||||
args.batch_size,
|
||||
)
|
||||
add_transcript_to_manifest(
|
||||
os.path.join(tmp_dir, low_score_manifest),
|
||||
os.path.join(manifests_dir, low_score_manifest),
|
||||
asr_model,
|
||||
args.batch_size,
|
||||
)
|
||||
print(f'High score files duration: {round(high_score_dur)}s or ~{round(high_score_dur/60)}min at {manifests_dir}')
|
||||
print(
|
||||
f'Low score files duration: {round(low_score_dur)}s or ~{round(low_score_dur/60)}min saved at {manifests_dir}'
|
||||
)
|
||||
|
||||
# save deleted segments along with manifest
|
||||
deleted = []
|
||||
del_duration = 0
|
||||
begin = 0
|
||||
i = 0
|
||||
with open(os.path.join(manifests_dir, del_manifest), 'w', encoding='utf8') as f:
|
||||
for i, (st, end, _) in enumerate(segments):
|
||||
if st - begin > 0.01:
|
||||
segment = signal[int(begin * sampling_rate) : int(st * sampling_rate)]
|
||||
audio_filepath = os.path.join(del_fragments, f'del_{base_name}_{i:04}.wav')
|
||||
wavfile.write(audio_filepath, sampling_rate, segment)
|
||||
duration = len(segment) / sampling_rate
|
||||
del_duration += duration
|
||||
deleted.append((begin, st))
|
||||
transcript = 'n/a' # asr_model.transcribe(paths2audio_files=[audio_filepath], batch_size=1)[0]
|
||||
info = {'audio_filepath': audio_filepath, 'duration': duration, 'text': transcript}
|
||||
json.dump(info, f, ensure_ascii=False)
|
||||
f.write('\n')
|
||||
begin = end
|
||||
|
||||
segment = signal[int(begin * sampling_rate) :]
|
||||
audio_filepath = os.path.join(del_fragments, f'del_{base_name}_{i+1:04}.wav')
|
||||
wavfile.write(audio_filepath, sampling_rate, segment)
|
||||
duration = len(segment) / sampling_rate
|
||||
del_duration += duration
|
||||
deleted.append((begin, original_duration))
|
||||
|
||||
info = {'audio_filepath': audio_filepath, 'duration': duration, 'text': 'n/a'}
|
||||
json.dump(info, f)
|
||||
f.write('\n')
|
||||
|
||||
print(f'Saved DEL files duration: {round(del_duration)}s or ~ {round(del_duration/60)}min at {del_fragments}')
|
||||
missing_audio = original_duration - high_score_dur - del_duration - low_score_dur
|
||||
if missing_audio > 15:
|
||||
raise ValueError(f'{round(missing_audio)}s or ~ {round(missing_audio/60)}min is missing. Check the args')
|
||||
|
||||
stats = (
|
||||
args.output_dir,
|
||||
base_name,
|
||||
round(original_duration),
|
||||
round(high_score_dur),
|
||||
round(low_score_dur),
|
||||
round(del_duration),
|
||||
)
|
||||
return stats
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
start_time = time.time()
|
||||
args = parser.parse_args()
|
||||
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
if os.path.exists(args.model):
|
||||
asr_model = nemo_asr.models.EncDecCTCModel.restore_from(args.model, strict=False)
|
||||
elif args.model in nemo_asr.models.EncDecCTCModel.get_available_model_names():
|
||||
asr_model = nemo_asr.models.EncDecCTCModel.from_pretrained(args.model, strict=False)
|
||||
else:
|
||||
raise ValueError(
|
||||
f'Provide path to the pretrained checkpoint or choose from {nemo_asr.models.EncDecCTCModel.list_available_models()}'
|
||||
)
|
||||
|
||||
alignment_files = Path(args.alignment)
|
||||
if os.path.isdir(args.alignment):
|
||||
alignment_files = alignment_files.glob("*.txt")
|
||||
else:
|
||||
alignment_files = [Path(alignment_files)]
|
||||
|
||||
stats_file = os.path.join(args.output_dir, 'stats.tsv')
|
||||
with open(stats_file, 'w') as f:
|
||||
f.write('Folder\tSegment\tOriginal dur (s)\tHigh quality dur (s)\tLow quality dur (s)\tDeleted dur (s)\n')
|
||||
|
||||
high_score_dur = 0
|
||||
low_score_dur = 0
|
||||
del_duration = 0
|
||||
|
||||
for alignment_file in alignment_files:
|
||||
stats = process_alignment(alignment_file, args)
|
||||
high_score_dur += stats[-3]
|
||||
low_score_dur += stats[-2]
|
||||
del_duration += stats[-1]
|
||||
stats = '\t'.join([str(t) for t in stats]) + '\n'
|
||||
f.write(stats)
|
||||
|
||||
f.write(f'Total\t\t{round(high_score_dur)}\t{round(low_score_dur)}\t{del_duration}')
|
||||
|
||||
total_time = time.time() - start_time
|
||||
print(f'High score segments duration: {round(high_score_dur)}')
|
||||
print(f'Low score segments duration: {round(low_score_dur)}')
|
||||
print(f'Deleted segments duration: {round(del_duration)}')
|
||||
print(f'Stats saved at {stats_file}')
|
||||
print(f'Total execution time: ~{round(total_time / 60)}min')
|
||||
@@ -0,0 +1,69 @@
|
||||
# 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.
|
||||
|
||||
__all__ = ['LATIN_TO_RU', 'RU_ABBREVIATIONS', 'NUMBERS_TO_ENG', 'NUMBERS_TO_RU']
|
||||
|
||||
LATIN_TO_RU = {
|
||||
'a': 'а',
|
||||
'b': 'б',
|
||||
'c': 'к',
|
||||
'd': 'д',
|
||||
'e': 'е',
|
||||
'f': 'ф',
|
||||
'g': 'г',
|
||||
'h': 'х',
|
||||
'i': 'и',
|
||||
'j': 'ж',
|
||||
'k': 'к',
|
||||
'l': 'л',
|
||||
'm': 'м',
|
||||
'n': 'н',
|
||||
'o': 'о',
|
||||
'p': 'п',
|
||||
'q': 'к',
|
||||
'r': 'р',
|
||||
's': 'с',
|
||||
't': 'т',
|
||||
'u': 'у',
|
||||
'v': 'в',
|
||||
'w': 'в',
|
||||
'x': 'к',
|
||||
'y': 'у',
|
||||
'z': 'з',
|
||||
'à': 'а',
|
||||
'è': 'е',
|
||||
'é': 'е',
|
||||
'ß': 'в',
|
||||
'ä': 'а',
|
||||
'ö': 'о',
|
||||
'ü': 'у',
|
||||
'є': 'е',
|
||||
'ç': 'с',
|
||||
'ê': 'е',
|
||||
'ó': 'о',
|
||||
}
|
||||
|
||||
RU_ABBREVIATIONS = {
|
||||
' р.': ' рублей',
|
||||
' к.': ' копеек',
|
||||
' коп.': ' копеек',
|
||||
' копек.': ' копеек',
|
||||
' т.д.': ' так далее',
|
||||
' т. д.': ' так далее',
|
||||
' т.п.': ' тому подобное',
|
||||
' т. п.': ' тому подобное',
|
||||
' т.е.': ' то есть',
|
||||
' т. е.': ' то есть',
|
||||
' стр. ': ' страница ',
|
||||
}
|
||||
@@ -0,0 +1,412 @@
|
||||
# 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 multiprocessing
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
import regex
|
||||
import scipy.io.wavfile as wav
|
||||
from normalization_helpers import LATIN_TO_RU, RU_ABBREVIATIONS
|
||||
from num2words import num2words
|
||||
|
||||
from nemo.collections import asr as nemo_asr
|
||||
|
||||
try:
|
||||
from nemo_text_processing.text_normalization.normalize import Normalizer
|
||||
|
||||
NEMO_NORMALIZATION_AVAILABLE = True
|
||||
except (ModuleNotFoundError, ImportError):
|
||||
NEMO_NORMALIZATION_AVAILABLE = False
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser(description="Prepares text and audio files for segmentation")
|
||||
parser.add_argument("--in_text", type=str, default=None, help='Path to a text file or a directory with .txt files')
|
||||
parser.add_argument("--output_dir", type=str, required=True, help='Path to output directory')
|
||||
parser.add_argument("--audio_dir", type=str, help='Path to folder with .mp3 or .wav audio files')
|
||||
parser.add_argument(
|
||||
"--audio_format", type=str, default='.mp3', choices=['.mp3', '.wav'], help='Audio files format in --audio_dir'
|
||||
)
|
||||
parser.add_argument('--sample_rate', type=int, default=16000, help='Sampling rate used during ASR model training')
|
||||
parser.add_argument(
|
||||
'--language', type=str, default='eng', choices=['eng', 'ru', 'add other languages supported by num2words.']
|
||||
)
|
||||
parser.add_argument(
|
||||
'--cut_prefix', type=int, default=0, help='Number of seconds to cut from the beginning of the audio files.',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--model', type=str, default='QuartzNet15x5Base-En', help='Pre-trained model name or path to model checkpoint'
|
||||
)
|
||||
parser.add_argument('--min_length', type=int, default=0, help='Min number of chars of the text segment for alignment.')
|
||||
parser.add_argument(
|
||||
'--max_length', type=int, default=100, help='Max number of chars of the text segment for alignment.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--additional_split_symbols',
|
||||
type=str,
|
||||
default='',
|
||||
help='Additional symbols to use for \
|
||||
sentence split if eos sentence split resulted in sequence longer than --max_length. '
|
||||
'Use "|" as a separator between symbols, for example: ";|:|" ',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--use_nemo_normalization',
|
||||
action='store_true',
|
||||
help='Set to True to use NeMo Normalization tool to convert numbers from written to spoken format.',
|
||||
)
|
||||
|
||||
|
||||
def convert_audio(in_file: str, wav_file: str = None, sample_rate: int = 16000) -> str:
|
||||
"""
|
||||
Convert .mp3 to .wav and/or change sample rate if needed
|
||||
|
||||
Args:
|
||||
in_file: Path to .mp3 or .wav file
|
||||
sample_rate: Desired sample rate
|
||||
|
||||
Returns:
|
||||
path to .wav file
|
||||
"""
|
||||
print(f"Converting {in_file} to .wav format with sample rate {sample_rate}")
|
||||
if not os.path.exists(in_file):
|
||||
raise ValueError(f'{in_file} not found')
|
||||
if wav_file is None:
|
||||
wav_file = in_file.replace(os.path.splitext(in_file)[-1], f"_{sample_rate}.wav")
|
||||
|
||||
os.system(
|
||||
f'ffmpeg -i {in_file} -acodec pcm_s16le -ac 1 -af aresample=resampler=soxr -ar {sample_rate} {wav_file} -y'
|
||||
)
|
||||
return wav_file
|
||||
|
||||
|
||||
def process_audio(in_file: str, wav_file: str = None, cut_prefix: int = 0, sample_rate: int = 16000):
|
||||
"""Process audio file: .mp3 to .wav conversion and cut a few seconds from the beginning of the audio
|
||||
|
||||
Args:
|
||||
in_file: path to the .mp3 or .wav file for processing
|
||||
wav_file: path to the output .wav file
|
||||
cut_prefix: number of seconds to cut from the beginning of the audio file
|
||||
sample_rate: target sampling rate
|
||||
"""
|
||||
wav_audio = convert_audio(str(in_file), wav_file, sample_rate)
|
||||
|
||||
if cut_prefix > 0:
|
||||
# cut a few seconds of audio from the beginning
|
||||
sample_rate, signal = wav.read(wav_audio)
|
||||
wav.write(wav_audio, data=signal[cut_prefix * sample_rate :], rate=sample_rate)
|
||||
|
||||
|
||||
def split_text(
|
||||
in_file: str,
|
||||
out_file: str,
|
||||
vocabulary: List[str] = None,
|
||||
language='eng',
|
||||
remove_brackets=True,
|
||||
do_lower_case=True,
|
||||
min_length=0,
|
||||
max_length=100,
|
||||
additional_split_symbols=None,
|
||||
use_nemo_normalization=False,
|
||||
):
|
||||
"""
|
||||
Breaks down the in_file roughly into sentences. Each sentence will be on a separate line.
|
||||
Written form of the numbers will be converted to its spoken equivalent, OOV punctuation will be removed.
|
||||
|
||||
Args:
|
||||
in_file: path to original transcript
|
||||
out_file: path to the output file
|
||||
vocabulary: ASR model vocabulary
|
||||
language: text language
|
||||
remove_brackets: Set to True if square [] and curly {} brackets should be removed from text.
|
||||
Text in square/curly brackets often contains inaudible fragments like notes or translations
|
||||
do_lower_case: flag that determines whether to apply lower case to the in_file text
|
||||
min_length: Min number of chars of the text segment for alignment. Short segments will be combined to be
|
||||
at least min_length (not recommended for multi speaker data).
|
||||
max_length: Max number of chars of the text segment for alignment
|
||||
additional_split_symbols: Additional symbols to use for sentence split if eos sentence split resulted in
|
||||
segments longer than --max_length
|
||||
use_nemo_normalization: Set to True to use NeMo normalization tool to convert numbers from written to spoken
|
||||
format. Normalization using num2words will be applied afterwards to make sure there are no numbers present
|
||||
in the text, otherwise they will be replaced with a space and that could deteriorate segmentation results.
|
||||
"""
|
||||
|
||||
print(f'Splitting text in {in_file} into sentences.')
|
||||
with open(in_file, "r") as f:
|
||||
transcript = f.read()
|
||||
|
||||
# remove some symbols for better split into sentences
|
||||
transcript = (
|
||||
transcript.replace("\n", " ")
|
||||
.replace("\t", " ")
|
||||
.replace("…", "...")
|
||||
.replace("\\", " ")
|
||||
.replace("--", " -- ")
|
||||
.replace(". . .", "...")
|
||||
.replace("‘", "’")
|
||||
)
|
||||
# remove extra space
|
||||
transcript = re.sub(r' +', ' ', transcript)
|
||||
transcript = re.sub(r'(\.+)', '. ', transcript)
|
||||
|
||||
if remove_brackets:
|
||||
transcript = re.sub(r'(\[.*?\])', ' ', transcript)
|
||||
# remove text in curly brackets
|
||||
transcript = re.sub(r'(\{.*?\})', ' ', transcript)
|
||||
|
||||
lower_case_unicode = ''
|
||||
upper_case_unicode = ''
|
||||
if language == 'ru':
|
||||
lower_case_unicode = '\u0430-\u04FF'
|
||||
upper_case_unicode = '\u0410-\u042F'
|
||||
elif language not in ['ru', 'eng']:
|
||||
print(f'Consider using {language} unicode letters for better sentence split.')
|
||||
|
||||
# remove space in the middle of the lower case abbreviation to avoid splitting into separate sentences
|
||||
matches = re.findall(r'[a-z' + lower_case_unicode + ']\.\s[a-z' + lower_case_unicode + ']\.', transcript)
|
||||
for match in matches:
|
||||
transcript = transcript.replace(match, match.replace('. ', '.'))
|
||||
|
||||
# find phrases in quotes
|
||||
with_quotes = re.finditer(r'“[A-Za-z ?]+.*?”', transcript)
|
||||
sentences = []
|
||||
last_idx = 0
|
||||
for m in with_quotes:
|
||||
match = m.group()
|
||||
match_idx = m.start()
|
||||
if last_idx < match_idx:
|
||||
sentences.append(transcript[last_idx:match_idx])
|
||||
sentences.append(match)
|
||||
last_idx = m.end()
|
||||
sentences.append(transcript[last_idx:])
|
||||
sentences = [s.strip() for s in sentences if s.strip()]
|
||||
|
||||
# Read and split transcript by utterance (roughly, sentences)
|
||||
split_pattern = f"(?<!\w\.\w.)(?<![A-Z{upper_case_unicode}][a-z{lower_case_unicode}]\.)(?<![A-Z{upper_case_unicode}]\.)(?<=\.|\?|\!|\.”|\?”\!”)\s"
|
||||
|
||||
new_sentences = []
|
||||
for sent in sentences:
|
||||
new_sentences.extend(regex.split(split_pattern, sent))
|
||||
sentences = [s.strip() for s in new_sentences if s.strip()]
|
||||
|
||||
def additional_split(sentences, split_on_symbols, max_length):
|
||||
if len(split_on_symbols) == 0:
|
||||
return sentences
|
||||
|
||||
split_on_symbols = split_on_symbols.split('|')
|
||||
|
||||
def _split(sentences, delimiter, max_length):
|
||||
result = []
|
||||
for s in sentences:
|
||||
if len(s) <= max_length:
|
||||
result.append(s)
|
||||
else:
|
||||
split_sent = s.split(delimiter)
|
||||
result.extend([s + delimiter for s in split_sent[:-1]] + [split_sent[-1]])
|
||||
return result
|
||||
|
||||
another_sent_split = []
|
||||
for sent in sentences:
|
||||
split_sent = [sent]
|
||||
for delimiter in split_on_symbols:
|
||||
split_sent = _split(split_sent, delimiter + ' ', max_length)
|
||||
another_sent_split.extend(split_sent)
|
||||
|
||||
sentences = [s.strip() for s in another_sent_split if s.strip()]
|
||||
return sentences
|
||||
|
||||
sentences = additional_split(sentences, additional_split_symbols, max_length)
|
||||
|
||||
# check to make sure there will be no utterances for segmentation with only OOV symbols
|
||||
vocab_no_space_with_digits = set(vocabulary + [i for i in range(10)])
|
||||
vocab_no_space_with_digits.remove(' ')
|
||||
sentences = [s for s in sentences if len(vocab_no_space_with_digits.intersection(set(s))) > 0]
|
||||
|
||||
if min_length > 0:
|
||||
sentences_comb = []
|
||||
sentences_comb.append(sentences[0])
|
||||
# combines short sentence
|
||||
for i in range(1, len(sentences)):
|
||||
if len(sentences_comb[-1]) < min_length or len(sentences[i]) < min_length:
|
||||
sentences_comb[-1] += ' ' + sentences[i].strip()
|
||||
else:
|
||||
sentences_comb.append(sentences[i].strip())
|
||||
sentences = sentences_comb
|
||||
|
||||
sentences = [s.strip() for s in sentences if s.strip()]
|
||||
|
||||
# save split text with original punctuation and case
|
||||
out_dir, out_file_name = os.path.split(out_file)
|
||||
with open(os.path.join(out_dir, out_file_name[:-4] + '_with_punct.txt'), "w") as f:
|
||||
f.write("\n".join(sentences))
|
||||
|
||||
# substitute common abbreviations before applying lower case
|
||||
if language == 'ru':
|
||||
for k, v in RU_ABBREVIATIONS.items():
|
||||
sentences = [s.replace(k, v) for s in sentences]
|
||||
|
||||
if language == 'ru':
|
||||
# replace Latin characters with Russian
|
||||
for k, v in LATIN_TO_RU.items():
|
||||
sentences = [s.replace(k, v) for s in sentences]
|
||||
|
||||
if language == 'eng' and use_nemo_normalization:
|
||||
if not NEMO_NORMALIZATION_AVAILABLE:
|
||||
raise ValueError(f'NeMo normalization tool is not installed.')
|
||||
|
||||
print('Using NeMo normalization tool...')
|
||||
normalizer = Normalizer(input_case='cased')
|
||||
sentences_norm = normalizer.normalize_list(sentences, verbose=False)
|
||||
if len(sentences_norm) != len(sentences):
|
||||
raise ValueError(f'Normalization failed, number of sentences does not match.')
|
||||
|
||||
sentences = '\n'.join(sentences)
|
||||
|
||||
# replace numbers with num2words
|
||||
try:
|
||||
p = re.compile("\d+")
|
||||
new_text = ''
|
||||
match_end = 0
|
||||
for i, m in enumerate(p.finditer(sentences)):
|
||||
match = m.group()
|
||||
match_start = m.start()
|
||||
if i == 0:
|
||||
new_text = sentences[:match_start]
|
||||
else:
|
||||
new_text += sentences[match_end:match_start]
|
||||
match_end = m.end()
|
||||
new_text += sentences[match_start:match_end].replace(match, num2words(match, lang=language))
|
||||
new_text += sentences[match_end:]
|
||||
sentences = new_text
|
||||
except NotImplementedError:
|
||||
print(
|
||||
f'{language} might be missing in "num2words" package. Add required language to the choices for the'
|
||||
f'--language argument.'
|
||||
)
|
||||
raise
|
||||
|
||||
sentences = (
|
||||
sentences.replace("’", "'")
|
||||
.replace("»", '"')
|
||||
.replace("«", '"')
|
||||
.replace("\\", "")
|
||||
.replace("”", '"')
|
||||
.replace("„", '"')
|
||||
.replace("´", "'")
|
||||
.replace("-- --", "--")
|
||||
.replace("--", " -- ")
|
||||
.replace("’", "'")
|
||||
.replace('“', '"')
|
||||
.replace('“', '"')
|
||||
.replace("‘", "'")
|
||||
.replace('—', '-')
|
||||
.replace("- -", "--")
|
||||
.replace('`', "'")
|
||||
.replace(' !', '!')
|
||||
.replace(' ?', '?')
|
||||
.replace(' ,', ',')
|
||||
.replace(' .', '.')
|
||||
.replace(' ;', ';')
|
||||
.replace(' :', ':')
|
||||
.replace('!!', '!')
|
||||
.replace('--', '-')
|
||||
.replace('“', '"')
|
||||
.replace(', , ', ', ')
|
||||
.replace('=', '')
|
||||
)
|
||||
|
||||
allowed_punct = [',', '.', '?', '!', ':', ';', '-', '"', '(', ')']
|
||||
# clean up normalized text and keep only allowed_punct and ASR vocabulary (lower and upper case)
|
||||
symbols_to_remove = ''.join(
|
||||
set(sentences).difference(set(vocabulary + [s.upper() for s in vocabulary] + ['\n'] + allowed_punct))
|
||||
)
|
||||
sentences_norm = sentences.translate(''.maketrans(symbols_to_remove, len(symbols_to_remove) * ' '))
|
||||
|
||||
with open(os.path.join(out_dir, out_file_name[:-4] + '_with_punct_normalized.txt'), "w") as f:
|
||||
f.write(sentences_norm)
|
||||
|
||||
if do_lower_case:
|
||||
sentences = sentences.lower()
|
||||
|
||||
# remove all OOV symbols
|
||||
symbols_to_remove = ''.join(set(sentences).difference(set(vocabulary + ['\n'])))
|
||||
sentences = sentences.translate(''.maketrans(symbols_to_remove, len(symbols_to_remove) * ' '))
|
||||
|
||||
# remove extra space
|
||||
sentences = re.sub(r' +', ' ', sentences)
|
||||
with open(out_file, "w") as f:
|
||||
f.write(sentences)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = parser.parse_args()
|
||||
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
text_files = []
|
||||
if args.in_text:
|
||||
vocabulary = None
|
||||
if args.model is None:
|
||||
print(f"No model provided, vocabulary won't be used")
|
||||
elif os.path.exists(args.model):
|
||||
asr_model = nemo_asr.models.EncDecCTCModel.restore_from(args.model)
|
||||
vocabulary = asr_model.cfg.decoder.vocabulary
|
||||
elif args.model in nemo_asr.models.EncDecCTCModel.get_available_model_names():
|
||||
asr_model = nemo_asr.models.EncDecCTCModel.from_pretrained(args.model)
|
||||
vocabulary = asr_model.cfg.decoder.vocabulary
|
||||
else:
|
||||
raise ValueError(
|
||||
f'Provide path to the pretrained checkpoint or choose from {nemo_asr.models.EncDecCTCModel.get_available_model_names()}'
|
||||
)
|
||||
|
||||
if os.path.isdir(args.in_text):
|
||||
text_files = Path(args.in_text).glob(("*.txt"))
|
||||
else:
|
||||
text_files.append(Path(args.in_text))
|
||||
for text in text_files:
|
||||
base_name = os.path.basename(text)[:-4]
|
||||
out_text_file = os.path.join(args.output_dir, base_name + '.txt')
|
||||
|
||||
split_text(
|
||||
text,
|
||||
out_text_file,
|
||||
vocabulary=vocabulary,
|
||||
language=args.language,
|
||||
min_length=args.min_length,
|
||||
max_length=args.max_length,
|
||||
additional_split_symbols=args.additional_split_symbols,
|
||||
use_nemo_normalization=args.use_nemo_normalization,
|
||||
)
|
||||
print(f'Processed text saved at {args.output_dir}')
|
||||
|
||||
if args.audio_dir:
|
||||
if not os.path.exists(args.audio_dir):
|
||||
raise ValueError(f'{args.audio_dir} not found. "--audio_dir" should contain .mp3 or .wav files.')
|
||||
|
||||
audio_paths = list(Path(args.audio_dir).glob(f"*{args.audio_format}"))
|
||||
|
||||
workers = []
|
||||
for i in range(len(audio_paths)):
|
||||
wav_file = os.path.join(args.output_dir, audio_paths[i].name.replace(args.audio_format, ".wav"))
|
||||
worker = multiprocessing.Process(
|
||||
target=process_audio, args=(audio_paths[i], wav_file, args.cut_prefix, args.sample_rate),
|
||||
)
|
||||
workers.append(worker)
|
||||
worker.start()
|
||||
for w in workers:
|
||||
w.join()
|
||||
|
||||
print('Done.')
|
||||
@@ -0,0 +1,111 @@
|
||||
# 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.
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from random import randint
|
||||
from typing import List
|
||||
|
||||
parser = argparse.ArgumentParser(description="Combine manifests and extract samples")
|
||||
parser.add_argument(
|
||||
"--manifests_dir",
|
||||
default='_high_score_manifest.json',
|
||||
type=str,
|
||||
required=True,
|
||||
help='Path to directory with *_high_score_manifest.json files',
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_dir", default='output', required=True, help='Path to the directory to store final manifest'
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num_samples",
|
||||
default=0,
|
||||
type=int,
|
||||
help='The number of samples to get from the beginning, end and randomly from the middle of highly scored samples',
|
||||
)
|
||||
|
||||
|
||||
def get_samples(manifest: str, num_samples: int) -> List[str]:
|
||||
"""
|
||||
Samples num_samples from the beginning, end and randomly from the middle of the manifest
|
||||
|
||||
Args:
|
||||
manifest: path to manifest
|
||||
num_samples: number of samples to extract from the beginning, end and randomly from the middle of the manifest
|
||||
|
||||
Returns:
|
||||
num_samples: manifest lines
|
||||
"""
|
||||
samples = []
|
||||
with open(manifest, 'r', encoding='utf8') as f:
|
||||
lines = f.readlines()
|
||||
|
||||
# first n samples
|
||||
samples.extend(lines[:num_samples])
|
||||
# random in the middle
|
||||
for i in range(num_samples):
|
||||
samples.append(lines[randint(num_samples, len(lines) - num_samples)])
|
||||
# last n samples
|
||||
samples.extend(lines[-num_samples:])
|
||||
return samples
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = parser.parse_args()
|
||||
manifest_dir = Path(args.manifests_dir)
|
||||
manifests = sorted(list(manifest_dir.glob('*_high_score_manifest.json')))
|
||||
sample_manifest = os.path.join(args.output_dir, 'sample_manifest.json')
|
||||
all_manifest = os.path.join(args.output_dir, 'all_manifest.json')
|
||||
|
||||
if args.num_samples > 0:
|
||||
with open(sample_manifest, 'w', encoding='utf8') as f:
|
||||
for manifest in manifests:
|
||||
samples = get_samples(manifest, args.num_samples)
|
||||
for sample in samples:
|
||||
f.write(sample)
|
||||
|
||||
print(f'Sample manifest is saved at {sample_manifest}')
|
||||
|
||||
total_duration = 0
|
||||
with open(all_manifest, 'w') as out_f:
|
||||
for manifest in manifests:
|
||||
with open(manifest, 'r', encoding='utf8') as f:
|
||||
for line in f:
|
||||
info = json.loads(line)
|
||||
total_duration += info['duration']
|
||||
out_f.write(line)
|
||||
|
||||
print(f'Aggregated manifest is saved at {all_manifest}')
|
||||
print(f'Total files duration: ~{round(total_duration/60)} min or ~{round(total_duration/60/60)} hr')
|
||||
|
||||
sets = {'_high_score_manifest.json': 0, '_low_score_manifest.json': 0, '_del_manifest.json': 0}
|
||||
|
||||
for k in sets:
|
||||
files = manifest_dir.glob(f'*{k}')
|
||||
for manifest_f in files:
|
||||
with open(manifest_f, 'r') as f:
|
||||
for line in f:
|
||||
line = json.loads(line)
|
||||
sets[k] += line['duration']
|
||||
|
||||
total = 0
|
||||
for k in sets:
|
||||
sets[k] = round(sets[k] / 3600, 1)
|
||||
total += sets[k]
|
||||
name = k.split('_')[1].upper()
|
||||
print(f'Total {name} score duration: ~{sets[k]} hr')
|
||||
print(f'Total duration: ~{total} hr')
|
||||
print('Done.')
|
||||
@@ -0,0 +1,180 @@
|
||||
# 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
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import scipy.io.wavfile as wav
|
||||
import torch
|
||||
from utils import get_segments, listener_configurer, listener_process, worker_configurer, worker_process
|
||||
|
||||
import nemo.collections.asr as nemo_asr
|
||||
|
||||
parser = argparse.ArgumentParser(description="CTC Segmentation")
|
||||
parser.add_argument("--output_dir", default='output', type=str, help='Path to output directory')
|
||||
parser.add_argument(
|
||||
"--data",
|
||||
type=str,
|
||||
required=True,
|
||||
help='Path to directory with audio files and associated transcripts (same respective names only formats are '
|
||||
'different or path to wav file (transcript should have the same base name and be located in the same folder'
|
||||
'as the wav file.',
|
||||
)
|
||||
parser.add_argument('--window_len', type=int, default=8000, help='Window size for ctc segmentation algorithm')
|
||||
parser.add_argument('--no_parallel', action='store_true', help='Flag to disable parallel segmentation')
|
||||
parser.add_argument('--sample_rate', type=int, default=16000, help='Sampling rate')
|
||||
parser.add_argument(
|
||||
'--model', type=str, default='QuartzNet15x5Base-En', help='Path to model checkpoint or pre-trained model name',
|
||||
)
|
||||
parser.add_argument('--debug', action='store_true', help='Flag to enable debugging messages')
|
||||
|
||||
logger = logging.getLogger('ctc_segmentation') # use module name
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# setup logger
|
||||
log_dir = os.path.join(args.output_dir, 'logs')
|
||||
os.makedirs(log_dir, exist_ok=True)
|
||||
log_file = os.path.join(log_dir, f'ctc_segmentation_{args.window_len}.log')
|
||||
level = 'DEBUG' if args.debug else 'INFO'
|
||||
|
||||
logger = logging.getLogger('CTC')
|
||||
file_handler = logging.FileHandler(filename=log_file)
|
||||
stdout_handler = logging.StreamHandler(sys.stdout)
|
||||
handlers = [file_handler, stdout_handler]
|
||||
logging.basicConfig(handlers=handlers, level=level)
|
||||
|
||||
if os.path.exists(args.model):
|
||||
asr_model = nemo_asr.models.EncDecCTCModel.restore_from(args.model)
|
||||
elif args.model in nemo_asr.models.EncDecCTCModel.get_available_model_names():
|
||||
asr_model = nemo_asr.models.EncDecCTCModel.from_pretrained(args.model, strict=False)
|
||||
else:
|
||||
raise ValueError(
|
||||
f'{args.model} not a valid model name or path. Provide path to the pre-trained checkpoint '
|
||||
f'or choose from {nemo_asr.models.EncDecCTCModel.list_available_models()}'
|
||||
)
|
||||
|
||||
# extract ASR vocabulary and add blank symbol
|
||||
vocabulary = asr_model.cfg.decoder.vocabulary
|
||||
odim = len(vocabulary) + 1
|
||||
logging.debug(f'ASR Model vocabulary: {vocabulary}')
|
||||
|
||||
# add blank to vocab
|
||||
vocabulary = ["ε"] + list(vocabulary)
|
||||
data = Path(args.data)
|
||||
output_dir = Path(args.output_dir)
|
||||
|
||||
if os.path.isdir(data):
|
||||
audio_paths = data.glob("*.wav")
|
||||
data_dir = data
|
||||
else:
|
||||
audio_paths = [Path(data)]
|
||||
data_dir = Path(os.path.dirname(data))
|
||||
|
||||
all_log_probs = []
|
||||
all_transcript_file = []
|
||||
all_segment_file = []
|
||||
all_wav_paths = []
|
||||
segments_dir = os.path.join(args.output_dir, 'segments')
|
||||
os.makedirs(segments_dir, exist_ok=True)
|
||||
for path_audio in audio_paths:
|
||||
transcript_file = os.path.join(data_dir, path_audio.name.replace(".wav", ".txt"))
|
||||
segment_file = os.path.join(
|
||||
segments_dir, f"{args.window_len}_" + path_audio.name.replace(".wav", "_segments.txt")
|
||||
)
|
||||
try:
|
||||
sample_rate, signal = wav.read(path_audio)
|
||||
if sample_rate != args.sample_rate:
|
||||
raise ValueError(
|
||||
f'Sampling rate of the audio file {path_audio} doesn\'t match --sample_rate={args.sample_rate}'
|
||||
)
|
||||
except ValueError:
|
||||
logging.error(
|
||||
f"{path_audio} should be a .wav mono file with the sampling rate used for the ASR model training"
|
||||
f"specified with {args.sample_rate}."
|
||||
)
|
||||
raise
|
||||
|
||||
original_duration = len(signal) / sample_rate
|
||||
logging.debug(f'Duration: {original_duration}s, file_name: {path_audio}')
|
||||
log_probs = asr_model.transcribe(paths2audio_files=[str(path_audio)], batch_size=1, logprobs=True)[0]
|
||||
|
||||
# move blank values to the first column
|
||||
blank_col = log_probs[:, -1].reshape((log_probs.shape[0], 1))
|
||||
log_probs = np.concatenate((blank_col, log_probs[:, :-1]), axis=1)
|
||||
|
||||
all_log_probs.append(log_probs)
|
||||
all_segment_file.append(str(segment_file))
|
||||
all_transcript_file.append(str(transcript_file))
|
||||
all_wav_paths.append(path_audio)
|
||||
|
||||
del asr_model
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
if len(all_log_probs) == 0:
|
||||
raise ValueError(f'No valid audio files found at {args.data}')
|
||||
start_time = time.time()
|
||||
if args.no_parallel:
|
||||
for i in range(len(all_log_probs)):
|
||||
get_segments(
|
||||
all_log_probs[i],
|
||||
all_wav_paths[i],
|
||||
all_transcript_file[i],
|
||||
all_segment_file[i],
|
||||
vocabulary,
|
||||
args.window_len,
|
||||
)
|
||||
else:
|
||||
queue = multiprocessing.Queue(-1)
|
||||
|
||||
listener = multiprocessing.Process(target=listener_process, args=(queue, listener_configurer, log_file, level))
|
||||
listener.start()
|
||||
workers = []
|
||||
for i in range(len(all_log_probs)):
|
||||
worker = multiprocessing.Process(
|
||||
target=worker_process,
|
||||
args=(
|
||||
queue,
|
||||
worker_configurer,
|
||||
level,
|
||||
all_log_probs[i],
|
||||
all_wav_paths[i],
|
||||
all_transcript_file[i],
|
||||
all_segment_file[i],
|
||||
vocabulary,
|
||||
args.window_len,
|
||||
),
|
||||
)
|
||||
workers.append(worker)
|
||||
worker.start()
|
||||
for w in workers:
|
||||
w.join()
|
||||
queue.put_nowait(None)
|
||||
listener.join()
|
||||
|
||||
total_time = time.time() - start_time
|
||||
logger.info(f'Total execution time: ~{round(total_time/60)}min')
|
||||
logger.info(f'Saving logs to {log_file}')
|
||||
|
||||
if os.path.exists(log_file):
|
||||
with open(log_file, 'r') as f:
|
||||
lines = f.readlines()
|
||||
@@ -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 logging
|
||||
import logging.handlers
|
||||
import multiprocessing
|
||||
import os
|
||||
from pathlib import PosixPath
|
||||
from typing import List, Tuple, Union
|
||||
|
||||
import ctc_segmentation as cs
|
||||
import numpy as np
|
||||
|
||||
|
||||
def get_segments(
|
||||
log_probs: np.ndarray,
|
||||
path_wav: Union[PosixPath, str],
|
||||
transcript_file: Union[PosixPath, str],
|
||||
output_file: str,
|
||||
vocabulary: List[str],
|
||||
window_size: int = 8000,
|
||||
frame_duration_ms: int = 20,
|
||||
) -> None:
|
||||
"""
|
||||
Segments the audio into segments and saves segments timings to a file
|
||||
|
||||
Args:
|
||||
log_probs: Log probabilities for the original audio from an ASR model, shape T * |vocabulary|.
|
||||
values for blank should be at position 0
|
||||
path_wav: path to the audio .wav file
|
||||
transcript_file: path to
|
||||
output_file: path to the file to save timings for segments
|
||||
vocabulary: vocabulary used to train the ASR model, note blank is at position 0
|
||||
window_size: the length of each utterance (in terms of frames of the CTC outputs) fits into that window.
|
||||
frame_duration_ms: frame duration in ms
|
||||
"""
|
||||
config = cs.CtcSegmentationParameters()
|
||||
config.char_list = vocabulary
|
||||
config.min_window_size = window_size
|
||||
config.frame_duration_ms = frame_duration_ms
|
||||
config.blank = config.space
|
||||
config.subsampling_factor = 2
|
||||
|
||||
with open(transcript_file, "r") as f:
|
||||
text = f.readlines()
|
||||
text = [t.strip() for t in text if t.strip()]
|
||||
|
||||
# add corresponding original text without pre-processing
|
||||
transcript_file_no_preprocessing = transcript_file.replace('.txt', '_with_punct.txt')
|
||||
if not os.path.exists(transcript_file_no_preprocessing):
|
||||
raise ValueError(f'{transcript_file_no_preprocessing} not found.')
|
||||
|
||||
with open(transcript_file_no_preprocessing, "r") as f:
|
||||
text_no_preprocessing = f.readlines()
|
||||
text_no_preprocessing = [t.strip() for t in text_no_preprocessing if t.strip()]
|
||||
|
||||
# add corresponding normalized original text
|
||||
transcript_file_normalized = transcript_file.replace('.txt', '_with_punct_normalized.txt')
|
||||
if not os.path.exists(transcript_file_normalized):
|
||||
raise ValueError(f'{transcript_file_normalized} not found.')
|
||||
|
||||
with open(transcript_file_normalized, "r") as f:
|
||||
text_normalized = f.readlines()
|
||||
text_normalized = [t.strip() for t in text_normalized if t.strip()]
|
||||
|
||||
if len(text_no_preprocessing) != len(text):
|
||||
raise ValueError(f'{transcript_file} and {transcript_file_no_preprocessing} do not match')
|
||||
|
||||
if len(text_normalized) != len(text):
|
||||
raise ValueError(f'{transcript_file} and {transcript_file_normalized} do not match')
|
||||
|
||||
ground_truth_mat, utt_begin_indices = cs.prepare_text(config, text)
|
||||
logging.debug(f"Syncing {transcript_file}")
|
||||
logging.debug(
|
||||
f"Audio length {os.path.basename(path_wav)}: {log_probs.shape[0]}. "
|
||||
f"Text length {os.path.basename(transcript_file)}: {len(ground_truth_mat)}"
|
||||
)
|
||||
|
||||
timings, char_probs, char_list = cs.ctc_segmentation(config, log_probs, ground_truth_mat)
|
||||
segments = cs.determine_utterance_segments(config, utt_begin_indices, char_probs, timings, text)
|
||||
write_output(output_file, path_wav, segments, text, text_no_preprocessing, text_normalized)
|
||||
|
||||
|
||||
def write_output(
|
||||
out_path: str,
|
||||
path_wav: str,
|
||||
segments: List[Tuple[float]],
|
||||
text: str,
|
||||
text_no_preprocessing: str,
|
||||
text_normalized: str,
|
||||
stride: int = 2,
|
||||
):
|
||||
"""
|
||||
Write the segmentation output to a file
|
||||
|
||||
out_path: Path to output file
|
||||
path_wav: Path to the original audio file
|
||||
segments: Segments include start, end and alignment score
|
||||
text: Text used for alignment
|
||||
text_no_preprocessing: Reference txt without any pre-processing
|
||||
text_normalized: Reference text normalized
|
||||
stride: Stride applied to an ASR input
|
||||
"""
|
||||
# Uses char-wise alignments to get utterance-wise alignments and writes them into the given file
|
||||
with open(str(out_path), "w") as outfile:
|
||||
outfile.write(str(path_wav) + "\n")
|
||||
|
||||
for i, (start, end, score) in enumerate(segments):
|
||||
outfile.write(
|
||||
f'{start/stride} {end/stride} {score} | {text[i]} | {text_no_preprocessing[i]} | {text_normalized[i]}\n'
|
||||
)
|
||||
|
||||
|
||||
#####################
|
||||
# logging utils
|
||||
#####################
|
||||
def listener_configurer(log_file, level):
|
||||
root = logging.getLogger()
|
||||
h = logging.handlers.RotatingFileHandler(log_file, 'w')
|
||||
f = logging.Formatter('%(asctime)s %(processName)-10s %(name)s %(levelname)-8s %(message)s')
|
||||
h.setFormatter(f)
|
||||
ch = logging.StreamHandler()
|
||||
root.addHandler(h)
|
||||
root.setLevel(level)
|
||||
root.addHandler(ch)
|
||||
|
||||
|
||||
def listener_process(queue, configurer, log_file, level):
|
||||
configurer(log_file, level)
|
||||
while True:
|
||||
try:
|
||||
record = queue.get()
|
||||
if record is None: # We send this as a sentinel to tell the listener to quit.
|
||||
break
|
||||
logger = logging.getLogger(record.name)
|
||||
logger.setLevel(logging.INFO)
|
||||
logger.handle(record) # No level or filter logic applied - just do it!
|
||||
|
||||
except Exception:
|
||||
import sys
|
||||
import traceback
|
||||
|
||||
print('Problem:', file=sys.stderr)
|
||||
traceback.print_exc(file=sys.stderr)
|
||||
|
||||
|
||||
def worker_configurer(queue, level):
|
||||
h = logging.handlers.QueueHandler(queue) # Just the one handler needed
|
||||
root = logging.getLogger()
|
||||
root.addHandler(h)
|
||||
root.setLevel(level)
|
||||
|
||||
|
||||
def worker_process(
|
||||
queue, configurer, level, log_probs, path_wav, transcript_file, output_file, vocabulary, window_len
|
||||
):
|
||||
configurer(queue, level)
|
||||
name = multiprocessing.current_process().name
|
||||
innerlogger = logging.getLogger('worker')
|
||||
innerlogger.info(f'{name} is processing {path_wav}, window_len={window_len}')
|
||||
get_segments(log_probs, path_wav, transcript_file, output_file, vocabulary, window_len)
|
||||
innerlogger.info(f'{name} completed segmentation of {path_wav}, segments saved to {output_file}')
|
||||
@@ -0,0 +1,143 @@
|
||||
# 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.
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import pandas as pd
|
||||
|
||||
parser = argparse.ArgumentParser(description="Compare alignment segments generated with different window sizes")
|
||||
parser.add_argument(
|
||||
"--base_dir",
|
||||
default='output',
|
||||
type=str,
|
||||
required=True,
|
||||
help='Path to directory with "logs" and "segments" folders generated during the segmentation step',
|
||||
)
|
||||
|
||||
|
||||
def _process_log_file(log_file: str):
|
||||
processes = {}
|
||||
with open(log_file, 'r') as f:
|
||||
for line in f:
|
||||
if line[0].isdigit() and 'Process-' in line:
|
||||
process_id = re.search(r'Process-\d+', line)[0]
|
||||
window_size = re.findall(r'\d+', line)[-1]
|
||||
path = re.search(r'\/.*?\.[\w:]+', line)
|
||||
if 'is processing' in line:
|
||||
if process_id not in processes:
|
||||
processes[process_id] = {}
|
||||
processes[process_id]['Window Increased'] = False
|
||||
processes[process_id]['Start Window'] = window_size
|
||||
processes[process_id]['Final Window'] = window_size
|
||||
base_dir, base_name = os.path.split(path[0])
|
||||
processes[process_id]['Audio Dir'] = base_dir
|
||||
processes[process_id]['Audio File'] = base_name
|
||||
elif 'Increasing' in line:
|
||||
processes[process_id]['Final Window'] = window_size
|
||||
processes[process_id]['Window Increased'] = True
|
||||
elif 'completed' in line:
|
||||
processes[process_id]['Completed'] = True
|
||||
|
||||
df = pd.DataFrame.from_dict(
|
||||
processes,
|
||||
columns=['Audio Dir', 'Audio File', 'Start Window', 'Final Window', 'Window Increased', 'Completed'],
|
||||
orient='index',
|
||||
).reset_index()
|
||||
return df
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = parser.parse_args()
|
||||
log_dir = os.path.join(args.base_dir, 'logs')
|
||||
if not os.path.exists(log_dir):
|
||||
raise ValueError(f'"log" directory was not found at {args.base_dir}.')
|
||||
|
||||
log_files = Path(log_dir).glob('*.log')
|
||||
dfs = []
|
||||
for log in log_files:
|
||||
dfs.append(_process_log_file(log))
|
||||
|
||||
dfs = pd.concat(dfs)
|
||||
summary_file = os.path.join(log_dir, 'log_summary.csv')
|
||||
dfs.to_csv(summary_file, index=False)
|
||||
print(f'Log summary saved to {summary_file}')
|
||||
|
||||
segments_dir = os.path.join(args.base_dir, 'segments')
|
||||
if not os.path.exists(segments_dir):
|
||||
raise ValueError(f'"segments" directory was not found at {args.base_dir}.')
|
||||
|
||||
all_files = Path(segments_dir).glob('*_segments.txt')
|
||||
all_alignment_files = {}
|
||||
for file in all_files:
|
||||
base_name = re.sub(r'^\d+_', '', file.name)
|
||||
if base_name not in all_alignment_files:
|
||||
all_alignment_files[base_name] = []
|
||||
all_alignment_files[base_name].append(file)
|
||||
|
||||
verified_dir = os.path.join(args.base_dir, 'verified_segments')
|
||||
os.makedirs(verified_dir, exist_ok=True)
|
||||
|
||||
def readlines(file):
|
||||
with open(file, 'r') as f:
|
||||
lines = f.readlines()
|
||||
return lines
|
||||
|
||||
stats = {}
|
||||
for part, alignment_files in all_alignment_files.items():
|
||||
stats[part] = {}
|
||||
num_alignment_files = len(alignment_files)
|
||||
all_alignments = []
|
||||
for alignment in alignment_files:
|
||||
all_alignments.append(readlines(alignment))
|
||||
|
||||
with open(os.path.join(verified_dir, part), 'w') as f:
|
||||
num_segments = len(all_alignments[0])
|
||||
stats[part]['Original number of segments'] = num_segments
|
||||
stats[part]['Verified segments'] = 0
|
||||
stats[part]['Original Duration, min'] = 0
|
||||
stats[part]['Verified Duration, min'] = 0
|
||||
|
||||
for i in range(num_segments):
|
||||
line = all_alignments[0][i]
|
||||
valid_line = True
|
||||
if i == 0:
|
||||
duration = 0
|
||||
else:
|
||||
info = line.split('|')[0].split()
|
||||
duration = (float(info[1]) - float(info[0])) / 60
|
||||
stats[part]['Original Duration, min'] += duration
|
||||
for alignment in all_alignments:
|
||||
if line != alignment[i]:
|
||||
valid_line = False
|
||||
if valid_line:
|
||||
f.write(line)
|
||||
stats[part]['Verified segments'] += 1
|
||||
stats[part]['Verified Duration, min'] += duration
|
||||
|
||||
stats = pd.DataFrame.from_dict(stats, orient='index').reset_index()
|
||||
stats['Number dropped'] = stats['Original number of segments'] - stats['Verified segments']
|
||||
stats['Duration of dropped, min'] = round(stats['Original Duration, min'] - stats['Verified Duration, min'])
|
||||
stats['% dropped, min'] = round(stats['Duration of dropped, min'] / stats['Original number of segments'] * 100)
|
||||
stats['Misalignment present'] = stats['Number dropped'] > 0
|
||||
stats['Original Duration, min'] = round(stats['Original Duration, min'])
|
||||
stats['Verified Duration, min'] = round(stats['Verified Duration, min'])
|
||||
stats.loc["Total"] = stats.sum()
|
||||
|
||||
stats_file = os.path.join(args.base_dir, 'alignment_summary.csv')
|
||||
stats.to_csv(stats_file, index=False)
|
||||
print(stats)
|
||||
print(f'Alignment summary saved to {stats_file}')
|
||||
@@ -0,0 +1,26 @@
|
||||
Speech Data Explorer
|
||||
--------------------
|
||||
|
||||
[Dash](https://plotly.com/dash/)-based tool for interactive exploration of ASR/TTS datasets.
|
||||
|
||||
Features:
|
||||
- dataset's statistics (alphabet, vocabulary, duration-based histograms)
|
||||
- navigation across dataset (sorting, filtering)
|
||||
- inspection of individual utterances (waveform, spectrogram, audio player)
|
||||
- errors' analysis (Word Error Rate, Character Error Rate, Word Match Rate, Mean Word Accuracy, diff)
|
||||
|
||||
Please make sure that requirements are installed. Then run:
|
||||
```
|
||||
python data_explorer.py path_to_manifest.json
|
||||
```
|
||||
|
||||
JSON manifest file should contain the following fields:
|
||||
- "audio_filepath" (path to audio file)
|
||||
- "duration" (duration of the audio file in seconds)
|
||||
- "text" (reference transcript)
|
||||
|
||||
Errors' analysis requires "pred_text" (ASR transcript) for all utterances.
|
||||
|
||||
Any additional field will be parsed and displayed in 'Samples' tab.
|
||||
|
||||

|
||||
+717
@@ -0,0 +1,717 @@
|
||||
# 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 base64
|
||||
import datetime
|
||||
import difflib
|
||||
import io
|
||||
import json
|
||||
import math
|
||||
import operator
|
||||
import os
|
||||
import pickle
|
||||
from collections import defaultdict
|
||||
|
||||
import dash
|
||||
import dash_bootstrap_components as dbc
|
||||
import dash_core_components as dcc
|
||||
import dash_html_components as html
|
||||
import dash_table
|
||||
import diff_match_patch
|
||||
import editdistance
|
||||
import librosa
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import tqdm
|
||||
from dash.dependencies import Input, Output
|
||||
from dash.exceptions import PreventUpdate
|
||||
from plotly import express as px
|
||||
from plotly import graph_objects as go
|
||||
from plotly.subplots import make_subplots
|
||||
|
||||
# number of items in a table per page
|
||||
DATA_PAGE_SIZE = 10
|
||||
|
||||
# operators for filtering items
|
||||
filter_operators = [
|
||||
['ge ', '>='],
|
||||
['le ', '<='],
|
||||
['lt ', '<'],
|
||||
['gt ', '>'],
|
||||
['ne ', '!='],
|
||||
['eq ', '='],
|
||||
['contains '],
|
||||
]
|
||||
|
||||
# parse table filter queries
|
||||
def split_filter_part(filter_part):
|
||||
for operator_type in filter_operators:
|
||||
for op in operator_type:
|
||||
if op in filter_part:
|
||||
name_part, value_part = filter_part.split(op, 1)
|
||||
name = name_part[name_part.find('{') + 1 : name_part.rfind('}')]
|
||||
value_part = value_part.strip()
|
||||
v0 = value_part[0]
|
||||
if v0 == value_part[-1] and v0 in ("'", '"', '`'):
|
||||
value = value_part[1:-1].replace('\\' + v0, v0)
|
||||
else:
|
||||
try:
|
||||
value = float(value_part)
|
||||
except ValueError:
|
||||
value = value_part
|
||||
return name, operator_type[0].strip(), value
|
||||
return [None] * 3
|
||||
|
||||
|
||||
# standard command-line arguments parser
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description='Speech Data Explorer')
|
||||
parser.add_argument(
|
||||
'manifest', help='path to JSON manifest file',
|
||||
)
|
||||
parser.add_argument('--vocab', help='optional vocabulary to highlight OOV words')
|
||||
parser.add_argument('--port', default='8050', help='serving port for establishing connection')
|
||||
parser.add_argument(
|
||||
'--disable-caching-metrics', action='store_true', help='disable caching metrics for errors analysis'
|
||||
)
|
||||
parser.add_argument('--debug', '-d', action='store_true', help='enable debug mode')
|
||||
args = parser.parse_args()
|
||||
print(args)
|
||||
return args
|
||||
|
||||
|
||||
# load data from JSON manifest file
|
||||
def load_data(data_filename, disable_caching, vocab=None):
|
||||
if not disable_caching:
|
||||
pickle_filename = data_filename.split('.json')[0]
|
||||
json_mtime = datetime.datetime.fromtimestamp(os.path.getmtime(data_filename))
|
||||
timestamp = json_mtime.strftime('%Y%m%d_%H%M')
|
||||
pickle_filename += '_' + timestamp + '.pkl'
|
||||
if os.path.exists(pickle_filename):
|
||||
with open(pickle_filename, 'rb') as f:
|
||||
data, wer, cer, wmr, mwa, num_hours, vocabulary_data, alphabet, metrics_available = pickle.load(f)
|
||||
return data, wer, cer, wmr, mwa, num_hours, vocabulary_data, alphabet, metrics_available
|
||||
|
||||
if vocab is not None:
|
||||
# load external vocab
|
||||
vocabulary_ext = {}
|
||||
with open(vocab, 'r') as f:
|
||||
for line in f:
|
||||
if '\t' in line:
|
||||
# parse word from TSV file
|
||||
word = line.split('\t')[0]
|
||||
else:
|
||||
# assume each line contains just a single word
|
||||
word = line.strip()
|
||||
vocabulary_ext[word] = 1
|
||||
|
||||
data = []
|
||||
wer_dist = 0.0
|
||||
wer_count = 0
|
||||
cer_dist = 0.0
|
||||
cer_count = 0
|
||||
wmr_count = 0
|
||||
wer = 0
|
||||
cer = 0
|
||||
wmr = 0
|
||||
mwa = 0
|
||||
num_hours = 0
|
||||
vocabulary = defaultdict(lambda: 0)
|
||||
alphabet = set()
|
||||
match_vocab = defaultdict(lambda: 0)
|
||||
|
||||
sm = difflib.SequenceMatcher()
|
||||
metrics_available = False
|
||||
with open(data_filename, 'r', encoding='utf8') as f:
|
||||
for line in tqdm.tqdm(f):
|
||||
item = json.loads(line)
|
||||
num_words = len(item['text'].split())
|
||||
num_chars = len(item['text'])
|
||||
orig = item['text'].split()
|
||||
for word in orig:
|
||||
vocabulary[word] += 1
|
||||
for char in item['text']:
|
||||
alphabet.add(char)
|
||||
num_hours += item['duration']
|
||||
|
||||
if 'pred_text' in item:
|
||||
metrics_available = True
|
||||
pred = item['pred_text'].split()
|
||||
word_dist = editdistance.eval(orig, pred)
|
||||
char_dist = editdistance.eval(item['text'], item['pred_text'])
|
||||
wer_dist += word_dist
|
||||
cer_dist += char_dist
|
||||
wer_count += num_words
|
||||
cer_count += num_chars
|
||||
|
||||
sm.set_seqs(orig, pred)
|
||||
num_matches = 0
|
||||
for m in sm.get_matching_blocks():
|
||||
for word_idx in range(m[0], m[0] + m[2]):
|
||||
match_vocab[orig[word_idx]] += 1
|
||||
num_matches += 1
|
||||
wmr_count += num_matches
|
||||
|
||||
data.append(
|
||||
{
|
||||
'audio_filepath': item['audio_filepath'],
|
||||
'duration': round(item['duration'], 2),
|
||||
'num_words': num_words,
|
||||
'num_chars': num_chars,
|
||||
'word_rate': round(num_words / item['duration'], 2),
|
||||
'char_rate': round(num_chars / item['duration'], 2),
|
||||
'text': item['text'],
|
||||
}
|
||||
)
|
||||
if metrics_available:
|
||||
data[-1]['pred_text'] = item['pred_text']
|
||||
data[-1]['WER'] = round(word_dist / num_words * 100.0, 2)
|
||||
data[-1]['CER'] = round(char_dist / num_chars * 100.0, 2)
|
||||
data[-1]['WMR'] = round(num_matches / len(orig) * 100.0, 2)
|
||||
for k in item:
|
||||
if k not in data[-1]:
|
||||
data[-1][k] = item[k]
|
||||
|
||||
vocabulary_data = [{'word': word, 'count': vocabulary[word]} for word in vocabulary]
|
||||
if vocab is not None:
|
||||
for item in vocabulary_data:
|
||||
item['OOV'] = item['word'] not in vocabulary_ext
|
||||
|
||||
if metrics_available:
|
||||
wer = wer_dist / wer_count * 100.0
|
||||
cer = cer_dist / cer_count * 100.0
|
||||
wmr = wmr_count / wer_count * 100.0
|
||||
|
||||
acc_sum = 0
|
||||
for item in vocabulary_data:
|
||||
w = item['word']
|
||||
word_accuracy = match_vocab[w] / vocabulary[w] * 100.0
|
||||
acc_sum += word_accuracy
|
||||
item['accuracy'] = round(word_accuracy, 1)
|
||||
mwa = acc_sum / len(vocabulary_data)
|
||||
|
||||
num_hours /= 3600.0
|
||||
|
||||
if not disable_caching:
|
||||
with open(pickle_filename, 'wb') as f:
|
||||
pickle.dump(
|
||||
[data, wer, cer, wmr, mwa, num_hours, vocabulary_data, alphabet, metrics_available],
|
||||
f,
|
||||
pickle.HIGHEST_PROTOCOL,
|
||||
)
|
||||
|
||||
return data, wer, cer, wmr, mwa, num_hours, vocabulary_data, alphabet, metrics_available
|
||||
|
||||
|
||||
# plot histogram of specified field in data list
|
||||
def plot_histogram(data, key, label):
|
||||
fig = px.histogram(
|
||||
data_frame=[item[key] for item in data],
|
||||
nbins=50,
|
||||
log_y=True,
|
||||
labels={'value': label},
|
||||
opacity=0.5,
|
||||
color_discrete_sequence=['green'],
|
||||
height=200,
|
||||
)
|
||||
fig.update_layout(showlegend=False, margin=dict(l=0, r=0, t=0, b=0, pad=0))
|
||||
return fig
|
||||
|
||||
|
||||
def plot_word_accuracy(vocabulary_data):
|
||||
labels = ['Unrecognized', 'Sometimes recognized', 'Always recognized']
|
||||
counts = [0, 0, 0]
|
||||
for word in vocabulary_data:
|
||||
if word['accuracy'] == 0:
|
||||
counts[0] += 1
|
||||
elif word['accuracy'] < 100:
|
||||
counts[1] += 1
|
||||
else:
|
||||
counts[2] += 1
|
||||
colors = ['red', 'orange', 'green']
|
||||
|
||||
fig = go.Figure(
|
||||
data=[
|
||||
go.Bar(
|
||||
x=labels,
|
||||
y=counts,
|
||||
marker_color=colors,
|
||||
text=['{:.2%}'.format(count / sum(counts)) for count in counts],
|
||||
textposition='auto',
|
||||
)
|
||||
]
|
||||
)
|
||||
fig.update_layout(
|
||||
showlegend=False, margin=dict(l=0, r=0, t=0, b=0, pad=0), height=200, yaxis={'title_text': '#words'}
|
||||
)
|
||||
|
||||
return fig
|
||||
|
||||
|
||||
args = parse_args()
|
||||
print('Loading data...')
|
||||
data, wer, cer, wmr, mwa, num_hours, vocabulary, alphabet, metrics_available = load_data(
|
||||
args.manifest, args.disable_caching_metrics, args.vocab
|
||||
)
|
||||
print('Starting server...')
|
||||
app = dash.Dash(
|
||||
__name__,
|
||||
suppress_callback_exceptions=True,
|
||||
external_stylesheets=[dbc.themes.BOOTSTRAP],
|
||||
title=os.path.basename(args.manifest),
|
||||
)
|
||||
|
||||
|
||||
figure_duration = plot_histogram(data, 'duration', 'Duration (sec)')
|
||||
figure_num_words = plot_histogram(data, 'num_words', '#words')
|
||||
figure_num_chars = plot_histogram(data, 'num_chars', '#chars')
|
||||
figure_word_rate = plot_histogram(data, 'word_rate', '#words/sec')
|
||||
figure_char_rate = plot_histogram(data, 'char_rate', '#chars/sec')
|
||||
|
||||
if metrics_available:
|
||||
figure_wer = plot_histogram(data, 'WER', 'WER, %')
|
||||
figure_cer = plot_histogram(data, 'CER', 'CER, %')
|
||||
figure_wmr = plot_histogram(data, 'WMR', 'WMR, %')
|
||||
figure_word_acc = plot_word_accuracy(vocabulary)
|
||||
|
||||
stats_layout = [
|
||||
dbc.Row(dbc.Col(html.H5(children='Global Statistics'), className='text-secondary'), className='mt-3'),
|
||||
dbc.Row(
|
||||
[
|
||||
dbc.Col(html.Div('Number of hours', className='text-secondary'), width=3, className='border-right'),
|
||||
dbc.Col(html.Div('Number of utterances', className='text-secondary'), width=3, className='border-right'),
|
||||
dbc.Col(html.Div('Vocabulary size', className='text-secondary'), width=3, className='border-right'),
|
||||
dbc.Col(html.Div('Alphabet size', className='text-secondary'), width=3),
|
||||
],
|
||||
className='bg-light mt-2 rounded-top border-top border-left border-right',
|
||||
),
|
||||
dbc.Row(
|
||||
[
|
||||
dbc.Col(
|
||||
html.H5(
|
||||
'{:.2f} hours'.format(num_hours),
|
||||
className='text-center p-1',
|
||||
style={'color': 'green', 'opacity': 0.7},
|
||||
),
|
||||
width=3,
|
||||
className='border-right',
|
||||
),
|
||||
dbc.Col(
|
||||
html.H5(len(data), className='text-center p-1', style={'color': 'green', 'opacity': 0.7}),
|
||||
width=3,
|
||||
className='border-right',
|
||||
),
|
||||
dbc.Col(
|
||||
html.H5(
|
||||
'{} words'.format(len(vocabulary)),
|
||||
className='text-center p-1',
|
||||
style={'color': 'green', 'opacity': 0.7},
|
||||
),
|
||||
width=3,
|
||||
className='border-right',
|
||||
),
|
||||
dbc.Col(
|
||||
html.H5(
|
||||
'{} chars'.format(len(alphabet)),
|
||||
className='text-center p-1',
|
||||
style={'color': 'green', 'opacity': 0.7},
|
||||
),
|
||||
width=3,
|
||||
),
|
||||
],
|
||||
className='bg-light rounded-bottom border-bottom border-left border-right',
|
||||
),
|
||||
]
|
||||
if metrics_available:
|
||||
stats_layout += [
|
||||
dbc.Row(
|
||||
[
|
||||
dbc.Col(
|
||||
html.Div('Word Error Rate (WER), %', className='text-secondary'), width=3, className='border-right'
|
||||
),
|
||||
dbc.Col(
|
||||
html.Div('Character Error Rate (CER), %', className='text-secondary'),
|
||||
width=3,
|
||||
className='border-right',
|
||||
),
|
||||
dbc.Col(
|
||||
html.Div('Word Match Rate (WMR), %', className='text-secondary'),
|
||||
width=3,
|
||||
className='border-right',
|
||||
),
|
||||
dbc.Col(html.Div('Mean Word Accuracy, %', className='text-secondary'), width=3),
|
||||
],
|
||||
className='bg-light mt-2 rounded-top border-top border-left border-right',
|
||||
),
|
||||
dbc.Row(
|
||||
[
|
||||
dbc.Col(
|
||||
html.H5(
|
||||
'{:.2f}'.format(wer), className='text-center p-1', style={'color': 'green', 'opacity': 0.7},
|
||||
),
|
||||
width=3,
|
||||
className='border-right',
|
||||
),
|
||||
dbc.Col(
|
||||
html.H5(
|
||||
'{:.2f}'.format(cer), className='text-center p-1', style={'color': 'green', 'opacity': 0.7}
|
||||
),
|
||||
width=3,
|
||||
className='border-right',
|
||||
),
|
||||
dbc.Col(
|
||||
html.H5(
|
||||
'{:.2f}'.format(wmr), className='text-center p-1', style={'color': 'green', 'opacity': 0.7},
|
||||
),
|
||||
width=3,
|
||||
className='border-right',
|
||||
),
|
||||
dbc.Col(
|
||||
html.H5(
|
||||
'{:.2f}'.format(mwa), className='text-center p-1', style={'color': 'green', 'opacity': 0.7},
|
||||
),
|
||||
width=3,
|
||||
),
|
||||
],
|
||||
className='bg-light rounded-bottom border-bottom border-left border-right',
|
||||
),
|
||||
]
|
||||
stats_layout += [
|
||||
dbc.Row(dbc.Col(html.H5(children='Alphabet'), className='text-secondary'), className='mt-3'),
|
||||
dbc.Row(
|
||||
dbc.Col(html.Div('{}'.format(sorted(alphabet))),), className='mt-2 bg-light text-monospace rounded border'
|
||||
),
|
||||
dbc.Row(dbc.Col(html.H5('Duration (per utterance)'), className='text-secondary'), className='mt-3'),
|
||||
dbc.Row(dbc.Col(dcc.Graph(id='duration-graph', figure=figure_duration),),),
|
||||
dbc.Row(dbc.Col(html.H5('Number of words (per utterance)'), className='text-secondary'), className='mt-3'),
|
||||
dbc.Row(dbc.Col(dcc.Graph(id='num-words-graph', figure=figure_num_words),),),
|
||||
dbc.Row(dbc.Col(html.H5('Number of characters (per utterance)'), className='text-secondary'), className='mt-3'),
|
||||
dbc.Row(dbc.Col(dcc.Graph(id='num-chars-graph', figure=figure_num_chars),),),
|
||||
dbc.Row(dbc.Col(html.H5('Word rate (per utterance)'), className='text-secondary'), className='mt-3'),
|
||||
dbc.Row(dbc.Col(dcc.Graph(id='word-rate-graph', figure=figure_word_rate),),),
|
||||
dbc.Row(dbc.Col(html.H5('Character rate (per utterance)'), className='text-secondary'), className='mt-3'),
|
||||
dbc.Row(dbc.Col(dcc.Graph(id='char-rate-graph', figure=figure_char_rate),),),
|
||||
]
|
||||
|
||||
if metrics_available:
|
||||
stats_layout += [
|
||||
dbc.Row(dbc.Col(html.H5('WER (per utterance)'), className='text-secondary'), className='mt-3'),
|
||||
dbc.Row(dbc.Col(dcc.Graph(id='wer-graph', figure=figure_wer),),),
|
||||
dbc.Row(dbc.Col(html.H5('CER (per utterance)'), className='text-secondary'), className='mt-3'),
|
||||
dbc.Row(dbc.Col(dcc.Graph(id='cer-graph', figure=figure_cer),),),
|
||||
dbc.Row(dbc.Col(html.H5('WMR (per utterance)'), className='text-secondary'), className='mt-3'),
|
||||
dbc.Row(dbc.Col(dcc.Graph(id='wmr-graph', figure=figure_wmr),),),
|
||||
dbc.Row(dbc.Col(html.H5('Word accuracy distribution'), className='text-secondary'), className='mt-3'),
|
||||
dbc.Row(dbc.Col(dcc.Graph(id='word-acc-graph', figure=figure_word_acc),),),
|
||||
]
|
||||
|
||||
wordstable_columns = [{'name': 'Word', 'id': 'word'}, {'name': 'Count', 'id': 'count'}]
|
||||
if args.vocab is not None:
|
||||
wordstable_columns.append({'name': 'OOV', 'id': 'OOV'})
|
||||
if metrics_available:
|
||||
wordstable_columns.append({'name': 'Accuracy, %', 'id': 'accuracy'})
|
||||
|
||||
stats_layout += [
|
||||
dbc.Row(dbc.Col(html.H5('Vocabulary'), className='text-secondary'), className='mt-3'),
|
||||
dbc.Row(
|
||||
dbc.Col(
|
||||
dash_table.DataTable(
|
||||
id='wordstable',
|
||||
columns=wordstable_columns,
|
||||
filter_action='custom',
|
||||
filter_query='',
|
||||
sort_action='custom',
|
||||
sort_mode='single',
|
||||
page_action='custom',
|
||||
page_current=0,
|
||||
page_size=DATA_PAGE_SIZE,
|
||||
cell_selectable=False,
|
||||
page_count=math.ceil(len(vocabulary) / DATA_PAGE_SIZE),
|
||||
sort_by=[{'column_id': 'word', 'direction': 'asc'}],
|
||||
style_cell={'maxWidth': 0, 'textAlign': 'left'},
|
||||
style_header={'color': 'text-primary'},
|
||||
),
|
||||
),
|
||||
className='m-2',
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@app.callback(
|
||||
[Output('wordstable', 'data'), Output('wordstable', 'page_count')],
|
||||
[Input('wordstable', 'page_current'), Input('wordstable', 'sort_by'), Input('wordstable', 'filter_query')],
|
||||
)
|
||||
def update_wordstable(page_current, sort_by, filter_query):
|
||||
vocabulary_view = vocabulary
|
||||
filtering_expressions = filter_query.split(' && ')
|
||||
for filter_part in filtering_expressions:
|
||||
col_name, op, filter_value = split_filter_part(filter_part)
|
||||
|
||||
if op in ('eq', 'ne', 'lt', 'le', 'gt', 'ge'):
|
||||
vocabulary_view = [x for x in vocabulary_view if getattr(operator, op)(x[col_name], filter_value)]
|
||||
print(len(vocabulary_view), len(vocabulary))
|
||||
elif op == 'contains':
|
||||
vocabulary_view = [x for x in vocabulary_view if filter_value in str(x[col_name])]
|
||||
print(len(vocabulary_view), len(vocabulary))
|
||||
|
||||
if len(sort_by):
|
||||
col = sort_by[0]['column_id']
|
||||
descending = sort_by[0]['direction'] == 'desc'
|
||||
vocabulary_view = sorted(vocabulary_view, key=lambda x: x[col], reverse=descending)
|
||||
if page_current * DATA_PAGE_SIZE >= len(vocabulary_view):
|
||||
page_current = len(vocabulary_view) // DATA_PAGE_SIZE
|
||||
return [
|
||||
vocabulary_view[page_current * DATA_PAGE_SIZE : (page_current + 1) * DATA_PAGE_SIZE],
|
||||
math.ceil(len(vocabulary_view) / DATA_PAGE_SIZE),
|
||||
]
|
||||
|
||||
|
||||
samples_layout = [
|
||||
dbc.Row(dbc.Col(html.H5('Data'), className='text-secondary'), className='mt-3'),
|
||||
dbc.Row(
|
||||
dbc.Col(
|
||||
dash_table.DataTable(
|
||||
id='datatable',
|
||||
columns=[{'name': k.replace('_', ' '), 'id': k, 'hideable': True} for k in data[0]],
|
||||
filter_action='custom',
|
||||
filter_query='',
|
||||
sort_action='custom',
|
||||
sort_mode='single',
|
||||
sort_by=[],
|
||||
row_selectable='single',
|
||||
selected_rows=[0],
|
||||
page_action='custom',
|
||||
page_current=0,
|
||||
page_size=DATA_PAGE_SIZE,
|
||||
page_count=math.ceil(len(data) / DATA_PAGE_SIZE),
|
||||
style_cell={'overflow': 'hidden', 'textOverflow': 'ellipsis', 'maxWidth': 0, 'textAlign': 'center'},
|
||||
style_header={
|
||||
'color': 'text-primary',
|
||||
'text_align': 'center',
|
||||
'height': 'auto',
|
||||
'whiteSpace': 'normal',
|
||||
},
|
||||
css=[{'selector': '.dash-spreadsheet-menu', 'rule': 'position:absolute; bottom:-36px'}],
|
||||
),
|
||||
)
|
||||
),
|
||||
] + [
|
||||
dbc.Row(
|
||||
[
|
||||
dbc.Col(
|
||||
html.Div(children=k.replace('_', ' ')),
|
||||
width=2,
|
||||
className='mt-1 bg-light text-monospace text-break small rounded border',
|
||||
),
|
||||
dbc.Col(html.Div(id='_' + k), className='mt-1 bg-light text-monospace text-break small rounded border'),
|
||||
]
|
||||
)
|
||||
for k in data[0]
|
||||
]
|
||||
|
||||
if metrics_available:
|
||||
samples_layout += [
|
||||
dbc.Row(
|
||||
[
|
||||
dbc.Col(
|
||||
html.Div(children='diff'),
|
||||
width=2,
|
||||
className='mt-1 bg-light text-monospace text-break small rounded border',
|
||||
),
|
||||
dbc.Col(
|
||||
html.Iframe(
|
||||
id='_diff',
|
||||
sandbox='',
|
||||
srcDoc='',
|
||||
style={'border': 'none', 'width': '100%', 'height': '100%'},
|
||||
className='bg-light text-monospace text-break small',
|
||||
),
|
||||
className='mt-1 bg-light text-monospace text-break small rounded border',
|
||||
),
|
||||
]
|
||||
)
|
||||
]
|
||||
samples_layout += [
|
||||
dbc.Row(dbc.Col(html.Audio(id='player', controls=True),), className='mt-3 '),
|
||||
dbc.Row(dbc.Col(dcc.Graph(id='signal-graph')), className='mt-3'),
|
||||
]
|
||||
|
||||
|
||||
@app.callback(
|
||||
[Output('datatable', 'data'), Output('datatable', 'page_count')],
|
||||
[Input('datatable', 'page_current'), Input('datatable', 'sort_by'), Input('datatable', 'filter_query')],
|
||||
)
|
||||
def update_datatable(page_current, sort_by, filter_query):
|
||||
data_view = data
|
||||
filtering_expressions = filter_query.split(' && ')
|
||||
for filter_part in filtering_expressions:
|
||||
col_name, op, filter_value = split_filter_part(filter_part)
|
||||
|
||||
if op in ('eq', 'ne', 'lt', 'le', 'gt', 'ge'):
|
||||
data_view = [x for x in data_view if getattr(operator, op)(x[col_name], filter_value)]
|
||||
elif op == 'contains':
|
||||
data_view = [x for x in data_view if filter_value in str(x[col_name])]
|
||||
|
||||
if len(sort_by):
|
||||
col = sort_by[0]['column_id']
|
||||
descending = sort_by[0]['direction'] == 'desc'
|
||||
data_view = sorted(data_view, key=lambda x: x[col], reverse=descending)
|
||||
if page_current * DATA_PAGE_SIZE >= len(data_view):
|
||||
page_current = len(data_view) // DATA_PAGE_SIZE
|
||||
return [
|
||||
data_view[page_current * DATA_PAGE_SIZE : (page_current + 1) * DATA_PAGE_SIZE],
|
||||
math.ceil(len(data_view) / DATA_PAGE_SIZE),
|
||||
]
|
||||
|
||||
|
||||
app.layout = html.Div(
|
||||
[
|
||||
dcc.Location(id='url', refresh=False),
|
||||
dbc.NavbarSimple(
|
||||
children=[
|
||||
dbc.NavItem(dbc.NavLink('Statistics', id='stats_link', href='/', active=True)),
|
||||
dbc.NavItem(dbc.NavLink('Samples', id='samples_link', href='/samples')),
|
||||
],
|
||||
brand='Speech Data Explorer',
|
||||
sticky='top',
|
||||
color='green',
|
||||
dark=True,
|
||||
),
|
||||
dbc.Container(id='page-content'),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@app.callback(
|
||||
[Output('page-content', 'children'), Output('stats_link', 'active'), Output('samples_link', 'active')],
|
||||
[Input('url', 'pathname')],
|
||||
)
|
||||
def nav_click(url):
|
||||
if url == '/samples':
|
||||
return [samples_layout, False, True]
|
||||
else:
|
||||
return [stats_layout, True, False]
|
||||
|
||||
|
||||
@app.callback(
|
||||
[Output('_' + k, 'children') for k in data[0]], [Input('datatable', 'selected_rows'), Input('datatable', 'data')]
|
||||
)
|
||||
def show_item(idx, data):
|
||||
if len(idx) == 0:
|
||||
raise PreventUpdate
|
||||
return [data[idx[0]][k] for k in data[0]]
|
||||
|
||||
|
||||
@app.callback(Output('_diff', 'srcDoc'), [Input('datatable', 'selected_rows'), Input('datatable', 'data')])
|
||||
def show_diff(idx, data):
|
||||
if len(idx) == 0:
|
||||
raise PreventUpdate
|
||||
|
||||
orig_words = data[idx[0]]['text']
|
||||
orig_words = '\n'.join(orig_words.split()) + '\n'
|
||||
|
||||
pred_words = data[idx[0]]['pred_text']
|
||||
pred_words = '\n'.join(pred_words.split()) + '\n'
|
||||
|
||||
diff = diff_match_patch.diff_match_patch()
|
||||
diff.Diff_Timeout = 0
|
||||
orig_enc, pred_enc, enc = diff.diff_linesToChars(orig_words, pred_words)
|
||||
diffs = diff.diff_main(orig_enc, pred_enc, False)
|
||||
diff.diff_charsToLines(diffs, enc)
|
||||
diffs_post = []
|
||||
for d in diffs:
|
||||
diffs_post.append((d[0], d[1].replace('\n', ' ')))
|
||||
|
||||
diff_html = diff.diff_prettyHtml(diffs_post)
|
||||
|
||||
return diff_html
|
||||
|
||||
|
||||
@app.callback(Output('signal-graph', 'figure'), [Input('datatable', 'selected_rows'), Input('datatable', 'data')])
|
||||
def plot_signal(idx, data):
|
||||
if len(idx) == 0:
|
||||
raise PreventUpdate
|
||||
figs = make_subplots(rows=2, cols=1, subplot_titles=('Waveform', 'Spectrogram'))
|
||||
try:
|
||||
filename = data[idx[0]]['audio_filepath']
|
||||
audio, fs = librosa.load(filename, sr=None)
|
||||
if 'offset' in data[idx[0]]:
|
||||
audio = audio[
|
||||
int(data[idx[0]]['offset'] * fs) : int((data[idx[0]]['offset'] + data[idx[0]]['duration']) * fs)
|
||||
]
|
||||
time_stride = 0.01
|
||||
hop_length = int(fs * time_stride)
|
||||
n_fft = 512
|
||||
# linear scale spectrogram
|
||||
s = librosa.stft(y=audio, n_fft=n_fft, hop_length=hop_length)
|
||||
s_db = librosa.power_to_db(np.abs(s) ** 2, ref=np.max, top_db=100)
|
||||
figs.add_trace(
|
||||
go.Scatter(
|
||||
x=np.arange(audio.shape[0]) / fs,
|
||||
y=audio,
|
||||
line={'color': 'green'},
|
||||
name='Waveform',
|
||||
hovertemplate='Time: %{x:.2f} s<br>Amplitude: %{y:.2f}<br><extra></extra>',
|
||||
),
|
||||
row=1,
|
||||
col=1,
|
||||
)
|
||||
figs.add_trace(
|
||||
go.Heatmap(
|
||||
z=s_db,
|
||||
colorscale=[[0, 'rgb(30,62,62)'], [0.5, 'rgb(30,128,128)'], [1, 'rgb(30,255,30)'],],
|
||||
colorbar=dict(yanchor='middle', lenmode='fraction', y=0.2, len=0.5, ticksuffix=' dB'),
|
||||
dx=time_stride,
|
||||
dy=fs / n_fft / 1000,
|
||||
name='Spectrogram',
|
||||
hovertemplate='Time: %{x:.2f} s<br>Frequency: %{y:.2f} kHz<br>Magnitude: %{z:.2f} dB<extra></extra>',
|
||||
),
|
||||
row=2,
|
||||
col=1,
|
||||
)
|
||||
figs.update_layout({'margin': dict(l=0, r=0, t=20, b=0, pad=0), 'height': 500})
|
||||
figs.update_xaxes(title_text='Time, s', row=1, col=1)
|
||||
figs.update_yaxes(title_text='Amplitude', row=1, col=1)
|
||||
figs.update_xaxes(title_text='Time, s', row=2, col=1)
|
||||
figs.update_yaxes(title_text='Frequency, kHz', row=2, col=1)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return figs
|
||||
|
||||
|
||||
@app.callback(Output('player', 'src'), [Input('datatable', 'selected_rows'), Input('datatable', 'data')])
|
||||
def update_player(idx, data):
|
||||
if len(idx) == 0:
|
||||
raise PreventUpdate
|
||||
try:
|
||||
filename = data[idx[0]]['audio_filepath']
|
||||
signal, sr = librosa.load(filename, sr=None)
|
||||
if 'offset' in data[idx[0]]:
|
||||
signal = signal[
|
||||
int(data[idx[0]]['offset'] * sr) : int((data[idx[0]]['offset'] + data[idx[0]]['duration']) * sr)
|
||||
]
|
||||
with io.BytesIO() as buf:
|
||||
# convert to PCM .wav
|
||||
sf.write(buf, signal, sr, format='WAV')
|
||||
buf.seek(0)
|
||||
encoded = base64.b64encode(buf.read())
|
||||
return 'data:audio/wav;base64,{}'.format(encoded.decode())
|
||||
except Exception:
|
||||
return ''
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
app.run_server(host='0.0.0.0', port=args.port, debug=args.debug)
|
||||
@@ -0,0 +1,10 @@
|
||||
librosa
|
||||
dash
|
||||
dash_html_components
|
||||
dash_bootstrap_components
|
||||
plotly
|
||||
dash_core_components
|
||||
dash_table
|
||||
numpy
|
||||
diff_match_patch
|
||||
editdistance
|
||||
BIN
Binary file not shown.
|
After Width: | Height: | Size: 1.1 MiB |
@@ -0,0 +1,38 @@
|
||||
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
|
||||
# Copyright 2015 and onwards Google, Inc.
|
||||
#
|
||||
# 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.
|
||||
|
||||
# Dockerfile for C++ (inverse) text normalization backend Sparrowhawk https://github.com/google/sparrowhawk
|
||||
|
||||
# set base image (host OS)
|
||||
FROM conda/miniconda3
|
||||
|
||||
# set the working directory in the container
|
||||
WORKDIR /workspace
|
||||
|
||||
# install dependencies
|
||||
RUN conda update -n base -c defaults conda -y
|
||||
RUN conda install conda-build -y
|
||||
RUN apt-get update && apt-get install -y --reinstall build-essential && apt-get upgrade -y && apt-get install -y git && apt-get install make
|
||||
RUN git clone https://github.com/google/re2
|
||||
RUN cd re2 && make && make install
|
||||
RUN apt-get install build-essential -y && apt-get install wget -y
|
||||
RUN wget https://github.com/protocolbuffers/protobuf/releases/download/v2.5.0/protobuf-2.5.0.tar.gz
|
||||
RUN tar xzvf protobuf-2.5.0.tar.gz
|
||||
RUN cd protobuf-2.5.0 && ./configure && make && make install && ldconfig
|
||||
RUN conda install -c conda-forge thrax=1.3.4 -y
|
||||
RUN git clone https://github.com/yzhang123/sparrowhawk.git
|
||||
RUN cd sparrowhawk && git checkout test && apt-get install -y autoconf && bash autoreconf && ./configure && make && make install && ldconfig
|
||||
RUN git clone https://github.com/kward/shunit2.git
|
||||
RUN echo "DONE"
|
||||
@@ -0,0 +1,10 @@
|
||||
**Text Processing Deployment**
|
||||
=========================================
|
||||
|
||||
Introduction
|
||||
------------
|
||||
|
||||
This folder provides scripts to deploy WFST-based grammars in `NeMo Text Processing <https://github.com/NVIDIA/NeMo/blob/main/nemo_text_processing>`_ for
|
||||
for production.
|
||||
|
||||
See `documentation <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/>`_ for details.
|
||||
@@ -0,0 +1,18 @@
|
||||
# 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.
|
||||
|
||||
|
||||
#!/bin/bash
|
||||
|
||||
docker build . --rm -t sparrowhawk
|
||||
@@ -0,0 +1,44 @@
|
||||
# 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.
|
||||
|
||||
#!/bin/bash
|
||||
|
||||
MODE=${1:-"export"}
|
||||
SCRIPT_DIR=$(cd $(dirname $0); pwd)
|
||||
: ${CLASSIFY_DIR:="$SCRIPT_DIR/../classify"}
|
||||
: ${VERBALIZE_DIR:="$SCRIPT_DIR/../verbalize"}
|
||||
: ${CMD:=${2:-"/bin/bash"}}
|
||||
|
||||
MOUNTS=""
|
||||
MOUNTS+=" -v $CLASSIFY_DIR:/workspace/sparrowhawk/documentation/grammars/en_toy/classify"
|
||||
MOUNTS+=" -v $VERBALIZE_DIR:/workspace/sparrowhawk/documentation/grammars/en_toy/verbalize"
|
||||
|
||||
WORK_DIR="/workspace/sparrowhawk/documentation/grammars"
|
||||
if [[ $MODE == "test_tn_grammars" ]]; then
|
||||
CMD="bash test_sparrowhawk_normalization.sh"
|
||||
WORK_DIR="/workspace/tests"
|
||||
elif [[ $MODE == "test_itn_grammars" ]]; then
|
||||
CMD="bash test_sparrowhawk_inverse_text_normalization.sh"
|
||||
WORK_DIR="/workspace/tests"
|
||||
fi
|
||||
|
||||
echo $MOUNTS
|
||||
docker run -it --rm \
|
||||
--shm-size=4g \
|
||||
--ulimit memlock=-1 \
|
||||
--ulimit stack=67108864 \
|
||||
$MOUNTS \
|
||||
-v $SCRIPT_DIR/../../../tests/nemo_text_processing/:/workspace/tests \
|
||||
-w $WORK_DIR \
|
||||
sparrowhawk:latest $CMD
|
||||
@@ -0,0 +1,59 @@
|
||||
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
|
||||
# Copyright 2015 and onwards Google, Inc.
|
||||
#
|
||||
# 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 script compiles and exports WFST-grammars from nemo_text_processing, builds C++ production backend Sparrowhawk (https://github.com/google/sparrowhawk) in docker,
|
||||
# plugs grammars into Sparrowhawk and returns prompt inside docker.
|
||||
# For inverse text normalization run:
|
||||
# bash export_grammars.sh --GRAMMARS=itn_grammars
|
||||
# echo "two dollars fifty" | ../../src/bin/normalizer_main --config=sparrowhawk_configuration.ascii_proto
|
||||
# For text normalization run:
|
||||
# bash export_grammars.sh --GRAMMARS=tn_grammars
|
||||
# echo "\$2.5" | ../../src/bin/normalizer_main --config=sparrowhawk_configuration.ascii_proto
|
||||
#
|
||||
# To test TN grammars, run:
|
||||
# bash export_grammars.sh --GRAMMARS=tn_grammars --MODE=test
|
||||
#
|
||||
# To test ITN grammars, run:
|
||||
# bash export_grammars.sh --GRAMMARS=itn_grammars --MODE=test
|
||||
|
||||
GRAMMARS="itn_grammars" # tn_grammars
|
||||
INPUT_CASE="cased" # lower_cased, only for tn_grammars
|
||||
MODE=""
|
||||
|
||||
for ARG in "$@"
|
||||
do
|
||||
key=$(echo $ARG | cut -f1 -d=)
|
||||
value=$(echo $ARG | cut -f2 -d=)
|
||||
|
||||
if [[ $key == *"--"* ]]; then
|
||||
v="${key/--/}"
|
||||
declare $v="${value}"
|
||||
fi
|
||||
done
|
||||
|
||||
echo "GRAMMARS = $GRAMMARS"
|
||||
echo "MODE = $MODE"
|
||||
|
||||
python pynini_export.py --output_dir=. --grammars=${GRAMMARS} --input_case=${INPUT_CASE} || exit 1
|
||||
find . -name "Makefile" -type f -delete
|
||||
bash docker/build.sh
|
||||
|
||||
if [[ $MODE == "test" ]]; then
|
||||
MODE=${MODE}_${GRAMMARS}
|
||||
fi
|
||||
|
||||
bash docker/launch.sh $MODE
|
||||
@@ -0,0 +1,111 @@
|
||||
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
|
||||
# Copyright 2015 and onwards Google, Inc.
|
||||
#
|
||||
# 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 os
|
||||
import time
|
||||
from argparse import ArgumentParser
|
||||
from typing import Dict
|
||||
|
||||
from nemo_text_processing.inverse_text_normalization.taggers.tokenize_and_classify import ClassifyFst as ITNClassifyFst
|
||||
from nemo_text_processing.inverse_text_normalization.verbalizers.verbalize import VerbalizeFst as ITNVerbalizeFst
|
||||
from nemo_text_processing.text_normalization.taggers.tokenize_and_classify import ClassifyFst as TNClassifyFst
|
||||
from nemo_text_processing.text_normalization.verbalizers.verbalize import VerbalizeFst as TNVerbalizeFst
|
||||
|
||||
from nemo.utils import logging
|
||||
|
||||
try:
|
||||
import pynini
|
||||
from pynini.export import export
|
||||
|
||||
PYNINI_AVAILABLE = True
|
||||
except (ModuleNotFoundError, ImportError):
|
||||
|
||||
logging.warning(
|
||||
"`pynini` is not installed ! \n"
|
||||
"Please run the `nemo_text_processing/setup.sh` script"
|
||||
"prior to usage of this toolkit."
|
||||
)
|
||||
|
||||
PYNINI_AVAILABLE = False
|
||||
|
||||
|
||||
# This script exports compiled grammars inside nemo_text_processing into OpenFst finite state archive files tokenize_and_classify.far and verbalize.far for production purposes
|
||||
|
||||
|
||||
def _generator_main(file_name: str, graphs: Dict[str, pynini.FstLike]):
|
||||
"""
|
||||
Exports graph as OpenFst finite state archive (FAR) file with given file name and rule name.
|
||||
|
||||
Args:
|
||||
file_name: exported file name
|
||||
graph: Pynini WFST graph to be exported
|
||||
rule_name: rule name for graph in created FAR file
|
||||
|
||||
"""
|
||||
exporter = export.Exporter(file_name)
|
||||
for rule, graph in graphs.items():
|
||||
exporter[rule] = graph.optimize()
|
||||
exporter.close()
|
||||
print(f'Created {file_name}')
|
||||
|
||||
|
||||
def itn_grammars(**kwargs):
|
||||
d = {}
|
||||
d['classify'] = {'TOKENIZE_AND_CLASSIFY': ITNClassifyFst().fst}
|
||||
d['verbalize'] = {'ALL': ITNVerbalizeFst().fst, 'REDUP': pynini.accep("REDUP")}
|
||||
return d
|
||||
|
||||
|
||||
def tn_grammars(**kwargs):
|
||||
d = {}
|
||||
d['classify'] = {'TOKENIZE_AND_CLASSIFY': TNClassifyFst(input_case=kwargs["input_case"], deterministic=True).fst}
|
||||
d['verbalize'] = {'ALL': TNVerbalizeFst(deterministic=True).fst, 'REDUP': pynini.accep("REDUP")}
|
||||
return d
|
||||
|
||||
|
||||
def export_grammars(output_dir, grammars):
|
||||
"""
|
||||
Exports tokenizer_and_classify and verbalize Fsts as OpenFst finite state archive (FAR) files.
|
||||
|
||||
Args:
|
||||
output_dir: directory to export FAR files to. Subdirectories will be created for tagger and verbalizer respectively.
|
||||
"""
|
||||
|
||||
for category, graphs in grammars.items():
|
||||
out_dir = os.path.join(output_dir, category)
|
||||
if not os.path.exists(out_dir):
|
||||
os.makedirs(out_dir)
|
||||
time.sleep(1)
|
||||
if category == "classify":
|
||||
category = "tokenize_and_classify"
|
||||
_generator_main(f"{out_dir}/{category}.far", graphs)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument("--output_dir", help="output directory for grammars", required=True, type=str)
|
||||
parser.add_argument(
|
||||
"--grammars", help="grammars to be exported", choices=["tn_grammars", "itn_grammars"], type=str, required=True
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input_case", help="input capitalization", choices=["lower_cased", "cased"], default="cased", type=str
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = parse_args()
|
||||
export_grammars(output_dir=args.output_dir, grammars=locals()[args.grammars](input_case=args.input_case))
|
||||
Reference in New Issue
Block a user