espnet/forced-alignment

1"""Forced alignment: when each line was said.2 3Published as https://huggingface.co/spaces/espnet/forced-alignment; the source4lives in espnet, at egs2/owsm_ctc_v4/s2t1/demo_align.5 6The other demos are one input and one output - audio in, text out - and the7page for them is built from the checkpoint by espnet2.bin.demo. This one takes8audio *and* the text that goes with it, and answers with a table, so it has a9page of its own.10 11The model is only a CTC head here: espnet2.bin.align aligns on whatever12checkpoint it is handed, and this Space happens to hand it OWSM-CTC v4.13"""14 15# The ZeroGPU package patches torch as it is imported, so it has to come16# first - before torch, and before anything that imports torch.17try:  # only Hugging Face's runners have it18    import spaces  # isort: skip19except ImportError:  # running elsewhere: the decorator does nothing20 21    class spaces:  # noqa: N801 - stands in for the module22        @staticmethod23        def GPU(func=None, **kwargs):24            return func if func is not None else (lambda f: f)25 26 27import os  # noqa: E40228 29import gradio as gr  # noqa: E40230import librosa  # noqa: E40231import matplotlib  # noqa: E40232import torch  # noqa: E40233 34from espnet2.bin.align import ForcedAligner  # noqa: E40235 36matplotlib.use("Agg")  # a Space has no display, and gradio wants the figure37import matplotlib.pyplot as plt  # noqa: E40238 39# ZeroGPU gives a decorated call a fixed slice of GPU time and kills it at the40# end, so the demo asks for a slice and refuses audio it could not finish in41# one. The same two minutes the other demos take.42MAX_SECS = 12043GPU_SECONDS = 12044MODEL_TAG = os.environ.get("ALIGN_MODEL_TAG", "espnet/owsm_ctc_v4_1B")45# A score is a probability under the checkpoint that produced it, so this is46# a heuristic rather than a calibrated confidence: 0.3 is where a line that47# was not said sits, on OWSM-CTC v4 and read English. ALIGN_WARN_BELOW moves48# it for a checkpoint that scores differently.49WARN_BELOW = float(os.environ.get("ALIGN_WARN_BELOW", "0.3"))50# ZeroGPU attaches the GPU only while a @spaces.GPU function runs, so51# torch.cuda.is_available() is False here and asking it would pin the models to52# the CPU on the very hardware bought to run them. SPACES_ZERO_GPU is the53# runtime's own marker; `spaces` being importable is not, since anyone can54# install it.55ZERO_GPU = bool(os.environ.get("SPACES_ZERO_GPU"))56if os.environ.get("DEVICE"):57    DEVICE = os.environ["DEVICE"]58elif ZERO_GPU or torch.cuda.is_available():59    DEVICE = "cuda"60else:61    DEVICE = "cpu"62 63EXAMPLE_WAV = (64    "https://github.com/espnet/espnet/raw/master/test_utils/ctc_align_test.wav"65)66EXAMPLE_TEXT = """The sale of the hotels67is part of Holiday's strategy68to sell off assets69and concentrate on property management"""70 71TITLE = "Forced alignment"72DESCRIPTION = """# Forced alignment73 74You have a recording and the text of what was said. This says *when*: a start75and an end for every line, and a score for how well the line and the audio76agree. Subtitles come from the times; data cleaning comes from the score.77 78Nothing is trained here. The alignment is a Viterbi path through the CTC head79of a model that already exists - [OWSM-CTC80v4](https://huggingface.co/espnet/owsm_ctc_v4_1B) - which is why it needs no81model of its own and works in any language that one covers.82 83**Write the text the way the model writes it.** The score is a probability84under this model: the same words in capitals, which a reference transcript85often uses, score near zero while still landing in roughly the right place.86The app says so when it sees it.87"""88ARTICLE = """One line an utterance, in the order they were said - the times are89what you are asking for, so they are not needed. The score is the mean90probability of the line's tokens: 1.0 is a perfect match, and a line that does91not belong to this audio scores near zero, which is what alignment-score92filtering uses.93 94`espnet align audio.wav --text "..."` is the same thing from a terminal, and95`espnet2.bin.align.ForcedAligner` is the class behind both. Source of this96Space: [`egs2/owsm_ctc_v4/s2t1/demo_align`](https://github.com/espnet/espnet/tree/master/egs2/owsm_ctc_v4/s2t1/demo_align).97"""  # noqa: E501 - one markdown link, and breaking a URL breaks the link98 99aligner = ForcedAligner.from_pretrained(MODEL_TAG, device=DEVICE)100# the rate the checkpoint wants, not this file's idea of it: ALIGN_MODEL_TAG101# can point at a model trained at another rate102SAMPLE_RATE = aligner.sample_rate103 104 105def _read(path):106    if path is None:107        raise gr.Error("Record or upload some audio first.")108    speech, _ = librosa.load(path, sr=SAMPLE_RATE)109    seconds = len(speech) / SAMPLE_RATE110    if seconds > MAX_SECS:111        raise gr.Error(112            f"This demo takes up to {MAX_SECS} s; that file is {seconds:.0f} s. "113            "Run the app yourself for longer audio - the model has no such limit."114        )115    return speech116 117 118def _figure(speech, segments):119    """The waveform with each segment marked, which is the answer to look at."""120    figure, axes = plt.subplots(figsize=(11, 2.6))121    seconds = [i / SAMPLE_RATE for i in range(len(speech))]122    axes.plot(seconds, speech, linewidth=0.4, color="#888")123    for index, segment in enumerate(segments):124        axes.axvspan(segment.start, segment.end, color=f"C{index % 10}", alpha=0.25)125        axes.text(126            (segment.start + segment.end) / 2,127            0.85 * max(abs(speech.max()), 1e-6),128            segment.text.split()[0] if segment.text.split() else "",129            ha="center",130            fontsize=8,131        )132    axes.set_xlabel("seconds")133    axes.set_yticks([])134    figure.tight_layout()135    return figure136 137 138@spaces.GPU(duration=GPU_SECONDS)139def predict(audio_path, text):140    speech = _read(audio_path)141    lines = [line.strip() for line in (text or "").splitlines() if line.strip()]142    if not lines:143        raise gr.Error("Type the lines that were said, one to a line.")144 145    try:146        segments = aligner(speech, lines)147    except ValueError as error:148        # "this text cannot fit in this recording", and the like149        raise gr.Error(str(error)) from error150 151    rows = [152        [f"{s.start:.2f}", f"{s.end:.2f}", f"{s.score:.3f}", s.text] for s in segments153    ]154    worst = min(s.score for s in segments)155    if worst < WARN_BELOW:156        gr.Warning(157            f"The weakest line scores {worst:.2f}, under {WARN_BELOW:.2f}. "158            "Either it was not said, or the text is spelled a way this model "159            "does not use."160        )161    return rows, _figure(speech, segments)162 163 164with gr.Blocks(title=TITLE) as demo:165    gr.Markdown(DESCRIPTION)166    with gr.Row():167        with gr.Column():168            audio = gr.Audio(169                sources=["microphone", "upload"], type="filepath", label="Speech"170            )171            text = gr.Textbox(172                label="What was said", lines=6, placeholder="One utterance a line"173            )174            button = gr.Button("Align", variant="primary")175        with gr.Column():176            table = gr.Dataframe(177                headers=["start", "end", "score", "text"],178                label="Segments",179                wrap=True,180            )181            drawing = gr.Plot(label="Where they fall")182    button.click(predict, [audio, text], [table, drawing])183    gr.Examples(184        [[EXAMPLE_WAV, EXAMPLE_TEXT]],185        inputs=[audio, text],186        outputs=[table, drawing],187        fn=predict,188        cache_examples=False,189    )190    gr.Markdown(ARTICLE)191 192 193if __name__ == "__main__":194    demo.launch()195