chore: import upstream snapshot with attribution

tools/asr_evaluator/README.md

@@ -0,0 +1,44 @@
+ASR evaluator
+--------------------
+
+A tool for thoroughly evaluating the performance of ASR models and other features such as Voice Activity Detection. 
+
+Features:
+   - Simple step to evaluate a model in all three modes currently supported by NeMo: offline, chunked, and offline_by_chunked.
+   - On-the-fly data augmentation (such as silence, noise, etc.,) for ASR robustness evaluation. 
+   - Investigate the model's performance by detailed insertion, deletion, and substitution error rates for each and all samples.
+   - Evaluate models' reliability on different target groups such as gender, and audio length if metadata is presented.
+
+
+ASR evaluator contains two main parts: 
+- **ENGINE**. To conduct ASR inference.
+- **ANALYST**. To evaluate model performance based on predictions. 
+
+In Analyst, we can evaluate on metadata (such as duration, emotion, etc.) if it presents in manifest. For example, with the following config, we can calculate WERs for audios in different interval groups, where each group (in seconds) is defined by [[0,2],[2,5],[5,10],[10,20],[20,100000]]. Also, we can calculate the WERs for three groups of emotions, where each group is defined by [['happy','laugh'],['neutral'],['sad']]. Moreover, if we set save_wer_per_class=True, it will calculate WERs for audios in all classes presented in the data (i.e. above 5 classes + 'cry' which presented in data but not in the slot). 
+
+```
+analyst:   
+   metadata:
+        duration: 
+            enable: True
+            slot: [[0,2],[2,5],[5,10],[10,20],[20,100000]] 
+            save_wer_per_class: False # whether to save wer for each presented class.
+
+        emotion: 
+            enable: True
+            slot: [['happy','laugh'],['neutral'],['sad']] # we could have 'cry' in data but not in slot we focus on.
+            save_wer_per_class: False
+ ```           
+            
+            
+Check `./conf/eval.yaml` for the supported configuration. 
+
+If you plan to evaluate/add new tasks such as Punctuation and Capitalization, add it to the engine.
+
+Run
+```
+python asr_evaluator.py \
+engine.pretrained_name="stt_en_conformer_transducer_large" \
+engine.inference.mode="offline" \
+engine.test_ds.augmentor.noise.manifest_path=<manifest file for noise data>
+```

tools/asr_evaluator/asr_evaluator.py

@@ -0,0 +1,132 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import json
+
+import git
+from omegaconf import OmegaConf, open_dict
+from utils import cal_target_metadata_wer, run_asr_inference
+
+from nemo.collections.asr.parts.utils.eval_utils import cal_write_text_metric, cal_write_wer
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+
+"""
+This script serves as evaluator of ASR models
+Usage:
+python asr_evaluator.py \
+engine.pretrained_name="stt_en_conformer_transducer_large" \
+engine.inference.mode="offline" \
+engine.test_ds.augmentor.noise.manifest_path=<manifest file for noise data> \
+.....
+
+Check out parameters in ./conf/eval.yaml
+"""
+
+
+@hydra_runner(config_path="conf", config_name="eval.yaml")
+def main(cfg):
+    report = {}
+    logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}')
+
+    # Store git hash for reproducibility
+    if cfg.env.save_git_hash:
+        repo = git.Repo(search_parent_directories=True)
+        report['git_hash'] = repo.head.object.hexsha
+
+    ## Engine
+    # Could skip run_asr_inference and use the generated manifest by
+    # specifying analyst.metric_calculator.exist_pred_manifest
+    if cfg.analyst.metric_calculator.exist_pred_manifest is None:
+        # If need to change more parameters for ASR inference, change it in
+        # 1) shell script in utils.py
+        # 2) TranscriptionConfig on top of the executed scripts such as transcribe_speech.py in examples/asr
+        # Note we SKIP calculating wer during asr_inference stage with calculate_wer=False and calculate wer for each sample below
+        # for more flexibility and reducing possible redundant inference cost.
+        cfg.engine = run_asr_inference(cfg=cfg.engine)
+
+    else:
+        logging.info(
+            f"Use generated prediction manifest {cfg.analyst.metric_calculator.exist_pred_manifest} and skip enigneer"
+        )
+        with open_dict(cfg):
+            cfg.engine.output_filename = cfg.analyst.metric_calculator.exist_pred_manifest
+
+    ## Analyst
+    if cfg.analyst.metric_calculator.get("metric", "wer") == "wer":
+        output_manifest_w_wer, total_res, eval_metric = cal_write_wer(
+            pred_manifest=cfg.engine.output_filename,
+            gt_text_attr_name=cfg.analyst.metric_calculator.get("gt_text_attr_name", "text"),
+            pred_text_attr_name=cfg.analyst.metric_calculator.get("pred_text_attr_name", "pred_text"),
+            clean_groundtruth_text=cfg.analyst.metric_calculator.clean_groundtruth_text,
+            langid=cfg.analyst.metric_calculator.langid,
+            use_cer=cfg.analyst.metric_calculator.use_cer,
+            output_filename=cfg.analyst.metric_calculator.output_filename,
+            ignore_capitalization=cfg.analyst.metric_calculator.get("ignore_capitalization", False),
+            ignore_punctuation=cfg.analyst.metric_calculator.get("ignore_punctuation", False),
+            punctuations=cfg.analyst.metric_calculator.get("punctuations", None),
+            strip_punc_space=cfg.analyst.metric_calculator.get("strip_punc_space", False),
+        )
+    else:
+        output_manifest_w_wer, total_res, eval_metric = cal_write_text_metric(
+            pred_manifest=cfg.engine.output_filename,
+            gt_text_attr_name=cfg.analyst.metric_calculator.get("gt_text_attr_name", "text"),
+            pred_text_attr_name=cfg.analyst.metric_calculator.get("pred_text_attr_name", "pred_text"),
+            output_filename=cfg.analyst.metric_calculator.output_filename,
+            ignore_capitalization=cfg.analyst.metric_calculator.get("ignore_capitalization", False),
+            ignore_punctuation=cfg.analyst.metric_calculator.get("ignore_punctuation", False),
+            punctuations=cfg.analyst.metric_calculator.get("punctuations", None),
+            metric=cfg.analyst.metric_calculator.get("metric", "bleu"),
+            metric_args=cfg.analyst.metric_calculator.get("metric_args", None),
+            strip_punc_space=cfg.analyst.metric_calculator.get("strip_punc_space", False),
+        )
+
+    with open_dict(cfg):
+        cfg.analyst.metric_calculator.output_filename = output_manifest_w_wer
+
+    report.update({"res": total_res})
+
+    for target in cfg.analyst.metadata:
+        if cfg.analyst.metadata[target].enable:
+            occ_avg_wer = cal_target_metadata_wer(
+                manifest=cfg.analyst.metric_calculator.output_filename,
+                target=target,
+                meta_cfg=cfg.analyst.metadata[target],
+                eval_metric=eval_metric,
+            )
+            report[target] = occ_avg_wer
+
+    config_engine = OmegaConf.to_object(cfg.engine)
+    report.update(config_engine)
+
+    config_metric_calculator = OmegaConf.to_object(cfg.analyst.metric_calculator)
+    report.update(config_metric_calculator)
+
+    pretty = json.dumps(report, indent=4)
+    res = "%.3f" % (report["res"][eval_metric] * 100)
+    logging.info(pretty)
+    logging.info(f"Overall {eval_metric} is {res} %")
+
+    ## Writer
+    report_file = "report.json"
+    if "report_filename" in cfg.writer and cfg.writer.report_filename:
+        report_file = cfg.writer.report_filename
+
+    with open(report_file, "a") as fout:
+        json.dump(report, fout)
+        fout.write('\n')
+        fout.flush()
+
+
+if __name__ == "__main__":
+    main()

tools/asr_evaluator/conf/eval.yaml

@@ -0,0 +1,84 @@
+env:
+    save_git_hash: True
+
+engine:    
+    model_path: null  
+    pretrained_name: null
+    
+    output_filename: null
+    random_seed: &random_seed 42
+
+    inference: 
+        mode: offline # choose from offline, chunked or offline_by_chunked
+        chunk_len_in_secs: 1.6 #null # Need to specify if use buffered inference (default for offline_by_chunked is 20)
+        total_buffer_in_secs: 4 #null # Need to specify if use buffered inference (default for offline_by_chunked is 22)
+        model_stride: 8 # Model downsampling factor, 8 for Citrinet and FastConformer models, and 4 for Conformer models
+        decoder_type: null # Used for hybrid CTC RNNT model only. Specify decoder_type *ctc* or *rnnt* for hybrid CTC RNNT model.
+   
+    test_ds:
+        manifest_filepath: null
+        sample_rate: 16000
+        batch_size: 32
+        num_workers: 4
+        
+        augmentor:
+          silence:
+            prob: 0.8
+            min_start_silence_secs: 0
+            max_start_silence_secs: 5
+            min_end_silence_secs: 0
+            max_end_silence_secs: 5
+            rng: *random_seed
+            
+          noise:
+            manifest_path: null
+            prob: 0.8
+            min_snr_db: 0
+            max_snr_db: 15
+            rng: *random_seed
+
+    transcribe_params:
+        # Put additional overrides for params in TranscriptionConfig used by transcribe_speech.py here
+        # Don't put the following fields here: 'calculate_wer', 'model_path', 'pretrained_name', 'dataset_manifest', 
+        # 'output_filename', 'batch_size', 'num_workers', 'random_seed', 'eval_config_yaml', 'decoder_type'
+        allow_partial_transcribe: False  # only set True if your audio is too long and have 'offset' in manifest
+
+analyst:
+    metric_calculator:
+        exist_pred_manifest: null # specify the previously generated manifest will skip engine
+        clean_groundtruth_text: True
+        langid: "en" # speciify language to clean text. Note use text normalization in NeMo for better performancce
+        output_filename: null
+        use_cer: False
+        ignore_capitalization: False
+        ignore_punctuation: False
+        punctuations: null  # a string of punctuations to remove when ignore_punctuation=True. if not set, default to '!"#$%&()*+,-./:;<=>?@[\\]^_`{|}~'
+        strip_punc_space: False # strip spaces before punctuations. e.g., "I do ." -> "I do."
+
+    metadata:
+        duration: 
+            enable: True
+            slot: [[0,2],[2,5],[5,10],[10,20],[20,100000]] # a slot accepts List[List[str]] or List[List[float]]. i.e. 1.8s belongs to slot [0,2]
+            save_wer_per_class: False # whether to save wer for each presented class.
+            
+        gender:
+            enable: False
+            slot: [["female"]] # One could also report only one group/class though there are multiple classes in the data. 
+            save_wer_per_class: True
+            
+        speaker:
+            enable: True
+            save_wer_per_class: False
+            
+        age:
+            enable: False
+            slot: null
+            save_wer_per_class: False
+            
+        emotion: 
+            enable: True
+            slot: [['happy','laugh'],['neutral'],['sad']] 
+            save_wer_per_class: False
+                 
+writer:
+    report_filename: null

tools/asr_evaluator/utils.py

@@ -0,0 +1,401 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import json
+import subprocess
+import tempfile
+from pathlib import Path
+
+from omegaconf import DictConfig, OmegaConf, open_dict
+from nemo.collections.asr.parts.utils.eval_utils import get_hydra_override_from_config
+from nemo.utils import logging
+
+
+def model_name_contains(model_name: str, *keywords) -> bool:
+    """
+    Check if any of the given keywords appear (case-insensitive) in the model name.
+    Args:
+        model_name (str): Model name.
+        *keywords: Variable length argument list of keywords to check.
+    Returns:
+        bool: True if any of the keywords are found in the model name, False otherwise.
+    """
+    model_name_lower = model_name.lower()
+    return any(kw.lower() in model_name_lower for kw in keywords)
+
+
+def run_asr_inference(cfg: DictConfig) -> DictConfig:
+    """
+    Execute ASR inference based on input mode and parameters.
+    """
+    if (cfg.model_path and cfg.pretrained_name) or (not cfg.model_path and not cfg.pretrained_name):
+        raise ValueError("Please specify either cfg.model_path or cfg.pretrained_name!")
+
+    if cfg.inference.decoder_type not in [None, 'ctc', 'rnnt', 'aed']:
+        raise ValueError("decoder_type could only be null, ctc, rnnt or aed")
+
+    if cfg.inference.mode == "offline":
+        cfg = run_offline_inference(cfg)
+
+    elif cfg.inference.mode == "chunked":
+        if (
+            "total_buffer_in_secs" not in cfg.inference
+            or "chunk_len_in_secs" not in cfg.inference
+            or not cfg.inference.total_buffer_in_secs
+            or not cfg.inference.chunk_len_in_secs
+        ):
+            raise ValueError(f"Please specify both total_buffer_in_secs and chunk_len_in_secs for chunked inference")
+        cfg = run_chunked_inference(cfg)
+
+    elif cfg.inference.mode == "offline_by_chunked":
+        # When use Conformer to transcribe long audio sample, we could probably encounter CUDA out of memory issue.
+        # Here we use offline_by_chunked mode to simulate offline mode for Conformer.
+        # And we specify default total_buffer_in_secs=22 and chunk_len_in_secs=20 to avoid above problem.
+        OmegaConf.set_struct(cfg, True)
+        if 'total_buffer_in_secs' not in cfg.inference or not cfg.inference.total_buffer_in_secs:
+            with open_dict(cfg):
+                cfg.inference.total_buffer_in_secs = 22
+                logging.info(
+                    f"Does not provide total_buffer_in_secs required by {cfg.inference.mode} mode. Using default value {cfg.inference.total_buffer_in_secs}"
+                )
+        if 'chunk_len_in_secs' not in cfg.inference or not cfg.inference.chunk_len_in_secs:
+            with open_dict(cfg):
+                cfg.inference.chunk_len_in_secs = 20
+                logging.info(
+                    f"Does not provide total_buffer_in_secs required by {cfg.inference.mode} mode. Using default value {cfg.inference.chunk_len_in_secs}"
+                )
+        cfg = run_chunked_inference(cfg)
+
+    else:
+        raise ValueError(f"inference could only be offline or chunked, but got {cfg.inference.mode}")
+
+    return cfg
+
+
+def run_chunked_inference(cfg: DictConfig) -> DictConfig:
+
+    if cfg.model_path:
+        model_name = Path(cfg.model_path).stem
+    else:
+        model_name = cfg.pretrained_name
+
+    if "output_filename" not in cfg or not cfg.output_filename:
+        dataset_name = Path(cfg.test_ds.manifest_filepath).stem
+        mode_name = (
+            cfg.inference.mode
+            + "B"
+            + str(cfg.inference.total_buffer_in_secs)
+            + "C"
+            + str(cfg.inference.chunk_len_in_secs)
+        )
+
+        OmegaConf.set_struct(cfg, True)
+        with open_dict(cfg):
+            cfg.output_filename = f"{model_name}-{dataset_name}-{mode_name}.json"
+
+    use_ctc_script = False
+    use_rnnt_scrpit = False
+    use_aed_script = False
+
+    # hybrid model
+    if model_name_contains(model_name, "hybrid"):
+        if cfg.inference.decoder_type:
+            if cfg.inference.decoder_type == "rnnt":
+                use_rnnt_scrpit = True
+            elif cfg.inference.decoder_type == "ctc":
+                use_ctc_script = True
+            else:
+                raise ValueError(
+                    f"Hybrid models only support rnnt or ctc decoding! Current decoder_type: {cfg.inference.decoder_type}! Change it to null, rnnt or ctc for hybrid models"
+                )
+        else:
+            # By default, use RNNT for hybrid models
+            use_rnnt_scrpit = True
+
+    # rnnt model
+    elif model_name_contains(model_name, "rnnt", "transducer"):
+        if cfg.inference.decoder_type and cfg.inference.decoder_type != 'rnnt':
+            raise ValueError(
+                f"rnnt models only support rnnt decoding! Current decoder_type: {cfg.inference.decoder_type}! Change it to null or rnnt for rnnt models"
+            )
+        use_rnnt_scrpit = True
+
+    # ctc model
+    elif model_name_contains(model_name, "ctc"):
+        if cfg.inference.decoder_type and cfg.inference.decoder_type != 'ctc':
+            raise ValueError(
+                f"ctc models only support ctc decoding! Current decoder_type: {cfg.inference.decoder_type}! Change it to null or ctc for ctc models"
+            )
+        use_ctc_script = True
+
+    # aed model
+    elif model_name_contains(model_name, "canary"):
+        if cfg.inference.decoder_type and cfg.inference.decoder_type != 'aed':
+            raise ValueError(
+                f"Canary models only support aed decoding! Current decoder_type: {cfg.inference.decoder_type}! Change it to null or aed for aed models"
+            )
+        use_aed_script = True
+
+    else:
+        raise ValueError(
+            "Please make sure your pretrained_name or model_path contains \n\
+            'hybrid' for EncDecHybridRNNTCTCModel model, \n\
+            'transducer/rnnt' for EncDecRNNTModel model, \n\
+            'ctc' for EncDecCTCModel, or \n\
+            'aed' for EncDecMultiTaskModel."
+        )
+
+    script_path = None
+    if use_rnnt_scrpit:
+        script_path = (
+            Path(__file__).parents[2]
+            / "examples"
+            / "asr"
+            / "asr_chunked_inference"
+            / "rnnt"
+            / "speech_to_text_buffered_infer_rnnt.py"
+        )
+    elif use_aed_script:
+        script_path = (
+            Path(__file__).parents[2]
+            / "examples"
+            / "asr"
+            / "asr_chunked_inference"
+            / "aed"
+            / "speech_to_text_aed_chunked_infer.py"
+        )
+    elif use_ctc_script:
+        raise ValueError("Evaluation of CTC models with chunked inference is not supported")
+    else:
+        raise ValueError(f"Unsupported model: {model_name}")
+
+    # If need to change other config such as decoding strategy, could either:
+    # 1) change TranscriptionConfig on top of the executed scripts such as speech_to_text_buffered_infer_rnnt.py, or
+    # 2) add command as "decoding.strategy=greedy_batch " to below script
+
+    base_cmd = [
+        "python",
+        str(script_path),
+        "calculate_wer=False",
+        f"model_path={cfg.model_path}",
+        f"pretrained_name={cfg.pretrained_name}",
+        f"dataset_manifest={cfg.test_ds.manifest_filepath}",
+        f"output_filename={cfg.output_filename}",
+        f"random_seed={cfg.random_seed}",
+        f"batch_size={cfg.test_ds.batch_size}",
+        f"++num_workers={cfg.test_ds.num_workers}",
+        f"chunk_len_in_secs={cfg.inference.chunk_len_in_secs}",
+        f"++total_buffer_in_secs={cfg.inference.total_buffer_in_secs}",
+        f"model_stride={cfg.inference.model_stride}",
+        f"++timestamps={cfg.inference.timestamps}",
+    ]
+
+    subprocess.run(
+        base_cmd,
+        shell=False,
+        check=True,
+    )
+    return cfg
+
+
+def run_offline_inference(cfg: DictConfig) -> DictConfig:
+    if "output_filename" not in cfg or not cfg.output_filename:
+        if cfg.model_path:
+            model_name = Path(cfg.model_path).stem
+        else:
+            model_name = cfg.pretrained_name
+        dataset_name = Path(cfg.test_ds.manifest_filepath).stem
+        mode_name = cfg.inference.mode
+
+        OmegaConf.set_struct(cfg, True)
+        with open_dict(cfg):
+            cfg.output_filename = f"{model_name}-{dataset_name}-{mode_name}.json"
+
+    with tempfile.NamedTemporaryFile(mode='w', encoding='utf-8') as f:
+        OmegaConf.save(cfg, f)
+        f.seek(0)  # reset file pointer
+        script_path = Path(__file__).parents[2] / "examples" / "asr" / "transcribe_speech.py"
+        # some keys to ingore when generating hydra overrides
+        exclude_keys = [
+            'calculate_wer',
+            'model_path',
+            'pretrained_name',
+            'dataset_manifest',
+            'output_filename',
+            'batch_size',
+            'num_workers',
+            'random_seed',
+            'eval_config_yaml',
+            'decoder_type',
+        ]
+        hydra_overrides = get_hydra_override_from_config(cfg.get("transcribe_params", None), exclude_keys=exclude_keys)
+        # If need to change other config such as decoding strategy, could either:
+        # 1) change TranscriptionConfig on top of the executed scripts such as transcribe_speech.py in examples/asr, or
+        # 2) add command as "rnnt_decoding.strategy=greedy_batch " to below script
+        base_cmd = [
+            "python",
+            str(script_path),
+            "calculate_wer=False",
+            f"model_path={cfg.model_path}",
+            f"pretrained_name={cfg.pretrained_name}",
+            f"dataset_manifest={cfg.test_ds.manifest_filepath}",
+            f"output_filename={cfg.output_filename}",
+            f"batch_size={cfg.test_ds.batch_size}",
+            f"num_workers={cfg.test_ds.num_workers}",
+            f"random_seed={cfg.random_seed}",
+            f"eval_config_yaml={f.name}",
+            f"decoder_type={cfg.inference.decoder_type}",
+        ]
+        if hydra_overrides:
+            base_cmd.extend(hydra_overrides.split())
+        subprocess.run(
+            base_cmd,
+            shell=False,
+            check=True,
+        )
+
+    return cfg
+
+
+def cal_target_metadata_wer(
+    manifest: str,
+    target: str,
+    meta_cfg: DictConfig,
+    eval_metric: str = "wer",
+) -> dict:
+    """
+    Caculating number of samples (samples), number of words/characters/tokens (tokens),
+    wer/cer, insertion error rate (ins_rate), deletion error rate (del_rate), substitution error rate (sub_rate) of the group/slot of target metadata.
+
+    The group could be [female, male] or slot group like [0-2s, 2-5s, >5s audios]
+
+
+    Args:
+        manifest (str): Filepath of the generated manifest which contains prediction and eval result for each samples.
+        target (str): Target metadata. Execute the target metadata if field presents in manifest.
+            such as 'duration', 'speaker', 'emotion', etc.
+        meta_cfg (DictConfig): Config for calculating group eval_metric for the target metadata.
+        eval_metric: (str): Supported evaluation metrics. Currently support 'wer' and 'cer'.
+
+    Return:
+        ret (dict): Generated dictionary containing all results regarding the target metadata.
+    """
+    if eval_metric not in ['wer', 'cer']:
+        raise ValueError(
+            "Currently support wer and cer as eval_metric. Please implement it in cal_target_metadata_wer if using different eval_metric"
+        )
+
+    wer_per_class = {}
+    with open(manifest, 'r') as fp:
+        for line in fp:
+            sample = json.loads(line)
+            if target in sample:
+                target_class = sample[target]
+                if target_class not in wer_per_class:
+                    wer_per_class[target_class] = {
+                        'samples': 0,
+                        'tokens': 0,
+                        "errors": 0,
+                        "inss": 0,
+                        "dels": 0,
+                        "subs": 0,
+                    }
+                wer_per_class[target_class]['samples'] += 1
+
+                tokens = sample["tokens"]
+                wer_per_class[target_class]["tokens"] += tokens
+                wer_per_class[target_class]["errors"] += tokens * sample[eval_metric]
+                wer_per_class[target_class]["inss"] += tokens * sample["ins_rate"]
+                wer_per_class[target_class]["dels"] += tokens * sample["del_rate"]
+                wer_per_class[target_class]["subs"] += tokens * sample["sub_rate"]
+
+    if len(wer_per_class) > 0:
+        res_wer_per_class = {}
+        for target_class in wer_per_class:
+            res_wer_per_class[target_class] = {}
+            res_wer_per_class[target_class]["samples"] = wer_per_class[target_class]["samples"]
+            res_wer_per_class[target_class][eval_metric] = (
+                wer_per_class[target_class]["errors"] / wer_per_class[target_class]["tokens"]
+            )
+            res_wer_per_class[target_class]["tokens"] = wer_per_class[target_class]["tokens"]
+            res_wer_per_class[target_class]["ins_rate"] = (
+                wer_per_class[target_class]["inss"] / wer_per_class[target_class]["tokens"]
+            )
+            res_wer_per_class[target_class]["del_rate"] = (
+                wer_per_class[target_class]["dels"] / wer_per_class[target_class]["tokens"]
+            )
+            res_wer_per_class[target_class]["sub_rate"] = (
+                wer_per_class[target_class]["subs"] / wer_per_class[target_class]["tokens"]
+            )
+    else:
+        logging.info(f"metadata '{target}' does not present in manifest. Skipping! ")
+        return None
+
+    values = ['samples', 'tokens', 'errors', 'inss', 'dels', 'subs']
+    slot_wer = {}
+    if 'slot' in meta_cfg and meta_cfg.slot:
+        for target_class in wer_per_class:
+            for s in meta_cfg.slot:
+                if isinstance(s[0], float) or isinstance(s[0], int):
+                    if s[0] <= target_class < s[1]:
+                        slot_key = "slot-" + ",".join(str(i) for i in s)
+                        if slot_key not in slot_wer:
+                            slot_wer[slot_key] = {
+                                'samples': 0,
+                                'tokens': 0,
+                                "errors": 0,
+                                "inss": 0,
+                                "dels": 0,
+                                "subs": 0,
+                            }
+
+                        for v in values:
+                            slot_wer[slot_key][v] += wer_per_class[target_class][v]
+                        break
+
+                elif isinstance(s[0], str):
+                    if target_class in s:
+                        slot_key = "slot-" + ",".join(s)
+                        if slot_key not in slot_wer:
+                            slot_wer[slot_key] = {
+                                'samples': 0,
+                                'tokens': 0,
+                                "errors": 0,
+                                "inss": 0,
+                                "dels": 0,
+                                "subs": 0,
+                            }
+
+                        for v in values:
+                            slot_wer[slot_key][v] += wer_per_class[target_class][v]
+                        break
+                else:
+                    raise ValueError("Current only support target metadata belongs to numeric or string ")
+
+        for slot_key in slot_wer:
+            slot_wer[slot_key][eval_metric] = slot_wer[slot_key]['errors'] / slot_wer[slot_key]['tokens']
+            slot_wer[slot_key]['ins_rate'] = slot_wer[slot_key]['inss'] / slot_wer[slot_key]['tokens']
+            slot_wer[slot_key]['del_rate'] = slot_wer[slot_key]['dels'] / slot_wer[slot_key]['tokens']
+            slot_wer[slot_key]['sub_rate'] = slot_wer[slot_key]['subs'] / slot_wer[slot_key]['tokens']
+            slot_wer[slot_key].pop('errors')
+            slot_wer[slot_key].pop('inss')
+            slot_wer[slot_key].pop('dels')
+            slot_wer[slot_key].pop('subs')
+        res_wer_per_class.update(slot_wer)
+
+    ret = None
+    if meta_cfg.save_wer_per_class:
+        ret = res_wer_per_class
+    if (not meta_cfg.save_wer_per_class) and ('slot' in meta_cfg and meta_cfg.slot):
+        ret = slot_wer
+    return ret

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

tools/customization_dataset_preparation/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2023, 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.

tools/customization_dataset_preparation/customization_dataset_preparation.py

@@ -0,0 +1,472 @@
+# Copyright (c) 2023, 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.
+
+"""
+
+NeMo LLM Customization service requires data to be in the form of .jsonl file with each line having only two fields (namely prompt and completion).
+
+However, you might not have your data readily in this format (or even filetype).
+
+This script will help you to convert from what you have to what you will need quickly and easily.
+
+You will need your datafile (in the form of a .jsonl, .json, .csv, .tsv or .xlsx). 
+Each row should contain one sample. 
+Make sure that the directory your file is in is readable and writeable.
+Otherwise, please change it using chmod. Don't worry, we will not overwrite your existing file.
+
+With close to a dozen consideration factors that makes training optimal, there might just be something you overlook (we all do!). 
+To check if dataset has been prepared correctly
+
+!python customization_dataset_preparation.py --filename <filename> 
+
+To format dataset from an alternative jsonl/json/csv/tsv/xlsx column structure (example here for Question Answering task)
+
+For instances, if you are working on a Question Answering Task, you would typically have the columns `context`, `question` and `answer`
+
+!python customization_dataset_preparation.py --filename <filename> --prompt_template "Context: {context} Question: {question} Answer:" --completion_template "{answer}"
+
+Other flags that can be set
+
+1.   `--drop_duplicates` : Use this flag to drop rows that are exactly the same for both prompt and completion
+2.   `--split_train_validation` : Use this flag to split one file into separate train and validation files.
+3.   `--val_proportion 0.1`: Use a float (default 0.1) between 0 and 1 to control how much of the dataset to allocate to the validation set and the remaining for the train dataset.
+4.   `--short_context_model`: Use this flag to prepare data for use with models that have shorter context length of 2048 tokens (e.g. 5B and 20B models)
+
+What to expect
+
+After running this code, you see a list of suggestions to use under ACTIONABLE MESSAGES as well as some insights into your dataset under INFORMATIONAL MESSAGES.
+
+We suggest you prioritize changes suggested under ACTIONABLE MESSAGES but also have a look at the INFORMATIONAL MESSAGES to ensure that changes are done in an expected manner.
+
+"""
+import argparse
+import math
+import os
+import pathlib
+from collections import Counter
+
+import numpy as np
+import pandas as pd
+
+
+def load_file_into_df(filename):
+    message = None
+    if not os.path.isfile(filename):
+        raise ValueError(f"File {filename} does not exist")
+    if filename.lower().endswith(".jsonl"):
+        df = pd.read_json(filename, lines=True, dtype=str).fillna("")
+    elif filename.lower().endswith(".json"):
+        df = pd.read_json(filename, dtype=str).fillna("")
+    elif filename.lower().endswith(".xlsx"):
+        df = pd.read_excel(filename, dtype=str).fillna("")
+        message = "Note only the first sheet in your Excel file will be read."
+    elif filename.lower().endswith(".csv"):
+        df = pd.read_csv(filename, sep=",", dtype=str).fillna("")
+    elif filename.lower().endswith(".tsv"):
+        df = pd.read_csv(filename, sep="\t", dtype=str).fillna("")
+    else:
+        raise ValueError(
+            f"Filename {filename} does not have the acceptable extension of .jsonl, .json, .xlsx, .csv or .tsv"
+        )
+    return df, message
+
+
+def recommend_hyperparameters_human_readable(recommended_hyperparameters):
+    message = 'TODO: Recommended hyperparameters\n'
+    for param, param_value in recommended_hyperparameters.items():
+        message += f'{param}: {param_value}\n'
+    return message
+
+
+def recommend_hyperparameters(df, model=None):
+    """
+    Makes recommendations on the batch_size to use for training, based on the dataset size
+
+    """
+    potential_batch_sizes = [2, 4, 8, 12, 16, 32, 64, 128]
+
+    max_bs = 128
+    if len(df) < 128:
+        max_bs = 2
+        for potential_bs in potential_batch_sizes:
+            if potential_bs < len(df) * 0.9:
+                max_bs = potential_bs
+
+    bs = min(max_bs, 32)
+
+    df_char_length = df.apply(lambda x: len(x.prompt) + len(x.completion), axis=1)
+    length_by_chars = sorted(list(df_char_length))
+    n_samples_under_99p5_limit = math.ceil(len(df_char_length) * 0.995)
+    char_length_99p5 = length_by_chars[n_samples_under_99p5_limit - 1]
+    mean_char_length = np.mean(length_by_chars)
+    std_char_length = np.std(length_by_chars)
+
+    # filter out only outliers that are >2 std above mean
+    max_char_length = max(min(mean_char_length + 2 * std_char_length, length_by_chars[-1]), char_length_99p5)
+
+    # every token is around 4 chars + 100 for extra capacity
+    max_seq_length = max_char_length // 4 + 100
+
+    if len(df) <= 100:
+        encoder_hidden_size = 1024
+    elif len(df) <= 1000:
+        encoder_hidden_size = 2048
+    else:
+        encoder_hidden_size = 4096
+
+    if len(df) <= 100:
+        lr = 5e-3
+    elif len(df) <= 1000:
+        lr = 1e-3
+    elif len(df) <= 10000:
+        lr = 5e-4
+    else:
+        lr = 1e-4
+
+    return {
+        'batch_size': bs,
+        'max_batch_size': max_bs,
+        'num_virtual_tokens': 10,
+        'lr': lr,
+        'epochs': 10,
+        'max_seq_length': max_seq_length,
+        'encoder_hidden_size': encoder_hidden_size,
+    }
+
+
+def estimating_customization_job_time(df, recommended_hyperparameters):
+    recommended_batch_size = recommended_hyperparameters['batch_size']
+
+    size = df.memory_usage(index=True, deep=True).sum()
+    time_in_seconds_per_epoch = size / recommended_batch_size * 0.0025
+
+    if time_in_seconds_per_epoch < 60:
+        time_per_epoch = f"{round(time_in_seconds_per_epoch, 2)} seconds"
+    elif time_in_seconds_per_epoch < 3600:
+        time_per_epoch = f"{round(time_in_seconds_per_epoch/60, 2)} minutes"
+    else:
+        time_per_epoch = f"{round(time_in_seconds_per_epoch/3600, 2)} hours"
+
+    message = f"TODO: Training will take around {time_per_epoch} for each epoch for gpt20b model and around half of that for gpt5b. Please set no. of epochs accordingly to ensure that the limit of 8h total is not exceeded."
+    return message
+
+
+def warn_completion_is_not_empty(df):
+    message = None
+    field = "completion"
+    empty_rows = (df[field] == "") | (df[field].isnull())
+    empty_indexes = df.reset_index().index[empty_rows].tolist()
+    if len(empty_indexes) == len(df):
+        message = (
+            "TODO: Note all completion fields are empty. This is possibly expected for inference but not for training"
+        )
+    elif len(empty_indexes) != 0:
+        message = f"""TODO: completion contains {len(empty_indexes)} empty values at rows ({empty_indexes})
+                Please check the original file that the fields for prompt template are 
+                not empty and rerun dataset validation"""
+    return message
+
+
+def warn_imbalanced_completion(df):
+    completions = df["completion"].tolist()
+    completions_counter = Counter(completions)
+    message = None
+    # low variety of unique completions relative to completions
+    # suggesting it is a classification set up
+    if len(completions_counter) < len(completions) / 3:
+        message = f"There are {len(completions_counter)} unique completions over {len(completions)} samples.\nThe five most common completions are:"
+        for completion, n in completions_counter.most_common(5):
+            message += f"\n {n} samples ({round(100*n/len(completions),0)}%) with completion: {completion}"
+    return message
+
+
+def get_common_suffix(series):
+    common_suffix = ""
+    while True:
+        candidate_common_suffixes = series.str[-(len(common_suffix) + 1) :]
+        if candidate_common_suffixes.nunique() != 1:
+            # candidate_common_suffixes contains more than one value
+            # therefore, it is no longer a common suffix
+            break
+        elif common_suffix == candidate_common_suffixes.values[0]:
+            # candidate is the same as previous common_suffix
+            # therefore values in series are too short to move back by one char
+            break
+        else:
+            common_suffix = candidate_common_suffixes.values[0]
+    return common_suffix
+
+
+def warn_missing_suffix(df):
+    message = ''
+    for field in ["prompt", "completion"]:
+        if not get_common_suffix(df[field]):
+            message += f"TODO: {field} does not have common suffix, please add one (e.g. \\n) at the end of {field}_template\n"
+    return message if message else None
+
+
+def validate_template(template):
+    template_with_only_brackets = [i for i in template if i in ["{", "}"]]
+    error_msg = (
+        "Your template ("
+        + template
+        + ") is not in the correct format.\
+                Template must be in the format contains zero or more fields, \
+                each field specified by {field}\
+                For instance, it can be 'Context: {context} Question: {question}:"
+    )
+    if len(template_with_only_brackets) % 2 != 0:
+        raise ValueError(error_msg)
+    for i in range(0, len(template_with_only_brackets), 2):
+        if not (template_with_only_brackets[i] == "{" and template_with_only_brackets[i + 1] == "}"):
+            raise ValueError(error_msg)
+    return None
+
+
+def parse_template(template):
+    field_names = []
+    i = 0
+    in_field = False
+    while i < len(template):
+        if template[i] == "{":
+            field_names.append("")
+            in_field = True
+        elif template[i] == "}":
+            in_field = False
+        elif in_field:
+            field_names[-1] += template[i]
+        else:
+            pass
+        i += 1
+    return field_names
+
+
+def warn_duplicated_rows(df):
+    message = None
+    duplicated_rows = df.duplicated()
+    duplicated_indices = df.reset_index().index[duplicated_rows].tolist()
+    if len(duplicated_indices) > 0:
+        message = f"TODO: There are {len(duplicated_indices)} duplicated rows "
+        message += f"at rows ({duplicated_indices}) \n"
+        message += "Please check the original file to make sure that is expected\n"
+        message += "If it is not, please add the argument --drop_duplicate"
+    return message
+
+
+def drop_duplicated_rows(df):
+    duplicated_rows = df.duplicated()
+    duplicated_indices = df.reset_index().index[duplicated_rows].tolist()
+    message = None
+    if len(duplicated_indices) > 0:
+        df = df.drop_duplicates()
+        message = f"There are {len(duplicated_indices)} duplicated rows\n"
+        message += f"Removed {len(duplicated_indices)} duplicate rows"
+    return df, message
+
+
+def template_mapper(row, field_names, template):
+    for field_name in field_names:
+        template = template.replace("{" + field_name + "}", row[field_name])
+    return template
+
+
+def drop_unrequired_fields(df, required_fields=["prompt", "completion"]):
+    for column in df.columns:
+        if column not in required_fields:
+            df = df.drop(column, axis=1)
+    return df
+
+
+def convert_into_template(df, template, prompt_or_completion="prompt"):
+    validate_template(template)
+    template = template.replace("\\n", "\n")
+    field_names = parse_template(template)
+    for field_name in field_names:
+        if field_name not in df.columns:
+            raise ValueError(
+                f"Field {field_name} requested in {prompt_or_completion}_template ({template}) but not found in file columns, which contains {list(df.columns)}"
+            )
+    df[prompt_or_completion] = df.apply(lambda row: template_mapper(row, field_names, template), axis=1)
+    return df
+
+
+def convert_into_prompt_completion_only(df, prompt_template="{prompt}", completion_template="{completion}"):
+    df = convert_into_template(df, prompt_template, prompt_or_completion="prompt")
+    df = convert_into_template(df, completion_template, prompt_or_completion="completion")
+    df = drop_unrequired_fields(df)
+    return df
+
+
+def warn_and_drop_long_samples(df, max_total_char_length):
+    long_examples = df.apply(lambda x: len(x.prompt) + len(x.completion) > max_total_char_length, axis=1)
+    indices_of_long_examples = df.reset_index().index[long_examples].tolist()
+    message = None
+    if len(indices_of_long_examples) > 0:
+        message = f"""TODO: There are {len(indices_of_long_examples)} / {len(df)} 
+        samples that have its prompt and completion too long 
+        (over {max_total_char_length} chars), which have been dropped."""
+        df = df.drop(indices_of_long_examples).reset_index()
+        df = df.drop('index', axis=1)
+    return df, message
+
+
+def warn_low_n_samples(df, min_samples=64):
+    if len(df) < min_samples:
+        return f"""TODO: We would recommend having more samples (>{min_samples}) if possible but current_file only contains {len(df)} samples. """
+    return None
+
+
+def show_first_example_in_df(df):
+    message = ''
+    for column in df.columns:
+        # prints \n instead of an a newline
+        column_value = df[column][0].replace('\n', '\\n')
+        message += f"-->Column {column}:\n{column_value}\n"
+    return message
+
+
+def get_prepared_filename(filename, split_train_validation=False):
+    message = ""
+
+    file_extension = pathlib.Path(filename).suffix
+    if not split_train_validation:
+        new_filename = filename.replace(file_extension, "_prepared.jsonl")
+        retry = 0
+        while os.path.isfile(new_filename):
+            message += f"File {new_filename} exists. Trying next available filename increment\n"
+            retry += 1
+            new_filename = filename.replace(file_extension, f"_prepared{retry}.jsonl")
+        return new_filename, message if message else None
+    else:
+        train_filename = filename.replace(file_extension, "_prepared_train.jsonl")
+        val_filename = filename.replace(file_extension, "_prepared_val.jsonl")
+        retry = 0
+        while os.path.isfile(train_filename) or os.path.isfile(val_filename):
+            message += f"File {train_filename} or {val_filename} exists. Trying next available filename increment\n"
+            retry += 1
+            train_filename = filename.replace(file_extension, f"_prepared_train{retry}.jsonl")
+            val_filename = filename.replace(file_extension, f"_prepared_val{retry}.jsonl")
+        return [train_filename, val_filename], message if message else None
+
+
+def split_into_train_validation(df, val_proportion=0.1):
+    n_val = int(val_proportion * len(df))
+    df_val = df.sample(n=n_val, random_state=42)
+    df_train = df.drop(df_val.index)
+    return df_train, df_val
+
+
+def write_df_to_jsonl(df, filename):
+    df.to_json(filename, lines=True, orient="records", force_ascii=False)
+    return f"File {filename} written"
+
+
+def print_select_messages(title, select_messages):
+    print("*" * 40)
+    print(title)
+    print("*" * 40)
+    for idx, message in enumerate(select_messages):
+        print(f"{idx+1}.")
+        print(message)
+
+
+def print_all_messages(messages):
+    messages = [message for message in messages if message]
+    info_messages = [message for message in messages if not message.startswith("TODO")]
+    to_do_messages = [message for message in messages if message.startswith("TODO")]
+
+    print_select_messages("ACTIONABLE MESSAGES", to_do_messages)
+    print_select_messages("INFORMATIONAL MESSAGES", info_messages)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Prepares data for NeMoLLM Customization Service")
+    parser.add_argument("--filename", "-f", required=True)
+    parser.add_argument("--prompt_template", "-pt", default="{prompt}")
+    parser.add_argument("--completion_template", "-ct", default="{completion}")
+    parser.add_argument("--drop_duplicates", "-dd", action="store_true")
+    parser.add_argument("--split_train_validation", "-stv", action="store_true")
+    parser.add_argument(
+        "--short_context_model",
+        "-scm",
+        action="store_true",
+        help="Specifies if using models with shorter context length of 2048 tokens e.g. 5B and 20B models",
+    )
+    parser.add_argument(
+        "--val_proportion",
+        "-vp",
+        default=0.1,
+        type=float,
+        help="Give a number between 0 to 1, \
+                        representing proportion of samples to go into the validation set\
+                        only use when --split_train_validation is set",
+    )
+    args = parser.parse_args()
+    messages = []
+    messages.append(str(args))
+
+    if args.short_context_model:
+        MAX_TOKEN_LENGTH = 2048
+    else:
+        MAX_TOKEN_LENGTH = 4096
+
+    # every token is around 4 chars
+    MAX_TOTAL_CHAR_LENGTH = 4 * MAX_TOKEN_LENGTH
+
+    df, message = load_file_into_df(args.filename)
+    messages.append(message)
+
+    messages.append("-------Before converting into prompt and completion template------ \n")
+    messages[-1] += show_first_example_in_df(df)
+
+    df = convert_into_prompt_completion_only(
+        df, prompt_template=args.prompt_template, completion_template=args.completion_template
+    )
+
+    messages.append("-------After converting into prompt and completion template------ \n")
+    messages[-1] += show_first_example_in_df(df)
+
+    if args.drop_duplicates:
+        df, message = drop_duplicated_rows(df)
+        messages.append(message)
+    else:
+        messages.append(warn_duplicated_rows(df))
+
+    messages.append(warn_missing_suffix(df))
+
+    messages.append(warn_completion_is_not_empty(df))
+    messages.append(warn_imbalanced_completion(df))
+    messages.append(warn_low_n_samples(df))
+
+    df, message = warn_and_drop_long_samples(df, MAX_TOTAL_CHAR_LENGTH)
+    messages.append(message)
+
+    recommended_hyperparameters = recommend_hyperparameters(df)
+    recommend_hyperparameters_message = recommend_hyperparameters_human_readable(recommended_hyperparameters)
+    messages.append(recommend_hyperparameters_message)
+
+    messages.append(estimating_customization_job_time(df, recommended_hyperparameters))
+
+    prepared_filename, message = get_prepared_filename(
+        args.filename, split_train_validation=args.split_train_validation
+    )
+    messages.append(message)
+    if args.split_train_validation:
+        df_train, df_val = split_into_train_validation(df, val_proportion=args.val_proportion)
+        messages.append(write_df_to_jsonl(df_train, prepared_filename[0]))
+        messages.append(write_df_to_jsonl(df_val, prepared_filename[1]))
+    else:
+        messages.append(write_df_to_jsonl(df, prepared_filename))
+
+    print_all_messages(messages)

tools/customization_dataset_preparation/tests/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2023, 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.

tools/customization_dataset_preparation/tests/test_customization_dataset_preparation.py

@@ -0,0 +1,382 @@
+# Copyright (c) 2023, 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 pandas as pd
+import pytest
+from ..customization_dataset_preparation import (
+    convert_into_prompt_completion_only,
+    convert_into_template,
+    drop_duplicated_rows,
+    drop_unrequired_fields,
+    get_common_suffix,
+    get_prepared_filename,
+    parse_template,
+    recommend_hyperparameters,
+    show_first_example_in_df,
+    split_into_train_validation,
+    template_mapper,
+    validate_template,
+    warn_and_drop_long_samples,
+    warn_completion_is_not_empty,
+    warn_duplicated_rows,
+    warn_imbalanced_completion,
+    warn_low_n_samples,
+    warn_missing_suffix,
+)
+
+
+def test_recommend_hyperparameters():
+    df_100 = pd.DataFrame({'prompt': ['prompt'] * 100, 'completion': ['completion'] * 100})
+    assert recommend_hyperparameters(df_100) == {
+        'batch_size': 32,
+        'max_batch_size': 64,
+        'num_virtual_tokens': 10,
+        'encoder_hidden_size': 1024,
+        'lr': 0.005,
+        'epochs': 10,
+        'max_seq_length': 104,
+    }
+
+    df_1000 = pd.DataFrame({'prompt': ['prompt'] * 1000, 'completion': ['completion'] * 1000})
+    assert recommend_hyperparameters(df_1000) == {
+        'batch_size': 32,
+        'max_batch_size': 128,
+        'num_virtual_tokens': 10,
+        'encoder_hidden_size': 2048,
+        'lr': 0.001,
+        'epochs': 10,
+        'max_seq_length': 104,
+    }
+    df_10000 = pd.DataFrame({'prompt': ['prompt'] * 10000, 'completion': ['completion'] * 10000})
+    assert recommend_hyperparameters(df_10000) == {
+        'batch_size': 32,
+        'max_batch_size': 128,
+        'num_virtual_tokens': 10,
+        'encoder_hidden_size': 4096,
+        'lr': 0.0005,
+        'epochs': 10,
+        'max_seq_length': 104,
+    }
+    df_100000 = pd.DataFrame({'prompt': ['prompt'] * 100000, 'completion': ['completion'] * 100000})
+    assert recommend_hyperparameters(df_100000) == {
+        'batch_size': 32,
+        'max_batch_size': 128,
+        'num_virtual_tokens': 10,
+        'encoder_hidden_size': 4096,
+        'lr': 0.0001,
+        'epochs': 10,
+        'max_seq_length': 104,
+    }
+
+
+def test_warn_completion_is_not_empty():
+
+    df_all_empty = pd.DataFrame({'prompt': ['prompt'] * 2, 'completion': [''] * 2})
+
+    msg_all_empty = (
+        "TODO: Note all completion fields are empty. This is possibly expected for inference but not for training"
+    )
+
+    assert warn_completion_is_not_empty(df_all_empty) == msg_all_empty
+
+    df_some_empty = pd.DataFrame({'prompt': ['prompt'] * 2, 'completion': ['', 'completion']})
+
+    msg_some_empty = f"""TODO: completion contains {1} empty values at rows ({[0]})
+                Please check the original file that the fields for prompt template are 
+                not empty and rerun dataset validation"""
+
+    assert warn_completion_is_not_empty(df_some_empty) == msg_some_empty
+
+    df_no_empty = pd.DataFrame({'prompt': ['prompt'] * 2, 'completion': ['completion'] * 2})
+
+    assert warn_completion_is_not_empty(df_no_empty) is None
+
+
+def test_warn_imbalanced_completion():
+    df_generation = pd.DataFrame(
+        {'prompt': [f'prompt{i}' for i in range(100)], 'completion': [f'completion{i}' for i in range(100)]}
+    )
+    assert warn_imbalanced_completion(df_generation) is None
+
+    df_classification_balanced = pd.DataFrame(
+        {'prompt': [f'prompt{i}' for i in range(100)], 'completion': [f'completion{i}' for i in range(5)] * 20}
+    )
+
+    msg_classification_balanced = (
+        f"There are {5} unique completions over {100} samples.\nThe five most common completions are:"
+    )
+    for i in range(5):
+        msg_classification_balanced += f"\n {20} samples ({20.0}%) with completion: completion{i}"
+
+    assert warn_imbalanced_completion(df_classification_balanced) == msg_classification_balanced
+
+    df_classification_imbalanced = pd.DataFrame(
+        {
+            'prompt': [f'prompt{i}' for i in range(100)],
+            'completion': ['completion0'] * 95 + [f'completion{i}' for i in range(5)],
+        }
+    )
+
+    msg_classification_imbalanced = (
+        f"There are {5} unique completions over {100} samples.\nThe five most common completions are:"
+    )
+    msg_classification_imbalanced += f"\n {96} samples ({96.0}%) with completion: completion0"
+    for i in range(1, 5):
+        msg_classification_imbalanced += f"\n {1} samples ({1.0}%) with completion: completion{i}"
+
+    assert warn_imbalanced_completion(df_classification_imbalanced) == msg_classification_imbalanced
+
+
+def test_get_common_suffix():
+    df = pd.DataFrame(
+        {
+            'prompt': [f'prompt{i} answer:' for i in range(100)],
+            'completion': [f'completion{i}' for i in range(100)],
+            'empty_completion': [''] * 100,
+            'some_empty_completion': ['', 'completion'] * 50,
+        }
+    )
+    assert get_common_suffix(df.prompt) == " answer:"
+    assert get_common_suffix(df.completion) == ""
+    assert get_common_suffix(df.empty_completion) == ""
+    assert get_common_suffix(df.some_empty_completion) == ""
+
+
+def test_warn_missing_suffix():
+    df_no_common = pd.DataFrame(
+        {
+            'prompt': [f'prompt{i}' for i in range(100)],
+            'completion': [f'completion{i}' for i in range(100)],
+        }
+    )
+    message = f"TODO: prompt does not have common suffix, please add one (e.g. \\n) at the end of prompt_template\n"
+    message += (
+        f"TODO: completion does not have common suffix, please add one (e.g. \\n) at the end of completion_template\n"
+    )
+
+    assert warn_missing_suffix(df_no_common) == message
+    df_common = pd.DataFrame(
+        {
+            'prompt': [f'prompt{i} answer:' for i in range(100)],
+            'completion': [f'completion{i}\n' for i in range(100)],
+        }
+    )
+    assert warn_missing_suffix(df_common) is None
+
+
+def test_parse_template():
+    template_qa_prompt = "Context: {context}, Question: {question} Answer:"
+    template_qa_completion = "{answer}"
+    template_prompt = "{prompt}"
+    template_completion = "{completion}"
+    assert parse_template(template_qa_prompt) == ['context', 'question']
+    assert parse_template(template_qa_completion) == ['answer']
+    assert parse_template(template_prompt) == ['prompt']
+    assert parse_template(template_completion) == ['completion']
+
+
+def test_validate_template():
+    template = "{prompt}"
+    template_missing_left = "prompt}"
+    template_missing_right = "{prompt"
+    template_twice = "{{prompt}}"
+    template_enclosed = "{prompt{enclosed}}"
+    assert validate_template(template) is None
+    with pytest.raises(ValueError):
+        validate_template(template_missing_left)
+    with pytest.raises(ValueError):
+        validate_template(template_missing_right)
+    with pytest.raises(ValueError):
+        validate_template(template_twice)
+    with pytest.raises(ValueError):
+        validate_template(template_enclosed)
+
+
+def test_warn_duplicated_rows():
+    df_duplicated = pd.DataFrame({'prompt': ['prompt'] * 2, 'completion': ['completion'] * 2})
+
+    message_duplicated = f"TODO: There are {1} duplicated rows "
+    message_duplicated += f"at rows ([1]) \n"
+    message_duplicated += "Please check the original file to make sure that is expected\n"
+    message_duplicated += "If it is not, please add the argument --drop_duplicate"
+
+    assert warn_duplicated_rows(df_duplicated) == message_duplicated
+
+    df_unique = pd.DataFrame({'prompt': ['prompt', 'prompt1'], 'completion': ['completion', 'completion1']})
+    assert warn_duplicated_rows(df_unique) is None
+
+    df_only_prompt_duplicated = pd.DataFrame({'prompt': ['prompt'] * 2, 'completion': ['completion', 'completion1']})
+    assert warn_duplicated_rows(df_only_prompt_duplicated) is None
+
+
+def test_drop_duplicated_rows():
+    df_deduplicated = pd.DataFrame({'prompt': ['prompt'], 'completion': ['completion']})
+
+    df_duplicated = pd.DataFrame({'prompt': ['prompt'] * 2, 'completion': ['completion'] * 2})
+    message_duplicated = "There are 1 duplicated rows\n"
+    message_duplicated += "Removed 1 duplicate rows"
+
+    assert drop_duplicated_rows(df_duplicated)[0].equals(df_deduplicated)
+    assert drop_duplicated_rows(df_duplicated)[1] == message_duplicated
+
+    df_unique = pd.DataFrame({'prompt': ['prompt', 'prompt1'], 'completion': ['completion', 'completion1']})
+    assert drop_duplicated_rows(df_unique) == (df_unique, None)
+
+    df_only_prompt_duplicated = pd.DataFrame({'prompt': ['prompt'] * 2, 'completion': ['completion', 'completion1']})
+    assert drop_duplicated_rows(df_only_prompt_duplicated) == (df_only_prompt_duplicated, None)
+
+
+def test_template_mapper():
+
+    df = pd.DataFrame(
+        {
+            'prompt': ['prompt sample'],
+        }
+    )
+
+    template = "{prompt}"
+    field_names = ['prompt']
+    assert template_mapper(df.iloc[0], field_names, template) == 'prompt sample'
+
+    df_qa = pd.DataFrame({'question': ['question sample'], 'context': ['context sample']})
+
+    template_qa = "Context: {context} Question: {question} Answer:"
+    field_names_qa = ['context', 'question']
+    assert (
+        template_mapper(df_qa.iloc[0], field_names_qa, template_qa)
+        == "Context: context sample Question: question sample Answer:"
+    )
+
+
+def test_drop_unrequired_fields():
+    df = pd.DataFrame(
+        {'question': ['question'], 'context': ['context'], 'prompt': ['prompt'], 'completion': ['completion']}
+    )
+
+    df_dropped_unnecessary_fields = pd.DataFrame({'prompt': ['prompt'], 'completion': ['completion']})
+    assert df_dropped_unnecessary_fields.equals(drop_unrequired_fields(df))
+
+
+def test_convert_into_template():
+    df_non_existant_field_name = pd.DataFrame({'question': ['question']})
+
+    template = "Context: {context} Question: {question} Answer:"
+    with pytest.raises(ValueError):
+        convert_into_template(df_non_existant_field_name, template)
+
+    df = pd.DataFrame(
+        {
+            'question': ['question sample'],
+            'context': ['context sample'],
+        }
+    )
+
+    df_prompt = pd.DataFrame(
+        {
+            'question': ['question sample'],
+            'context': ['context sample'],
+            'prompt': ["Context: context sample Question: question sample Answer:"],
+        }
+    )
+    assert convert_into_template(df, template).equals(df_prompt)
+
+
+def test_convert_into_prompt_completion_only():
+    df = pd.DataFrame({'question': ['question sample'], 'context': ['context sample'], 'answer': ['answer sample']})
+
+    df_prompt = pd.DataFrame(
+        {'prompt': ["Context: context sample Question: question sample Answer:"], 'completion': ["answer sample"]}
+    )
+
+    prompt_template = "Context: {context} Question: {question} Answer:"
+    completion_template = "{answer}"
+
+    assert df_prompt.equals(
+        convert_into_prompt_completion_only(
+            df, prompt_template=prompt_template, completion_template=completion_template
+        )
+    )
+    assert df_prompt.equals(convert_into_prompt_completion_only(df_prompt))
+
+
+def get_indexes_of_long_examples(df, max_total_char_length):
+    long_examples = df.apply(lambda x: len(x.prompt) + len(x.completion) > max_total_char_length, axis=1)
+    return df.reset_index().index[long_examples].tolist()
+
+
+def test_warn_and_drop_long_samples():
+    df = pd.DataFrame({'prompt': ['a' * 12000, 'a' * 9000, 'a'], 'completion': ['b' * 12000, 'b' * 2000, 'b']})
+
+    expected_df = pd.DataFrame({'prompt': ['a'], 'completion': ['b']})
+    message = f"""TODO: There are {2} / {3} 
+        samples that have its prompt and completion too long 
+        (over {10000} chars), which have been dropped."""
+
+    assert expected_df.equals(warn_and_drop_long_samples(df, 10000)[0])
+    assert warn_and_drop_long_samples(df, 10000)[1] == message
+
+    df_short = pd.DataFrame({'prompt': ['a'] * 2, 'completion': ['b'] * 2})
+
+    assert warn_and_drop_long_samples(df_short, 10000) == (df_short, None)
+
+
+def test_warn_low_n_samples():
+    df_low = pd.DataFrame({'prompt': ['a'] * 10, 'completion': ['b'] * 10})
+
+    df_high = pd.DataFrame({'prompt': ['a'] * 100, 'completion': ['b'] * 100})
+
+    message = (
+        "TODO: We would recommend having more samples (>64) if possible but current_file only contains 10 samples. "
+    )
+    assert warn_low_n_samples(df_low) == message
+    assert warn_low_n_samples(df_high) is None
+
+
+def test_show_first_example_in_df():
+    df = pd.DataFrame({'question': ['question sample'], 'context': ['context sample'], 'answer': ['answer sample']})
+
+    message = f"-->Column question:\nquestion sample\n"
+    message += f"-->Column context:\ncontext sample\n"
+    message += f"-->Column answer:\nanswer sample\n"
+
+    assert message == show_first_example_in_df(df)
+
+
+def test_get_prepared_filename():
+    filename = "tmp/sample.jsonl"
+    prepared_filename = "tmp/sample_prepared.jsonl"
+    prepared_train_filename = "tmp/sample_prepared_train.jsonl"
+    prepared_val_filename = "tmp/sample_prepared_val.jsonl"
+    assert get_prepared_filename(filename) == (prepared_filename, None)
+    assert get_prepared_filename(filename, split_train_validation=True) == (
+        [
+            prepared_train_filename,
+            prepared_val_filename,
+        ],
+        None,
+    )
+    csv_filename = "tmp/sample.csv"
+    prepared_filename = "tmp/sample_prepared.jsonl"
+    assert get_prepared_filename(csv_filename) == (prepared_filename, None)
+
+
+def test_split_into_train_validation():
+    df = pd.DataFrame({'prompt': ['a'] * 10, 'completion': ['b'] * 10})
+    df_train, df_val = split_into_train_validation(df, val_proportion=0.1)
+    assert len(df_train) == 9
+    assert len(df_val) == 1
+
+    df_train, df_val = split_into_train_validation(df, val_proportion=0.2)
+    assert len(df_train) == 8
+    assert len(df_val) == 2

tools/nemo_forced_aligner/README.md

@@ -0,0 +1,30 @@
+# NeMo Forced Aligner (NFA)
+
+<p align="center">
+Try it out: <a href="https://huggingface.co/spaces/erastorgueva-nv/NeMo-Forced-Aligner">HuggingFace Space 🎤</a> | Tutorial: <a href="https://colab.research.google.com/github/NVIDIA/NeMo/blob/main/tutorials/tools/NeMo_Forced_Aligner_Tutorial.ipynb">"How to use NFA?" 🚀</a> | Blog post: <a href="https://nvidia.github.io/NeMo/blogs/2023/2023-08-forced-alignment/">"How does forced alignment work?" 📚</a>
+</p>
+
+<p align="center">
+<img width="80%" src="https://github.com/NVIDIA/NeMo/releases/download/v1.20.0/nfa_forced_alignment_pipeline.png">
+</p>
+
+NFA is a tool for generating token-, word- and segment-level timestamps of speech in audio using NeMo's CTC-based Automatic Speech Recognition models. You can provide your own reference text, or use ASR-generated transcription. You can use NeMo's ASR Model checkpoints out of the box in [14+ languages](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/results.html#speech-recognition-languages), or train your own model. NFA can be used on long audio files of 1+ hours duration (subject to your hardware and the ASR model used).
+
+
+## Quickstart
+1. Install [NeMo](https://docs.nvidia.com/nemo/speech/nightly/starthere/install.html) with the ASR collection.
+2. Prepare a NeMo-style manifest containing the paths of audio files you would like to process, and (optionally) their text.
+3. Run NFA's `align.py` script with the desired config, e.g.:
+    ``` bash
+    python <path_to_NeMo>/tools/nemo_forced_aligner/align.py \
+	    pretrained_name="stt_en_fastconformer_hybrid_large_pc" \
+	    manifest_filepath=<path to manifest of utterances you want to align> \
+	    output_dir=<path to where your output files will be saved>
+    ```
+
+<p align="center">
+	<img src="https://github.com/NVIDIA/NeMo/releases/download/v1.20.0/nfa_run.png">
+</p>
+
+## Documentation 
+More documentation is available [here](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/tools/nemo_forced_aligner.html).

tools/nemo_forced_aligner/align.py

@@ -0,0 +1,387 @@
+# Copyright (c) 2023, 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 copy
+import math
+import os
+from dataclasses import dataclass, field, is_dataclass
+from pathlib import Path
+from typing import List, Optional
+
+import torch
+from omegaconf import OmegaConf
+from utils.data_prep import (
+    get_batch_starts_ends,
+    get_manifest_lines_batch,
+    is_entry_in_all_lines,
+    is_entry_in_any_lines,
+)
+from utils.make_ass_files import make_ass_files
+from utils.make_ctm_files import make_ctm_files
+from utils.make_output_manifest import write_manifest_out_line
+
+from nemo.collections.asr.models.ctc_models import EncDecCTCModel
+from nemo.collections.asr.models.hybrid_rnnt_ctc_models import EncDecHybridRNNTCTCModel
+from nemo.collections.asr.parts.utils.streaming_utils import FrameBatchASR
+from nemo.collections.asr.parts.utils.transcribe_utils import setup_model
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+
+try:
+    from nemo.collections.asr.parts.utils.aligner_utils import (
+        add_t_start_end_to_utt_obj,
+        get_batch_variables,
+        viterbi_decoding,
+    )
+except ImportError:
+    raise ImportError(
+        "Missing required dependency for NFA. "
+        "Install NeMo with NFA utilities support:\n"
+        "  pip install 'nemo-toolkit[all]>=2.5.0'\n"
+        "Or install the latest development version:\n"
+        "  pip install git+https://github.com/NVIDIA-NeMo/NeMo.git"
+    )
+"""
+Align the utterances in manifest_filepath. 
+Results are saved in ctm files in output_dir.
+
+Arguments:
+    pretrained_name: string specifying the name of a CTC NeMo ASR model which will be automatically downloaded
+        from NGC and used for generating the log-probs which we will use to do alignment.
+        Note: NFA can only use CTC models (not Transducer models) at the moment.
+    model_path: string specifying the local filepath to a CTC NeMo ASR model which will be used to generate the
+        log-probs which we will use to do alignment.
+        Note: NFA can only use CTC models (not Transducer models) at the moment.
+        Note: if a model_path is provided, it will override the pretrained_name.
+    manifest_filepath: filepath to the manifest of the data you want to align,
+        containing 'audio_filepath' and 'text' fields.
+    output_dir: the folder where output CTM files and new JSON manifest will be saved.
+    align_using_pred_text: if True, will transcribe the audio using the specified model and then use that transcription 
+        as the reference text for the forced alignment. 
+    transcribe_device: None, or a string specifying the device that will be used for generating log-probs (i.e. "transcribing").
+        The string needs to be in a format recognized by torch.device(). If None, NFA will set it to 'cuda' if it is available 
+        (otherwise will set it to 'cpu').
+    viterbi_device: None, or string specifying the device that will be used for doing Viterbi decoding. 
+        The string needs to be in a format recognized by torch.device(). If None, NFA will set it to 'cuda' if it is available 
+        (otherwise will set it to 'cpu').
+    batch_size: int specifying batch size that will be used for generating log-probs and doing Viterbi decoding.
+    use_local_attention: boolean flag specifying whether to try to use local attention for the ASR Model (will only
+        work if the ASR Model is a Conformer model). If local attention is used, we will set the local attention context 
+        size to [64,64].
+    additional_segment_grouping_separator: an optional string or list of strings used to separate the text into smaller segments. 
+        If this is not specified, then the whole text will be treated as a single segment. 
+    remove_blank_tokens_from_ctm:  a boolean denoting whether to remove <blank> tokens from token-level output CTMs. 
+    audio_filepath_parts_in_utt_id: int specifying how many of the 'parts' of the audio_filepath
+        we will use (starting from the final part of the audio_filepath) to determine the 
+        utt_id that will be used in the CTM files. Note also that any spaces that are present in the audio_filepath 
+        will be replaced with dashes, so as not to change the number of space-separated elements in the 
+        CTM files.
+        e.g. if audio_filepath is "/a/b/c/d/e 1.wav" and audio_filepath_parts_in_utt_id is 1 => utt_id will be "e1"
+        e.g. if audio_filepath is "/a/b/c/d/e 1.wav" and audio_filepath_parts_in_utt_id is 2 => utt_id will be "d_e1"
+        e.g. if audio_filepath is "/a/b/c/d/e 1.wav" and audio_filepath_parts_in_utt_id is 3 => utt_id will be "c_d_e1"
+    use_buffered_infer: False, if set True, using streaming to do get the logits for alignment
+                        This flag is useful when aligning large audio file.
+                        However, currently the chunk streaming inference does not support batch inference,
+                        which means even you set batch_size > 1, it will only infer one by one instead of doing
+                        the whole batch inference together.
+    chunk_len_in_secs: float chunk length in seconds
+    total_buffer_in_secs: float  Length of buffer (chunk + left and right padding) in seconds
+    chunk_batch_size: int batch size for buffered chunk inference,
+                      which will cut one audio into segments and do inference on chunk_batch_size segments at a time
+
+    simulate_cache_aware_streaming: False, if set True, using cache aware streaming to do get the logits for alignment
+
+    save_output_file_formats: List of strings specifying what type of output files to save (default: ["ctm", "ass"])
+    ctm_file_config: CTMFileConfig to specify the configuration of the output CTM files
+    ass_file_config: ASSFileConfig to specify the configuration of the output ASS files
+"""
+
+
+@dataclass
+class CTMFileConfig:
+    remove_blank_tokens: bool = False
+    # minimum duration (in seconds) for timestamps in the CTM.If any line in the CTM has a
+    # duration lower than this, it will be enlarged from the middle outwards until it
+    # meets the minimum_timestamp_duration, or reaches the beginning or end of the audio file.
+    # Note that this may cause timestamps to overlap.
+    minimum_timestamp_duration: float = 0
+
+
+@dataclass
+class ASSFileConfig:
+    fontsize: int = 20
+    vertical_alignment: str = "center"
+    # if resegment_text_to_fill_space is True, the ASS files will use new segments
+    # such that each segment will not take up more than (approximately) max_lines_per_segment
+    # when the ASS file is applied to a video
+    resegment_text_to_fill_space: bool = False
+    max_lines_per_segment: int = 2
+    text_already_spoken_rgb: List[int] = field(default_factory=lambda: [49, 46, 61])  # dark gray
+    text_being_spoken_rgb: List[int] = field(default_factory=lambda: [57, 171, 9])  # dark green
+    text_not_yet_spoken_rgb: List[int] = field(default_factory=lambda: [194, 193, 199])  # light gray
+
+
+@dataclass
+class AlignmentConfig:
+    # Required configs
+    pretrained_name: Optional[str] = None
+    model_path: Optional[str] = None
+    manifest_filepath: Optional[str] = None
+    output_dir: Optional[str] = None
+
+    # General configs
+    align_using_pred_text: bool = False
+    transcribe_device: Optional[str] = None
+    viterbi_device: Optional[str] = None
+    batch_size: int = 1
+    use_local_attention: bool = True
+    additional_segment_grouping_separator: Optional[List[str]] = field(default_factory=lambda: ['.', '?', '!', '...'])
+    audio_filepath_parts_in_utt_id: int = 1
+
+    # Buffered chunked streaming configs
+    use_buffered_chunked_streaming: bool = False
+    chunk_len_in_secs: float = 1.6
+    total_buffer_in_secs: float = 4.0
+    chunk_batch_size: int = 32
+
+    # Cache aware streaming configs
+    simulate_cache_aware_streaming: Optional[bool] = False
+
+    # Output file configs
+    save_output_file_formats: List[str] = field(default_factory=lambda: ["ctm", "ass"])
+    ctm_file_config: CTMFileConfig = field(default_factory=lambda: CTMFileConfig())
+    ass_file_config: ASSFileConfig = field(default_factory=lambda: ASSFileConfig())
+
+
+@hydra_runner(config_name="AlignmentConfig", schema=AlignmentConfig)
+def main(cfg: AlignmentConfig):
+
+    logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}')
+
+    if is_dataclass(cfg):
+        cfg = OmegaConf.structured(cfg)
+
+    # Validate config
+    if cfg.model_path is None and cfg.pretrained_name is None:
+        raise ValueError("Both cfg.model_path and cfg.pretrained_name cannot be None")
+
+    if cfg.model_path is not None and cfg.pretrained_name is not None:
+        raise ValueError("One of cfg.model_path and cfg.pretrained_name must be None")
+
+    if cfg.manifest_filepath is None:
+        raise ValueError("cfg.manifest_filepath must be specified")
+
+    if cfg.output_dir is None:
+        raise ValueError("cfg.output_dir must be specified")
+
+    if cfg.batch_size < 1:
+        raise ValueError("cfg.batch_size cannot be zero or a negative number")
+
+    if cfg.additional_segment_grouping_separator == "" or cfg.additional_segment_grouping_separator == " ":
+        raise ValueError("cfg.additional_grouping_separator cannot be empty string or space character")
+    elif cfg.additional_segment_grouping_separator is not None and cfg.additional_segment_grouping_separator != []:
+        logging.warning(
+            f"`additional_segment_grouping_separator` is set to {cfg.additional_segment_grouping_separator}. "
+            "BEHAVIOR CHANGE: Starting in NeMo 2.5.0, separators are preserved in segment text after splitting. "
+            "In previous versions, separators were removed. This affects the behavior of NFA."
+        )
+
+    if cfg.ctm_file_config.minimum_timestamp_duration < 0:
+        raise ValueError("cfg.minimum_timestamp_duration cannot be a negative number")
+
+    if cfg.ass_file_config.vertical_alignment not in ["top", "center", "bottom"]:
+        raise ValueError("cfg.ass_file_config.vertical_alignment must be one of 'top', 'center' or 'bottom'")
+
+    for rgb_list in [
+        cfg.ass_file_config.text_already_spoken_rgb,
+        cfg.ass_file_config.text_already_spoken_rgb,
+        cfg.ass_file_config.text_already_spoken_rgb,
+    ]:
+        if len(rgb_list) != 3:
+            raise ValueError(
+                "cfg.ass_file_config.text_already_spoken_rgb,"
+                " cfg.ass_file_config.text_being_spoken_rgb,"
+                " and cfg.ass_file_config.text_already_spoken_rgb all need to contain"
+                " exactly 3 elements."
+            )
+
+    # Validate manifest contents
+    if not is_entry_in_all_lines(cfg.manifest_filepath, "audio_filepath"):
+        raise RuntimeError(
+            "At least one line in cfg.manifest_filepath does not contain an 'audio_filepath' entry. "
+            "All lines must contain an 'audio_filepath' entry."
+        )
+
+    if cfg.align_using_pred_text:
+        if is_entry_in_any_lines(cfg.manifest_filepath, "pred_text"):
+            raise RuntimeError(
+                "Cannot specify cfg.align_using_pred_text=True when the manifest at cfg.manifest_filepath "
+                "contains 'pred_text' entries. This is because the audio will be transcribed and may produce "
+                "a different 'pred_text'. This may cause confusion."
+            )
+    else:
+        if not is_entry_in_all_lines(cfg.manifest_filepath, "text"):
+            raise RuntimeError(
+                "At least one line in cfg.manifest_filepath does not contain a 'text' entry. "
+                "NFA requires all lines to contain a 'text' entry when cfg.align_using_pred_text=False."
+            )
+
+    # init devices
+    if cfg.transcribe_device is None:
+        transcribe_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        transcribe_device = torch.device(cfg.transcribe_device)
+    logging.info(f"Device to be used for transcription step (`transcribe_device`) is {transcribe_device}")
+
+    if cfg.viterbi_device is None:
+        viterbi_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        viterbi_device = torch.device(cfg.viterbi_device)
+    logging.info(f"Device to be used for viterbi step (`viterbi_device`) is {viterbi_device}")
+
+    if transcribe_device.type == 'cuda' or viterbi_device.type == 'cuda':
+        logging.warning(
+            'One or both of transcribe_device and viterbi_device are GPUs. If you run into OOM errors '
+            'it may help to change both devices to be the CPU.'
+        )
+
+    # load model
+    model, _ = setup_model(cfg, transcribe_device)
+    model.eval()
+
+    if isinstance(model, EncDecHybridRNNTCTCModel):
+        model.change_decoding_strategy(decoder_type="ctc")
+
+    if cfg.use_local_attention:
+        logging.info(
+            "Flag use_local_attention is set to True => will try to use local attention for model if it allows it"
+        )
+        model.change_attention_model(self_attention_model="rel_pos_local_attn", att_context_size=[64, 64])
+
+    if not (isinstance(model, EncDecCTCModel) or isinstance(model, EncDecHybridRNNTCTCModel)):
+        raise NotImplementedError(
+            f"Model is not an instance of NeMo EncDecCTCModel or ENCDecHybridRNNTCTCModel."
+            " Currently only instances of these models are supported"
+        )
+
+    if cfg.ctm_file_config.minimum_timestamp_duration > 0:
+        logging.warning(
+            f"cfg.ctm_file_config.minimum_timestamp_duration has been set to {cfg.ctm_file_config.minimum_timestamp_duration} seconds. "
+            "This may cause the alignments for some tokens/words/additional segments to be overlapping."
+        )
+
+    buffered_chunk_params = {}
+    if cfg.use_buffered_chunked_streaming:
+        model_cfg = copy.deepcopy(model._cfg)
+
+        OmegaConf.set_struct(model_cfg.preprocessor, False)
+        # some changes for streaming scenario
+        model_cfg.preprocessor.dither = 0.0
+        model_cfg.preprocessor.pad_to = 0
+
+        if model_cfg.preprocessor.normalize != "per_feature":
+            logging.error(
+                "Only EncDecCTCModelBPE models trained with per_feature normalization are supported currently"
+            )
+        # Disable config overwriting
+        OmegaConf.set_struct(model_cfg.preprocessor, True)
+
+        feature_stride = model_cfg.preprocessor['window_stride']
+        model_stride_in_secs = feature_stride * cfg.model_downsample_factor
+        total_buffer = cfg.total_buffer_in_secs
+        chunk_len = float(cfg.chunk_len_in_secs)
+        tokens_per_chunk = math.ceil(chunk_len / model_stride_in_secs)
+        mid_delay = math.ceil((chunk_len + (total_buffer - chunk_len) / 2) / model_stride_in_secs)
+        logging.info(f"tokens_per_chunk is {tokens_per_chunk}, mid_delay is {mid_delay}")
+
+        model = FrameBatchASR(
+            asr_model=model,
+            frame_len=chunk_len,
+            total_buffer=cfg.total_buffer_in_secs,
+            batch_size=cfg.chunk_batch_size,
+        )
+        buffered_chunk_params = {
+            "delay": mid_delay,
+            "model_stride_in_secs": model_stride_in_secs,
+            "tokens_per_chunk": tokens_per_chunk,
+        }
+    # get start and end line IDs of batches
+    starts, ends = get_batch_starts_ends(cfg.manifest_filepath, cfg.batch_size)
+
+    # init output_timestep_duration = None and we will calculate and update it during the first batch
+    output_timestep_duration = None
+
+    # init f_manifest_out
+    os.makedirs(cfg.output_dir, exist_ok=True)
+    tgt_manifest_name = str(Path(cfg.manifest_filepath).stem) + "_with_output_file_paths.json"
+    tgt_manifest_filepath = str(Path(cfg.output_dir) / tgt_manifest_name)
+    f_manifest_out = open(tgt_manifest_filepath, 'w')
+
+    # get alignment and save in CTM batch-by-batch
+    for start, end in zip(starts, ends):
+        manifest_lines_batch = get_manifest_lines_batch(cfg.manifest_filepath, start, end)
+
+        if not cfg.align_using_pred_text:
+            gt_text_batch = [line.get('text', '') for line in manifest_lines_batch]
+        else:
+            gt_text_batch = None
+
+        (
+            log_probs_batch,
+            y_batch,
+            T_batch,
+            U_batch,
+            utt_obj_batch,
+            output_timestep_duration,
+        ) = get_batch_variables(
+            audio=[line['audio_filepath'] for line in manifest_lines_batch],
+            model=model,
+            segment_separators=cfg.additional_segment_grouping_separator,
+            align_using_pred_text=cfg.align_using_pred_text,
+            audio_filepath_parts_in_utt_id=cfg.audio_filepath_parts_in_utt_id,
+            gt_text_batch=gt_text_batch,
+            output_timestep_duration=output_timestep_duration,
+            simulate_cache_aware_streaming=cfg.simulate_cache_aware_streaming,
+            use_buffered_chunked_streaming=cfg.use_buffered_chunked_streaming,
+            buffered_chunk_params=buffered_chunk_params,
+        )
+
+        alignments_batch = viterbi_decoding(log_probs_batch, y_batch, T_batch, U_batch, viterbi_device)
+
+        for utt_obj, alignment_utt in zip(utt_obj_batch, alignments_batch):
+
+            utt_obj = add_t_start_end_to_utt_obj(utt_obj, alignment_utt, output_timestep_duration)
+
+            if "ctm" in cfg.save_output_file_formats:
+                utt_obj = make_ctm_files(
+                    utt_obj,
+                    cfg.output_dir,
+                    cfg.ctm_file_config,
+                )
+
+            if "ass" in cfg.save_output_file_formats:
+                utt_obj = make_ass_files(utt_obj, cfg.output_dir, cfg.ass_file_config)
+
+            write_manifest_out_line(
+                f_manifest_out,
+                utt_obj,
+            )
+
+    f_manifest_out.close()
+
+    return None
+
+
+if __name__ == "__main__":
+    main()

tools/nemo_forced_aligner/align_eou.py

@@ -0,0 +1,589 @@
+# Copyright (c) 2025, 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 copy
+import json
+import math
+import os
+import shutil
+import unicodedata
+import uuid
+from dataclasses import dataclass, field, is_dataclass
+from pathlib import Path
+from string import punctuation
+from typing import List, Optional
+
+import torch
+from omegaconf import OmegaConf
+from utils.data_prep import (
+    get_batch_starts_ends,
+    get_manifest_lines_batch,
+    is_entry_in_all_lines,
+    is_entry_in_any_lines,
+)
+from utils.make_ass_files import make_ass_files
+from utils.make_ctm_files import make_ctm_files
+from utils.make_output_manifest import write_manifest_out_line
+
+from nemo.collections.asr.models.ctc_models import EncDecCTCModel
+from nemo.collections.asr.models.hybrid_rnnt_ctc_models import EncDecHybridRNNTCTCModel
+from nemo.collections.asr.parts.utils.streaming_utils import FrameBatchASR
+from nemo.collections.asr.parts.utils.transcribe_utils import setup_model
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+
+try:
+    from nemo.collections.asr.parts.utils.aligner_utils import (
+        add_t_start_end_to_utt_obj,
+        get_batch_variables,
+        viterbi_decoding,
+    )
+except ImportError:
+    raise ImportError(
+        "Missing required dependency for NFA. "
+        "Install NeMo with NFA utilities support:\n"
+        "  pip install 'nemo-toolkit[all]>=2.5.0'\n"
+        "Or install the latest development version:\n"
+        "  pip install git+https://github.com/NVIDIA-NeMo/NeMo.git"
+    )
+
+"""
+Align the utterances in manifest_filepath. 
+Results are saved in ctm files in output_dir as well as json manifest in output_manifest_filepath.
+If no output_manifest_filepath is specified, it will save the results in the same parent directory as 
+the input manifest_filepath.
+
+Arguments:
+    pretrained_name: string specifying the name of a CTC NeMo ASR model which will be automatically downloaded
+        from NGC and used for generating the log-probs which we will use to do alignment.
+        Note: NFA can only use CTC models (not Transducer models) at the moment.
+    model_path: string specifying the local filepath to a CTC NeMo ASR model which will be used to generate the
+        log-probs which we will use to do alignment.
+        Note: NFA can only use CTC models (not Transducer models) at the moment.
+        Note: if a model_path is provided, it will override the pretrained_name.
+    manifest_filepath: filepath to the manifest of the data you want to align,
+        containing 'audio_filepath' and 'text' fields.
+    output_dir: the folder where output CTM files and new JSON manifest will be saved.
+    output_manifest_filepath: Optional[str] = None  # output of manfiest with sou_time and eou_time
+    manifest_pattern: Optional[str] = None  # pattern used in Path.glob() for finding manifests
+
+    align_using_pred_text: if True, will transcribe the audio using the specified model and then use that transcription 
+        as the reference text for the forced alignment. 
+    transcribe_device: None, or a string specifying the device that will be used for generating log-probs (i.e. "transcribing").
+        The string needs to be in a format recognized by torch.device(). If None, NFA will set it to 'cuda' if it is available 
+        (otherwise will set it to 'cpu').
+    viterbi_device: None, or string specifying the device that will be used for doing Viterbi decoding. 
+        The string needs to be in a format recognized by torch.device(). If None, NFA will set it to 'cuda' if it is available 
+        (otherwise will set it to 'cpu').
+    batch_size: int specifying batch size that will be used for generating log-probs and doing Viterbi decoding.
+    use_local_attention: boolean flag specifying whether to try to use local attention for the ASR Model (will only
+        work if the ASR Model is a Conformer model). If local attention is used, we will set the local attention context 
+        size to [64,64].
+    additional_segment_grouping_separator: an optional string used to separate the text into smaller segments. 
+        If this is not specified, then the whole text will be treated as a single segment. 
+    remove_blank_tokens_from_ctm:  a boolean denoting whether to remove <blank> tokens from token-level output CTMs. 
+    audio_filepath_parts_in_utt_id: int specifying how many of the 'parts' of the audio_filepath
+        we will use (starting from the final part of the audio_filepath) to determine the 
+        utt_id that will be used in the CTM files. Note also that any spaces that are present in the audio_filepath 
+        will be replaced with dashes, so as not to change the number of space-separated elements in the 
+        CTM files.
+        e.g. if audio_filepath is "/a/b/c/d/e 1.wav" and audio_filepath_parts_in_utt_id is 1 => utt_id will be "e1"
+        e.g. if audio_filepath is "/a/b/c/d/e 1.wav" and audio_filepath_parts_in_utt_id is 2 => utt_id will be "d_e1"
+        e.g. if audio_filepath is "/a/b/c/d/e 1.wav" and audio_filepath_parts_in_utt_id is 3 => utt_id will be "c_d_e1"
+    use_buffered_infer: False, if set True, using streaming to do get the logits for alignment
+                        This flag is useful when aligning large audio file.
+                        However, currently the chunk streaming inference does not support batch inference,
+                        which means even you set batch_size > 1, it will only infer one by one instead of doing
+                        the whole batch inference together.
+    chunk_len_in_secs: float chunk length in seconds
+    total_buffer_in_secs: float  Length of buffer (chunk + left and right padding) in seconds
+    chunk_batch_size: int batch size for buffered chunk inference,
+                      which will cut one audio into segments and do inference on chunk_batch_size segments at a time
+
+    simulate_cache_aware_streaming: False, if set True, using cache aware streaming to do get the logits for alignment
+
+    save_output_file_formats: List of strings specifying what type of output files to save (default: ["ctm", "ass"])
+    ctm_file_config: CTMFileConfig to specify the configuration of the output CTM files
+    ass_file_config: ASSFileConfig to specify the configuration of the output ASS files
+"""
+
+
+@dataclass
+class CTMFileConfig:
+    remove_blank_tokens: bool = False
+    # minimum duration (in seconds) for timestamps in the CTM.If any line in the CTM has a
+    # duration lower than this, it will be enlarged from the middle outwards until it
+    # meets the minimum_timestamp_duration, or reaches the beginning or end of the audio file.
+    # Note that this may cause timestamps to overlap.
+    minimum_timestamp_duration: float = 0
+
+
+@dataclass
+class ASSFileConfig:
+    fontsize: int = 20
+    vertical_alignment: str = "center"
+    # if resegment_text_to_fill_space is True, the ASS files will use new segments
+    # such that each segment will not take up more than (approximately) max_lines_per_segment
+    # when the ASS file is applied to a video
+    resegment_text_to_fill_space: bool = False
+    max_lines_per_segment: int = 2
+    text_already_spoken_rgb: List[int] = field(default_factory=lambda: [49, 46, 61])  # dark gray
+    text_being_spoken_rgb: List[int] = field(default_factory=lambda: [57, 171, 9])  # dark green
+    text_not_yet_spoken_rgb: List[int] = field(default_factory=lambda: [194, 193, 199])  # light gray
+
+
+@dataclass
+class AlignmentConfig:
+    # Required configs
+    pretrained_name: Optional[str] = None
+    model_path: Optional[str] = None
+    manifest_filepath: Optional[str] = None  # path to manifest file or directory
+    output_dir: Optional[str] = '.tmp'  # set it to .tmp and will be removed after alignment
+    output_manifest_filepath: Optional[str] = None  # output of manfiest with sou_time and eou_time
+    manifest_pattern: Optional[str] = None  # pattern used in Path.glob() for finding manifests
+
+    # General configs
+    align_using_pred_text: bool = False
+    transcribe_device: Optional[str] = None
+    viterbi_device: Optional[str] = None
+    batch_size: int = 1
+    use_local_attention: bool = True
+    additional_segment_grouping_separator: Optional[str] = None
+    audio_filepath_parts_in_utt_id: int = 4
+
+    # Buffered chunked streaming configs
+    use_buffered_chunked_streaming: bool = False
+    chunk_len_in_secs: float = 1.6
+    total_buffer_in_secs: float = 4.0
+    chunk_batch_size: int = 32
+
+    # Cache aware streaming configs
+    simulate_cache_aware_streaming: Optional[bool] = False
+
+    # Output file configs
+    save_output_file_formats: List[str] = field(default_factory=lambda: ["ctm", "ass"])
+    ctm_file_config: CTMFileConfig = field(default_factory=lambda: CTMFileConfig())
+    ass_file_config: ASSFileConfig = field(default_factory=lambda: ASSFileConfig())
+
+    # remove tmp dir after alignment
+    remove_tmp_dir: bool = False
+    clean_text: bool = True
+
+    # For multi-node multi-gpu processing
+    num_nodes: int = 1  # total num of nodes/machines
+    num_gpus: int = 1  # num of GPUs per node/machine
+    node_idx: int = 0  # current node index
+    gpu_idx: int = 0  # current GPU index
+
+
+def unicode_to_ascii(text: str) -> str:
+    """
+    Converts text with accented or special Latin characters (e.g., ó, ñ, ū, ō)
+    into their closest ASCII equivalents.
+    """
+    # Normalize the string to NFKD to separate base characters from diacritics
+    normalized = unicodedata.normalize('NFKD', text)
+
+    # Encode to ASCII bytes, ignoring characters that can't be converted
+    ascii_bytes = normalized.encode('ascii', 'ignore')
+
+    # Decode back to string
+    ascii_text = ascii_bytes.decode('ascii')
+
+    return ascii_text
+
+
+def drop_pnc(text):
+    """
+    Clean the text by removing invalid characters and converting to lowercase.
+
+    :param text: Input text.
+    :return: Cleaned text.
+    """
+    valid_chars = "abcdefghijklmnopqrstuvwxyz'"
+    text = text.lower()
+    text = unicode_to_ascii(text)
+    text = text.replace(":", " ")
+    text = ''.join([c for c in text if c in valid_chars or c.isspace() or c == "'"])
+    return " ".join(text.split()).strip()
+
+
+def clean_text(manifest: List[dict]):
+    """
+    Clean the text in the manifest.
+    Args:
+        manifest: List of dictionaries with the text to clean.
+
+    Returns:
+        manifest: List of dictionaries with the cleaned text.
+    """
+    punctuations = punctuation.replace("'", "")
+    # replace_with_space = [char for char in '/?*\",.:=?_{|}~¨«·»¡¿„…‧‹›≪≫!:;ː→']
+    replace_with_blank = [char for char in '`¨´‘’“”`ʻ‘’“"‘”']
+    replace_with_apos = [char for char in '‘’ʻ‘’‘']
+
+    for i in range(len(manifest)):
+        text = manifest[i]["text"].strip().lower()  # type: str
+        text = text.translate(str.maketrans("", "", punctuations))
+        text = drop_pnc(text)
+        for c in replace_with_blank:
+            text = text.replace(c, "")
+        for c in replace_with_apos:
+            text = text.replace(c, "'")
+        manifest[i]["text"] = text
+    return manifest
+
+
+def get_manifests_for_this_rank(manifest_list, num_nodes, num_gpus, node_idx, gpu_idx):
+    """
+    Get the manifest files for this rank.
+    """
+    if len(manifest_list) == 0:
+        return manifest_list
+
+    assert num_nodes > 0, "num_nodes must be greater than 0"
+    assert num_gpus > 0, "num_gpus must be greater than 0"
+    assert 0 <= node_idx < num_nodes, f"node_idx {node_idx} must be between 0 and {num_nodes - 1}"
+    assert 0 <= gpu_idx < num_gpus, f"gpu_idx {gpu_idx} must be between 0 and {num_gpus - 1}"
+
+    manifests_this_node = []
+    for i, manifest_file in enumerate(manifest_list):
+        if num_nodes > 1:
+            if i % num_nodes == node_idx:
+                manifests_this_node.append(manifest_file)
+        else:
+            manifests_this_node.append(manifest_file)
+
+    manifests_this_gpu = []
+    for i, manifest_file in enumerate(manifests_this_node):
+        if num_gpus > 1:
+            if i % num_gpus == gpu_idx:
+                manifests_this_gpu.append(manifest_file)
+        else:
+            manifests_this_gpu.append(manifest_file)
+    return manifests_this_gpu
+
+
+@hydra_runner(config_name="AlignmentConfig", schema=AlignmentConfig)
+def main(cfg: AlignmentConfig):
+
+    logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}')
+
+    if is_dataclass(cfg):
+        cfg = OmegaConf.structured(cfg)
+
+    # Validate config
+    if cfg.model_path is None and cfg.pretrained_name is None:
+        raise ValueError("Both cfg.model_path and cfg.pretrained_name cannot be None")
+
+    if cfg.model_path is not None and cfg.pretrained_name is not None:
+        raise ValueError("One of cfg.model_path and cfg.pretrained_name must be None")
+
+    if cfg.manifest_filepath is None:
+        raise ValueError("cfg.manifest_filepath must be specified")
+
+    if cfg.output_dir is None and not cfg.remove_tmp_dir:
+        raise ValueError("cfg.output_dir must be specified if cfg.remove_tmp_dir is False")
+
+    if cfg.batch_size < 1:
+        raise ValueError("cfg.batch_size cannot be zero or a negative number")
+
+    if cfg.additional_segment_grouping_separator == "" or cfg.additional_segment_grouping_separator == " ":
+        raise ValueError("cfg.additional_grouping_separator cannot be empty string or space character")
+
+    if cfg.ctm_file_config.minimum_timestamp_duration < 0:
+        raise ValueError("cfg.minimum_timestamp_duration cannot be a negative number")
+
+    if cfg.ass_file_config.vertical_alignment not in ["top", "center", "bottom"]:
+        raise ValueError("cfg.ass_file_config.vertical_alignment must be one of 'top', 'center' or 'bottom'")
+
+    for rgb_list in [
+        cfg.ass_file_config.text_already_spoken_rgb,
+        cfg.ass_file_config.text_already_spoken_rgb,
+        cfg.ass_file_config.text_already_spoken_rgb,
+    ]:
+        if len(rgb_list) != 3:
+            raise ValueError(
+                "cfg.ass_file_config.text_already_spoken_rgb,"
+                " cfg.ass_file_config.text_being_spoken_rgb,"
+                " and cfg.ass_file_config.text_already_spoken_rgb all need to contain"
+                " exactly 3 elements."
+            )
+
+    # init devices
+    if cfg.transcribe_device is None:
+        transcribe_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        transcribe_device = torch.device(cfg.transcribe_device)
+    logging.info(f"Device to be used for transcription step (`transcribe_device`) is {transcribe_device}")
+
+    if cfg.viterbi_device is None:
+        viterbi_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        viterbi_device = torch.device(cfg.viterbi_device)
+    logging.info(f"Device to be used for viterbi step (`viterbi_device`) is {viterbi_device}")
+
+    if transcribe_device.type == 'cuda' or viterbi_device.type == 'cuda':
+        logging.warning(
+            'One or both of transcribe_device and viterbi_device are GPUs. If you run into OOM errors '
+            'it may help to change both devices to be the CPU.'
+        )
+
+    # load model
+    model, _ = setup_model(cfg, transcribe_device)
+    model.eval()
+
+    if isinstance(model, EncDecHybridRNNTCTCModel):
+        model.change_decoding_strategy(decoder_type="ctc")
+
+    if cfg.use_local_attention:
+        logging.info(
+            "Flag use_local_attention is set to True => will try to use local attention for model if it allows it"
+        )
+        model.change_attention_model(self_attention_model="rel_pos_local_attn", att_context_size=[64, 64])
+
+    if not (isinstance(model, EncDecCTCModel) or isinstance(model, EncDecHybridRNNTCTCModel)):
+        raise NotImplementedError(
+            f"Model is not an instance of NeMo EncDecCTCModel or ENCDecHybridRNNTCTCModel."
+            " Currently only instances of these models are supported"
+        )
+
+    if cfg.ctm_file_config.minimum_timestamp_duration > 0:
+        logging.warning(
+            f"cfg.ctm_file_config.minimum_timestamp_duration has been set to {cfg.ctm_file_config.minimum_timestamp_duration} seconds. "
+            "This may cause the alignments for some tokens/words/additional segments to be overlapping."
+        )
+
+    buffered_chunk_params = {}
+    if cfg.use_buffered_chunked_streaming:
+        model_cfg = copy.deepcopy(model._cfg)
+
+        OmegaConf.set_struct(model_cfg.preprocessor, False)
+        # some changes for streaming scenario
+        model_cfg.preprocessor.dither = 0.0
+        model_cfg.preprocessor.pad_to = 0
+
+        if model_cfg.preprocessor.normalize != "per_feature":
+            logging.error(
+                "Only EncDecCTCModelBPE models trained with per_feature normalization are supported currently"
+            )
+        # Disable config overwriting
+        OmegaConf.set_struct(model_cfg.preprocessor, True)
+
+        feature_stride = model_cfg.preprocessor['window_stride']
+        model_stride_in_secs = feature_stride * cfg.model_downsample_factor
+        total_buffer = cfg.total_buffer_in_secs
+        chunk_len = float(cfg.chunk_len_in_secs)
+        tokens_per_chunk = math.ceil(chunk_len / model_stride_in_secs)
+        mid_delay = math.ceil((chunk_len + (total_buffer - chunk_len) / 2) / model_stride_in_secs)
+        logging.info(f"tokens_per_chunk is {tokens_per_chunk}, mid_delay is {mid_delay}")
+
+        model = FrameBatchASR(
+            asr_model=model,
+            frame_len=chunk_len,
+            total_buffer=cfg.total_buffer_in_secs,
+            batch_size=cfg.chunk_batch_size,
+        )
+        buffered_chunk_params = {
+            "delay": mid_delay,
+            "model_stride_in_secs": model_stride_in_secs,
+            "tokens_per_chunk": tokens_per_chunk,
+        }
+
+    if Path(cfg.manifest_filepath).is_file():
+        manifest_list = [cfg.manifest_filepath]
+    elif Path(cfg.manifest_filepath).is_dir():
+        if cfg.manifest_pattern is not None:
+            manifest_list = list(Path(cfg.manifest_filepath).glob(cfg.manifest_pattern))
+        else:
+            manifest_list = list(Path(cfg.manifest_filepath).glob("*.json"))
+    else:
+        raise ValueError(
+            f"cfg.manifest_filepath is not a valid file or directory. "
+            f"Please check the path: {cfg.manifest_filepath}"
+        )
+
+    origin_output_manifest_filepath = cfg.output_manifest_filepath
+
+    manifest_list = get_manifests_for_this_rank(manifest_list, cfg.num_nodes, cfg.num_gpus, cfg.node_idx, cfg.gpu_idx)
+    logging.info(f"Found {len(manifest_list)} manifest files to process.")
+    # process each manifest file
+    for manifest_filepath in manifest_list:
+        logging.info(f"Processing manifest file: {manifest_filepath}")
+        cfg.manifest_filepath = str(manifest_filepath)
+
+        if origin_output_manifest_filepath is None:
+            manifest_stem = Path(manifest_filepath).stem
+            cfg.output_manifest_filepath = str(Path(manifest_filepath).parent / f"{manifest_stem}-aligned.json")
+        elif len(manifest_list) > 1 and origin_output_manifest_filepath is not None:
+            raise ValueError(
+                "cfg.output_manifest_filepath must be None when processing multiple manifest files. "
+                "Please set it to None."
+            )
+
+        if not cfg.remove_tmp_dir and len(manifest_list) > 1:
+            # if keep alignment files, then we need to set output_dir to be different for each manifest
+            cfg.output_dir = str(Path(manifest_filepath).parent / f"{Path(manifest_filepath).stem}_alignment")
+
+        process_single_manifest(cfg, model, buffered_chunk_params, viterbi_device)
+        logging.info(f"Output manifest saved to: {cfg.output_manifest_filepath}")
+
+    logging.info("All manifest files processed successfully.")
+
+
+def process_single_manifest(cfg: AlignmentConfig, model, buffered_chunk_params, viterbi_device):
+    # Validate manifest contents
+    if not is_entry_in_all_lines(cfg.manifest_filepath, "audio_filepath"):
+        raise RuntimeError(
+            "At least one line in cfg.manifest_filepath does not contain an 'audio_filepath' entry. "
+            "All lines must contain an 'audio_filepath' entry."
+        )
+
+    if cfg.align_using_pred_text:
+        if is_entry_in_any_lines(cfg.manifest_filepath, "pred_text"):
+            raise RuntimeError(
+                "Cannot specify cfg.align_using_pred_text=True when the manifest at cfg.manifest_filepath "
+                "contains 'pred_text' entries. This is because the audio will be transcribed and may produce "
+                "a different 'pred_text'. This may cause confusion."
+            )
+    else:
+        if not is_entry_in_all_lines(cfg.manifest_filepath, "text"):
+            raise RuntimeError(
+                "At least one line in cfg.manifest_filepath does not contain a 'text' entry. "
+                "NFA requires all lines to contain a 'text' entry when cfg.align_using_pred_text=False."
+            )
+
+    # get start and end line IDs of batches
+    starts, ends = get_batch_starts_ends(cfg.manifest_filepath, cfg.batch_size)
+
+    # init output_timestep_duration = None and we will calculate and update it during the first batch
+    output_timestep_duration = None
+
+    if cfg.remove_tmp_dir and cfg.output_dir is None:
+        cfg.output_dir = f"alignment-{uuid.uuid4()}"
+
+    # init f_manifest_out
+    os.makedirs(cfg.output_dir, exist_ok=True)
+    tgt_manifest_name = str(Path(cfg.manifest_filepath).stem) + "_with_output_file_paths.json"
+    tgt_manifest_filepath = str(Path(cfg.output_dir) / tgt_manifest_name)
+    f_manifest_out = open(tgt_manifest_filepath, 'w')
+
+    # get alignment and save in CTM batch-by-batch
+    for start, end in zip(starts, ends):
+        manifest_lines_batch = get_manifest_lines_batch(cfg.manifest_filepath, start, end)
+
+        if cfg.clean_text:
+            manifest_lines_batch = clean_text(manifest_lines_batch)
+
+        if not cfg.align_using_pred_text:
+            gt_text_batch = [line.get('text', '') for line in manifest_lines_batch]
+        else:
+            gt_text_batch = None
+
+        (
+            log_probs_batch,
+            y_batch,
+            T_batch,
+            U_batch,
+            utt_obj_batch,
+            output_timestep_duration,
+        ) = get_batch_variables(
+            audio=[line["audio_filepath"] for line in manifest_lines_batch],
+            model=model,
+            segment_separators=cfg.additional_segment_grouping_separator,
+            align_using_pred_text=cfg.align_using_pred_text,
+            audio_filepath_parts_in_utt_id=cfg.audio_filepath_parts_in_utt_id,
+            gt_text_batch=gt_text_batch,
+            output_timestep_duration=output_timestep_duration,
+            simulate_cache_aware_streaming=cfg.simulate_cache_aware_streaming,
+            use_buffered_chunked_streaming=cfg.use_buffered_chunked_streaming,
+            buffered_chunk_params=buffered_chunk_params,
+        )
+
+        alignments_batch = viterbi_decoding(log_probs_batch, y_batch, T_batch, U_batch, viterbi_device)
+
+        for utt_obj, alignment_utt in zip(utt_obj_batch, alignments_batch):
+
+            utt_obj = add_t_start_end_to_utt_obj(utt_obj, alignment_utt, output_timestep_duration)
+
+            if "ctm" in cfg.save_output_file_formats:
+                utt_obj = make_ctm_files(
+                    utt_obj,
+                    cfg.output_dir,
+                    cfg.ctm_file_config,
+                )
+
+            if "ass" in cfg.save_output_file_formats:
+                utt_obj = make_ass_files(utt_obj, cfg.output_dir, cfg.ass_file_config)
+
+            write_manifest_out_line(
+                f_manifest_out,
+                utt_obj,
+            )
+
+    f_manifest_out.close()
+
+    # adding eou processing here
+    input_manifest_lines = []
+    with open(cfg.manifest_filepath, 'r') as f:
+        for line in f.readlines():
+            if line.strip():
+                input_manifest_lines.append(json.loads(line))
+
+    output_manifest_lines = []
+    with open(tgt_manifest_filepath, 'r') as f:
+        for i, line in enumerate(f.readlines()):
+            item = json.loads(line)
+            assert os.path.basename(input_manifest_lines[i]['audio_filepath']) == os.path.basename(
+                item['audio_filepath']
+            )
+
+            if 'segments_level_ctm_filepath' not in item:
+                print(
+                    f"`segments_level_ctm_filepath` not found for {input_manifest_lines[i]['audio_filepath']}, skipping"
+                )
+                continue
+
+            # get sou/eou time
+            with open(item['segments_level_ctm_filepath']) as f:
+                lines = [line.split() for line in f]
+                start_time = min([float(line[2]) for line in lines])
+                end_time = max([float(line[2]) + float(line[3]) for line in lines])
+                input_manifest_lines[i]['sou_time'] = start_time
+                input_manifest_lines[i]['eou_time'] = end_time
+                output_manifest_lines.append(input_manifest_lines[i])
+
+    with open(cfg.output_manifest_filepath, 'w') as f:
+        for item in output_manifest_lines:
+            f.write(json.dumps(item) + '\n')
+
+    if cfg.remove_tmp_dir:  # safely removing tmp dir after alignment
+        for file_or_folder in [
+            tgt_manifest_filepath,
+            os.path.join(cfg.output_dir, 'ctm'),
+            os.path.join(cfg.output_dir, 'ass'),
+        ]:
+            if os.path.exists(file_or_folder):
+                if os.path.isfile(file_or_folder):
+                    os.remove(file_or_folder)
+                else:
+                    shutil.rmtree(file_or_folder)
+        if os.path.exists(cfg.output_dir) and len(os.listdir(cfg.output_dir)) == 0:
+            shutil.rmtree(cfg.output_dir)
+
+    return None
+
+
+if __name__ == "__main__":
+    main()

tools/nemo_forced_aligner/requirements.txt

@@ -0,0 +1,3 @@
+nemo-toolkit[all]
+prettyprinter # for testing
+pytest # for testing

tools/nemo_forced_aligner/tests/test_add_t_start_end_to_utt_obj.py

@@ -0,0 +1,290 @@
+# Copyright (c) 2023, 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 copy
+
+import pytest
+from nemo.collections.asr.parts.utils.aligner_utils import Segment, Token, Utterance, Word, add_t_start_end_to_utt_obj
+
+OUTPUT_TIMESTEP_DURATION = 0.04
+
+ALIGNMENT = [
+    1,
+    1,
+    3,
+    3,
+    4,
+    5,
+    7,
+    7,
+    9,
+    10,
+    11,
+    12,
+    13,
+    15,
+    17,
+    17,
+    19,
+    21,
+    23,
+    23,
+]
+
+EXPECTED_OUTPUT_UTTERANCE = Utterance(
+    text='hi world | hey',
+    token_ids_with_blanks=[
+        28,
+        8,
+        28,
+        9,
+        28,
+        0,
+        28,
+        23,
+        28,
+        15,
+        28,
+        18,
+        28,
+        12,
+        28,
+        4,
+        28,
+        0,
+        28,
+        8,
+        28,
+        5,
+        28,
+        25,
+        28,
+    ],
+    segments_and_tokens=[
+        Token(text='<b>', text_cased='<b>', s_start=0, s_end=0, t_start=-1, t_end=-1),
+        Segment(
+            text="hi world",
+            s_start=1,
+            s_end=15,
+            t_start=0 * OUTPUT_TIMESTEP_DURATION,
+            t_end=14 * OUTPUT_TIMESTEP_DURATION,
+            words_and_tokens=[
+                Word(
+                    text="hi",
+                    s_start=1,
+                    s_end=3,
+                    t_start=0 * OUTPUT_TIMESTEP_DURATION,
+                    t_end=4 * OUTPUT_TIMESTEP_DURATION,
+                    tokens=[
+                        Token(
+                            text='h',
+                            text_cased='h',
+                            s_start=1,
+                            s_end=1,
+                            t_start=0 * OUTPUT_TIMESTEP_DURATION,
+                            t_end=2 * OUTPUT_TIMESTEP_DURATION,
+                        ),
+                        Token(text='<b>', text_cased='<b>', s_start=2, s_end=2, t_start=-1, t_end=-1),
+                        Token(
+                            text='i',
+                            text_cased='i',
+                            s_start=3,
+                            s_end=3,
+                            t_start=2 * OUTPUT_TIMESTEP_DURATION,
+                            t_end=4 * OUTPUT_TIMESTEP_DURATION,
+                        ),
+                    ],
+                ),
+                Token(
+                    text='<b>',
+                    text_cased='<b>',
+                    s_start=4,
+                    s_end=4,
+                    t_start=4 * OUTPUT_TIMESTEP_DURATION,
+                    t_end=5 * OUTPUT_TIMESTEP_DURATION,
+                ),
+                Token(
+                    text='<space>',
+                    text_cased='<space>',
+                    s_start=5,
+                    s_end=5,
+                    t_start=5 * OUTPUT_TIMESTEP_DURATION,
+                    t_end=6 * OUTPUT_TIMESTEP_DURATION,
+                ),
+                Token(text='<b>', text_cased='<b>', s_start=6, s_end=6, t_start=-1, t_end=-1),
+                Word(
+                    text="world",
+                    s_start=7,
+                    s_end=15,
+                    t_start=6 * OUTPUT_TIMESTEP_DURATION,
+                    t_end=14 * OUTPUT_TIMESTEP_DURATION,
+                    tokens=[
+                        Token(
+                            text='w',
+                            text_cased='w',
+                            s_start=7,
+                            s_end=7,
+                            t_start=6 * OUTPUT_TIMESTEP_DURATION,
+                            t_end=8 * OUTPUT_TIMESTEP_DURATION,
+                        ),
+                        Token(text='<b>', text_cased='<b>', s_start=8, s_end=8, t_start=-1, t_end=-1),
+                        Token(
+                            text='o',
+                            text_cased='o',
+                            s_start=9,
+                            s_end=9,
+                            t_start=8 * OUTPUT_TIMESTEP_DURATION,
+                            t_end=9 * OUTPUT_TIMESTEP_DURATION,
+                        ),
+                        Token(
+                            text='<b>',
+                            text_cased='<b>',
+                            s_start=10,
+                            s_end=10,
+                            t_start=9 * OUTPUT_TIMESTEP_DURATION,
+                            t_end=10 * OUTPUT_TIMESTEP_DURATION,
+                        ),
+                        Token(
+                            text='r',
+                            text_cased='r',
+                            s_start=11,
+                            s_end=11,
+                            t_start=10 * OUTPUT_TIMESTEP_DURATION,
+                            t_end=11 * OUTPUT_TIMESTEP_DURATION,
+                        ),
+                        Token(
+                            text='<b>',
+                            text_cased='<b>',
+                            s_start=12,
+                            s_end=12,
+                            t_start=11 * OUTPUT_TIMESTEP_DURATION,
+                            t_end=12 * OUTPUT_TIMESTEP_DURATION,
+                        ),
+                        Token(
+                            text='l',
+                            text_cased='l',
+                            s_start=13,
+                            s_end=13,
+                            t_start=12 * OUTPUT_TIMESTEP_DURATION,
+                            t_end=13 * OUTPUT_TIMESTEP_DURATION,
+                        ),
+                        Token(text='<b>', text_cased='<b>', s_start=14, s_end=14, t_start=-1, t_end=-1),
+                        Token(
+                            text='d',
+                            text_cased='d',
+                            s_start=15,
+                            s_end=15,
+                            t_start=13 * OUTPUT_TIMESTEP_DURATION,
+                            t_end=14 * OUTPUT_TIMESTEP_DURATION,
+                        ),
+                    ],
+                ),
+            ],
+        ),
+        Token(text='<b>', text_cased='<b>', s_start=16, s_end=16, t_start=-1, t_end=-1),
+        Token(
+            text='<space>',
+            text_cased='<space>',
+            s_start=17,
+            s_end=17,
+            t_start=14 * OUTPUT_TIMESTEP_DURATION,
+            t_end=16 * OUTPUT_TIMESTEP_DURATION,
+        ),
+        Token(text='<b>', text_cased='<b>', s_start=18, s_end=18, t_start=-1, t_end=-1),
+        Segment(
+            text="hey",
+            s_start=19,
+            s_end=23,
+            t_start=16 * OUTPUT_TIMESTEP_DURATION,
+            t_end=20 * OUTPUT_TIMESTEP_DURATION,
+            words_and_tokens=[
+                Word(
+                    text="hey",
+                    s_start=19,
+                    s_end=23,
+                    t_start=16 * OUTPUT_TIMESTEP_DURATION,
+                    t_end=20 * OUTPUT_TIMESTEP_DURATION,
+                    tokens=[
+                        Token(
+                            text='h',
+                            text_cased='h',
+                            s_start=19,
+                            s_end=19,
+                            t_start=16 * OUTPUT_TIMESTEP_DURATION,
+                            t_end=17 * OUTPUT_TIMESTEP_DURATION,
+                        ),
+                        Token(text='<b>', text_cased='<b>', s_start=20, s_end=20, t_start=-1, t_end=-1),
+                        Token(
+                            text='e',
+                            text_cased='e',
+                            s_start=21,
+                            s_end=21,
+                            t_start=17 * OUTPUT_TIMESTEP_DURATION,
+                            t_end=18 * OUTPUT_TIMESTEP_DURATION,
+                        ),
+                        Token(text='<b>', text_cased='<b>', s_start=22, s_end=22, t_start=-1, t_end=-1),
+                        Token(
+                            text='y',
+                            text_cased='y',
+                            s_start=23,
+                            s_end=23,
+                            t_start=18 * OUTPUT_TIMESTEP_DURATION,
+                            t_end=20 * OUTPUT_TIMESTEP_DURATION,
+                        ),
+                    ],
+                )
+            ],
+        ),
+        Token(text='<b>', text_cased='<b>', s_start=24, s_end=24, t_start=-1, t_end=-1),
+    ],
+)
+
+
+@pytest.mark.parametrize(
+    "alignment,expected_output_utterance, output_timestep_duration",
+    [
+        (ALIGNMENT, EXPECTED_OUTPUT_UTTERANCE, OUTPUT_TIMESTEP_DURATION),
+    ],
+)
+def test_add_t_start_end_to_utt_obj(alignment, expected_output_utterance, output_timestep_duration):
+    input_utterance = copy.deepcopy(expected_output_utterance)
+
+    # set all t_start and t_end to None in input_utterance
+    for segment_or_token in input_utterance.segments_and_tokens:
+        if type(segment_or_token) is Segment:
+            segment = segment_or_token
+            segment.t_start = None
+            segment.t_end = None
+
+            for word_or_token in segment.words_and_tokens:
+                if type(word_or_token) is Word:
+                    word = word_or_token
+                    word.t_start = None
+                    word.t_end = None
+
+                    for token in word.tokens:
+                        token.t_start = None
+                        token.t_end = None
+                else:
+                    token = word_or_token
+                    token.t_start = None
+                    token.t_end = None
+
+        else:
+            token = segment_or_token
+            token.t_start = None
+            token.t_end = None
+
+    output_utterance = add_t_start_end_to_utt_obj(input_utterance, alignment, output_timestep_duration)
+    assert output_utterance == expected_output_utterance

tools/nemo_forced_aligner/tests/test_get_utt_obj.py

@@ -0,0 +1,801 @@
+# Copyright (c) 2023, 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 prettyprinter
+import pytest
+from prettyprinter import pretty_call, register_pretty
+
+from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.parts.utils.aligner_utils import Segment, Token, Utterance, Word, get_utt_obj
+
+
+def get_utt_obj_pp_string(utt_obj):
+    @register_pretty(Word)
+    def pretty_utterance(value, ctx):
+        return pretty_call(
+            ctx,
+            Word,
+            text=value.text,
+            s_start=value.s_start,
+            s_end=value.s_end,
+            t_start=value.t_start,
+            t_end=value.t_end,
+            tokens=value.tokens,
+        )
+
+    @register_pretty(Segment)
+    def pretty_utterance(value, ctx):
+        return pretty_call(
+            ctx,
+            Segment,
+            text=value.text,
+            s_start=value.s_start,
+            s_end=value.s_end,
+            t_start=value.t_start,
+            t_end=value.t_end,
+            words_and_tokens=value.words_and_tokens,
+        )
+
+    @register_pretty(Utterance)
+    def pretty_utterance(value, ctx):
+        return pretty_call(
+            ctx,
+            Utterance,
+            text=value.text,
+            token_ids_with_blanks=value.token_ids_with_blanks,
+            segments_and_tokens=value.segments_and_tokens,
+            audio_filepath=value.audio_filepath,
+            utt_id=value.utt_id,
+        )
+
+    return prettyprinter.pformat(utt_obj)
+
+
+T_FOR_TEST = 999
+AUDIO_FILEPATH_FOR_TEST = "arbitrary_string.wav"
+UTT_ID_FOR_TEST = "arbitrary_string"
+
+EN_TEXT = "hi world | hey"
+
+EN_CN_EXPECTED_UTTERANCE = Utterance(
+    text='hi world | hey',
+    token_ids_with_blanks=[1024, 317, 1024, 472, 1024, 105, 1024, 0, 1024, 25, 1024, 20, 1024],
+    segments_and_tokens=[
+        Token(text='<b>', text_cased='<b>', token_id=1024, token=None, s_start=0, s_end=0, t_start=None, t_end=None),
+        Segment(
+            text='hi world |',
+            s_start=1,
+            s_end=7,
+            t_start=None,
+            t_end=None,
+            words_and_tokens=[
+                Word(
+                    text='hi',
+                    s_start=1,
+                    s_end=1,
+                    t_start=None,
+                    t_end=None,
+                    tokens=[
+                        Token(
+                            text='▁hi',
+                            text_cased='▁hi',
+                            token_id=317,
+                            token=None,
+                            s_start=1,
+                            s_end=1,
+                            t_start=None,
+                            t_end=None,
+                        )
+                    ],
+                ),
+                Token(
+                    text='<b>',
+                    text_cased='<b>',
+                    token_id=1024,
+                    token=None,
+                    s_start=2,
+                    s_end=2,
+                    t_start=None,
+                    t_end=None,
+                ),
+                Word(
+                    text='world',
+                    s_start=3,
+                    s_end=3,
+                    t_start=None,
+                    t_end=None,
+                    tokens=[
+                        Token(
+                            text='▁world',
+                            text_cased='▁world',
+                            token_id=472,
+                            token=None,
+                            s_start=3,
+                            s_end=3,
+                            t_start=None,
+                            t_end=None,
+                        )
+                    ],
+                ),
+                Token(
+                    text='<b>',
+                    text_cased='<b>',
+                    token_id=1024,
+                    token=None,
+                    s_start=4,
+                    s_end=4,
+                    t_start=None,
+                    t_end=None,
+                ),
+                Word(
+                    text='|',
+                    s_start=5,
+                    s_end=7,
+                    t_start=None,
+                    t_end=None,
+                    tokens=[
+                        Token(
+                            text='▁',
+                            text_cased='▁',
+                            token_id=105,
+                            token=None,
+                            s_start=5,
+                            s_end=5,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='<b>',
+                            text_cased='<b>',
+                            token_id=1024,
+                            token=None,
+                            s_start=6,
+                            s_end=6,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='|',
+                            text_cased='|',
+                            token_id=0,
+                            token=None,
+                            s_start=7,
+                            s_end=7,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                    ],
+                ),
+            ],
+        ),
+        Token(text='<b>', text_cased='<b>', token_id=1024, token=None, s_start=8, s_end=8, t_start=None, t_end=None),
+        Segment(
+            text='hey',
+            s_start=9,
+            s_end=11,
+            t_start=None,
+            t_end=None,
+            words_and_tokens=[
+                Word(
+                    text='hey',
+                    s_start=9,
+                    s_end=11,
+                    t_start=None,
+                    t_end=None,
+                    tokens=[
+                        Token(
+                            text='▁he',
+                            text_cased='▁he',
+                            token_id=25,
+                            token=None,
+                            s_start=9,
+                            s_end=9,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='<b>',
+                            text_cased='<b>',
+                            token_id=1024,
+                            token=None,
+                            s_start=10,
+                            s_end=10,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='y',
+                            text_cased='y',
+                            token_id=20,
+                            token=None,
+                            s_start=11,
+                            s_end=11,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                    ],
+                )
+            ],
+        ),
+        Token(text='<b>', text_cased='<b>', token_id=1024, token=None, s_start=12, s_end=12, t_start=None, t_end=None),
+    ],
+    audio_filepath=AUDIO_FILEPATH_FOR_TEST,
+    utt_id=UTT_ID_FOR_TEST,
+)
+
+EN_QN_EXPECTED_UTTERANCE = Utterance(
+    text='hi world | hey',
+    token_ids_with_blanks=[
+        28,
+        8,
+        28,
+        9,
+        28,
+        0,
+        28,
+        23,
+        28,
+        15,
+        28,
+        18,
+        28,
+        12,
+        28,
+        4,
+        28,
+        0,
+        28,
+        28,
+        28,
+        0,
+        28,
+        8,
+        28,
+        5,
+        28,
+        25,
+        28,
+    ],
+    segments_and_tokens=[
+        Token(text='<b>', text_cased='<b>', token_id=28, token=None, s_start=0, s_end=0, t_start=None, t_end=None),
+        Segment(
+            text='hi world |',
+            s_start=1,
+            s_end=19,
+            t_start=None,
+            t_end=None,
+            words_and_tokens=[
+                Word(
+                    text='hi',
+                    s_start=1,
+                    s_end=3,
+                    t_start=None,
+                    t_end=None,
+                    tokens=[
+                        Token(
+                            text='h',
+                            text_cased='h',
+                            token_id=8,
+                            token=None,
+                            s_start=1,
+                            s_end=1,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='<b>',
+                            text_cased='<b>',
+                            token_id=28,
+                            token=None,
+                            s_start=2,
+                            s_end=2,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='i',
+                            text_cased='i',
+                            token_id=9,
+                            token=None,
+                            s_start=3,
+                            s_end=3,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                    ],
+                ),
+                Token(
+                    text='<b>', text_cased='<b>', token_id=28, token=None, s_start=4, s_end=4, t_start=None, t_end=None
+                ),
+                Token(
+                    text='<space>',
+                    text_cased='<space>',
+                    token_id=0,
+                    token=None,
+                    s_start=5,
+                    s_end=5,
+                    t_start=None,
+                    t_end=None,
+                ),
+                Token(
+                    text='<b>', text_cased='<b>', token_id=28, token=None, s_start=6, s_end=6, t_start=None, t_end=None
+                ),
+                Word(
+                    text='world',
+                    s_start=7,
+                    s_end=15,
+                    t_start=None,
+                    t_end=None,
+                    tokens=[
+                        Token(
+                            text='w',
+                            text_cased='w',
+                            token_id=23,
+                            token=None,
+                            s_start=7,
+                            s_end=7,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='<b>',
+                            text_cased='<b>',
+                            token_id=28,
+                            token=None,
+                            s_start=8,
+                            s_end=8,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='o',
+                            text_cased='o',
+                            token_id=15,
+                            token=None,
+                            s_start=9,
+                            s_end=9,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='<b>',
+                            text_cased='<b>',
+                            token_id=28,
+                            token=None,
+                            s_start=10,
+                            s_end=10,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='r',
+                            text_cased='r',
+                            token_id=18,
+                            token=None,
+                            s_start=11,
+                            s_end=11,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='<b>',
+                            text_cased='<b>',
+                            token_id=28,
+                            token=None,
+                            s_start=12,
+                            s_end=12,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='l',
+                            text_cased='l',
+                            token_id=12,
+                            token=None,
+                            s_start=13,
+                            s_end=13,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='<b>',
+                            text_cased='<b>',
+                            token_id=28,
+                            token=None,
+                            s_start=14,
+                            s_end=14,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='d',
+                            text_cased='d',
+                            token_id=4,
+                            token=None,
+                            s_start=15,
+                            s_end=15,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                    ],
+                ),
+                Token(
+                    text='<b>',
+                    text_cased='<b>',
+                    token_id=28,
+                    token=None,
+                    s_start=16,
+                    s_end=16,
+                    t_start=None,
+                    t_end=None,
+                ),
+                Token(
+                    text='<space>',
+                    text_cased='<space>',
+                    token_id=0,
+                    token=None,
+                    s_start=17,
+                    s_end=17,
+                    t_start=None,
+                    t_end=None,
+                ),
+                Token(
+                    text='<b>',
+                    text_cased='<b>',
+                    token_id=28,
+                    token=None,
+                    s_start=18,
+                    s_end=18,
+                    t_start=None,
+                    t_end=None,
+                ),
+                Word(
+                    text='|',
+                    s_start=19,
+                    s_end=19,
+                    t_start=None,
+                    t_end=None,
+                    tokens=[
+                        Token(
+                            text='|',
+                            text_cased='|',
+                            token_id=28,
+                            token=None,
+                            s_start=19,
+                            s_end=19,
+                            t_start=None,
+                            t_end=None,
+                        )
+                    ],
+                ),
+            ],
+        ),
+        Token(text='<b>', text_cased='<b>', token_id=28, token=None, s_start=20, s_end=20, t_start=None, t_end=None),
+        Token(
+            text='<space>',
+            text_cased='<space>',
+            token_id=0,
+            token=None,
+            s_start=21,
+            s_end=21,
+            t_start=None,
+            t_end=None,
+        ),
+        Token(text='<b>', text_cased='<b>', token_id=28, token=None, s_start=22, s_end=22, t_start=None, t_end=None),
+        Segment(
+            text='hey',
+            s_start=23,
+            s_end=27,
+            t_start=None,
+            t_end=None,
+            words_and_tokens=[
+                Word(
+                    text='hey',
+                    s_start=23,
+                    s_end=27,
+                    t_start=None,
+                    t_end=None,
+                    tokens=[
+                        Token(
+                            text='h',
+                            text_cased='h',
+                            token_id=8,
+                            token=None,
+                            s_start=23,
+                            s_end=23,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='<b>',
+                            text_cased='<b>',
+                            token_id=28,
+                            token=None,
+                            s_start=24,
+                            s_end=24,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='e',
+                            text_cased='e',
+                            token_id=5,
+                            token=None,
+                            s_start=25,
+                            s_end=25,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='<b>',
+                            text_cased='<b>',
+                            token_id=28,
+                            token=None,
+                            s_start=26,
+                            s_end=26,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='y',
+                            text_cased='y',
+                            token_id=25,
+                            token=None,
+                            s_start=27,
+                            s_end=27,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                    ],
+                )
+            ],
+        ),
+        Token(text='<b>', text_cased='<b>', token_id=28, token=None, s_start=28, s_end=28, t_start=None, t_end=None),
+    ],
+    audio_filepath=AUDIO_FILEPATH_FOR_TEST,
+    utt_id=UTT_ID_FOR_TEST,
+)
+
+
+ZH_TEXT = "人工 智能|技术"
+
+ZH_CN_EXPECTED_UTTERANCE = Utterance(
+    text='人工 智能|技术',
+    token_ids_with_blanks=[
+        5206,
+        125,
+        5206,
+        1329,
+        5206,
+        0,
+        5206,
+        2029,
+        5206,
+        3668,
+        5206,
+        5206,
+        5206,
+        1695,
+        5206,
+        2075,
+        5206,
+    ],
+    segments_and_tokens=[
+        Token(text='<b>', text_cased='<b>', token_id=5206, token=None, s_start=0, s_end=0, t_start=None, t_end=None),
+        Segment(
+            text='人工 智能|技术',
+            s_start=1,
+            s_end=15,
+            t_start=None,
+            t_end=None,
+            words_and_tokens=[
+                Word(
+                    text='人工',
+                    s_start=1,
+                    s_end=3,
+                    t_start=None,
+                    t_end=None,
+                    tokens=[
+                        Token(
+                            text='人',
+                            text_cased='人',
+                            token_id=125,
+                            token=None,
+                            s_start=1,
+                            s_end=1,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='<b>',
+                            text_cased='<b>',
+                            token_id=5206,
+                            token=None,
+                            s_start=2,
+                            s_end=2,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='工',
+                            text_cased='工',
+                            token_id=1329,
+                            token=None,
+                            s_start=3,
+                            s_end=3,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                    ],
+                ),
+                Token(
+                    text='<b>',
+                    text_cased='<b>',
+                    token_id=5206,
+                    token=None,
+                    s_start=4,
+                    s_end=4,
+                    t_start=None,
+                    t_end=None,
+                ),
+                Token(
+                    text='<space>',
+                    text_cased='<space>',
+                    token_id=0,
+                    token=None,
+                    s_start=5,
+                    s_end=5,
+                    t_start=None,
+                    t_end=None,
+                ),
+                Token(
+                    text='<b>',
+                    text_cased='<b>',
+                    token_id=5206,
+                    token=None,
+                    s_start=6,
+                    s_end=6,
+                    t_start=None,
+                    t_end=None,
+                ),
+                Word(
+                    text='智能|技术',
+                    s_start=7,
+                    s_end=15,
+                    t_start=None,
+                    t_end=None,
+                    tokens=[
+                        Token(
+                            text='智',
+                            text_cased='智',
+                            token_id=2029,
+                            token=None,
+                            s_start=7,
+                            s_end=7,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='<b>',
+                            text_cased='<b>',
+                            token_id=5206,
+                            token=None,
+                            s_start=8,
+                            s_end=8,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='能',
+                            text_cased='能',
+                            token_id=3668,
+                            token=None,
+                            s_start=9,
+                            s_end=9,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='<b>',
+                            text_cased='<b>',
+                            token_id=5206,
+                            token=None,
+                            s_start=10,
+                            s_end=10,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='|',
+                            text_cased='|',
+                            token_id=5206,
+                            token=None,
+                            s_start=11,
+                            s_end=11,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='<b>',
+                            text_cased='<b>',
+                            token_id=5206,
+                            token=None,
+                            s_start=12,
+                            s_end=12,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='技',
+                            text_cased='技',
+                            token_id=1695,
+                            token=None,
+                            s_start=13,
+                            s_end=13,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='<b>',
+                            text_cased='<b>',
+                            token_id=5206,
+                            token=None,
+                            s_start=14,
+                            s_end=14,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                        Token(
+                            text='术',
+                            text_cased='术',
+                            token_id=2075,
+                            token=None,
+                            s_start=15,
+                            s_end=15,
+                            t_start=None,
+                            t_end=None,
+                        ),
+                    ],
+                ),
+            ],
+        ),
+        Token(text='<b>', text_cased='<b>', token_id=5206, token=None, s_start=16, s_end=16, t_start=None, t_end=None),
+    ],
+    audio_filepath=AUDIO_FILEPATH_FOR_TEST,
+    utt_id=UTT_ID_FOR_TEST,
+)
+
+
+@pytest.mark.parametrize(
+    "text,model_pretrained_name,separator,expected_utterance",
+    [
+        (EN_TEXT, "stt_en_citrinet_256_gamma_0_25", "|", EN_CN_EXPECTED_UTTERANCE),
+        (EN_TEXT, "stt_en_quartznet15x5", "|", EN_QN_EXPECTED_UTTERANCE),
+        (ZH_TEXT, "stt_zh_citrinet_512", "|", ZH_CN_EXPECTED_UTTERANCE),
+    ],
+)
+def test_token_info(text, model_pretrained_name, separator, expected_utterance):
+    model = ASRModel.from_pretrained(model_pretrained_name)
+    utt_obj = get_utt_obj(
+        text=text,
+        model=model,
+        segment_separators=separator,
+        T=T_FOR_TEST,
+        audio_filepath=AUDIO_FILEPATH_FOR_TEST,
+        utt_id=UTT_ID_FOR_TEST,
+    )
+    print(f"expected utterance object: {get_utt_obj_pp_string(expected_utterance)}\n")
+    print(f"output utterance object in test: {get_utt_obj_pp_string(utt_obj)}\n")
+
+    assert utt_obj == expected_utterance

tools/nemo_forced_aligner/tests/test_restore_token_case.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2023, 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 pytest
+
+from nemo.collections.asr.parts.utils.aligner_utils import restore_token_case
+
+
+@pytest.mark.parametrize(
+    "word,word_tokens,expected_word_tokens_cased",
+    [
+        ("HEY!", ['▁he', 'y', '!'], ['▁HE', 'Y', '!']),
+        ("BabABa▁", ['▁b', 'a', 'b', 'a', 'b', 'a'], ['▁B', 'a', 'b', 'A', 'B', 'a']),
+        ("BabAB▁a", ['▁b', 'a', 'b', 'a', 'b', '_a'], ['▁B', 'a', 'b', 'A', 'B', '_a']),
+        ("Bab▁AB▁a", ['▁b', 'a', 'b', '▁a', 'b', '▁a'], ['▁B', 'a', 'b', '▁A', 'B', '▁a']),
+        ("▁Bab▁AB▁a", ['▁b', 'a', 'b', '▁a', 'b', '▁a'], ['▁B', 'a', 'b', '▁A', 'B', '▁a']),
+        ("▁Bab▁AB▁▁a", ['▁b', 'a', 'b', '▁a', 'b', '▁a'], ['▁B', 'a', 'b', '▁A', 'B', '▁a']),
+        ("▁▁BabAB▁a", ['▁b', 'a', 'b', 'a', 'b', '▁a'], ['▁B', 'a', 'b', 'A', 'B', '▁a']),
+        ("m²", ['▁', 'm', '2'], ['▁', 'm', '2']),
+        ("²", ['▁', '2'], ['▁', '2']),
+    ],
+)
+def test_restore_token_case(word, word_tokens, expected_word_tokens_cased):
+    word_tokens_cased = restore_token_case(word, word_tokens)
+    assert word_tokens_cased == expected_word_tokens_cased

tools/nemo_forced_aligner/utils/constants.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2023, 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.
+
+BLANK_TOKEN = "<b>"
+
+SPACE_TOKEN = "<space>"

tools/nemo_forced_aligner/utils/data_prep.py

@@ -0,0 +1,100 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+from pathlib import Path
+
+from nemo.collections.common.parts.preprocessing.manifest import get_full_path
+
+
+def get_batch_starts_ends(manifest_filepath, batch_size):
+    """
+    Get the start and end ids of the lines we will use for each 'batch'.
+    """
+
+    with open(manifest_filepath, 'r') as f:
+        num_lines_in_manifest = sum(1 for _ in f)
+
+    starts = [x for x in range(0, num_lines_in_manifest, batch_size)]
+    ends = [x - 1 for x in starts]
+    ends.pop(0)
+    ends.append(num_lines_in_manifest)
+
+    return starts, ends
+
+
+def is_entry_in_any_lines(manifest_filepath, entry):
+    """
+    Returns True if entry is a key in any of the JSON lines in manifest_filepath
+    """
+
+    entry_in_manifest = False
+
+    with open(manifest_filepath, 'r') as f:
+        for line in f:
+            data = json.loads(line)
+
+            if entry in data:
+                entry_in_manifest = True
+
+    return entry_in_manifest
+
+
+def is_entry_in_all_lines(manifest_filepath, entry):
+    """
+    Returns True is entry is a key in all of the JSON lines in manifest_filepath.
+    """
+    with open(manifest_filepath, 'r') as f:
+        for line in f:
+            data = json.loads(line)
+
+            if entry not in data:
+                return False
+
+    return True
+
+
+def get_manifest_lines_batch(manifest_filepath, start, end):
+    manifest_lines_batch = []
+    with open(manifest_filepath, "r", encoding="utf-8-sig") as f:
+        for line_i, line in enumerate(f):
+            if line_i >= start and line_i <= end:
+                data = json.loads(line)
+                data["audio_filepath"] = get_full_path(data["audio_filepath"], manifest_filepath)
+                if "text" in data:
+                    # remove any BOM, any duplicated spaces, convert any
+                    # newline chars to spaces
+                    data["text"] = data["text"].replace("\ufeff", "")
+                    data["text"] = " ".join(data["text"].split())
+
+                    # Replace any horizontal ellipses with 3 separate periods.
+                    # The tokenizer will do this anyway. But making this replacement
+                    # now helps avoid errors when restoring punctuation when saving
+                    # the output files
+                    data["text"] = data["text"].replace("\u2026", "...")
+
+                if not Path(data['audio_filepath']).exists():
+                    extended_path = Path(Path(manifest_filepath).parent, data['audio_filepath'])
+                    if extended_path.exists():
+                        data['audio_filepath'] = str(extended_path)
+                    else:
+                        raise FileNotFoundError(
+                            f"Audio file {data['audio_filepath']} not found in {manifest_filepath}"
+                        )
+
+                manifest_lines_batch.append(data)
+
+            if line_i == end:
+                break
+    return manifest_lines_batch

tools/nemo_forced_aligner/utils/make_ass_files.py

@@ -0,0 +1,525 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This file contains functions for make ASS-format subtitle files based on the generated alignment.
+ASS files can be generated highlighting token-level alignments or word-level alignments.
+In both cases, 'segment' boundaries will be used to determine which parts of the text will appear
+at the same time. 
+For the token-level ASS files, the text will be highlighted token-by-token, with the timings determined 
+by the NFA alignments. 
+For the word-level ASS files, the text will be highlighted word-by-word, with the timings determined
+by the NFA alignemtns.
+"""
+
+import math
+import os
+
+import soundfile as sf
+from utils.constants import BLANK_TOKEN, SPACE_TOKEN
+
+from nemo.collections.asr.parts.utils.aligner_utils import Segment, Token, Word
+
+PLAYERRESX = 384
+PLAYERRESY = 288
+MARGINL = 10
+MARGINR = 10
+MARGINV = 20
+
+
+def seconds_to_ass_format(seconds_float):
+    seconds_float = float(seconds_float)
+    mm, ss_decimals = divmod(seconds_float, 60)
+    hh, mm = divmod(mm, 60)
+
+    hh = str(round(hh))
+    if len(hh) == 1:
+        hh = '0' + hh
+
+    mm = str(round(mm))
+    if len(mm) == 1:
+        mm = '0' + mm
+
+    ss_decimals = f"{ss_decimals:.2f}"
+    if len(ss_decimals.split(".")[0]) == 1:
+        ss_decimals = "0" + ss_decimals
+
+    srt_format_time = f"{hh}:{mm}:{ss_decimals}"
+
+    return srt_format_time
+
+
+def rgb_list_to_hex_bgr(rgb_list):
+    r, g, b = rgb_list
+    return f"{b:x}{g:x}{r:x}"
+
+
+def make_ass_files(
+    utt_obj,
+    output_dir_root,
+    ass_file_config,
+):
+
+    # don't try to make files if utt_obj.segments_and_tokens is empty, which will happen
+    # in the case of the ground truth text being empty or the number of tokens being too large vs audio duration
+    if not utt_obj.segments_and_tokens:
+        return utt_obj
+
+    if ass_file_config.resegment_text_to_fill_space:
+        utt_obj = resegment_utt_obj(utt_obj, ass_file_config)
+
+    # get duration of the utterance, so we know the final timestamp of the final set of subtitles,
+    # which we will keep showing until the end
+    with sf.SoundFile(utt_obj.audio_filepath) as f:
+        audio_dur = f.frames / f.samplerate
+
+    utt_obj = make_word_level_ass_file(utt_obj, output_dir_root, ass_file_config, audio_dur)
+    utt_obj = make_token_level_ass_file(utt_obj, output_dir_root, ass_file_config, audio_dur)
+
+    return utt_obj
+
+
+def _get_word_n_chars(word):
+    n_chars = 0
+    for token in word.tokens:
+        if token.text != BLANK_TOKEN:
+            n_chars += len(token.text)
+    return n_chars
+
+
+def _get_segment_n_chars(segment):
+    n_chars = 0
+    for word_or_token in segment.words_and_tokens:
+        if word_or_token.text == SPACE_TOKEN:
+            n_chars += 1
+        elif word_or_token.text != BLANK_TOKEN:
+            n_chars += len(word_or_token.text)
+    return n_chars
+
+
+def resegment_utt_obj(utt_obj, ass_file_config):
+
+    # get list of just all words and tokens
+    all_words_and_tokens = []
+    for segment_or_token in utt_obj.segments_and_tokens:
+        if type(segment_or_token) is Segment:
+            all_words_and_tokens.extend(segment_or_token.words_and_tokens)
+        else:
+            all_words_and_tokens.append(segment_or_token)
+
+    # figure out how many chars will fit into one 'slide' and thus should be the max
+    # size of a segment
+    approx_chars_per_line = (PLAYERRESX - MARGINL - MARGINR) / (
+        ass_file_config.fontsize * 0.6
+    )  # assume chars 0.6 as wide as they are tall
+    approx_lines_per_segment = (PLAYERRESY - MARGINV) / (
+        ass_file_config.fontsize * 1.15
+    )  # assume line spacing is 1.15
+    if approx_lines_per_segment > ass_file_config.max_lines_per_segment:
+        approx_lines_per_segment = ass_file_config.max_lines_per_segment
+
+    max_chars_per_segment = int(approx_chars_per_line * approx_lines_per_segment)
+
+    new_segments_and_tokens = []
+    all_words_and_tokens_pointer = 0
+    for word_or_token in all_words_and_tokens:
+        if type(word_or_token) is Token:
+            new_segments_and_tokens.append(word_or_token)
+            all_words_and_tokens_pointer += 1
+        else:
+            break
+
+    new_segments_and_tokens.append(Segment())
+
+    while all_words_and_tokens_pointer < len(all_words_and_tokens):
+        word_or_token = all_words_and_tokens[all_words_and_tokens_pointer]
+        if type(word_or_token) is Word:
+
+            # if this is going to be the first word in the segment, we definitely want
+            # to add it to the segment
+            if not new_segments_and_tokens[-1].words_and_tokens:
+                new_segments_and_tokens[-1].words_and_tokens.append(word_or_token)
+
+            else:
+                # if not the first word, check what the new length of the segment will be
+                # if short enough - add this word to this segment;
+                # if too long - add to a new segment
+                this_word_n_chars = _get_word_n_chars(word_or_token)
+                segment_so_far_n_chars = _get_segment_n_chars(new_segments_and_tokens[-1])
+                if this_word_n_chars + segment_so_far_n_chars < max_chars_per_segment:
+                    new_segments_and_tokens[-1].words_and_tokens.append(word_or_token)
+                else:
+                    new_segments_and_tokens.append(Segment())
+                    new_segments_and_tokens[-1].words_and_tokens.append(word_or_token)
+
+        else:  # i.e. word_or_token is a token
+            # currently this breaks the convention of tokens at the end/beginning
+            # of segments being listed as separate tokens in segment.word_and_tokens
+            # TODO: change code so we follow this convention
+            new_segments_and_tokens[-1].words_and_tokens.append(word_or_token)
+
+        all_words_and_tokens_pointer += 1
+
+    utt_obj.segments_and_tokens = new_segments_and_tokens
+
+    return utt_obj
+
+
+def make_word_level_ass_file(utt_obj, output_dir_root, ass_file_config, audio_dur):
+
+    default_style_dict = {
+        "Name": "Default",
+        "Fontname": "Arial",
+        "Fontsize": str(ass_file_config.fontsize),
+        "PrimaryColour": "&Hffffff",
+        "SecondaryColour": "&Hffffff",
+        "OutlineColour": "&H0",
+        "BackColour": "&H0",
+        "Bold": "0",
+        "Italic": "0",
+        "Underline": "0",
+        "StrikeOut": "0",
+        "ScaleX": "100",
+        "ScaleY": "100",
+        "Spacing": "0",
+        "Angle": "0",
+        "BorderStyle": "1",
+        "Outline": "1",
+        "Shadow": "0",
+        "Alignment": None,  # will specify below
+        "MarginL": str(MARGINL),
+        "MarginR": str(MARGINR),
+        "MarginV": str(MARGINV),
+        "Encoding": "0",
+    }
+
+    if ass_file_config.vertical_alignment == "top":
+        default_style_dict["Alignment"] = "8"  # text will be 'center-justified' and in the top of the screen
+    elif ass_file_config.vertical_alignment == "center":
+        default_style_dict["Alignment"] = "5"  # text will be 'center-justified' and in the middle of the screen
+    elif ass_file_config.vertical_alignment == "bottom":
+        default_style_dict["Alignment"] = "2"  # text will be 'center-justified' and in the bottom of the screen
+    else:
+        raise ValueError(f"got an unexpected value for ass_file_config.vertical_alignment")
+
+    output_dir = os.path.join(output_dir_root, "ass", "words")
+    os.makedirs(output_dir, exist_ok=True)
+    output_file = os.path.join(output_dir, f"{utt_obj.utt_id}.ass")
+
+    already_spoken_color_code = r"{\c&H" + rgb_list_to_hex_bgr(ass_file_config.text_already_spoken_rgb) + r"&}"
+    being_spoken_color_code = r"{\c&H" + rgb_list_to_hex_bgr(ass_file_config.text_being_spoken_rgb) + r"&}"
+    not_yet_spoken_color_code = r"{\c&H" + rgb_list_to_hex_bgr(ass_file_config.text_not_yet_spoken_rgb) + r"&}"
+
+    with open(output_file, 'w') as f:
+        default_style_top_line = "Format: " + ", ".join(default_style_dict.keys())
+        default_style_bottom_line = "Style: " + ",".join(default_style_dict.values())
+
+        f.write(
+            (
+                "[Script Info]\n"
+                "ScriptType: v4.00+\n"
+                f"PlayResX: {PLAYERRESX}\n"
+                f"PlayResY: {PLAYERRESY}\n"
+                "\n"
+                "[V4+ Styles]\n"
+                f"{default_style_top_line}\n"
+                f"{default_style_bottom_line}\n"
+                "\n"
+                "[Events]\n"
+                "Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text\n\n"
+            )
+        )
+
+        # write first set of subtitles for text before speech starts to be spoken
+        words_in_first_segment = []
+        for segment_or_token in utt_obj.segments_and_tokens:
+            if type(segment_or_token) is Segment:
+                first_segment = segment_or_token
+
+                for word_or_token in first_segment.words_and_tokens:
+                    if type(word_or_token) is Word:
+                        words_in_first_segment.append(word_or_token)
+                break
+
+        text_before_speech = not_yet_spoken_color_code + " ".join([x.text for x in words_in_first_segment]) + r"{\r}"
+        subtitle_text = (
+            f"Dialogue: 0,{seconds_to_ass_format(0)},{seconds_to_ass_format(words_in_first_segment[0].t_start)},Default,,0,0,0,,"
+            + text_before_speech.rstrip()
+        )
+
+        f.write(subtitle_text + '\n')
+
+        for segment_or_token in utt_obj.segments_and_tokens:
+            if type(segment_or_token) is Segment:
+                segment = segment_or_token
+
+                words_in_segment = []
+                for word_or_token in segment.words_and_tokens:
+                    if type(word_or_token) is Word:
+                        words_in_segment.append(word_or_token)
+
+                for word_i, word in enumerate(words_in_segment):
+
+                    text_before = " ".join([x.text for x in words_in_segment[:word_i]])
+                    if text_before != "":
+                        text_before += " "
+                    text_before = already_spoken_color_code + text_before + r"{\r}"
+
+                    if word_i < len(words_in_segment) - 1:
+                        text_after = " " + " ".join([x.text for x in words_in_segment[word_i + 1 :]])
+                    else:
+                        text_after = ""
+                    text_after = not_yet_spoken_color_code + text_after + r"{\r}"
+
+                    aligned_text = being_spoken_color_code + word.text + r"{\r}"
+                    aligned_text_off = already_spoken_color_code + word.text + r"{\r}"
+
+                    subtitle_text = (
+                        f"Dialogue: 0,{seconds_to_ass_format(word.t_start)},{seconds_to_ass_format(word.t_end)},Default,,0,0,0,,"
+                        + text_before
+                        + aligned_text
+                        + text_after.rstrip()
+                    )
+                    f.write(subtitle_text + '\n')
+
+                    # add subtitles without word-highlighting for when words are not being spoken
+                    if word_i < len(words_in_segment) - 1:
+                        last_word_end = float(words_in_segment[word_i].t_end)
+                        next_word_start = float(words_in_segment[word_i + 1].t_start)
+                        if next_word_start - last_word_end > 0.001:
+                            subtitle_text = (
+                                f"Dialogue: 0,{seconds_to_ass_format(last_word_end)},{seconds_to_ass_format(next_word_start)},Default,,0,0,0,,"
+                                + text_before
+                                + aligned_text_off
+                                + text_after.rstrip()
+                            )
+                            f.write(subtitle_text + '\n')
+
+        # write final set of subtitles for text after speech has been spoken
+        words_in_final_segment = []
+        for segment_or_token in utt_obj.segments_and_tokens[::-1]:
+            if type(segment_or_token) is Segment:
+                final_segment = segment_or_token
+
+                for word_or_token in final_segment.words_and_tokens:
+                    if type(word_or_token) is Word:
+                        words_in_final_segment.append(word_or_token)
+                break
+
+        text_after_speech = already_spoken_color_code + " ".join([x.text for x in words_in_final_segment]) + r"{\r}"
+        # note: for now doing some extra padding with math.ceil(audio_dur)+1) to account for the fact that the video with subtitles can become
+        # longer than the original audio during the MP4 creation stage.
+        subtitle_text = (
+            f"Dialogue: 0,{seconds_to_ass_format(words_in_final_segment[-1].t_end)},{seconds_to_ass_format(math.ceil(audio_dur)+1)},Default,,0,0,0,,"
+            + text_after_speech.rstrip()
+        )
+
+        f.write(subtitle_text + '\n')
+
+    utt_obj.saved_output_files[f"words_level_ass_filepath"] = output_file
+
+    return utt_obj
+
+
+def make_token_level_ass_file(utt_obj, output_dir_root, ass_file_config, audio_dur):
+
+    default_style_dict = {
+        "Name": "Default",
+        "Fontname": "Arial",
+        "Fontsize": str(ass_file_config.fontsize),
+        "PrimaryColour": "&Hffffff",
+        "SecondaryColour": "&Hffffff",
+        "OutlineColour": "&H0",
+        "BackColour": "&H0",
+        "Bold": "0",
+        "Italic": "0",
+        "Underline": "0",
+        "StrikeOut": "0",
+        "ScaleX": "100",
+        "ScaleY": "100",
+        "Spacing": "0",
+        "Angle": "0",
+        "BorderStyle": "1",
+        "Outline": "1",
+        "Shadow": "0",
+        "Alignment": None,  # will specify below
+        "MarginL": str(MARGINL),
+        "MarginR": str(MARGINR),
+        "MarginV": str(MARGINV),
+        "Encoding": "0",
+    }
+
+    if ass_file_config.vertical_alignment == "top":
+        default_style_dict["Alignment"] = "8"  # text will be 'center-justified' and in the top of the screen
+    elif ass_file_config.vertical_alignment == "center":
+        default_style_dict["Alignment"] = "5"  # text will be 'center-justified' and in the middle of the screen
+    elif ass_file_config.vertical_alignment == "bottom":
+        default_style_dict["Alignment"] = "2"  # text will be 'center-justified' and in the bottom of the screen
+    else:
+        raise ValueError(f"got an unexpected value for ass_file_config.vertical_alignment")
+
+    output_dir = os.path.join(output_dir_root, "ass", "tokens")
+    os.makedirs(output_dir, exist_ok=True)
+    output_file = os.path.join(output_dir, f"{utt_obj.utt_id}.ass")
+
+    already_spoken_color_code = r"{\c&H" + rgb_list_to_hex_bgr(ass_file_config.text_already_spoken_rgb) + r"&}"
+    being_spoken_color_code = r"{\c&H" + rgb_list_to_hex_bgr(ass_file_config.text_being_spoken_rgb) + r"&}"
+    not_yet_spoken_color_code = r"{\c&H" + rgb_list_to_hex_bgr(ass_file_config.text_not_yet_spoken_rgb) + r"&}"
+
+    with open(output_file, 'w') as f:
+        default_style_top_line = "Format: " + ", ".join(default_style_dict.keys())
+        default_style_bottom_line = "Style: " + ",".join(default_style_dict.values())
+
+        f.write(
+            (
+                "[Script Info]\n"
+                "ScriptType: v4.00+\n"
+                f"PlayResX: {PLAYERRESX}\n"
+                f"PlayResY: {PLAYERRESY}\n"
+                "ScaledBorderAndShadow: yes\n"
+                "\n"
+                "[V4+ Styles]\n"
+                f"{default_style_top_line}\n"
+                f"{default_style_bottom_line}\n"
+                "\n"
+                "[Events]\n"
+                "Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text\n\n"
+            )
+        )
+
+        # write first set of subtitles for text before speech starts to be spoken
+        tokens_in_first_segment = []
+        for segment_or_token in utt_obj.segments_and_tokens:
+            if type(segment_or_token) is Segment:
+                for word_or_token in segment_or_token.words_and_tokens:
+                    if type(word_or_token) is Token:
+                        if word_or_token.text != BLANK_TOKEN:
+                            tokens_in_first_segment.append(word_or_token)
+                    else:
+                        for token in word_or_token.tokens:
+                            if token.text != BLANK_TOKEN:
+                                tokens_in_first_segment.append(token)
+
+                break
+
+        for token in tokens_in_first_segment:
+            token.text_cased = token.text_cased.replace(
+                "▁", " "
+            )  # replace underscores used in subword tokens with spaces
+            token.text_cased = token.text_cased.replace(SPACE_TOKEN, " ")  # space token with actual space
+
+        text_before_speech = (
+            not_yet_spoken_color_code + "".join([x.text_cased for x in tokens_in_first_segment]) + r"{\r}"
+        )
+        subtitle_text = (
+            f"Dialogue: 0,{seconds_to_ass_format(0)},{seconds_to_ass_format(tokens_in_first_segment[0].t_start)},Default,,0,0,0,,"
+            + text_before_speech.rstrip()
+        )
+
+        f.write(subtitle_text + '\n')
+
+        for segment_or_token in utt_obj.segments_and_tokens:
+            if type(segment_or_token) is Segment:
+                segment = segment_or_token
+
+                tokens_in_segment = []  # make list of (non-blank) tokens
+                for word_or_token in segment.words_and_tokens:
+                    if type(word_or_token) is Token:
+                        if word_or_token.text != BLANK_TOKEN:
+                            tokens_in_segment.append(word_or_token)
+                    else:
+                        for token in word_or_token.tokens:
+                            if token.text != BLANK_TOKEN:
+                                tokens_in_segment.append(token)
+
+                for token in tokens_in_segment:
+                    token.text_cased = token.text_cased.replace(
+                        "▁", " "
+                    )  # replace underscores used in subword tokens with spaces
+                    token.text_cased = token.text_cased.replace(SPACE_TOKEN, " ")  # space token with actual space
+
+                for token_i, token in enumerate(tokens_in_segment):
+
+                    text_before = "".join([x.text_cased for x in tokens_in_segment[:token_i]])
+                    text_before = already_spoken_color_code + text_before + r"{\r}"
+
+                    if token_i < len(tokens_in_segment) - 1:
+                        text_after = "".join([x.text_cased for x in tokens_in_segment[token_i + 1 :]])
+                    else:
+                        text_after = ""
+                    text_after = not_yet_spoken_color_code + text_after + r"{\r}"
+
+                    aligned_text = being_spoken_color_code + token.text_cased + r"{\r}"
+                    aligned_text_off = already_spoken_color_code + token.text_cased + r"{\r}"
+
+                    subtitle_text = (
+                        f"Dialogue: 0,{seconds_to_ass_format(token.t_start)},{seconds_to_ass_format(token.t_end)},Default,,0,0,0,,"
+                        + text_before
+                        + aligned_text
+                        + text_after.rstrip()
+                    )
+                    f.write(subtitle_text + '\n')
+
+                    # add subtitles without word-highlighting for when words are not being spoken
+                    if token_i < len(tokens_in_segment) - 1:
+                        last_token_end = float(tokens_in_segment[token_i].t_end)
+                        next_token_start = float(tokens_in_segment[token_i + 1].t_start)
+                        if next_token_start - last_token_end > 0.001:
+                            subtitle_text = (
+                                f"Dialogue: 0,{seconds_to_ass_format(last_token_end)},{seconds_to_ass_format(next_token_start)},Default,,0,0,0,,"
+                                + text_before
+                                + aligned_text_off
+                                + text_after.rstrip()
+                            )
+                            f.write(subtitle_text + '\n')
+
+        # Write final set of subtitles for text after speech has been spoken.
+        # To do this, we need to collect 'tokens_in_final_segment' so that we know what the final line is.
+        tokens_in_final_segment = []
+        for segment_or_token in utt_obj.segments_and_tokens[::-1]:
+            # Collect tokens from final segment - will 'break' so we only look at the final one.
+            if type(segment_or_token) is Segment:
+                # 'segment_or_token' is known to be Segment, which has attribute 'words_and_tokens'
+                for word_or_token in segment_or_token.words_and_tokens:
+                    if type(word_or_token) is Token:
+                        if word_or_token.text != BLANK_TOKEN:
+                            tokens_in_final_segment.append(word_or_token)
+                    else:
+                        # 'word_or_token' is known to be a Word, which has attribute 'tokens'
+                        for token in word_or_token.tokens:
+                            if token.text != BLANK_TOKEN:
+                                tokens_in_final_segment.append(token)
+                break
+
+        for token in tokens_in_final_segment:
+            token.text_cased = token.text_cased.replace(
+                "▁", " "
+            )  # replace underscores used in subword tokens with spaces
+            token.text_cased = token.text_cased.replace(SPACE_TOKEN, " ")  # space token with actual space
+
+        text_after_speech = (
+            already_spoken_color_code + "".join([x.text_cased for x in tokens_in_final_segment]) + r"{\r}"
+        )
+        # note: for now doing some extra padding with math.ceil(audio_dur)+1) to account for the fact that the video with subtitles can become
+        # longer than the original audio during the MP4 creation stage.
+        subtitle_text = (
+            f"Dialogue: 0,{seconds_to_ass_format(tokens_in_final_segment[-1].t_end)},{seconds_to_ass_format(math.ceil(audio_dur)+1)},Default,,0,0,0,,"
+            + text_after_speech.rstrip()
+        )
+
+        f.write(subtitle_text + '\n')
+
+    utt_obj.saved_output_files[f"tokens_level_ass_filepath"] = output_file
+
+    return utt_obj

tools/nemo_forced_aligner/utils/make_ctm_files.py

@@ -0,0 +1,149 @@
+# Copyright (c) 2023, 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 os
+
+import soundfile as sf
+from utils.constants import BLANK_TOKEN, SPACE_TOKEN
+from nemo.collections.asr.parts.utils.aligner_utils import Segment, Word
+from nemo.collections.asr.parts.utils.manifest_utils import get_ctm_line
+
+
+def make_ctm_files(
+    utt_obj,
+    output_dir_root,
+    ctm_file_config,
+):
+    """
+    Function to save CTM files for all the utterances in the incoming batch.
+    """
+
+    # don't try to make files if utt_obj.segments_and_tokens is empty, which will happen
+    # in the case of the ground truth text being empty or the number of tokens being too large vs audio duration
+    if not utt_obj.segments_and_tokens:
+        return utt_obj
+
+    # get audio file duration if we will need it later
+    if ctm_file_config.minimum_timestamp_duration > 0:
+        with sf.SoundFile(utt_obj.audio_filepath) as f:
+            audio_file_duration = f.frames / f.samplerate
+    else:
+        audio_file_duration = None
+
+    utt_obj = make_ctm(
+        "tokens",
+        utt_obj,
+        output_dir_root,
+        audio_file_duration,
+        ctm_file_config,
+    )
+    utt_obj = make_ctm(
+        "words",
+        utt_obj,
+        output_dir_root,
+        audio_file_duration,
+        ctm_file_config,
+    )
+    utt_obj = make_ctm(
+        "segments",
+        utt_obj,
+        output_dir_root,
+        audio_file_duration,
+        ctm_file_config,
+    )
+
+    return utt_obj
+
+
+def make_ctm(
+    alignment_level,
+    utt_obj,
+    output_dir_root,
+    audio_file_duration,
+    ctm_file_config,
+):
+    output_dir = os.path.join(output_dir_root, "ctm", alignment_level)
+    os.makedirs(output_dir, exist_ok=True)
+
+    boundary_info_utt = []
+    for segment_or_token in utt_obj.segments_and_tokens:
+        if type(segment_or_token) is Segment:
+            segment = segment_or_token
+            if alignment_level == "segments":
+                boundary_info_utt.append(segment)
+
+            for word_or_token in segment.words_and_tokens:
+                if type(word_or_token) is Word:
+                    word = word_or_token
+                    if alignment_level == "words":
+                        boundary_info_utt.append(word)
+
+                    for token in word.tokens:
+                        if alignment_level == "tokens":
+                            boundary_info_utt.append(token)
+
+                else:
+                    token = word_or_token
+                    if alignment_level == "tokens":
+                        boundary_info_utt.append(token)
+
+        else:
+            token = segment_or_token
+            if alignment_level == "tokens":
+                boundary_info_utt.append(token)
+
+    with open(os.path.join(output_dir, f"{utt_obj.utt_id}.ctm"), "w") as f_ctm:
+        for boundary_info_ in boundary_info_utt:  # loop over every token/word/segment
+
+            # skip if t_start = t_end = negative number because we used it as a marker to skip some blank tokens
+            if not (boundary_info_.t_start < 0 or boundary_info_.t_end < 0):
+                text = boundary_info_.text
+                start_time = boundary_info_.t_start
+                end_time = boundary_info_.t_end
+
+                if (
+                    ctm_file_config.minimum_timestamp_duration > 0
+                    and ctm_file_config.minimum_timestamp_duration > end_time - start_time
+                ):
+                    # make the predicted duration of the token/word/segment longer, growing it outwards equal
+                    # amounts from the predicted center of the token/word/segment
+                    token_mid_point = (start_time + end_time) / 2
+                    start_time = max(token_mid_point - ctm_file_config.minimum_timestamp_duration / 2, 0.0)
+                    end_time = min(
+                        token_mid_point + ctm_file_config.minimum_timestamp_duration / 2, audio_file_duration
+                    )
+
+                if not (
+                    text == BLANK_TOKEN and ctm_file_config.remove_blank_tokens
+                ):  # don't save blanks if we don't want to
+                    # replace any spaces with <space> so we dont introduce extra space characters to our CTM files
+                    text = text.replace(" ", SPACE_TOKEN)
+
+                    ctm_line = get_ctm_line(
+                        source=utt_obj.utt_id,
+                        channel=1,
+                        start_time=start_time,
+                        duration=end_time - start_time,
+                        token=text,
+                        conf=None,
+                        type_of_token='lex',
+                        speaker=None,
+                    )
+                    f_ctm.write(ctm_line)
+
+    utt_obj.saved_output_files[f"{alignment_level}_level_ctm_filepath"] = os.path.join(
+        output_dir, f"{utt_obj.utt_id}.ctm"
+    )
+
+    return utt_obj

tools/nemo_forced_aligner/utils/make_output_manifest.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+
+
+def write_manifest_out_line(
+    f_manifest_out,
+    utt_obj,
+):
+
+    data = {"audio_filepath": utt_obj.audio_filepath}
+    if not utt_obj.text is None:
+        data["text"] = utt_obj.text
+
+    if not utt_obj.pred_text is None:
+        data["pred_text"] = utt_obj.pred_text
+
+    for key, val in utt_obj.saved_output_files.items():
+        data[key] = val
+
+    new_line = json.dumps(data)
+    f_manifest_out.write(f"{new_line}\n")
+
+    return None

tools/rir_corpus_generator/README.md

@@ -0,0 +1,205 @@
+# Table of Contents
+1. [RIR Corpus Generator](#rir-corpus-generator)
+2. [RIR Mix Generator](#rir-mix-generator)
+
+---
+
+# RIR Corpus Generator
+
+## Outline
+
+This tool can be used to generate a corpus of room impulse responses (RIR) for a set of randomly generated rooms and microphone array placements. The tool is using [Pyroomacoustics](https://github.com/LCAV/pyroomacoustics) to simulate RIRs.
+
+The generator loads configuration from a config file, generates RIRs, saves the metadata in a manifest file and RIRs in a HDF5 file.
+
+```mermaid
+flowchart TD
+	A[Hydra Config + Overrides] --> B[rir_corpus_generator.py]
+	B --> C[RIRCorpusGenerator]
+	C --> |generate| D[Initialize subset]
+	D --> E[Generate rooms, mic array, source placement]
+	E --> F[Simulate rooms]
+	F --> G[Save RIRs and metadata]
+	G --> |next subset| D
+	G --> |RIR| H[*.h5]
+	G --> |Metadata| I[*.json]
+	G --> |Statistics| J[*.png]
+```
+
+## Features
+
+The generator is easily reconfigurable and supports configuration of the following options
+
+- Number of rooms for each subset, range of dimensions and RT60 for randomization
+- Microphone array placement and orientation randomization
+- Number of sources per room and their placement inside the room
+
+## Parameters
+
+An example of a RIR corpus generator setup is provided in `conf/rir_corpus.yaml`
+
+## Example
+
+The RIR corpus generator with the example config can be used by running:
+
+```bash
+python rir_corpus_generator.py output_dir=OUTPUT_DIR
+```
+
+where `OUTPUT_DIR` is a path to the output directory.
+
+The output will be structured as
+
+```bash
+OUTPUT_DIR
++--{train, dev, test}
+|	+--*.h5
++--config.yaml
++--{train, dev, test}_manifest.json
++--{train, dev, test}_info.png
+```
+
+Each directory, e.g, `{train, dev, test}`, corresponds to a subset of data and contain the `*.h5` files with RIRs. Corresponding `*_manifest.json` files contain metadata for each subset. Each row corresponds to a single room/`*.h5` file and includes the following fields
+
+- `room_filepath`: path to `*.h5` file with RIRs
+- `sample_rate`: sample rate
+- `dim`: dimensions of the room `[width, length, height]`
+- `num_sources`: number of sources simulated in this room
+- `rir_rt60_theory`: Theoretically calculated RT60
+- `rir_rt60_measured`: RT60 calculated from the generated RIRs, list with `num_sources` elements
+- `mic_positions`: microphone positions in the room, list with `num_mics` elements
+- `mic_center`: center of the microphone array
+- `source_position`: position of each source, list with `num_source` elements
+- `source_distance`: distance of each source to microphone center, list with `num_source` elements
+- `source_azimuth`: azimuth of each source relative to microphone array, list with `num_source` elements
+- `source_elevation`: elevation of each source relative to microphone array, list with `num_source` elements
+
+## Loading generated data
+
+The following function can be used to load the RIR data from a simulated room file 
+
+```bash
+from nemo.collections.asr.data.data_simulation import load_rir_simulation
+
+# Specify the file
+filepath = 'OUTPUT_DIR/test/test_room_00000.h5'
+
+# Load RIRs for the first source
+mc_rir, sample_rate = load_rir_simulation(filepath=filepath, source=0)
+
+# Plot RIRs for all microphones
+import matplotlib.pyplot as plt
+plt.plot(mc_rir)
+```
+
+## Requirements
+
+Pyroomacoustics needs to be installed. If not available, it can be installed as
+
+```bash
+pip install pyroomacoustics
+```
+
+---
+
+# RIR Mix Generator
+
+## Outline
+
+This tool can be used to generate a corpus of signals by using a database of RIRs, speech, noise, and interfering sources.
+
+The generator loads configuration from a config file, configures target and interference RIRs and signals, noise signals, mix configuration, prepares the corresponding audio files, and saves the metadata in a manifest file.
+
+```mermaid
+flowchart TD
+	A[Hydra Config + Overrides] --> B[rir_mix_generator.py]
+	B --> C[RIRMixGenerator]
+	C --> |generate| D[Initialize subset]
+	D --> E[Generate target, noise,\ninterference and mix configuration]
+	E --> F[Simulate signals]
+	F --> G[Save signals and metadata]
+	G --> |next subset| D
+	G --> |Signals| H[*.wav]
+	G --> |Metadata| I[*.json]
+	G --> |Statistics| J[*.png]
+```
+
+Microphone signals are constructed by mixing target, background noise and interference signal. This is illustrated in the following diagram for an example with two interfering sources:
+
+```mermaid
+flowchart LR
+    N1[noise signal] ====> N2[scaling]
+    N2 ==> SUM
+    T2 -.-> |RSNR| N2
+    T2 ==> SUM((+))
+    T2 -.-> |RSIR| ISCALE
+
+	T1[target signal] ---> T2[RIR]
+
+    IA1[interference 1] --> IA2[RIR]
+    IA2 ==> ISUM((+))
+    IB1[interference 2] --> IB2[RIR]
+    IB2 ==> ISUM
+    ISUM ==> ISCALE[scaling]
+    ISCALE ==> SUM
+
+    SUM ==> RMS[scaling]
+    RMS ==> MIC[mic signal]
+```
+
+## Features
+
+The generator is easily reconfigurable and supports configuration of the following options
+
+- Manifest files for RIRs, signals target, noise, interference
+- Constrains for target positions relative to the microphone array (azimuth, elevation, distance)
+- Probability and number of interfering sources and separation to target
+- Reverberant-signal-to-noise ratio (RSNR) and reverberant-signal-to-interference ratio (RSIR) for scaling noise and interference
+
+
+## Parameters
+
+An example of a RIR corpus generator setup is provided in `conf/rir_mix.yaml`
+
+## Example
+
+The RIR mix generator with the example config can be used by running:
+
+```bash
+python rir_mix_generator.py output_dir=OUTPUT_DIR
+```
+
+where `OUTPUT_DIR` is a path to the output directory.
+
+The output will be structured as
+
+```bash
+OUTPUT_DIR
++--{train, dev, test}
+|	+--*_mic.wav
+|	+--*_target_reverberant.wav
+|	+--*_target_anechoic.wav
+|	+--*_noise.wav
+|	+--*_interference.wav
++--config.yaml
++--{train, dev, test}_manifest.json
++--{train, dev, test}_info.png
+```
+
+Each directory, e.g, `{train, dev, test}`, corresponds to a subset of data and contains the `*.wav` files with the generated audio signals. Corresponding `*_manifest.json` files contain metadata for each subset. Each row corresponds to a single example/set of `*.wav` files and includes the following fields
+
+- `audio_filepath`: path to the mic file
+- `{tag}_filepath`: path to the corresponding signal component, such as `noise` or `interference`
+- `text`: transcription of the target utterance (if available)
+- `duration`: duration of the signal
+- `{target/noise/interference/mix}_cfg`: dictionary with configuration used to generate the corresponding signal
+- `rt60`: RT60 measured from RIRs for the target source
+- `drr`: Direct-to-reverberant ratio calculated from the RIRs for the target source [2]
+
+---
+
+# References
+
+1. R. Scheibler, E. Bezzam, I. Dokmanić, Pyroomacoustics: A Python package for audio room simulations and array processing algorithms, Proc. IEEE ICASSP, Calgary, CA, 2018.
+
+2. J. Eaton, N. D. Gaubitch, A. H. Moore, P. A. Naylor, The ACE challenge: Corpus description and performance evaluation, Proc. IEEE Workshop on Applications of Signal Processing to Audio and Acoustics (WASPAA), New Paltz, NY, USA, 2015

tools/rir_corpus_generator/conf/rir_corpus.yaml

@@ -0,0 +1,56 @@
+output_dir: ${hydra:job.config_name}
+num_workers: 8
+random_seed: 42
+
+sample_rate: 16000
+
+room:
+  # number of generated rooms per subset
+  num:
+    train: 2000
+    dev: 200
+    test: 200
+  # room dimensions
+  dim:
+    width: [3, 7] # range (meters)
+    length: [3, 9] # range (meters)
+    height: [2.3, 3.5] # range (meters)
+  rt60: [0.15, 0.60] # range (seconds)
+
+mic_array:
+  # 3D Cartesian coordinates per mic [x, y, z],
+  # relative to the center of the array (meters)
+  positions:
+    - [-0.10,  0.095, 0.00]
+    - [ 0.00,  0.095, 0.00]
+    - [ 0.10,  0.095, 0.00]
+    - [-0.10, -0.095, 0.00]
+    - [ 0.00, -0.095, 0.00]
+    - [ 0.10, -0.095, 0.00]
+  placement:
+    # placement of the mic array center
+    x: null # constrained by room dimensions
+    y: null # constrained by room dimensions
+    height: [1.0, 1.5] # range (meters)
+    min_to_wall: 0.5 # min distance to wall (meters)
+  orientation:
+    # orientation of the mic array
+    yaw: [-180, 180] # randomize (degrees)
+    pitch: 0 # keep fixed (degrees)
+    roll: 0 # keep fixed (degrees)
+
+source:
+  # number of sources per room
+  # total number of RIRs is num_rooms x num_sources
+  num: 20
+  # placement of a source inside a room
+  placement:
+    x: null # constrained by room dimensions
+    y: null # constrained by room dimensions
+    height: [1.4, 1.8] # range (meters)
+    min_to_wall: 0.5 # min distance to wall (meters)
+
+anechoic:
+  # parameters for generating the ideal anechoic RIR
+  max_order: 0
+  absorption: 0.999

tools/rir_corpus_generator/conf/rir_mix.yaml

@@ -0,0 +1,52 @@
+output_dir: ${hydra:job.config_name}
+num_workers: 8
+random_seed: 42
+
+sample_rate: 16000
+
+room:
+  # Paths to RIR manifests
+  train: ??? # manifest
+  dev: ??? # manifest
+  test: ??? # manifest
+
+target:
+  # Target source setup
+  train: ??? # manifest
+  dev: ??? # manifest
+  test: ??? # manifest
+  azimuth: [-180, 180] # acceptable range
+  elevation: [-90, 90] # acceptable range
+  distance: [0, 10] # acceptable range
+
+noise:
+  # Background noise setup
+  train: ??? # manifest
+  dev: ??? # manifest
+  test: ??? # manifest
+
+interference:
+  # Directional interference setup
+  train: ??? # manifest
+  dev: ??? # manifest
+  test: ??? # manifest
+  interference_probability: 0.3 # probability an interfering source is present
+  max_num_interferers: 2 # max number of interfering sources, used number will be randomly selected between 0 and max
+  min_azimuth_to_target: 20 # min separation between the target source an an interfering source
+
+mix:
+  train:
+    num: 40000
+    rsnr: [-5, 20] # reverberant signal-to-noise ratio
+    rsir: [5, 25] # reverberant signal-to-interference ratio
+  dev:
+    num: 5000
+    rsnr: [-5, 0, 5, 10, 15, 20] # Two values: uniform distribution, otherwise: random choice from the list
+    rsir: [10, 20, 30]
+  test:
+    num: 5000
+    rsnr: [-5, 0, 5, 10, 15, 20] # Two values: uniform distribution, otherwise: random choice from the list
+    rsir: [10, 20, 30]
+  ref_mic: 0 # index of the ref microphone (zero-based)
+  ref_mic_rms: [-40, -20] # desired RMS for the ref mic signal. If necessary, may be reduced to prevent clipping.
+

tools/rir_corpus_generator/rir_corpus_generator.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2022, 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.
+
+from nemo.collections.audio.data.data_simulation import RIRCorpusGenerator
+from nemo.core.config import hydra_runner
+
+
+"""
+This script creates a corpus of room impulse responses.
+
+Usage:
+    python rir_corpus_generator.py --config-path PATH_TO_CONFIG_DIR --config-name CONFIG_NAME output_dir=OUTPUT_DIR
+
+Details of the configuration can be found in the example configuration files in conf/*
+"""
+
+
+@hydra_runner(config_path="conf", config_name="rir_corpus.yaml")
+def main(cfg):
+    room_corpus = RIRCorpusGenerator(cfg=cfg)
+    room_corpus.generate()
+
+
+if __name__ == "__main__":
+    main()

tools/rir_corpus_generator/rir_mix_generator.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2022, 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.
+
+from nemo.collections.audio.data.data_simulation import RIRMixGenerator
+from nemo.core.config import hydra_runner
+
+
+"""
+This script creates a corpus of signals using RIRs, speech and noise.
+
+Usage:
+    python rir_mix_generator.py --config-path PATH_TO_CONFIG_DIR --config-name CONFIG_NAME output_dir=OUTPUT_DIR
+
+Details of the configuration can be found in the example configuration files in conf/*
+"""
+
+
+@hydra_runner(config_path="conf", config_name="rir_mix.yaml")
+def main(cfg):
+    rir_mix = RIRMixGenerator(cfg=cfg)
+    rir_mix.generate()
+
+
+if __name__ == "__main__":
+    main()

tools/speech_data_explorer/README.md

@@ -0,0 +1,145 @@
+Speech Data Explorer
+--------------------
+
+[Dash](https://plotly.com/dash/)-based tool for interactive exploration of ASR/TTS datasets.
+
+Features:
+- dataset's statistics (alphabet, vocabulary, duration-based histograms)
+- navigation across dataset (sorting, filtering)
+- inspection of individual utterances (waveform, spectrogram, audio player)
+- errors' analysis (Word Error Rate, Character Error Rate, Word Match Rate, Mean Word Accuracy, diff)
+- comparison of two ASR models using interactive word-level accuracy plot
+- read manifests and audio directly from S3-compatible storage (including AIStore)
+- support for tarred audio datasets with efficient byte-range reads via DALI index files
+
+## Quick Start
+
+Install the requirements:
+```
+pip install -r requirements.txt
+```
+
+Run with a local manifest:
+```
+python data_explorer.py path_to_manifest.json
+```
+
+## S3 / AIStore Support
+
+Speech Data Explorer can read manifests and audio files directly from S3-compatible
+object storage, including NVIDIA AIStore (AIS).
+
+### Using an S3 config file
+
+```
+python data_explorer.py s3://bucket/manifest.json --s3cfg ~/.s3cfg[default]
+```
+
+### Using AIStore with environment variables
+
+```
+export AIS_ENDPOINT=http://ais-gateway:8080
+export AIS_AUTHN_TOKEN=your_token
+python data_explorer.py s3://bucket/manifest.json --s3cfg AIS
+```
+
+### Sharded paths (`_OP_/_CL_` syntax)
+
+Manifests and tar files are often split into numbered shards. Instead of listing
+every shard explicitly, use the `_OP_start..end_CL_` range pattern. The tool
+expands it into individual paths automatically:
+
+```
+s3://bucket/manifest__OP_0..255_CL_.json
+→  s3://bucket/manifest_0.json
+   s3://bucket/manifest_1.json
+   ...
+   s3://bucket/manifest_255.json
+```
+
+Multiple ranges in a single path produce a **cartesian product** — useful when
+shards are spread across several buckets or directories:
+
+```
+s3://store_OP_1..2_CL_/audio__OP_0..1_CL_.tar
+→  s3://store1/audio_0.tar
+   s3://store1/audio_1.tar
+   s3://store2/audio_0.tar
+   s3://store2/audio_1.tar
+```
+
+### Tarred audio
+
+When audio is stored in tar archives locally or on S3, use `--tar-base-path` to
+point to the tar files. DALI index files are used automatically (if available at
+`<tar_dir>/dali_index/`) for fast byte-range lookups:
+
+```
+python data_explorer.py /data/manifests/manifest.json \
+    --tar-base-path /data/tarred/audio.tar
+```
+
+```
+python data_explorer.py s3://bucket/manifests/manifest__OP_0..255_CL_.json \
+    --tar-base-path s3://bucket/tarred/audio__OP_0..255_CL_.tar \
+    --s3cfg ~/.s3cfg[default]
+```
+
+You can also specify a custom DALI index location:
+```
+python data_explorer.py s3://bucket/manifest.json \
+    --tar-base-path s3://bucket/tarred/audio__OP_0..255_CL_.tar \
+    --dali-index-base s3://bucket/tarred/dali_index/ \
+    --s3cfg ~/.s3cfg[default]
+```
+
+## Comparing Two ASR Models
+
+### Single manifest with two prediction fields
+
+If your manifest contains two `pred_text_*` fields (e.g. `pred_text_contextnet`
+and `pred_text_conformer`):
+
+```
+python data_explorer.py path_to_manifest.json \
+    -nc pred_text_contextnet pred_text_conformer
+```
+
+### Two separate manifests
+
+You can also pass two separate manifests (order-invariant). Each manifest must
+contain a plain `pred_text` field, and `-nc` names the models:
+
+```
+python data_explorer.py manifest_model_A.json manifest_model_B.json \
+    -nc pred_text_model_A pred_text_model_B
+```
+
+## Manifest Format
+
+JSON manifest file should contain the following fields:
+- `audio_filepath` — path to audio file (local path, or filename inside a tar archive when using `--tar-base-path`)
+- `duration` — duration of the audio file in seconds
+- `text` — reference transcript
+
+Errors' analysis requires `pred_text` (ASR transcript) for all utterances.
+
+Any additional field will be parsed and displayed in the Samples tab.
+
+## Additional Options
+
+| Flag | Description |
+|------|-------------|
+| `--vocab` | Vocabulary file to highlight OOV words |
+| `--port` | Serving port (default: 8050) |
+| `--estimate-audio-metrics` / `-a` | Estimate audio metrics |
+| `--base-path` | Base path for relative audio paths in the manifest |
+| `--tar-base-path` | Local or S3 path to tarred audio files (supports sharded `_OP_..._CL_` patterns) |
+| `--dali-index-base` | Local or S3 path to DALI index directory for fast tar lookups |
+| `--s3cfg` / `-s3c` | S3 config file and section, or `AIS` for AIStore env vars |
+| `--force` / `-f` | Tolerate manifest entries with missing required fields |
+| `-nc` / `--names_compared` | Two field names for model comparison |
+| `--show_statistics` / `-shst` | Field name to show statistics for |
+| `--debug` / `-d` | Enable debug mode |
+
+![Speech Data Explorer](screenshot.png)

tools/speech_data_explorer/requirements.txt

@@ -0,0 +1,13 @@
+boto3
+braceexpand
+dash>=2.1.0
+dash_bootstrap_components>=1.0.3
+diff_match_patch
+kaldialign
+librosa>=0.9.1
+numpy
+pandas
+plotly
+requests
+SoundFile
+tqdm

tools/speech_data_simulator/README.md

@@ -0,0 +1,126 @@
+**Speech Data Simulator**
+===============
+
+Outline
+------------
+
+The speech data simulator generates synthetic multispeaker audio sessions for training or evaluating models for multispeaker ASR or speaker diarization. This tool aims to address the lack of labelled multispeaker training data and to help models deal with overlapping speech.
+
+The simulator loads audio files from different speakers as well as forced alignments for each sentence and concatenates the audio files together to build a synthetic multispeaker audio session. The simulator uses the word alignments to segment the audio from each speaker to produce utterances of the desired length. The simulator also incorporates synthetic room impulse response (RIR) generation in order to simulate multi-microphone multispeaker sessions.
+
+Features
+------------
+
+The simulator is reconfigurable and has several options including:
+
+* Amount of overlapping speech 
+  - The percentage of overlapping speech out of the total speaker time.
+* Percentage of silence
+  - The percentage of the overall audio session that has no speakers talking.
+* Sentence length distribution
+  - The distribution of sentence lengths that is used for sampling (the parameters passed in are for a negative binomial distribution).
+* Number of speakers per session
+* Length of each session
+* Variance in speaker dominance 
+  - Determines what portion of the speaking time will be used by each speaker in a session. Increasing this value will make it more likely that a few speakers dominate the conversation. 
+* Turn taking 
+  - Determines how likely it is that a speaker keeps talking after completing an utterance.
+* Background noise 
+
+The simulator can be used in two modes: near field (no Room Impulse Response) as well as far field (including synthetic RIR). When using synthetic RIR generation, multiple microphones can be placed in the simulated room environment for multichannel simulations.
+
+The simulator also has a speaker enforcement mode which ensures that the correct number of speakers appear in each session (otherwise not guaranteed since speaker turns are stochastic). In speaker enforcement mode, the length of the session or speaker probabilities may be adjusted to ensure all speakers are present.
+
+Required Datasets
+------------
+
+* LibriSpeech (or another single-speaker dataset)
+* LibriSpeech word alignments from [here](https://github.com/CorentinJ/librispeech-alignments) (or alignments corresponding to another single-speaker dataset)
+
+Example alignment format from the LibriSpeech dataset (to be passed as input to the `scripts/speaker_tasks/create_librispeech_alignment_manifest.py` script):
+
+* Alignment files are stored at <Speaker ID>/<Chapter ID>/<Speaker ID>-<Chapter ID>.txt, and each line in the alignment file corresponds to a separate sentence
+* Example of a line in `dev-clean/1272/128104/1272-128104.txt': '1272-128104-0000 ",MISTER,QUILTER,IS,THE,APOSTLE,OF,THE,MIDDLE,CLASSES,,AND,WE,ARE,GLAD,TO,WELCOME,HIS,GOSPEL," "0.500,0.800,1.270,1.400,1.520,2.150,2.270,2.350,2.620,3.270,3.300,3.450,3.600,3.670,4.070,4.200,4.600,4.840,5.510,5.855"`
+
+Optional Datasets
+------------
+
+* Room Impulse Response and Noise Database from [here](https://www.openslr.org/resources/28/rirs_noises.zip) (or another background noise dataset)
+
+Installation (after installing NeMo)
+------------
+
+Note that only one of gpuRIR or pyroomacoustics is required for RIR simulation.
+
+```bash
+pip install cmake
+pip install https://github.com/DavidDiazGuerra/gpuRIR/zipball/master
+pip install pyroomacoustics
+```
+
+Parameters
+------------
+
+* Data simulator parameters are contained in `conf/data_simulator.yaml`
+
+Example Session
+------------
+
+Example multispeaker audio session (using LibriSpeech audio samples and word alignments). RTTM and CTM output labels are highlighted.
+
+![Example multispeaker audio session (using LibriSpeech audio samples and word alignments). RTTM and CTM output labels are highlighted](pictures/audio_session.png)
+
+Running the data simulator for the LibriSpeech dataset
+------------
+
+1. Download the LibriSpeech dataset
+
+```bash
+python scripts/dataset_processing/get_librispeech_data.py \
+  --data_root <path to download LibriSpeech dataset to> \
+  --data_sets ALL
+```
+
+2. Download LibriSpeech alignments from [here](https://drive.google.com/file/d/1WYfgr31T-PPwMcxuAq09XZfHQO5Mw8fE/view?usp=sharing) (the base directory is the LibriSpeech-Alignments directory)
+
+3. Create the manifest file with alignments
+
+```bash
+python <NeMo base path>/scripts/speaker_tasks/create_alignment_manifest.py \
+  --input_manifest_filepath <Path to train_clean_100.json manifest file> \
+  --base_alignment_path <Path to LibriSpeech_Alignments directory> \
+  --output_manifest_filepath train-clean-100-align.json \
+  --ctm_output_directory ./ctm_out \
+  --libri_dataset_split train-clean-100
+```
+
+4. (Optional) Create the background noise manifest file
+
+```bash
+python <NeMo base path>/scripts/speaker_tasks/pathfiles_to_diarize_manifest.py \
+--paths2audio_files <Path to noise list file> \
+--manifest_filepath bg_noise.json
+```
+
+5. Create audio sessions (near field)
+
+```bash
+python multispeaker_simulator.py --config-path='conf' --config-name='data_simulator.yaml' \
+  data_simulator.random_seed=42 \
+  data_simulator.manifest_filepath=./train-clean-100-align.json \
+  data_simulator.outputs.output_dir=./test \
+  data_simulator.background_noise.add_bg=True \
+  data_simulator.background_noise.background_manifest=./bg_noise.json
+```
+
+6. Create multi-microphone audio sessions (with synthetic RIR generation)
+
+```bash
+python multispeaker_simulator.py --config-path='conf' --config-name='data_simulator.yaml' \
+  data_simulator.random_seed=42 \
+  data_simulator.manifest_filepath=./train-clean-100-align.json \
+  data_simulator.outputs.output_dir=./test_rir \
+  data_simulator.background_noise.add_bg=True \
+  data_simulator.background_noise.background_manifest=./bg_noise.json
+  data_simulator.rir_generation.use_rir=True
+```

tools/speech_data_simulator/conf/data_simulator.yaml

@@ -0,0 +1,159 @@
+data_simulator:
+  manifest_filepath: ??? # Manifest file with paths to single speaker audio files
+
+  sr: 16000 # Sampling rate of the input audio files from the manifest
+  random_seed: 42
+  multiprocessing_chunksize: 10000 # Max number that multiprocessing can handle at once
+
+  session_config:
+    num_speakers: 4 # Number of unique speakers per multispeaker audio session
+    num_sessions: 60 # Number of sessions to simulate
+    session_length: 600 # Length of each simulated multispeaker audio session (seconds)
+
+  session_params:
+    max_audio_read_sec: 20.0 # The maximum audio length in second when loading an audio file. The bigger the number, the slower the reading speed. Should be greater than 2.5 second.
+    sentence_length_params: # k,p values for a negative_binomial distribution which is sampled to get the sentence length (in number of words)
+    - 0.4 # k (Number of successes until the experiment is stopped) value must be a positive integer.
+    - 0.05 # p (Success probability) must be in the range (0, 1]. The average sentence length will be k*(1-p)/p
+    dominance_var: 0.11 # Variance in speaker dominance (where each speaker's dominance is sampled from a normal distribution centered on 1/`num_speakers`, and then the dominance values are together normalized to 1)
+    min_dominance: 0.05 # Minimum percentage of speaking time per speaker (note that this can cause the dominance of the other speakers to be slightly reduced)
+    turn_prob: 0.875 # Probability of switching speakers after each utterance
+    min_turn_prob: 0.5 # Minimum turn probability when enforce mode is True to prevent from making excessive session length
+    mean_silence: 0.15 # Mean proportion of silence to speaking time in the audio session. Should be in range [0, 1).
+    mean_silence_var: 0.01 # var for mean silence in all audio sessions. This value should be 0 <= mean_silence_var < mean_silence * (1 - mean_silence)
+    per_silence_var: 900 # var for per silence in each session, set large values to de-correlate silence lengths with the latest speech segment lengths
+    per_silence_min: 0.0 # minimum per silence duration in seconds
+    per_silence_max: -1 # maximum per silence duration in seconds, set -1 for no maximum
+    mean_overlap: 0.1 # Mean proportion of overlap in the overall non-silence duration. Should be in range [0, 1) and recommend [0, 0.15] range.
+    mean_overlap_var: 0.01 # var for mean overlap in all audio sessions. This value should be 0 <= mean_overlap_var < mean_overlap * (1 - mean_overlap)
+    per_overlap_var: 900 # var for per overlap in each session, set large values to de-correlate silence lengths with the latest speech segment lengths
+    per_overlap_min: 0.0 # minimum per overlap duration in seconds
+    per_overlap_max: -1 # maximum per overlap duration in seconds, set -1 for no maximum
+    start_window: true # Window the start of sentences to smooth the audio signal (and remove silence at the start of the clip)
+    window_type: hamming # Type of windowing used when segmenting utterances ("hamming", "hann", "cosine")
+    window_size: 0.05 # Length of window at the start or the end of segmented utterance (seconds)
+    start_buffer: 0.1 # Buffer of silence before the start of the sentence (to avoid cutting off speech or starting abruptly)
+    split_buffer: 0.1 # Split RTTM labels if greater than twice this amount of silence (to avoid long gaps between utterances as being labelled as speech)
+    release_buffer: 0.1 # Buffer before window at end of sentence (to avoid cutting off speech or ending abruptly)
+    normalize: true # Normalize speaker volumes 
+    normalization_type: equal # Normalizing speakers ("equal" - same volume per speaker, "var" - variable volume per speaker)
+    normalization_var: 0.1 # Variance in speaker volume (sample from standard deviation centered at 1)
+    min_volume: 0.75 # Minimum speaker volume (only used when variable normalization is used)
+    max_volume: 1.25 # Maximum speaker volume (only used when variable normalization is used)
+    end_buffer: 0.5 # Buffer at the end of the session to leave blank
+
+  outputs:
+    output_dir: ??? # Output directory for audio sessions and corresponding label files
+    output_filename: multispeaker_session # Output filename for the wav and rttm files
+    overwrite_output: true # If true, delete the output directory if it exists
+    output_precision: 3 # Number of decimal places in output files
+
+  background_noise: # If bg noise is used, a noise source position must be passed for RIR mode
+    add_bg: false # Add ambient background noise if true
+    background_manifest: null # Path to background noise manifest file
+    num_noise_files: 10 # Number of randomly chosen noise source files to be potentially included in one session
+    snr: 60 # SNR for background noise (using average speaker power), set `snr_min` and `snr_max` values to enable random SNR
+    snr_min: null  # Min random SNR for background noise (using average speaker power), set `null` to use fixed SNR
+    snr_max: null  # Max random SNR for background noise (using average speaker power), set `null` to use fixed SNR
+
+  # Segment and session augmentations. Available augmentations are in nemo/collections/asr/parts/preprocessing/perturb.py
+  # See tutorial at https://github.com/NVIDIA/NeMo/blob/main/tutorials/asr/Online_Noise_Augmentation.ipynb
+  # Note that ImpulsePerturbation, NoisePerturbation, RirAndNoisePerturbation and other perturbations that uses `collections.ASRAudioText`
+  # cannot use multi-proccessing in simulation, due to non-pickable errors.
+  segment_augmentor:
+    add_seg_aug: False # Set True to enable augmentation on each speech segment
+    augmentor:
+      gain:  # Randomly perturb the gain of each speech segment
+        prob: 0.5  # Probability of applying gain augmentation
+        min_gain_dbfs: -10.0  # Min dB level to add
+        max_gain_dbfs: 10.0  # Max dB level to add
+
+  session_augmentor:
+    add_sess_aug: False # Set True to enable audio augmentation on the whole session
+    augmentor:
+      white_noise:  # Add random white noise to the whole session
+        prob: 1.0  # Probability of adding white noise
+        min_level: -90  # Min level of noise loudness (dB)
+        max_level: -46  # Max level of noise loudness (dB)
+
+  speaker_enforcement:
+    enforce_num_speakers: true # Enforce that all requested speakers are present in the output wav file
+    enforce_time:  # Percentage of the way through the audio session that enforcement mode is triggered (sampled between time 1 and 2)
+    - 0.25
+    - 0.75
+
+  segment_manifest: # Parameters for regenerating the segment manifest file 
+    window: 0.5 # Window length for segmentation 
+    shift: 0.25 # Shift length for segmentation
+    step_count: 50 # Number of the unit segments you want to create per utterance
+    deci: 3 # Rounding decimals for segment manifest file
+
+  rir_generation: # Using synthetic RIR augmentation
+    use_rir: false # Whether to generate synthetic RIR
+    toolkit: 'pyroomacoustics' # Which toolkit to use ("pyroomacoustics", "gpuRIR")
+    room_config:
+      room_sz: # Size of the shoebox room environment (1d array for specific, 2d array for random range to be sampled from)
+      - - 2
+        - 3
+      - - 2
+        - 3
+      - - 2
+        - 3
+      pos_src: # Positions of the speakers in the simulated room environment (2d array for specific, 3d array for random ranges to be sampled from)
+      - - - 0.5
+          - 1.5
+        - - 0.5
+          - 1.5
+        - - 0.5
+          - 1.5
+      - - - 0.5
+          - 1.5
+        - - 0.5
+          - 1.5
+        - - 0.5
+          - 1.5
+      - - - 0.5
+          - 1.5
+        - - 0.5
+          - 1.5
+        - - 0.5
+          - 1.5
+      - - - 0.5
+          - 1.5
+        - - 0.5
+          - 1.5
+        - - 0.5
+          - 1.5
+      noise_src_pos: # Position in room for the ambient background noise source
+      - 1.5
+      - 1.5
+      - 2
+    mic_config:
+      num_channels: 2 # Number of output audio channels
+      pos_rcv: # Microphone positions in the simulated room environment (1d/2d array for specific, 2d/3d array for range assuming num_channels is 1/2+)
+      - - - 0.5
+          - 1.5
+        - - 0.5
+          - 1.5
+        - - 0.5
+          - 1.5
+      - - - 0.5
+          - 1.5
+        - - 0.5
+          - 1.5
+        - - 0.5
+          - 1.5
+      orV_rcv: null # Microphone orientations (needed for non-omnidirectional microphones)
+      mic_pattern: omni # Microphone type ("omni" - omnidirectional) - currently only omnidirectional microphones are supported for pyroomacoustics
+
+    absorbtion_params: # Note: only `T60` is used for pyroomacoustics simulations
+      abs_weights: # Absorption coefficient ratios for each surface 
+      - 0.9
+      - 0.9
+      - 0.9
+      - 0.9
+      - 0.9
+      - 0.9
+      T60: 0.1 # Room reverberation time (`T60` is the time it takes for the RIR to decay by 60DB)
+      att_diff: 15.0 # Starting attenuation (if this is different than att_max, the diffuse reverberation model is used by gpuRIR)
+      att_max: 60.0 # End attenuation when using the diffuse reverberation model (gpuRIR)

tools/speech_data_simulator/multispeaker_simulator.py

@@ -0,0 +1,53 @@
+# Copyright (c) 2022, 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.
+
+from multiprocessing import set_start_method
+
+from nemo.collections.asr.data.data_simulation import MultiSpeakerSimulator, RIRMultiSpeakerSimulator
+from nemo.core.config import hydra_runner
+
+
+"""
+This script creates a synthetic diarization session using the provided audio dataset with ctm files.
+Usage:
+  python <NEMO_ROOT>/tools/speech_data_simulator/multispeaker_simulator.py  \
+    num_workers=10 \
+    data_simulator.random_seed=42 \
+    data_simulator.manifest_filepath=manifest_with_alignment_file.json \
+    data_simulator.outputs.output_dir=./simulated_data \
+    data_simulator.outputs.output_filename=sim_spk2_sess20 \
+    data_simulator.session_config.num_sessions=1000 \
+    data_simulator.session_config.num_speakers=2 \
+    data_simulator.session_config.session_length=20 \
+    data_simulator.background_noise.add_bg=False \
+    data_simulator.background_noise.background_manifest=background_noise.json \
+    data_simulator.background_noise.snr=40 \
+
+Check out parameters in ./conf/data_simulator.yaml.
+"""
+
+
+@hydra_runner(config_path="conf", config_name="data_simulator.yaml")
+def main(cfg):
+    if cfg.data_simulator.rir_generation.use_rir:
+        simulator = RIRMultiSpeakerSimulator(cfg=cfg)
+    else:
+        simulator = MultiSpeakerSimulator(cfg=cfg)
+
+    set_start_method('spawn', force=True)
+    simulator.generate_sessions()
+
+
+if __name__ == "__main__":
+    main()