chore: import upstream snapshot with attribution

tools/ctc_segmentation/README.md

@@ -0,0 +1,37 @@
+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"
+}
+```
+Requirements
+~~~~~~~~~~~~
+The tool requires:
+- packages listed in requirements.txt
+- NeMo ASR 
+- see pysox’s documentation (https://pysox.readthedocs.io/en/latest/) if you want support for mp3, flac and ogg files

tools/ctc_segmentation/requirements.txt

@@ -0,0 +1,5 @@
+ctc_segmentation==1.7.4
+num2words
+numpy<2.0.0
+plotly
+wget

tools/ctc_segmentation/run_filter.sh

@@ -0,0 +1,85 @@
+#!/bin/bash
+
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  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.
+
+SCRIPTS_DIR="scripts" # /<PATH TO>/NeMo/tools/ctc_segmentation/tools/scripts/ directory
+MODEL_NAME_OR_PATH="" # ASR model to use for transcribing the segmented audio files
+INPUT_AUDIO_DIR="" # Path to original directory with audio files
+MANIFEST=""
+BATCH_SIZE=4 # batch size for ASR transcribe
+NUM_JOBS=-2 # The maximum number of concurrently running jobs, `-2` - all CPUs but one are used
+
+# Thresholds for filtering
+CER_THRESHOLD=30
+WER_THRESHOLD=75
+CER_EDGE_THRESHOLD=60
+LEN_DIFF_RATIO_THRESHOLD=0.3
+MIN_DURATION=1 # in seconds
+MAX_DURATION=20 # in seconds
+EDGE_LEN=5 # number of characters for calculating edge cer
+
+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
+
+if [[ -z $MODEL_NAME_OR_PATH ]] || [[ -z $INPUT_AUDIO_DIR ]] || [[ -z $MANIFEST ]]; then
+  echo "Usage: $(basename "$0")
+  --MODEL_NAME_OR_PATH=[path to .nemo ASR model or a pre-trained model name to use for metrics calculation]
+  --INPUT_AUDIO_DIR=[path to original directory with audio files used for segmentation (for retention rate estimate)]
+  --MANIFEST=[path to manifest file generated during segmentation]"
+  exit 1
+fi
+
+echo "--- Adding transcripts to ${MANIFEST} using ${MODEL_NAME_OR_PATH} ---"
+if [[ ${MODEL_NAME_OR_PATH,,} == *".nemo" ]]; then
+  ARG_MODEL="model_path";
+else
+  ARG_MODEL="pretrained_name";
+fi
+
+OUT_MANIFEST="$(dirname ${MANIFEST})"
+OUT_MANIFEST=$OUT_MANIFEST/manifest_transcribed.json
+# Add transcripts to the manifest file, ASR model predictions will be stored under "pred_text" field
+python ${SCRIPTS_DIR}/../../../examples/asr/transcribe_speech.py \
+$ARG_MODEL=$MODEL_NAME_OR_PATH \
+dataset_manifest=$MANIFEST \
+output_filename=${OUT_MANIFEST} \
+batch_size=${BATCH_SIZE} \
+num_workers=0 || exit
+
+echo "--- Calculating metrics and filtering out samples based on thresholds ---"
+echo "CER_THRESHOLD = ${CER_THRESHOLD}"
+echo "WER_THRESHOLD = ${WER_THRESHOLD}"
+echo "CER_EDGE_THRESHOLD = ${CER_EDGE_THRESHOLD}"
+echo "LEN_DIFF_RATIO_THRESHOLD = ${LEN_DIFF_RATIO_THRESHOLD}"
+
+python ${SCRIPTS_DIR}/get_metrics_and_filter.py \
+--manifest=${OUT_MANIFEST} \
+--audio_dir=${INPUT_AUDIO_DIR} \
+--max_cer=${CER_THRESHOLD} \
+--max_wer=${WER_THRESHOLD} \
+--max_len_diff_ratio=${LEN_DIFF_RATIO_THRESHOLD} \
+--max_edge_cer=${CER_EDGE_THRESHOLD} \
+--min_duration=${MIN_DURATION} \
+--max_duration=${MAX_DURATION} \
+--edge_len=${EDGE_LEN}
+

tools/ctc_segmentation/run_segmentation.sh

@@ -0,0 +1,135 @@
+#!/bin/bash
+
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  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.
+
+# default values for optional arguments
+MIN_SCORE=-2
+CUT_PREFIX=0
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)
+
+SCRIPTS_DIR=$SCRIPT_DIR/"scripts" # /<PATH TO>/NeMo/tools/ctc_segmentation/tools/scripts/ directory
+OFFSET=0
+LANGUAGE='en' # 'en', 'es', 'ru'...
+MAX_SEGMENT_LEN=30
+ADDITIONAL_SPLIT_SYMBOLS=":|;"
+USE_NEMO_NORMALIZATION='True'
+NUM_JOBS=-2       # The maximum number of concurrently running jobs, `-2` - all CPUs but one are used
+SAMPLE_RATE=16000 # Target sample rate (default for ASR data - 16000 Hz)
+MAX_DURATION=20   # Maximum audio segment duration, in seconds. Samples that are longer will be dropped.
+
+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 "SAMPLE_RATE = $SAMPLE_RATE"
+echo "ADDITIONAL_SPLIT_SYMBOLS = $ADDITIONAL_SPLIT_SYMBOLS"
+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)]
+  --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)]
+  --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
+
+# check if num2words and ctc_segmentation are installed
+if ! command -v num2words &> /dev/null; then
+  echo "num2words could not be found"
+  echo "please install using tools/ctc_segmentation/requirements.txt"
+  exit 1
+fi
+
+if ! python -c "import ctc_segmentation" &> /dev/null; then
+  echo "ctc_segmentation could not be found"
+  echo "please install using tools/ctc_segmentation/requirements.txt"
+  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
+echo "TEXT AND AUDIO PREPROCESSING..."
+python $SCRIPTS_DIR/prepare_data.py \
+  --in_text=$DATA_DIR/text \
+  --audio_dir=$DATA_DIR/audio \
+  --output_dir=$OUTPUT_DIR/processed/ \
+  --language=$LANGUAGE \
+  --cut_prefix=$CUT_PREFIX \
+  --model=$MODEL_NAME_OR_PATH \
+  --max_length=$MAX_SEGMENT_LEN \
+  --sample_rate=$SAMPLE_RATE \
+  --additional_split_symbols=$ADDITIONAL_SPLIT_SYMBOLS $NEMO_NORMALIZATION || exit
+
+# STEP #2
+# Run CTC-segmentation. 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
+echo "SEGMENTATION STEP..."
+for WINDOW in 8000 12000; do
+  python $SCRIPTS_DIR/run_ctc_segmentation.py \
+    --output_dir=$OUTPUT_DIR \
+    --data=$OUTPUT_DIR/processed \
+    --sample_rate=$SAMPLE_RATE \
+    --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)
+echo "VERIFYING SEGMENTS..."
+python $SCRIPTS_DIR/verify_segments.py \
+  --base_dir=$OUTPUT_DIR || exit
+
+# STEP #4
+# Cut the original audio files based on the alignment score. Only segments with alignment confidence score
+# above the MIN_SCORE value will be saved to $OUTPUT_DIR/manifests/manifest.json
+echo "CUTTING AUDIO..."
+python $SCRIPTS_DIR/cut_audio.py \
+  --output_dir=$OUTPUT_DIR \
+  --alignment=$OUTPUT_DIR/verified_segments \
+  --threshold=$MIN_SCORE \
+  --offset=$OFFSET \
+  --sample_rate=$SAMPLE_RATE \
+  --max_duration=$MAX_DURATION || exit

tools/ctc_segmentation/scripts/cut_audio.py

@@ -0,0 +1,193 @@
+# 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
+from glob import glob
+
+import numpy as np
+from scipy.io import wavfile
+from tqdm import tqdm
+
+parser = argparse.ArgumentParser(description="Cut audio on the segments based on segments")
+parser.add_argument("--output_dir", type=str, help="Path to output directory", required=True)
+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("--offset", type=int, default=0, help="Offset, s")
+parser.add_argument("--batch_size", type=int, default=64, help="Batch size for inference")
+parser.add_argument(
+    "--edge_duration",
+    type=float,
+    help="Duration of audio for mean absolute value calculation at the edges, s",
+    default=0.05,
+)
+parser.add_argument("--sample_rate", type=int, help="Sample rate, Hz", default=16000)
+parser.add_argument(
+    "--max_duration",
+    type=int,
+    help="Maximum audio duration (seconds). Samples that are longer will be dropped",
+    default=60,
+)
+
+
+def process_alignment(alignment_file: str, manifest: str, clips_dir: 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
+        manifest: path to .json manifest to save segments metadata
+        clips_dir: path to a directory to save audio clips
+        args: main script args
+    """
+    if not os.path.exists(alignment_file):
+        raise ValueError(f"{alignment_file} not found")
+
+    base_name = os.path.basename(alignment_file).replace("_segments.txt", "")
+
+    # read the segments, note the first line contains the path to the original audio
+    segments = []
+    ref_text_processed = []
+    ref_text_no_preprocessing = []
+    ref_text_normalized = []
+    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())
+            ref_text_normalized.append(line[3].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 and save the final manifests at output_dir
+    sampling_rate, signal = wavfile.read(audio_file)
+    original_duration = len(signal) / sampling_rate
+
+    num_samples = int(args.edge_duration * args.sample_rate)
+    low_score_dur = 0
+    high_score_dur = 0
+    with open(manifest, "a", encoding="utf8") as 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 > args.max_duration:
+                continue
+            if duration > 0:
+                text_processed = ref_text_processed[i].strip()
+                text_no_preprocessing = ref_text_no_preprocessing[i].strip()
+                text_normalized = ref_text_normalized[i].strip()
+                if score >= args.threshold:
+                    high_score_dur += duration
+                    audio_filepath = os.path.join(clips_dir, f"{base_name}_{i:04}.wav")
+                    wavfile.write(audio_filepath, sampling_rate, segment)
+
+                    assert len(signal.shape) == 1 and sampling_rate == args.sample_rate, "check sampling rate"
+
+                    info = {
+                        "audio_filepath": audio_filepath,
+                        "duration": duration,
+                        "text": text_processed,
+                        "text_no_preprocessing": text_no_preprocessing,
+                        "text_normalized": text_normalized,
+                        "score": round(score, 2),
+                        "start_abs": float(np.mean(np.abs(segment[:num_samples]))),
+                        "end_abs": float(np.mean(np.abs(segment[-num_samples:]))),
+                    }
+                    json.dump(info, f, ensure_ascii=False)
+                    f.write("\n")
+                else:
+                    low_score_dur += duration
+
+    # keep track of duration of the deleted segments
+    del_duration = 0
+    begin = 0
+
+    for i, (st, end, _) in enumerate(segments):
+        if st - begin > 0.01:
+            segment = signal[int(begin * sampling_rate) : int(st * sampling_rate)]
+            duration = len(segment) / sampling_rate
+            del_duration += duration
+        begin = end
+
+    segment = signal[int(begin * sampling_rate) :]
+    duration = len(segment) / sampling_rate
+    del_duration += duration
+
+    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__":
+    args = parser.parse_args()
+    print("Splitting audio files into segments...")
+
+    if os.path.isdir(args.alignment):
+        alignment_files = glob(f"{args.alignment}/*_segments.txt")
+    else:
+        alignment_files = [args.alignment]
+
+    # create a directory to store segments with alignement confindence score avove the threshold
+    args.output_dir = os.path.abspath(args.output_dir)
+    clips_dir = os.path.join(args.output_dir, "clips")
+    manifest_dir = os.path.join(args.output_dir, "manifests")
+    os.makedirs(clips_dir, exist_ok=True)
+    os.makedirs(manifest_dir, exist_ok=True)
+
+    manifest = os.path.join(manifest_dir, "manifest.json")
+    if os.path.exists(manifest):
+        os.remove(manifest)
+
+    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
+        original_dur = 0
+
+        for alignment_file in tqdm(alignment_files):
+            stats = process_alignment(alignment_file, manifest, clips_dir, args)
+            original_dur += stats[-4]
+            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}")
+
+    print(f"Original duration  : {round(original_dur / 60)}min")
+    print(f"High score segments: {round(high_score_dur / 60)}min ({round(high_score_dur/original_dur*100)}%)")
+    print(f"Low score segments : {round(low_score_dur / 60)}min ({round(low_score_dur/original_dur*100)}%)")
+    print(f"Deleted segments   : {round(del_duration / 60)}min ({round(del_duration/original_dur*100)}%)")
+    print(f"Stats saved at {stats_file}")
+    print(f"Manifest saved at {manifest}")

tools/ctc_segmentation/scripts/get_metrics_and_filter.py

@@ -0,0 +1,212 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  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 glob import glob
+
+from joblib import Parallel, delayed
+from kaldialign import edit_distance
+from tqdm import tqdm
+
+from nemo.collections.asr.parts.preprocessing.segment import AudioSegment
+from nemo.utils import logging
+
+parser = argparse.ArgumentParser("Calculate metrics and filters out samples based on thresholds")
+parser.add_argument(
+    "--manifest",
+    required=True,
+    help="Path .json manifest file with ASR predictions saved at `pred_text` field.",
+)
+parser.add_argument(
+    "--edge_len", type=int, help="Number of characters to use for CER calculation at the edges", default=5
+)
+parser.add_argument("--audio_dir", type=str, help="Path to original .wav files", default=None)
+parser.add_argument("--max_cer", type=int, help="Threshold CER value, %", default=30)
+parser.add_argument("--max_wer", type=int, help="Threshold WER value, %", default=75)
+parser.add_argument(
+    "--max_len_diff_ratio",
+    type=float,
+    help="Threshold for len diff ratio between reference text "
+    "length and predicted text length with respect to "
+    "the reference text length (length measured "
+    "in number of characters)",
+    default=0.3,
+)
+parser.add_argument("--max_edge_cer", type=int, help="Threshold edge CER value, %", default=60)
+parser.add_argument("--max_duration", type=int, help="Max duration of a segment, seconds", default=-1)
+parser.add_argument("--min_duration", type=int, help="Min duration of a segment, seconds", default=1)
+parser.add_argument(
+    "--num_jobs",
+    default=-2,
+    type=int,
+    help="The maximum number of concurrently running jobs, `-2` - all CPUs but one are used",
+)
+parser.add_argument(
+    "--only_filter",
+    action="store_true",
+    help="Set to True to perform only filtering (when transcripts" "are already available)",
+)
+
+
+def _calculate(line: dict, edge_len: int):
+    """
+    Calculates metrics for every entry on manifest.json.
+
+    Args:
+        line - line of manifest.json (dict)
+        edge_len - number of characters for edge Character Error Rate (CER) calculations
+
+    Returns:
+        line - line of manifest.json (dict) with the following metrics added:
+        WER - word error rate
+        CER - character error rate
+        start_CER - CER at the beginning of the audio sample considering first 'edge_len' characters
+        end_CER - CER at the end of the audio sample considering last 'edge_len' characters
+        len_diff_ratio - ratio between reference text length and predicted text length with respect to
+            the reference text length (length measured in number of characters)
+    """
+    eps = 1e-9
+
+    text = line["text"].split()
+    pred_text = line["pred_text"].split()
+
+    num_words = max(len(text), eps)
+    word_dist = edit_distance(text, pred_text)['total']
+    line["WER"] = word_dist / num_words * 100.0
+    num_chars = max(len(line["text"]), eps)
+    char_dist = edit_distance(list(line["text"]), list(line["pred_text"]))['total']
+    line["CER"] = char_dist / num_chars * 100.0
+
+    line["start_CER"] = (
+        edit_distance(list(line["text"][:edge_len]), list(line["pred_text"][:edge_len]))['total'] / edge_len * 100
+    )
+    line["end_CER"] = (
+        edit_distance(list(line["text"][-edge_len:]), list(line["pred_text"][-edge_len:]))['total'] / edge_len * 100
+    )
+    line["len_diff_ratio"] = 1.0 * abs(len(text) - len(pred_text)) / max(len(text), eps)
+    return line
+
+
+def get_metrics(manifest, manifest_out):
+    """Calculate metrics for sample in manifest and saves the results to manifest_out"""
+    with open(manifest, "r") as f:
+        lines = f.readlines()
+
+    lines = Parallel(n_jobs=args.num_jobs)(
+        delayed(_calculate)(json.loads(line), edge_len=args.edge_len) for line in tqdm(lines)
+    )
+    with open(manifest_out, "w") as f_out:
+        for line in lines:
+            f_out.write(json.dumps(line) + "\n")
+    logging.info(f"Metrics save at {manifest_out}")
+
+
+def _apply_filters(
+    manifest,
+    manifest_out,
+    max_cer,
+    max_wer,
+    max_edge_cer,
+    max_len_diff_ratio,
+    max_dur=-1,
+    min_dur=1,
+    original_duration=0,
+):
+    """Filters out samples that do not satisfy specified threshold values and saves remaining samples to manifest_out"""
+    remaining_duration = 0
+    segmented_duration = 0
+    with open(manifest, "r") as f, open(manifest_out, "w") as f_out:
+        for line in f:
+            item = json.loads(line)
+            cer = item["CER"]
+            wer = item["WER"]
+            len_diff_ratio = item["len_diff_ratio"]
+            duration = item["duration"]
+            segmented_duration += duration
+            if (
+                cer <= max_cer
+                and wer <= max_wer
+                and len_diff_ratio <= max_len_diff_ratio
+                and item["end_CER"] <= max_edge_cer
+                and item["start_CER"] <= max_edge_cer
+                and (max_dur == -1 or (max_dur > -1 and duration < max_dur))
+                and duration > min_dur
+            ):
+                remaining_duration += duration
+                f_out.write(json.dumps(item) + "\n")
+
+    logging.info("-" * 50)
+    logging.info("Threshold values:")
+    logging.info(f"max WER, %: {max_wer}")
+    logging.info(f"max CER, %: {max_cer}")
+    logging.info(f"max edge CER, %: {max_edge_cer}")
+    logging.info(f"max Word len diff: {max_len_diff_ratio}")
+    logging.info(f"max Duration, s: {max_dur}")
+    logging.info("-" * 50)
+
+    remaining_duration = remaining_duration / 60
+    original_duration = original_duration / 60
+    segmented_duration = segmented_duration / 60
+    logging.info(f"Original audio dur: {round(original_duration, 2)} min")
+    logging.info(
+        f"Segmented duration: {round(segmented_duration, 2)} min ({round(100 * segmented_duration / original_duration, 2)}% of original audio)"
+    )
+    logging.info(
+        f"Retained {round(remaining_duration, 2)} min ({round(100*remaining_duration/original_duration, 2)}% of original or {round(100 * remaining_duration / segmented_duration, 2)}% of segmented audio)."
+    )
+    logging.info(f"Retained data saved to {manifest_out}")
+
+
+def filter(manifest):
+    """
+    Filters out samples that do not satisfy specified threshold values.
+
+    Args:
+        manifest: path to .json manifest
+    """
+    original_duration = 0
+    if args.audio_dir:
+        audio_files = glob(f"{os.path.abspath(args.audio_dir)}/*")
+        for audio in audio_files:
+            try:
+                audio_data = AudioSegment.from_file(audio)
+                duration = len(audio_data._samples) / audio_data._sample_rate
+                original_duration += duration
+            except Exception as e:
+                logging.info(f"Skipping {audio} -- {e}")
+
+    _apply_filters(
+        manifest=manifest,
+        manifest_out=manifest.replace(".json", "_filtered.json"),
+        max_cer=args.max_cer,
+        max_wer=args.max_wer,
+        max_edge_cer=args.max_edge_cer,
+        max_len_diff_ratio=args.max_len_diff_ratio,
+        max_dur=args.max_duration,
+        min_dur=args.min_duration,
+        original_duration=original_duration,
+    )
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    if not args.only_filter:
+        manifest_with_metrics = args.manifest.replace(".json", "_metrics.json")
+        get_metrics(args.manifest, manifest_with_metrics)
+    else:
+        manifest_with_metrics = args.manifest
+    filter(manifest_with_metrics)

tools/ctc_segmentation/scripts/normalization_helpers.py

@@ -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"]
+
+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 = {
+    " р.": " рублей",
+    " к.": " копеек",
+    " коп.": " копеек",
+    " копек.": " копеек",
+    " т.д.": " так далее",
+    " т. д.": " так далее",
+    " т.п.": " тому подобное",
+    " т. п.": " тому подобное",
+    " т.е.": " то есть",
+    " т. е.": " то есть",
+    " стр. ": " страница ",
+}

tools/ctc_segmentation/scripts/prepare_data.py

@@ -0,0 +1,430 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import os
+import re
+from glob import glob
+from typing import List, Optional
+
+import regex
+from joblib import Parallel, delayed
+from normalization_helpers import LATIN_TO_RU, RU_ABBREVIATIONS
+from num2words import num2words
+from tqdm import tqdm
+
+from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.models.ctc_models import EncDecCTCModel
+from nemo.collections.asr.models.hybrid_rnnt_ctc_models import EncDecHybridRNNTCTCModel
+from nemo.utils import model_utils
+
+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("--sample_rate", type=int, default=16000, help="Sampling rate used during ASR model training, Hz")
+parser.add_argument("--bit_depth", type=int, default=16, help="Bit depth to use for processed audio files")
+parser.add_argument("--n_jobs", default=-2, type=int, help="The maximum number of concurrently running jobs")
+parser.add_argument(
+    "--language",
+    type=str,
+    default="en",
+    choices=["en", "ru", "de", "es", 'other'],
+    help='Add target language based on the num2words list of supported languages',
+)
+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="stt_en_fastconformer_ctc_large",
+    help="Pre-trained model name or path to model checkpoint",
+)
+parser.add_argument(
+    "--max_length", type=int, default=40, help="Max number of words 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: ';|:'. Use '\s' to split by space.",
+)
+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.",
+)
+parser.add_argument(
+    "--batch_size",
+    type=int,
+    default=100,
+    help="Batch size for NeMo Normalization tool.",
+)
+
+
+def _load_sox_transformer():
+    try:
+        from sox import Transformer
+    except ImportError:
+        raise ImportError(
+            "Optional dependency 'sox' is required by this script. Install it with: pip install sox"
+        ) from None
+
+    return Transformer
+
+
+def process_audio(
+    in_file: str, wav_file: str = None, cut_prefix: int = 0, sample_rate: int = 16000, bit_depth: int = 16
+):
+    """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
+        bit_depth: target bit_depth
+    """
+    try:
+        if not os.path.exists(in_file):
+            raise ValueError(f'{in_file} not found')
+        Transformer = _load_sox_transformer()
+        tfm = Transformer()
+        tfm.convert(samplerate=sample_rate, n_channels=1, bitdepth=bit_depth)
+        tfm.trim(cut_prefix)
+        tfm.build(input_filepath=in_file, output_filepath=wav_file)
+    except Exception as e:
+        print(f'{in_file} skipped - {e}')
+
+
+def split_text(
+    in_file: str,
+    out_file: str,
+    vocabulary: List[str],
+    language="en",
+    remove_brackets: bool = True,
+    do_lower_case: bool = True,
+    max_length: bool = 100,
+    additional_split_symbols: bool = None,
+    use_nemo_normalization: bool = False,
+    n_jobs: Optional[int] = 1,
+    batch_size: Optional[int] = 1.0,
+):
+    """
+    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
+        max_length: Max number of words 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.
+        n_jobs (if use_nemo_normalization=True): the maximum number of concurrently running jobs. If -1 all CPUs are used. If 1 is given,
+                no parallel computing code is used at all, which is useful for debugging. For n_jobs below -1,
+                (n_cpus + 1 + n_jobs) are used. Thus for n_jobs = -2, all CPUs but one are used.
+        batch_size (if use_nemo_normalization=True): Number of examples for each process
+    """
+    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(". . .", "...")
+    )
+
+    # end of quoted speech - to be able to split sentences by full stop
+    transcript = re.sub(r"([\.\?\!])([\"\'”])", r"\g<2>\g<1> ", transcript)
+
+    # remove extra space
+    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", "en"]:
+        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):
+        if len(split_on_symbols) == 0:
+            return sentences
+
+        split_on_symbols = split_on_symbols.split("|")
+
+        def _split(sentences, delimiter):
+            result = []
+            for sent in sentences:
+                split_sent = sent.split(delimiter)
+                # keep the delimiter
+                split_sent = [(s + delimiter).strip() for s in split_sent[:-1]] + [split_sent[-1]]
+
+                if "," in delimiter:
+                    # split based on comma usually results in too short utterance, combine sentences
+                    # that result in a single word split. It's usually not recommended to do that for other delimiters.
+                    comb = []
+                    for s in split_sent:
+                        MIN_LEN = 2
+                        # if the previous sentence is too short, combine it with the current sentence
+                        if len(comb) > 0 and (len(comb[-1].split()) <= MIN_LEN or len(s.split()) <= MIN_LEN):
+                            comb[-1] = comb[-1] + " " + s
+                        else:
+                            comb.append(s)
+                    result.extend(comb)
+                else:
+                    result.extend(split_sent)
+            return result
+
+        another_sent_split = []
+        for sent in sentences:
+            split_sent = [sent]
+            for delimiter in split_on_symbols:
+                if len(delimiter) == 0:
+                    continue
+                split_sent = _split(split_sent, delimiter + " " if delimiter != " " else delimiter)
+            another_sent_split.extend(split_sent)
+
+        sentences = [s.strip() for s in another_sent_split if s.strip()]
+        return sentences
+
+    additional_split_symbols = additional_split_symbols.replace("/s", " ")
+    sentences = additional_split(sentences, additional_split_symbols)
+
+    vocabulary_symbols = []
+    for x in vocabulary:
+        if x != "<unk>":
+            # for BPE models
+            vocabulary_symbols.extend([x for x in x.replace("##", "").replace("▁", "")])
+    vocabulary_symbols = list(set(vocabulary_symbols))
+    vocabulary_symbols += [x.upper() for x in vocabulary_symbols]
+
+    # check to make sure there will be no utterances for segmentation with only OOV symbols
+    vocab_no_space_with_digits = set(vocabulary_symbols + [str(i) for i in range(10)])
+    if " " in vocab_no_space_with_digits:
+        vocab_no_space_with_digits.remove(" ")
+
+    sentences = [
+        s.strip() for s in sentences if len(vocab_no_space_with_digits.intersection(set(s.lower()))) > 0 and s.strip()
+    ]
+
+    # when no punctuation marks present in the input text, split based on max_length
+    if len(sentences) == 1:
+        sent = sentences[0].split()
+        sentences = []
+        for i in range(0, len(sent), max_length):
+            sentences.append(" ".join(sent[i : i + max_length]))
+    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(re.sub(r' +', ' ', "\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]
+        # replace Latin characters with Russian
+        for k, v in LATIN_TO_RU.items():
+            sentences = [s.replace(k, v) for s in sentences]
+
+    if language == "en" and use_nemo_normalization:
+        if not NEMO_NORMALIZATION_AVAILABLE:
+            raise ValueError("NeMo normalization tool is not installed.")
+
+        print("Using NeMo normalization tool...")
+        normalizer = Normalizer(input_case="cased", cache_dir=os.path.join(os.path.dirname(out_file), "en_grammars"))
+        sentences_norm = normalizer.normalize_list(
+            sentences, verbose=False, punct_post_process=True, n_jobs=n_jobs, batch_size=batch_size
+        )
+        if len(sentences_norm) != len(sentences):
+            raise ValueError("Normalization failed, number of sentences does not match.")
+        else:
+            sentences = sentences_norm
+
+    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 = re.sub(r' +', ' ', sentences)
+
+    with open(os.path.join(out_dir, out_file_name[:-4] + "_with_punct_normalized.txt"), "w") as f:
+        f.write(sentences)
+
+    if do_lower_case:
+        sentences = sentences.lower()
+
+    symbols_to_remove = ''.join(set(sentences).difference(set(vocabulary_symbols + ["\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:
+        if args.model is None:
+            raise ValueError(f"ASR model must be provided to extract vocabulary for text processing")
+        elif os.path.exists(args.model):
+            model_cfg = ASRModel.restore_from(restore_path=args.model, return_config=True)
+            classpath = model_cfg.target  # original class path
+            imported_class = model_utils.import_class_by_path(classpath)  # type: ASRModel
+            print(f"Restoring model : {imported_class.__name__}")
+            asr_model = imported_class.restore_from(restore_path=args.model)  # type: ASRModel
+            model_name = os.path.splitext(os.path.basename(args.model))[0]
+        else:
+            # restore model by name
+            asr_model = ASRModel.from_pretrained(model_name=args.model)  # type: ASRModel
+            model_name = args.model
+
+        if not (isinstance(asr_model, EncDecCTCModel) or isinstance(asr_model, EncDecHybridRNNTCTCModel)):
+            raise NotImplementedError(
+                f"Model is not an instance of NeMo EncDecCTCModel or ENCDecHybridRNNTCTCModel."
+                " Currently only instances of these models are supported"
+            )
+
+        # get vocabulary list
+        if hasattr(asr_model, 'tokenizer'):  # i.e. tokenization is BPE-based
+            vocabulary = asr_model.tokenizer.vocab
+        elif hasattr(asr_model.decoder, "vocabulary"):  # i.e. tokenization is character-based
+            vocabulary = asr_model.cfg.decoder.vocabulary
+        else:
+            raise ValueError("Unexpected model type. Vocabulary list not found.")
+
+        if os.path.isdir(args.in_text):
+            text_files = glob(f"{args.in_text}/*.txt")
+        else:
+            text_files.append(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,
+                max_length=args.max_length,
+                additional_split_symbols=args.additional_split_symbols,
+                use_nemo_normalization=args.use_nemo_normalization,
+                n_jobs=args.n_jobs,
+                batch_size=args.batch_size,
+            )
+        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 = glob(f"{args.audio_dir}/*")
+
+        Parallel(n_jobs=args.n_jobs)(
+            delayed(process_audio)(
+                audio_paths[i],
+                os.path.join(args.output_dir, os.path.splitext(os.path.basename(audio_paths[i]))[0] + ".wav"),
+                args.cut_prefix,
+                args.sample_rate,
+                args.bit_depth,
+            )
+            for i in tqdm(range(len(audio_paths)))
+        )
+    print("Data preparation is complete.")

tools/ctc_segmentation/scripts/run_ctc_segmentation.py

@@ -0,0 +1,210 @@
+# 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 os
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import scipy.io.wavfile as wav
+import torch
+from joblib import Parallel, delayed
+from tqdm import tqdm
+from utils import get_segments
+
+import nemo.collections.asr as nemo_asr
+from nemo.collections.asr.models.ctc_models import EncDecCTCModel
+from nemo.collections.asr.models.hybrid_rnnt_ctc_models import EncDecHybridRNNTCTCModel
+
+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("--sample_rate", type=int, default=16000, help="Sampling rate, Hz")
+parser.add_argument(
+    "--model",
+    type=str,
+    default="stt_en_fastconformer_ctc_large",
+    help="Path to model checkpoint or pre-trained model name",
+)
+parser.add_argument("--debug", action="store_true", help="Flag to enable debugging messages")
+parser.add_argument(
+    "--num_jobs",
+    default=-2,
+    type=int,
+    help="The maximum number of concurrently running jobs, `-2` - all CPUs but one are used",
+)
+
+logger = logging.getLogger("ctc_segmentation")  # use module name
+
+if __name__ == "__main__":
+
+    args = parser.parse_args()
+    logging.basicConfig(level=logging.INFO)
+    # 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")
+    if os.path.exists(log_file):
+        os.remove(log_file)
+    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.ASRModel.restore_from(args.model)
+    else:
+        asr_model = nemo_asr.models.ASRModel.from_pretrained(args.model, strict=False)
+
+    if not (isinstance(asr_model, EncDecCTCModel) or isinstance(asr_model, EncDecHybridRNNTCTCModel)):
+        raise NotImplementedError(
+            f"Model is not an instance of NeMo EncDecCTCModel or ENCDecHybridRNNTCTCModel."
+            " Currently only instances of these models are supported"
+        )
+
+    bpe_model = isinstance(asr_model, nemo_asr.models.EncDecCTCModelBPE) or isinstance(
+        asr_model, nemo_asr.models.EncDecHybridRNNTCTCBPEModel
+    )
+
+    # get tokenizer used during training, None for char based models
+    if bpe_model:
+        tokenizer = asr_model.tokenizer
+    else:
+        tokenizer = None
+
+    if isinstance(asr_model, EncDecHybridRNNTCTCModel):
+        asr_model.change_decoding_strategy(decoder_type="ctc")
+
+    # extract ASR vocabulary and add blank symbol
+    if hasattr(asr_model, 'tokenizer'):  # i.e. tokenization is BPE-based
+        vocabulary = asr_model.tokenizer.vocab
+    elif hasattr(asr_model.decoder, "vocabulary"):  # i.e. tokenization is character-based
+        vocabulary = asr_model.cfg.decoder.vocabulary
+    else:
+        raise ValueError("Unexpected model type. Vocabulary list not found.")
+
+    vocabulary = ["ε"] + list(vocabulary)
+    logging.debug(f"ASR Model vocabulary: {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)
+
+    index_duration = None
+    for path_audio in audio_paths:
+        logging.info(f"Processing {path_audio.name}...")
+        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")
+        )
+        if not os.path.exists(transcript_file):
+            logging.info(f"{transcript_file} not found. Skipping {path_audio.name}")
+            continue
+        try:
+            sample_rate, signal = wav.read(path_audio)
+            if len(signal) == 0:
+                logging.error(f"Skipping {path_audio.name}")
+                continue
+
+            assert (
+                sample_rate == args.sample_rate
+            ), f"Sampling rate of the audio file {path_audio} doesn't match --sample_rate={args.sample_rate}"
+
+            original_duration = len(signal) / sample_rate
+            logging.debug(f"len(signal): {len(signal)}, sr: {sample_rate}")
+            logging.debug(f"Duration: {original_duration}s, file_name: {path_audio}")
+
+            hypotheses = asr_model.transcribe([str(path_audio)], batch_size=1, return_hypotheses=True)
+            # if hypotheses form a tuple (from Hybrid model), extract just "best" hypothesis
+            if type(hypotheses) == tuple and len(hypotheses) == 2:
+                hypotheses = hypotheses[0]
+            log_probs = hypotheses[
+                0
+            ].alignments  # note: "[0]" is for batch dimension unpacking (and here batch size=1)
+
+            # move blank values to the first column (ctc-package compatibility)
+            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)
+
+            if index_duration is None:
+                index_duration = len(signal) / log_probs.shape[0] / sample_rate
+
+        except Exception as e:
+            logging.error(e)
+            logging.error(f"Skipping {path_audio.name}")
+            continue
+
+    asr_model_type = type(asr_model)
+    del asr_model
+    torch.cuda.empty_cache()
+
+    if len(all_log_probs) > 0:
+        start_time = time.time()
+
+        normalized_lines = Parallel(n_jobs=args.num_jobs)(
+            delayed(get_segments)(
+                all_log_probs[i],
+                all_wav_paths[i],
+                all_transcript_file[i],
+                all_segment_file[i],
+                vocabulary,
+                tokenizer,
+                bpe_model,
+                index_duration,
+                args.window_len,
+                log_file=log_file,
+                debug=args.debug,
+            )
+            for i in tqdm(range(len(all_log_probs)))
+        )
+
+        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()

tools/ctc_segmentation/scripts/utils.py

@@ -0,0 +1,328 @@
+# 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 math
+import os
+import sys
+from pathlib import PosixPath
+from typing import List, Tuple, Union
+
+import ctc_segmentation as cs
+import numpy as np
+from tqdm import tqdm
+
+from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer
+
+
+def get_segments(
+    log_probs: np.ndarray,
+    path_wav: Union[PosixPath, str],
+    transcript_file: Union[PosixPath, str],
+    output_file: str,
+    vocabulary: List[str],
+    tokenizer: SentencePieceTokenizer,
+    bpe_model: bool,
+    index_duration: float,
+    window_size: int = 8000,
+    log_file: str = "log.log",
+    debug: bool = False,
+) -> 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 len(vocabulary) - 1
+        tokenizer: ASR model tokenizer (for BPE models, None for char-based models)
+        bpe_model: Indicates whether the model uses BPE
+        window_size: the length of each utterance (in terms of frames of the CTC outputs) fits into that window.
+        index_duration: corresponding time duration of one CTC output index (in seconds)
+    """
+    level = "DEBUG" if debug else "INFO"
+    file_handler = logging.FileHandler(filename=log_file)
+    stdout_handler = logging.StreamHandler(sys.stdout)
+    handlers = [file_handler, stdout_handler]
+    logging.basicConfig(handlers=handlers, level=level)
+
+    try:
+        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")
+
+        config = cs.CtcSegmentationParameters()
+        config.char_list = vocabulary
+        config.min_window_size = window_size
+        config.index_duration = index_duration
+
+        if bpe_model:
+            ground_truth_mat, utt_begin_indices = _prepare_tokenized_text_for_bpe_model(text, tokenizer, vocabulary, 0)
+        else:
+            config.excluded_characters = ".,-?!:»«;'›‹()"
+            config.blank = vocabulary.index(" ")
+            ground_truth_mat, utt_begin_indices = cs.prepare_text(config, text)
+
+        _print(ground_truth_mat, config.char_list)
+
+        # set this after text prepare_text()
+        config.blank = 0
+        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)
+        _print(ground_truth_mat, vocabulary)
+        segments = determine_utterance_segments(config, utt_begin_indices, char_probs, timings, text, char_list)
+
+        write_output(output_file, path_wav, segments, text, text_no_preprocessing, text_normalized)
+
+        # Also writes labels in audacity format
+        output_file_audacity = output_file[:-4] + "_audacity.txt"
+        write_labels_for_audacity(output_file_audacity, segments, text_no_preprocessing)
+        logging.info(f"Label file for Audacity written to {output_file_audacity}.")
+
+        for i, (word, segment) in enumerate(zip(text, segments)):
+            if i < 5:
+                logging.debug(f"{segment[0]:.2f} {segment[1]:.2f} {segment[2]:3.4f} {word}")
+        logging.info(f"segmentation of {transcript_file} complete.")
+
+    except Exception as e:
+        logging.info(f"{e} -- segmentation of {transcript_file} failed")
+
+
+def _prepare_tokenized_text_for_bpe_model(text: List[str], tokenizer, vocabulary: List[str], blank_idx: int = 0):
+    """Creates a transition matrix for BPE-based models"""
+    space_idx = vocabulary.index("▁")
+    ground_truth_mat = [[-1, -1]]
+    utt_begin_indices = []
+    for uttr in text:
+        ground_truth_mat += [[blank_idx, space_idx]]
+        utt_begin_indices.append(len(ground_truth_mat))
+        token_ids = tokenizer.text_to_ids(uttr)
+        # blank token is moved from the last to the first (0) position in the vocabulary
+        token_ids = [idx + 1 for idx in token_ids]
+        ground_truth_mat += [[t, -1] for t in token_ids]
+
+    utt_begin_indices.append(len(ground_truth_mat))
+    ground_truth_mat += [[blank_idx, space_idx]]
+    ground_truth_mat = np.array(ground_truth_mat, np.int64)
+    return ground_truth_mat, utt_begin_indices
+
+
+def _print(ground_truth_mat, vocabulary, limit=20):
+    """Prints transition matrix"""
+    chars = []
+    for row in ground_truth_mat:
+        chars.append([])
+        for ch_id in row:
+            if ch_id != -1:
+                chars[-1].append(vocabulary[int(ch_id)])
+
+    for x in chars[:limit]:
+        logging.debug(x)
+
+
+def _get_blank_spans(char_list, blank="ε"):
+    """
+    Returns a list of tuples:
+        (start index, end index (exclusive), count)
+
+    ignores blank symbols at the beginning and end of the char_list
+    since they're not suitable for split in between
+    """
+    blanks = []
+    start = None
+    end = None
+    for i, ch in enumerate(char_list):
+        if ch == blank:
+            if start is None:
+                start, end = i, i
+            else:
+                end = i
+        else:
+            if start is not None:
+                # ignore blank tokens at the beginning
+                if start > 0:
+                    end += 1
+                    blanks.append((start, end, end - start))
+                start = None
+                end = None
+    return blanks
+
+
+def _compute_time(index, align_type, timings):
+    """Compute start and end time of utterance.
+    Adapted from https://github.com/lumaku/ctc-segmentation
+
+    Args:
+        index:  frame index value
+        align_type:  one of ["begin", "end"]
+
+    Return:
+        start/end time of utterance in seconds
+    """
+    middle = (timings[index] + timings[index - 1]) / 2
+    if align_type == "begin":
+        return max(timings[index + 1] - 0.5, middle)
+    elif align_type == "end":
+        return min(timings[index - 1] + 0.5, middle)
+
+
+def determine_utterance_segments(config, utt_begin_indices, char_probs, timings, text, char_list):
+    """Utterance-wise alignments from char-wise alignments.
+    Adapted from https://github.com/lumaku/ctc-segmentation
+
+    Args:
+        config: an instance of CtcSegmentationParameters
+        utt_begin_indices: list of time indices of utterance start
+        char_probs:  character positioned probabilities obtained from backtracking
+        timings: mapping of time indices to seconds
+        text: list of utterances
+    Return:
+        segments, a list of: utterance start and end [s], and its confidence score
+    """
+    segments = []
+    min_prob = np.float64(-10000000000.0)
+    for i in tqdm(range(len(text))):
+        start = _compute_time(utt_begin_indices[i], "begin", timings)
+        end = _compute_time(utt_begin_indices[i + 1], "end", timings)
+
+        start_t = start / config.index_duration_in_seconds
+        start_t_floor = math.floor(start_t)
+
+        # look for the left most blank symbol and split in the middle to fix start utterance segmentation
+        if char_list[start_t_floor] == config.char_list[config.blank]:
+            start_blank = None
+            j = start_t_floor - 1
+            while char_list[j] == config.char_list[config.blank] and j > start_t_floor - 20:
+                start_blank = j
+                j -= 1
+            if start_blank:
+                start_t = int(round(start_blank + (start_t_floor - start_blank) / 2))
+            else:
+                start_t = start_t_floor
+            start = start_t * config.index_duration_in_seconds
+
+        else:
+            start_t = int(round(start_t))
+
+        end_t = int(round(end / config.index_duration_in_seconds))
+
+        # Compute confidence score by using the min mean probability after splitting into segments of L frames
+        n = config.score_min_mean_over_L
+        if end_t <= start_t:
+            min_avg = min_prob
+        elif end_t - start_t <= n:
+            min_avg = char_probs[start_t:end_t].mean()
+        else:
+            min_avg = np.float64(0.0)
+            for t in range(start_t, end_t - n):
+                min_avg = min(min_avg, char_probs[t : t + n].mean())
+        segments.append((start, end, min_avg))
+    return segments
+
+
+def write_output(
+    out_path: str,
+    path_wav: str,
+    segments: List[Tuple[float]],
+    text: str,
+    text_no_preprocessing: str,
+    text_normalized: str,
+):
+    """
+    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
+    """
+    # 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, segment in enumerate(segments):
+            if isinstance(segment, list):
+                for j, x in enumerate(segment):
+                    start, end, score = x
+                    outfile.write(
+                        f"{start} {end} {score} | {text[i][j]} | {text_no_preprocessing[i][j]} | {text_normalized[i][j]}\n"
+                    )
+            else:
+                start, end, score = segment
+                outfile.write(
+                    f"{start} {end} {score} | {text[i]} | {text_no_preprocessing[i]} | {text_normalized[i]}\n"
+                )
+
+
+def write_labels_for_audacity(
+    out_path: str,
+    segments: List[Tuple[float]],
+    text_no_preprocessing: str,
+):
+    """
+    Write the segmentation output to a file ready to be imported in Audacity with the unprocessed text as labels
+
+    out_path: Path to output file
+    segments: Segments include start, end and alignment score
+    text_no_preprocessing: Reference txt without any pre-processing
+    """
+    # Audacity uses tab to separate each field (start end text)
+    TAB_CHAR = "	"
+
+    # Uses char-wise alignments to get utterance-wise alignments and writes them into the given file
+    with open(str(out_path), "w") as outfile:
+
+        for i, segment in enumerate(segments):
+            if isinstance(segment, list):
+                for j, x in enumerate(segment):
+                    start, end, _ = x
+                    outfile.write(f"{start}{TAB_CHAR}{end}{TAB_CHAR}{text_no_preprocessing[i][j]} \n")
+            else:
+                start, end, _ = segment
+                outfile.write(f"{start}{TAB_CHAR}{end}{TAB_CHAR}{text_no_preprocessing[i]} \n")

tools/ctc_segmentation/scripts/verify_segments.py

@@ -0,0 +1,98 @@
+# 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",
+)
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    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}")