Drop Your Feature Requests HERE

[MahmoudAshraf97]

Hi Everyone,
I'm planning for a refactor for this repo to make it more modular and ease maintenance especially in w.r.t dependency management
Also since it's hard to find a solution fits all, I'm thinking about making the repo modular so that we can integrate multiple options for each stage (transcription, diarization, etc.)
To sum it up this is the roadmap:

QoL Goals:

• Write a GitHub action to frequently verify that all requirements are compatible across different OSes Add a simple action to test requirements weekly #241
• Write Unit Tests
• Create Docker Image and Docker Compose Interface

Refactor Goals:

• Split everything to modules (Preprocessing, Transcription, Diarization, Postprocessing), this will allow integrating more options for each stage such as better vocal separation, different transcription provider, custom postprocessing and such.
• Convert this repo to a PyPI package

Happy to hear everyone's thoughts and ideas

[transcriptionstream]

Fantastic idea! I'd love to assist in the docker/compose aspect any way I can.

as for feature request, I'd love to see an option to choose/utilize other models than ctc-forced-aligner, if needed, based on license/use

[MahmoudAshraf97]

Can you check #237? I'm not familiar with docker compose, so I don't know what is the functionality needed here

as for ctc-forced-aligner the code itself can be used commercially after changing the default model which can be done by passing the model name to the loading function, alternatives will be nemo forced aligner or Montreal forced aligner

[SIlver--]

Hello. Requesting to have word_timestamps as a command line option included with the output results. I realize adding just the option to the whisper_model.transcribe isn't enough to get the results as they get filtered out during the diarization merging methods that happen afterwards.

In my local version of the project, I included helper functions to line up each word of the formatted_transcript with the word element from of the Iterable[Words] prop from faster_whisper return, but I'm not so sure this is the best way to do it. And how I have it done so far, there's no way to enable/disable this from the Command Line Options so I feel embarrassed 😞 to share my work so far.

[MahmoudAshraf97]

#187

[MahmoudAshraf97]

#289

[raivisdejus]

Switch to python version of the uroman tool for romanization.

The preprocess_text from ctc_forced_aligner uses system call to a perl script. There seems to be also an updated python version of the same script that would make integrating your great tool into different applications a bit easier.

P.S. I am currently working on integrating whisper-diarization into https://github.com/chidiwilliams/buzz
Thanks for this great and easy to use tool that gives excellent results.

[Arche151]

@MahmoudAshraf97 It would be great, if we could get word-level timestamps by calling diarize.py with an argument like -word_level_timestamps. faster-whisper can produce word-level timestamps, so I believe this should be possible but @SIlver-- did mention that they had trouble implementing it.

[MahmoudAshraf97]

Switch to python version of the uroman tool for romanization.

The preprocess_text from ctc_forced_aligner uses system call to a perl script. There seems to be also an updated python version of the same script that would make integrating your great tool into different applications a bit easier.

P.S. I am currently working on integrating whisper-diarization into https://github.com/chidiwilliams/buzz Thanks for this great and easy to use tool that gives excellent results.

This was implemented in the latest version of ctc-forced-aligner

[MahmoudAshraf97]

@MahmoudAshraf97 It would be great, if we could get word-level timestamps by calling diarize.py with an argument like -word_level_timestamps. faster-whisper can produce word-level timestamps, so I believe this should be possible but @SIlver-- did mention that they had trouble implementing it.

It's already possible to obtain them, just extract this output to any file

whisper-diarization/diarize.py

Line 187 in 7348644

word_timestamps = postprocess_results(text_starred, spans, stride, scores)

or do you mean word level diarization output?

[Arche151]

@MahmoudAshraf97 Ah, perfect, thanks!

No, I did mean word-level srt-files.

I ended up creating this modified diarize.py, which outputs srt-files with word-level timestamps when running diarize.py with the argument --word-timestamps:

import argparse
import json
import logging
import os
import re

import faster_whisper
import torch
import torchaudio

from ctc_forced_aligner import (
    generate_emissions,
    get_alignments,
    get_spans,
    load_alignment_model,
    postprocess_results,
    preprocess_text,
)
from deepmultilingualpunctuation import PunctuationModel
from nemo.collections.asr.models.msdd_models import NeuralDiarizer

from helpers import (
    cleanup,
    create_config,
    find_numeral_symbol_tokens,
    get_realigned_ws_mapping_with_punctuation,
    get_sentences_speaker_mapping,
    get_speaker_aware_transcript,
    get_words_speaker_mapping,
    langs_to_iso,
    process_language_arg,
    punct_model_langs,
    whisper_langs,
    write_srt,
    write_word_level_srt,  # New function
    write_word_level_json,  # New function
)

mtypes = {"cpu": "int8", "cuda": "float16"}

pid = os.getpid()
temp_outputs_dir = f"temp_outputs_{pid}"

# Initialize parser
parser = argparse.ArgumentParser()
parser.add_argument(
    "-a", "--audio", help="name of the target audio file", required=True
)
parser.add_argument(
    "--no-stem",
    action="store_false",
    dest="stemming",
    default=True,
    help="Disables source separation."
    "This helps with long files that don't contain a lot of music.",
)

parser.add_argument(
    "--suppress_numerals",
    action="store_true",
    dest="suppress_numerals",
    default=False,
    help="Suppresses Numerical Digits."
    "This helps the diarization accuracy but converts all digits into written text.",
)

parser.add_argument(
    "--whisper-model",
    dest="model_name",
    default="medium.en",
    help="name of the Whisper model to use",
)

parser.add_argument(
    "--batch-size",
    type=int,
    dest="batch_size",
    default=8,
    help="Batch size for batched inference, reduce if you run out of memory, "
    "set to 0 for original whisper longform inference",
)

parser.add_argument(
    "--language",
    type=str,
    default=None,
    choices=whisper_langs,
    help="Language spoken in the audio, specify None to perform language detection",
)

parser.add_argument(
    "--device",
    dest="device",
    default="cuda" if torch.cuda.is_available() else "cpu",
    help="if you have a GPU use 'cuda', otherwise 'cpu'",
)

# Add new argument for word-level timestamps
parser.add_argument(
    "--word-timestamps",
    action="store_true",
    dest="word_timestamps",
    default=False,
    help="Generate word-level timestamps in addition to sentence-level output",
)

args = parser.parse_args()
language = process_language_arg(args.language, args.model_name)

if args.stemming:
    # Isolate vocals from the rest of the audio

    return_code = os.system(
        f'python -m demucs.separate -n htdemucs --two-stems=vocals "{args.audio}" -o "{temp_outputs_dir}" --device "{args.device}"'
    )

    if return_code != 0:
        logging.warning(
            "Source splitting failed, using original audio file. "
            "Use --no-stem argument to disable it."
        )
        vocal_target = args.audio
    else:
        vocal_target = os.path.join(
            temp_outputs_dir,
            "htdemucs",
            os.path.splitext(os.path.basename(args.audio))[0],
            "vocals.wav",
        )
else:
    vocal_target = args.audio


# Transcribe the audio file

whisper_model = faster_whisper.WhisperModel(
    args.model_name, device=args.device, compute_type=mtypes[args.device]
)
whisper_pipeline = faster_whisper.BatchedInferencePipeline(whisper_model)
audio_waveform = faster_whisper.decode_audio(vocal_target)
suppress_tokens = (
    find_numeral_symbol_tokens(whisper_model.hf_tokenizer)
    if args.suppress_numerals
    else [-1]
)

if args.batch_size > 0:
    transcript_segments, info = whisper_pipeline.transcribe(
        audio_waveform,
        language,
        suppress_tokens=suppress_tokens,
        batch_size=args.batch_size,
    )
else:
    transcript_segments, info = whisper_model.transcribe(
        audio_waveform,
        language,
        suppress_tokens=suppress_tokens,
        vad_filter=True,
    )

full_transcript = "".join(segment.text for segment in transcript_segments)

# clear gpu vram
del whisper_model, whisper_pipeline
torch.cuda.empty_cache()

# Forced Alignment
alignment_model, alignment_tokenizer = load_alignment_model(
    args.device,
    dtype=torch.float16 if args.device == "cuda" else torch.float32,
)

emissions, stride = generate_emissions(
    alignment_model,
    torch.from_numpy(audio_waveform)
    .to(alignment_model.dtype)
    .to(alignment_model.device),
    batch_size=args.batch_size,
)

del alignment_model
torch.cuda.empty_cache()

tokens_starred, text_starred = preprocess_text(
    full_transcript,
    romanize=True,
    language=langs_to_iso[info.language],
)

segments, scores, blank_token = get_alignments(
    emissions,
    tokens_starred,
    alignment_tokenizer,
)

spans = get_spans(tokens_starred, segments, blank_token)

word_timestamps = postprocess_results(text_starred, spans, stride, scores)


# convert audio to mono for NeMo combatibility
ROOT = os.getcwd()
temp_path = os.path.join(ROOT, temp_outputs_dir)
os.makedirs(temp_path, exist_ok=True)
torchaudio.save(
    os.path.join(temp_path, "mono_file.wav"),
    torch.from_numpy(audio_waveform).unsqueeze(0).float(),
    16000,
    channels_first=True,
)


# Initialize NeMo MSDD diarization model
msdd_model = NeuralDiarizer(cfg=create_config(temp_path)).to(args.device)
msdd_model.diarize()

del msdd_model
torch.cuda.empty_cache()

# Reading timestamps <> Speaker Labels mapping


speaker_ts = []
with open(os.path.join(temp_path, "pred_rttms", "mono_file.rttm"), "r") as f:
    lines = f.readlines()
    for line in lines:
        line_list = line.split(" ")
        s = int(float(line_list[5]) * 1000)
        e = s + int(float(line_list[8]) * 1000)
        speaker_ts.append([s, e, int(line_list[11].split("_")[-1])])

wsm = get_words_speaker_mapping(word_timestamps, speaker_ts, "start")

if info.language in punct_model_langs:
    # restoring punctuation in the transcript to help realign the sentences
    punct_model = PunctuationModel(model="kredor/punctuate-all")

    words_list = list(map(lambda x: x["word"], wsm))

    labled_words = punct_model.predict(words_list, chunk_size=230)

    ending_puncts = ".?!"
    model_puncts = ".,;:!?"

    # We don't want to punctuate U.S.A. with a period. Right?
    is_acronym = lambda x: re.fullmatch(r"\b(?:[a-zA-Z]\.){2,}", x)

    for word_dict, labeled_tuple in zip(wsm, labled_words):
        word = word_dict["word"]
        if (
            word
            and labeled_tuple[1] in ending_puncts
            and (word[-1] not in model_puncts or is_acronym(word))
        ):
            word += labeled_tuple[1]
            if word.endswith(".."):
                word = word.rstrip(".")
            word_dict["word"] = word

else:
    logging.warning(
        f"Punctuation restoration is not available for {info.language} language."
        " Using the original punctuation."
    )

wsm = get_realigned_ws_mapping_with_punctuation(wsm)

# Save word-level timestamps if requested
if args.word_timestamps:
    # Write word-level SRT
    with open(f"{os.path.splitext(args.audio)[0]}_words.srt", "w", encoding="utf-8-sig") as srt:
        write_word_level_srt(wsm, srt)
    
    # Write word-level JSON for programmatic access
    with open(f"{os.path.splitext(args.audio)[0]}_words.json", "w", encoding="utf-8-sig") as f:
        write_word_level_json(wsm, f)
    
    print(f"Word-level timestamps saved to {os.path.splitext(args.audio)[0]}_words.srt and .json")

# Generate sentence-level output (existing functionality)
ssm = get_sentences_speaker_mapping(wsm, speaker_ts)

with open(f"{os.path.splitext(args.audio)[0]}.txt", "w", encoding="utf-8-sig") as f:
    get_speaker_aware_transcript(ssm, f)

with open(f"{os.path.splitext(args.audio)[0]}.srt", "w", encoding="utf-8-sig") as srt:
    write_srt(ssm, srt)

cleanup(temp_path)

I also had to update helpers.py to this: ```

import json
import os
import shutil

import nltk
import wget

from omegaconf import OmegaConf

punct_model_langs = [
    "en",
    "fr",
    "de",
    "es",
    "it",
    "nl",
    "pt",
    "bg",
    "pl",
    "cs",
    "sk",
    "sl",
]

LANGUAGES = {
    "en": "english",
    "zh": "chinese",
    "de": "german",
    "es": "spanish",
    "ru": "russian",
    "ko": "korean",
    "fr": "french",
    "ja": "japanese",
    "pt": "portuguese",
    "tr": "turkish",
    "pl": "polish",
    "ca": "catalan",
    "nl": "dutch",
    "ar": "arabic",
    "sv": "swedish",
    "it": "italian",
    "id": "indonesian",
    "hi": "hindi",
    "fi": "finnish",
    "vi": "vietnamese",
    "he": "hebrew",
    "uk": "ukrainian",
    "el": "greek",
    "ms": "malay",
    "cs": "czech",
    "ro": "romanian",
    "da": "danish",
    "hu": "hungarian",
    "ta": "tamil",
    "no": "norwegian",
    "th": "thai",
    "ur": "urdu",
    "hr": "croatian",
    "bg": "bulgarian",
    "lt": "lithuanian",
    "la": "latin",
    "mi": "maori",
    "ml": "malayalam",
    "cy": "welsh",
    "sk": "slovak",
    "te": "telugu",
    "fa": "persian",
    "lv": "latvian",
    "bn": "bengali",
    "sr": "serbian",
    "az": "azerbaijani",
    "sl": "slovenian",
    "kn": "kannada",
    "et": "estonian",
    "mk": "macedonian",
    "br": "breton",
    "eu": "basque",
    "is": "icelandic",
    "hy": "armenian",
    "ne": "nepali",
    "mn": "mongolian",
    "bs": "bosnian",
    "kk": "kazakh",
    "sq": "albanian",
    "sw": "swahili",
    "gl": "galician",
    "mr": "marathi",
    "pa": "punjabi",
    "si": "sinhala",
    "km": "khmer",
    "sn": "shona",
    "yo": "yoruba",
    "so": "somali",
    "af": "afrikaans",
    "oc": "occitan",
    "ka": "georgian",
    "be": "belarusian",
    "tg": "tajik",
    "sd": "sindhi",
    "gu": "gujarati",
    "am": "amharic",
    "yi": "yiddish",
    "lo": "lao",
    "uz": "uzbek",
    "fo": "faroese",
    "ht": "haitian creole",
    "ps": "pashto",
    "tk": "turkmen",
    "nn": "nynorsk",
    "mt": "maltese",
    "sa": "sanskrit",
    "lb": "luxembourgish",
    "my": "myanmar",
    "bo": "tibetan",
    "tl": "tagalog",
    "mg": "malagasy",
    "as": "assamese",
    "tt": "tatar",
    "haw": "hawaiian",
    "ln": "lingala",
    "ha": "hausa",
    "ba": "bashkir",
    "jw": "javanese",
    "su": "sundanese",
    "yue": "cantonese",
}

# language code lookup by name, with a few language aliases
TO_LANGUAGE_CODE = {
    **{language: code for code, language in LANGUAGES.items()},
    "burmese": "my",
    "valencian": "ca",
    "flemish": "nl",
    "haitian": "ht",
    "letzeburgesch": "lb",
    "pushto": "ps",
    "panjabi": "pa",
    "moldavian": "ro",
    "moldovan": "ro",
    "sinhalese": "si",
    "castilian": "es",
}

whisper_langs = sorted(LANGUAGES.keys()) + sorted(
    [k.title() for k in TO_LANGUAGE_CODE.keys()]
)

langs_to_iso = {
    "af": "afr",
    "am": "amh",
    "ar": "ara",
    "as": "asm",
    "az": "aze",
    "ba": "bak",
    "be": "bel",
    "bg": "bul",
    "bn": "ben",
    "bo": "tib",
    "br": "bre",
    "bs": "bos",
    "ca": "cat",
    "cs": "cze",
    "cy": "wel",
    "da": "dan",
    "de": "ger",
    "el": "gre",
    "en": "eng",
    "es": "spa",
    "et": "est",
    "eu": "baq",
    "fa": "per",
    "fi": "fin",
    "fo": "fao",
    "fr": "fre",
    "gl": "glg",
    "gu": "guj",
    "ha": "hau",
    "haw": "haw",
    "he": "heb",
    "hi": "hin",
    "hr": "hrv",
    "ht": "hat",
    "hu": "hun",
    "hy": "arm",
    "id": "ind",
    "is": "ice",
    "it": "ita",
    "ja": "jpn",
    "jw": "jav",
    "ka": "geo",
    "kk": "kaz",
    "km": "khm",
    "kn": "kan",
    "ko": "kor",
    "la": "lat",
    "lb": "ltz",
    "ln": "lin",
    "lo": "lao",
    "lt": "lit",
    "lv": "lav",
    "mg": "mlg",
    "mi": "mao",
    "mk": "mac",
    "ml": "mal",
    "mn": "mon",
    "mr": "mar",
    "ms": "may",
    "mt": "mlt",
    "my": "bur",
    "ne": "nep",
    "nl": "dut",
    "nn": "nno",
    "no": "nor",
    "oc": "oci",
    "pa": "pan",
    "pl": "pol",
    "ps": "pus",
    "pt": "por",
    "ro": "rum",
    "ru": "rus",
    "sa": "san",
    "sd": "snd",
    "si": "sin",
    "sk": "slo",
    "sl": "slv",
    "sn": "sna",
    "so": "som",
    "sq": "alb",
    "sr": "srp",
    "su": "sun",
    "sv": "swe",
    "sw": "swa",
    "ta": "tam",
    "te": "tel",
    "tg": "tgk",
    "th": "tha",
    "tk": "tuk",
    "tl": "tgl",
    "tr": "tur",
    "tt": "tat",
    "uk": "ukr",
    "ur": "urd",
    "uz": "uzb",
    "vi": "vie",
    "yi": "yid",
    "yo": "yor",
    "yue": "yue",
    "zh": "chi",
}


def create_config(output_dir):
    DOMAIN_TYPE = "telephonic"
    CONFIG_LOCAL_DIRECTORY = "nemo_msdd_configs"
    CONFIG_FILE_NAME = f"diar_infer_{DOMAIN_TYPE}.yaml"
    MODEL_CONFIG_PATH = os.path.join(CONFIG_LOCAL_DIRECTORY, CONFIG_FILE_NAME)
    if not os.path.exists(MODEL_CONFIG_PATH):
        os.makedirs(CONFIG_LOCAL_DIRECTORY, exist_ok=True)
        CONFIG_URL = f"https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/speaker_tasks/diarization/conf/inference/{CONFIG_FILE_NAME}"
        MODEL_CONFIG_PATH = wget.download(CONFIG_URL, MODEL_CONFIG_PATH)

    config = OmegaConf.load(MODEL_CONFIG_PATH)

    data_dir = os.path.join(output_dir, "data")
    os.makedirs(data_dir, exist_ok=True)

    meta = {
        "audio_filepath": os.path.join(output_dir, "mono_file.wav"),
        "offset": 0,
        "duration": None,
        "label": "infer",
        "text": "-",
        "rttm_filepath": None,
        "uem_filepath": None,
    }
    with open(os.path.join(data_dir, "input_manifest.json"), "w") as fp:
        json.dump(meta, fp)
        fp.write("\n")

    pretrained_vad = "vad_multilingual_marblenet"
    pretrained_speaker_model = "titanet_large"
    config.num_workers = 0
    config.diarizer.manifest_filepath = os.path.join(data_dir, "input_manifest.json")
    config.diarizer.out_dir = (
        output_dir  # Directory to store intermediate files and prediction outputs
    )

    config.diarizer.speaker_embeddings.model_path = pretrained_speaker_model
    config.diarizer.oracle_vad = (
        False  # compute VAD provided with model_path to vad config
    )
    config.diarizer.clustering.parameters.oracle_num_speakers = False

    # Here, we use our in-house pretrained NeMo VAD model
    config.diarizer.vad.model_path = pretrained_vad
    config.diarizer.vad.parameters.onset = 0.8
    config.diarizer.vad.parameters.offset = 0.6
    config.diarizer.vad.parameters.pad_offset = -0.05
    config.diarizer.msdd_model.model_path = (
        "diar_msdd_telephonic"  # Telephonic speaker diarization model
    )

    return config


def get_word_ts_anchor(s, e, option="start"):
    if option == "end":
        return e
    elif option == "mid":
        return (s + e) / 2
    return s


def get_words_speaker_mapping(wrd_ts, spk_ts, word_anchor_option="start"):
    s, e, sp = spk_ts[0]
    wrd_pos, turn_idx = 0, 0
    wrd_spk_mapping = []
    for wrd_dict in wrd_ts:
        ws, we, wrd = (
            int(wrd_dict["start"] * 1000),
            int(wrd_dict["end"] * 1000),
            wrd_dict["text"],
        )
        wrd_pos = get_word_ts_anchor(ws, we, word_anchor_option)
        while wrd_pos > float(e):
            turn_idx += 1
            turn_idx = min(turn_idx, len(spk_ts) - 1)
            s, e, sp = spk_ts[turn_idx]
            if turn_idx == len(spk_ts) - 1:
                e = get_word_ts_anchor(ws, we, option="end")
        wrd_spk_mapping.append(
            {"word": wrd, "start_time": ws, "end_time": we, "speaker": sp}
        )
    return wrd_spk_mapping


sentence_ending_punctuations = ".?!"


def get_first_word_idx_of_sentence(word_idx, word_list, speaker_list, max_words):
    is_word_sentence_end = (
        lambda x: x >= 0 and word_list[x][-1] in sentence_ending_punctuations
    )
    left_idx = word_idx
    while (
        left_idx > 0
        and word_idx - left_idx < max_words
        and speaker_list[left_idx - 1] == speaker_list[left_idx]
        and not is_word_sentence_end(left_idx - 1)
    ):
        left_idx -= 1

    return left_idx if left_idx == 0 or is_word_sentence_end(left_idx - 1) else -1


def get_last_word_idx_of_sentence(word_idx, word_list, max_words):
    is_word_sentence_end = (
        lambda x: x >= 0 and word_list[x][-1] in sentence_ending_punctuations
    )
    right_idx = word_idx
    while (
        right_idx < len(word_list) - 1
        and right_idx - word_idx < max_words
        and not is_word_sentence_end(right_idx)
    ):
        right_idx += 1

    return (
        right_idx
        if right_idx == len(word_list) - 1 or is_word_sentence_end(right_idx)
        else -1
    )


def get_realigned_ws_mapping_with_punctuation(
    word_speaker_mapping, max_words_in_sentence=50
):
    is_word_sentence_end = (
        lambda x: x >= 0
        and word_speaker_mapping[x]["word"][-1] in sentence_ending_punctuations
    )
    wsp_len = len(word_speaker_mapping)

    words_list, speaker_list = [], []
    for k, line_dict in enumerate(word_speaker_mapping):
        word, speaker = line_dict["word"], line_dict["speaker"]
        words_list.append(word)
        speaker_list.append(speaker)

    k = 0
    while k < len(word_speaker_mapping):
        line_dict = word_speaker_mapping[k]
        if (
            k < wsp_len - 1
            and speaker_list[k] != speaker_list[k + 1]
            and not is_word_sentence_end(k)
        ):
            left_idx = get_first_word_idx_of_sentence(
                k, words_list, speaker_list, max_words_in_sentence
            )
            right_idx = (
                get_last_word_idx_of_sentence(
                    k, words_list, max_words_in_sentence - k + left_idx - 1
                )
                if left_idx > -1
                else -1
            )
            if min(left_idx, right_idx) == -1:
                k += 1
                continue

            spk_labels = speaker_list[left_idx : right_idx + 1]
            mod_speaker = max(set(spk_labels), key=spk_labels.count)
            if spk_labels.count(mod_speaker) < len(spk_labels) // 2:
                k += 1
                continue

            speaker_list[left_idx : right_idx + 1] = [mod_speaker] * (
                right_idx - left_idx + 1
            )
            k = right_idx

        k += 1

    k, realigned_list = 0, []
    while k < len(word_speaker_mapping):
        line_dict = word_speaker_mapping[k].copy()
        line_dict["speaker"] = speaker_list[k]
        realigned_list.append(line_dict)
        k += 1

    return realigned_list


def get_sentences_speaker_mapping(word_speaker_mapping, spk_ts):
    sentence_checker = nltk.tokenize.PunktSentenceTokenizer().text_contains_sentbreak
    s, e, spk = spk_ts[0]
    prev_spk = spk

    snts = []
    snt = {"speaker": f"Speaker {spk}", "start_time": s, "end_time": e, "text": ""}

    for wrd_dict in word_speaker_mapping:
        wrd, spk = wrd_dict["word"], wrd_dict["speaker"]
        s, e = wrd_dict["start_time"], wrd_dict["end_time"]
        if spk != prev_spk or sentence_checker(snt["text"] + " " + wrd):
            snts.append(snt)
            snt = {
                "speaker": f"Speaker {spk}",
                "start_time": s,
                "end_time": e,
                "text": "",
            }
        else:
            snt["end_time"] = e
        snt["text"] += wrd + " "
        prev_spk = spk

    snts.append(snt)
    return snts


def get_speaker_aware_transcript(sentences_speaker_mapping, f):
    previous_speaker = sentences_speaker_mapping[0]["speaker"]
    f.write(f"{previous_speaker}: ")

    for sentence_dict in sentences_speaker_mapping:
        speaker = sentence_dict["speaker"]
        sentence = sentence_dict["text"]

        # If this speaker doesn't match the previous one, start a new paragraph
        if speaker != previous_speaker:
            f.write(f"\n\n{speaker}: ")
            previous_speaker = speaker

        # No matter what, write the current sentence
        f.write(sentence + " ")


def format_timestamp(
    milliseconds: float, always_include_hours: bool = False, decimal_marker: str = "."
):
    assert milliseconds >= 0, "non-negative timestamp expected"

    hours = milliseconds // 3_600_000
    milliseconds -= hours * 3_600_000

    minutes = milliseconds // 60_000
    milliseconds -= minutes * 60_000

    seconds = milliseconds // 1_000
    milliseconds -= seconds * 1_000

    hours_marker = f"{hours:02d}:" if always_include_hours or hours > 0 else ""
    return (
        f"{hours_marker}{minutes:02d}:{seconds:02d}{decimal_marker}{milliseconds:03d}"
    )


def write_srt(transcript, file):
    """
    Write a transcript to a file in SRT format.

    """
    for i, segment in enumerate(transcript, start=1):
        # write srt lines
        print(
            f"{i}\n"
            f"{format_timestamp(segment['start_time'], always_include_hours=True, decimal_marker=',')} --> "
            f"{format_timestamp(segment['end_time'], always_include_hours=True, decimal_marker=',')}\n"
            f"{segment['speaker']}: {segment['text'].strip().replace('-->', '->')}\n",
            file=file,
            flush=True,
        )


def find_numeral_symbol_tokens(tokenizer):
    numeral_symbol_tokens = [
        -1,
    ]
    for token, token_id in tokenizer.get_vocab().items():
        has_numeral_symbol = any(c in "0123456789%$£" for c in token)
        if has_numeral_symbol:
            numeral_symbol_tokens.append(token_id)
    return numeral_symbol_tokens


def _get_next_start_timestamp(word_timestamps, current_word_index, final_timestamp):
    # if current word is the last word
    if current_word_index == len(word_timestamps) - 1:
        return word_timestamps[current_word_index]["start"]

    next_word_index = current_word_index + 1
    while current_word_index < len(word_timestamps) - 1:
        if word_timestamps[next_word_index].get("start") is None:
            # if next word doesn't have a start timestamp
            # merge it with the current word and delete it
            word_timestamps[current_word_index]["word"] += (
                " " + word_timestamps[next_word_index]["word"]
            )

            word_timestamps[next_word_index]["word"] = None
            next_word_index += 1
            if next_word_index == len(word_timestamps):
                return final_timestamp

        else:
            return word_timestamps[next_word_index]["start"]


def filter_missing_timestamps(
    word_timestamps, initial_timestamp=0, final_timestamp=None
):
    # handle the first and last word
    if word_timestamps[0].get("start") is None:
        word_timestamps[0]["start"] = (
            initial_timestamp if initial_timestamp is not None else 0
        )
        word_timestamps[0]["end"] = _get_next_start_timestamp(
            word_timestamps, 0, final_timestamp
        )

    result = [
        word_timestamps[0],
    ]

    for i, ws in enumerate(word_timestamps[1:], start=1):
        # if ws doesn't have a start and end
        # use the previous end as start and next start as end
        if ws.get("start") is None and ws.get("word") is not None:
            ws["start"] = word_timestamps[i - 1]["end"]
            ws["end"] = _get_next_start_timestamp(word_timestamps, i, final_timestamp)

        if ws["word"] is not None:
            result.append(ws)
    return result


def cleanup(path: str):
    """path could either be relative or absolute."""
    # check if file or directory exists
    if os.path.isfile(path) or os.path.islink(path):
        # remove file
        os.remove(path)
    elif os.path.isdir(path):
        # remove directory and all its content
        shutil.rmtree(path)
    else:
        raise ValueError(f"Path {path} is not a file or dir.")


def process_language_arg(language: str, model_name: str):
    """
    Process the language argument to make sure it's valid
    and convert language names to language codes.
    """
    if language is not None:
        language = language.lower()
        if language not in LANGUAGES:
            if language in TO_LANGUAGE_CODE:
                language = TO_LANGUAGE_CODE[language]
            else:
                raise ValueError(f"Unsupported language: {language}")

        if model_name.endswith(".en") and language != "en":
            raise ValueError(
                f"{model_name} is an English-only model but choosen language is '{language}'"
            )

    return language


# ============================================
# NEW FUNCTIONS FOR WORD-LEVEL TIMESTAMPS
# ============================================

def write_word_level_srt(word_speaker_mapping, file):
    """
    Write word-level timestamps to SRT format.
    Each word gets its own subtitle entry with speaker information.
    """
    for i, word_dict in enumerate(word_speaker_mapping, start=1):
        start_time = word_dict['start_time']
        end_time = word_dict['end_time']
        word = word_dict['word']
        speaker = word_dict['speaker']
        
        # Write SRT entry for each word
        print(
            f"{i}\n"
            f"{format_timestamp(start_time, always_include_hours=True, decimal_marker=',')} --> "
            f"{format_timestamp(end_time, always_include_hours=True, decimal_marker=',')}\n"
            f"[Speaker {speaker}] {word}\n",
            file=file,
            flush=True,
        )


def write_word_level_json(word_speaker_mapping, file):
    """
    Write word-level timestamps to JSON format for easy programmatic access.
    """
    # Convert to a more structured format
    words_data = []
    for word_dict in word_speaker_mapping:
        words_data.append({
            "word": word_dict['word'],
            "speaker": f"Speaker {word_dict['speaker']}",
            "start": word_dict['start_time'] / 1000.0,  # Convert to seconds
            "end": word_dict['end_time'] / 1000.0,  # Convert to seconds
            "start_ms": word_dict['start_time'],  # Keep milliseconds too
            "end_ms": word_dict['end_time']
        })
    
    json.dump(words_data, file, indent=2, ensure_ascii=False)


def write_word_level_ass(word_speaker_mapping, file, video_width=1920, video_height=1080):
    """
    Write word-level timestamps to Advanced SubStation Alpha (ASS) format.
    This format supports more styling options and overlapping subtitles.
    """
    # Define colors for different speakers (in BGR format for ASS)
    speaker_colors = {
        0: "&H00FF00",  # Green
        1: "&H0000FF",  # Red  
        2: "&HFF0000",  # Blue
        3: "&HFF00FF",  # Magenta
        4: "&H00FFFF",  # Yellow
        5: "&HFFFF00",  # Cyan
    }
    
    # Write ASS header
    file.write("[Script Info]\n")
    file.write("Title: Word-level Transcription\n")
    file.write("ScriptType: v4.00+\n")
    file.write(f"PlayResX: {video_width}\n")
    file.write(f"PlayResY: {video_height}\n\n")
    
    file.write("[V4+ Styles]\n")
    file.write("Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding\n")
    
    # Create styles for each speaker
    unique_speakers = set(w['speaker'] for w in word_speaker_mapping)
    for speaker in unique_speakers:
        color = speaker_colors.get(speaker, "&HFFFFFF")
        file.write(f"Style: Speaker{speaker},Arial,48,{color},&H00FFFFFF,&H00000000,&H80000000,0,0,0,0,100,100,0,0,1,2,1,2,10,10,10,1\n")
    
    file.write("\n[Events]\n")
    file.write("Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text\n")
    
    # Write each word as a dialogue event
    for word_dict in word_speaker_mapping:
        start = format_ass_timestamp(word_dict['start_time'])
        end = format_ass_timestamp(word_dict['end_time'])
        speaker = word_dict['speaker']
        word = word_dict['word'].replace('\n', '\\N')  # Handle newlines
        
        file.write(f"Dialogue: 0,{start},{end},Speaker{speaker},,0,0,0,,{word}\n")


def format_ass_timestamp(milliseconds):
    """Format timestamp for ASS subtitle format (h:mm:ss.cc)"""
    hours = milliseconds // 3_600_000
    milliseconds -= hours * 3_600_000
    
    minutes = milliseconds // 60_000
    milliseconds -= minutes * 60_000
    
    seconds = milliseconds // 1_000
    centiseconds = (milliseconds % 1_000) // 10
    
    return f"{hours:01d}:{minutes:02d}:{seconds:02d}.{centiseconds:02d}"


def get_word_level_transcript_with_speakers(word_speaker_mapping, 
                                           words_per_line=10, 
                                           group_by_speaker=True):
    """
    Generate a formatted transcript with word-level speaker annotations.
    
    Args:
        word_speaker_mapping: List of word dictionaries with speaker info
        words_per_line: Maximum words per line (for readability)
        group_by_speaker: If True, group consecutive words by the same speaker
    
    Returns:
        Formatted string with word-level transcript
    """
    transcript_lines = []
    
    if group_by_speaker:
        # Group consecutive words by speaker
        current_speaker = None
        current_words = []
        current_start = None
        
        for word_dict in word_speaker_mapping:
            speaker = word_dict['speaker']
            word = word_dict['word']
            
            if speaker != current_speaker:
                # Output previous speaker's words
                if current_words:
                    timestamp = format_timestamp(current_start, always_include_hours=True, decimal_marker='.')
                    transcript_lines.append(f"[{timestamp}] Speaker {current_speaker}: {' '.join(current_words)}")
                
                # Start new speaker group
                current_speaker = speaker
                current_words = [word]
                current_start = word_dict['start_time']
            else:
                current_words.append(word)
                
                # Break line if too many words
                if len(current_words) >= words_per_line:
                    timestamp = format_timestamp(current_start, always_include_hours=True, decimal_marker='.')
                    transcript_lines.append(f"[{timestamp}] Speaker {current_speaker}: {' '.join(current_words)}")
                    current_words = []
                    current_start = word_dict['start_time']
        
        # Don't forget the last group
        if current_words:
            timestamp = format_timestamp(current_start, always_include_hours=True, decimal_marker='.')
            transcript_lines.append(f"[{timestamp}] Speaker {current_speaker}: {' '.join(current_words)}")
    
    else:
        # Simple word-by-word output
        for word_dict in word_speaker_mapping:
            timestamp = format_timestamp(word_dict['start_time'], always_include_hours=True, decimal_marker='.')
            transcript_lines.append(
                f"[{timestamp}] Speaker {word_dict['speaker']}: {word_dict['word']}"
            )
    
    return '\n'.join(transcript_lines)


def write_word_level_vtt(word_speaker_mapping, file):
    """
    Write word-level timestamps to WebVTT format.
    WebVTT is more modern than SRT and supports additional features.
    """
    file.write("WEBVTT\n\n")
    
    for i, word_dict in enumerate(word_speaker_mapping):
        start_time = format_vtt_timestamp(word_dict['start_time'])
        end_time = format_vtt_timestamp(word_dict['end_time'])
        word = word_dict['word']
        speaker = word_dict['speaker']
        
        # Write WebVTT cue
        file.write(f"{i+1}\n")
        file.write(f"{start_time} --> {end_time}\n")
        file.write(f"<v Speaker {speaker}>{word}</v>\n\n")


def format_vtt_timestamp(milliseconds):
    """Format timestamp for WebVTT (hh:mm:ss.mmm)"""
    hours = milliseconds // 3_600_000
    milliseconds -= hours * 3_600_000
    
    minutes = milliseconds // 60_000  
    milliseconds -= minutes * 60_000
    
    seconds = milliseconds // 1_000
    milliseconds = milliseconds % 1_000
    
    return f"{hours:02d}:{minutes:02d}:{seconds:02d}.{milliseconds:03d}"