chore: import upstream snapshot with attribution

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}")