chore: import upstream snapshot with attribution

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