Create app.py

app.py

@@ -0,0 +1,309 @@
+import os
+import tempfile
+import time
+
+import gradio as gr
+import numpy as np
+import torch
+import yt_dlp as youtube_dl
+from gradio_client import Client
+from pyannote.audio import Pipeline
+from transformers.pipelines.audio_utils import ffmpeg_read
+
+
+YT_LENGTH_LIMIT_S = 36000  # limit to 1 hour YouTube files
+SAMPLING_RATE = 16000
+
+API_URL = "https://sanchit-gandhi-whisper-jax.hf.space/"
+
+# set up the Gradio client
+client = Client(API_URL)
+
+# set up the diarization pipeline
+diarization_pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization", use_auth_token=True)
+
+
+def format_string(timestamp):
+    """
+    Reformat a timestamp string from (HH:)MM:SS to float seconds. Note that the hour column
+    is optional, and is appended within the function if not input.
+
+    Args:
+        timestamp (str):
+            Timestamp in string format, either MM:SS or HH:MM:SS.
+    Returns:
+        seconds (float):
+            Total seconds corresponding to the input timestamp.
+    """
+    split_time = timestamp.split(":")
+    split_time = [float(sub_time) for sub_time in split_time]
+
+    if len(split_time) == 2:
+        split_time.insert(0, 0)
+
+    seconds = split_time[0] * 3600 + split_time[1] * 60 + split_time[2]
+    return seconds
+
+
+# Adapted from https://github.com/openai/whisper/blob/c09a7ae299c4c34c5839a76380ae407e7d785914/whisper/utils.py#L50
+def format_timestamp(seconds: float, always_include_hours: bool = False, decimal_marker: str = "."):
+    """
+    Reformat a timestamp from a float of seconds to a string in format (HH:)MM:SS. Note that the hour
+    column is optional, and is appended in the function if the number of hours > 0.
+
+    Args:
+        seconds (float):
+            Total seconds corresponding to the input timestamp.
+    Returns:
+        timestamp (str):
+            Timestamp in string format, either MM:SS or HH:MM:SS.
+    """
+    if seconds is not None:
+        milliseconds = round(seconds * 1000.0)
+
+        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}"
+    else:
+        # we have a malformed timestamp so just return it as is
+        return seconds
+
+
+def format_as_transcription(raw_segments):
+    return "\n".join(
+        [
+            f"{chunk['speaker']} [{format_timestamp(chunk['timestamp'][0])} -> {format_timestamp(chunk['timestamp'][1])}] {chunk['text']}"
+            for chunk in raw_segments
+        ]
+    )
+
+
+def _return_yt_html_embed(yt_url):
+    video_id = yt_url.split("?v=")[-1]
+    HTML_str = (
+        f'<center> <iframe width="500" height="320" src="https://www.youtube.com/embed/{video_id}"> </iframe>'
+        " </center>"
+    )
+    return HTML_str
+
+
+def download_yt_audio(yt_url, filename):
+    info_loader = youtube_dl.YoutubeDL()
+    try:
+        info = info_loader.extract_info(yt_url, download=False)
+    except youtube_dl.utils.DownloadError as err:
+        raise gr.Error(str(err))
+
+    file_length = info["duration_string"]
+    file_length_s = format_string(file_length)
+
+    if file_length_s > YT_LENGTH_LIMIT_S:
+        yt_length_limit_hms = time.strftime("%HH:%MM:%SS", time.gmtime(YT_LENGTH_LIMIT_S))
+        file_length_hms = time.strftime("%HH:%MM:%SS", time.gmtime(file_length_s))
+        raise gr.Error(f"Maximum YouTube length is {yt_length_limit_hms}, got {file_length_hms} YouTube video.")
+
+    ydl_opts = {"outtmpl": filename, "format": "worstvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best"}
+    with youtube_dl.YoutubeDL(ydl_opts) as ydl:
+        try:
+            ydl.download([yt_url])
+        except youtube_dl.utils.ExtractorError as err:
+            raise gr.Error(str(err))
+
+
+def align(transcription, segments, group_by_speaker=True):
+    transcription_split = transcription.split("\n")
+
+    # re-format transcription from string to List[Dict]
+    transcript = []
+    for chunk in transcription_split:
+        start_end, transcription = chunk[1:].split("] ")
+        start, end = start_end.split("->")
+
+        transcript.append({"timestamp": (format_string(start), format_string(end)), "text": transcription})
+
+    # diarizer output may contain consecutive segments from the same speaker (e.g. {(0 -> 1, speaker_1), (1 -> 1.5, speaker_1), ...})
+    # we combine these segments to give overall timestamps for each speaker's turn (e.g. {(0 -> 1.5, speaker_1), ...})
+    new_segments = []
+    prev_segment = cur_segment = segments[0]
+
+    for i in range(1, len(segments)):
+        cur_segment = segments[i]
+
+        # check if we have changed speaker ("label")
+        if cur_segment["label"] != prev_segment["label"] and i < len(segments):
+            # add the start/end times for the super-segment to the new list
+            new_segments.append(
+                {
+                    "segment": {"start": prev_segment["segment"]["start"], "end": cur_segment["segment"]["start"]},
+                    "speaker": prev_segment["label"],
+                }
+            )
+            prev_segment = segments[i]
+
+    # add the last segment(s) if there was no speaker change
+    new_segments.append(
+        {
+            "segment": {"start": prev_segment["segment"]["start"], "end": cur_segment["segment"]["end"]},
+            "speaker": prev_segment["label"],
+        }
+    )
+
+    # get the end timestamps for each chunk from the ASR output
+    end_timestamps = np.array([chunk["timestamp"][-1] for chunk in transcript])
+    segmented_preds = []
+
+    # align the diarizer timestamps and the ASR timestamps
+    for segment in new_segments:
+        # get the diarizer end timestamp
+        end_time = segment["segment"]["end"]
+        # find the ASR end timestamp that is closest to the diarizer's end timestamp and cut the transcript to here
+        upto_idx = np.argmin(np.abs(end_timestamps - end_time))
+
+        if group_by_speaker:
+            segmented_preds.append(
+                {
+                    "speaker": segment["speaker"],
+                    "text": "".join([chunk["text"] for chunk in transcript[: upto_idx + 1]]),
+                    "timestamp": (transcript[0]["timestamp"][0], transcript[upto_idx]["timestamp"][1]),
+                }
+            )
+        else:
+            for i in range(upto_idx + 1):
+                segmented_preds.append({"speaker": segment["speaker"], **transcript[i]})
+
+        # crop the transcripts and timestamp lists according to the latest timestamp (for faster argmin)
+        transcript = transcript[upto_idx + 1 :]
+        end_timestamps = end_timestamps[upto_idx + 1 :]
+
+    # final post-processing
+    transcription = format_as_transcription(segmented_preds)
+    return transcription
+
+
+def transcribe(audio_path, group_by_speaker=True):
+    # run Whisper JAX asynchronously using Gradio client (endpoint)
+    job = client.submit(
+        audio_path,
+        "transcribe",
+        True,
+        api_name="/predict_1",
+    )
+
+    # run diarization while we wait for Whisper JAX
+    diarization = diarization_pipeline(audio_path)
+    segments = diarization.for_json()["content"]
+
+    # only fetch the transcription result after performing diarization
+    transcription, _ = job.result()
+
+    # align the ASR transcriptions and diarization timestamps
+    transcription = align(transcription, segments, group_by_speaker=group_by_speaker)
+
+    return transcription
+
+
+def transcribe_yt(yt_url, group_by_speaker=True):
+    # run Whisper JAX asynchronously using Gradio client (endpoint)
+    job = client.submit(
+        yt_url,
+        "transcribe",
+        True,
+        api_name="/predict_2",
+    )
+
+    _return_yt_html_embed(yt_url)
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        filepath = os.path.join(tmpdirname, "video.mp4")
+        download_yt_audio(yt_url, filepath)
+
+        with open(filepath, "rb") as f:
+            inputs = f.read()
+
+    inputs = ffmpeg_read(inputs, SAMPLING_RATE)
+    inputs = torch.from_numpy(inputs).float()
+    inputs = inputs.unsqueeze(0)
+
+    diarization = diarization_pipeline(
+        {"waveform": inputs, "sample_rate": SAMPLING_RATE},
+    )
+    segments = diarization.for_json()["content"]
+
+    # only fetch the transcription result after performing diarization
+    transcription, _ = job.result()
+
+    # align the ASR transcriptions and diarization timestamps
+    transcription = align(transcription, segments, group_by_speaker=group_by_speaker)
+
+    return transcription
+
+
+title = "Whisper JAX + Speaker Diarization ⚡️"
+
+description = """Combine the speed of Whisper JAX with pyannote speaker diarization to transcribe meetings in super fast time.
+"""
+
+article = "Whisper large-v2 model by OpenAI. Speaker diarization model by pyannote. Whisper JAX backend running JAX on a TPU v4-8 through the generous support of the [TRC](https://sites.research.google/trc/about/) programme. Whisper JAX [code](https://github.com/sanchit-gandhi/whisper-jax) and Gradio demo by 🤗 Hugging Face."
+
+microphone = gr.Interface(
+    fn=transcribe,
+    inputs=[
+        gr.inputs.Audio(source="microphone", optional=True, type="filepath"),
+        gr.inputs.Checkbox(default=True, label="Group by speaker"),
+    ],
+    outputs=[
+        gr.outputs.Textbox(label="Transcription").style(show_copy_button=True),
+    ],
+    allow_flagging="never",
+    title=title,
+    description=description,
+    article=article,
+)
+
+audio_file = gr.Interface(
+    fn=transcribe,
+    inputs=[
+        gr.inputs.Audio(source="upload", optional=True, label="Audio file", type="filepath"),
+        gr.inputs.Checkbox(default=True, label="Group by speaker"),
+    ],
+    outputs=[
+        gr.outputs.Textbox(label="Transcription").style(show_copy_button=True),
+    ],
+    allow_flagging="never",
+    title=title,
+    description=description,
+    article=article,
+)
+
+youtube = gr.Interface(
+    fn=transcribe_yt,
+    inputs=[
+        gr.inputs.Textbox(lines=1, placeholder="Paste the URL to a YouTube video here", label="YouTube URL"),
+        gr.inputs.Checkbox(default=True, label="Group by speaker"),
+    ],
+    outputs=[
+        gr.outputs.HTML(label="Video"),
+        gr.outputs.Textbox(label="Transcription").style(show_copy_button=True),
+    ],
+    allow_flagging="never",
+    title=title,
+    examples=[["https://www.youtube.com/watch?v=m8u-18Q0s7I", True]],
+    cache_examples=False,
+    description=description,
+    article=article,
+)
+
+demo = gr.Blocks()
+
+with demo:
+    gr.TabbedInterface([microphone, audio_file, youtube], ["Microphone", "Audio File", "YouTube"])
+
+demo.queue(concurrency_count=1, max_size=5)
+demo.launch()