Update app.py

app.py

@@ -3,7 +3,7 @@ import os
 
 import gradio as gr
 import pytube as pt
-from asr_diarizer import ASRDiarizationPipeline
+from asr_diarize import ASRDiarizationPipeline  # TODO: speechbox import
 
 MODEL_NAME = "openai/whisper-tiny"
 
@@ -16,22 +16,21 @@ pipe = ASRDiarizationPipeline.from_pretrained(
     use_auth_token=HF_TOKEN,
 )
 
-def transcribe(microphone, file_upload):
-    warn_output = ""
-    if (microphone is not None) and (file_upload is not None):
-        warn_output = (
-            "WARNING: You've uploaded an audio file and used the microphone. "
-            "The recorded file from the microphone will be used and the uploaded audio will be discarded.\n"
-        )
+def tuple_to_string(start_end_tuple, ndigits=1):
+    return str((round(start_end_tuple[0], ndigits), round(start_end_tuple[1], ndigits)))
 
-    elif (microphone is None) and (file_upload is None):
-        return "ERROR: You have to either use the microphone or upload an audio file"
 
-    file = microphone if microphone is not None else file_upload
+def format_as_transcription(raw_segments, with_timestamps=False):
+    if with_timestamps:
+        return "\n\n".join([chunk["speaker"] + " " + tuple_to_string(chunk["timestamp"]) +  chunk["text"] for chunk in raw_segments])
+    else:
+        return "\n\n".join([chunk["speaker"] + chunk["text"] for chunk in raw_segments])
 
-    text = pipe(file)
 
-    return warn_output + "\n\n".join([chunk["speaker"] + chunk["text"] for chunk in text])
+def transcribe(file_upload, with_timestamps):
+    raw_segments = pipe(file_upload)
+    transcription = format_as_transcription(raw_segments, with_timestamps=with_timestamps)
+    return transcription
 
 
 def _return_yt_html_embed(yt_url):
@@ -43,7 +42,7 @@ def _return_yt_html_embed(yt_url):
     return HTML_str
 
 
-def yt_transcribe(yt_url):
+def yt_transcribe(yt_url, with_timestamps):
     yt = pt.YouTube(yt_url)
     html_embed_str = _return_yt_html_embed(yt_url)
     stream = yt.streams.filter(only_audio=True)[0]
@@ -51,7 +50,7 @@ def yt_transcribe(yt_url):
 
     text = pipe("audio.mp3")
 
-    return html_embed_str, "\n\n".join([chunk["speaker"] + chunk["text"] for chunk in text])
+    return html_embed_str, format_as_transcription(text, with_timestamps=with_timestamps)
 
 
 demo = gr.Blocks()
@@ -59,37 +58,43 @@ demo = gr.Blocks()
 mf_transcribe = gr.Interface(
     fn=transcribe,
     inputs=[
-        gr.inputs.Audio(source="microphone", type="filepath", optional=True),
         gr.inputs.Audio(source="upload", type="filepath", optional=True),
+        gr.Checkbox(label="With timestamps?", value=True),
     ],
     outputs="text",
     layout="horizontal",
     theme="huggingface",
     title="Whisper Demo: Transcribe Audio",
     description=(
-        "Transcribe long-form microphone or audio inputs with the click of a button! Demo uses the the fine-tuned"
-        f" checkpoint [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}) and 🤗 Transformers to transcribe audio files"
-        " of arbitrary length."
-    ),
+        "Transcribe audio files with speaker diarization using 🤗 Speechbox. Demo uses the pre-trained checkpoint"
+        f" [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}) for ASR transcriptions and [PyAnnote Audio](https://huggingface.co/pyannote/speaker-diarization)"
+        " to label the speakers."
+    )
     allow_flagging="never",
 )
 
 yt_transcribe = gr.Interface(
     fn=yt_transcribe,
-    inputs=[gr.inputs.Textbox(lines=1, placeholder="Paste the URL to a YouTube video here", label="YouTube URL")],
+    inputs=[
+        gr.inputs.Textbox(lines=1, placeholder="Paste the URL to a YouTube video here", label="YouTube URL"),
+        gr.Checkbox(label="With timestamps?", value=True),
+    ],
     outputs=["html", "text"],
     layout="horizontal",
     theme="huggingface",
-    title="Whisper Demo: Transcribe YouTube",
+    title="Whisper Speaker Diarization Demo: Transcribe YouTube",
     description=(
-        "Transcribe long-form YouTube videos with the click of a button! Demo uses the the fine-tuned checkpoint:"
-        f" [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}) and 🤗 Transformers to transcribe audio files of"
-        " arbitrary length."
-    ),
+        "Transcribe YouTube videos with speaker diarization using 🤗 Speechbox. Demo uses the pre-trained checkpoint"
+        f" [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}) for ASR transcriptions and [PyAnnote Audio](https://huggingface.co/pyannote/speaker-diarization)"
+        " to label the speakers."
+    )
+    examples=[
+        ["https://www.youtube.com/watch?v=9dAWIPixYxc", True],
+    ],
     allow_flagging="never",
 )
 
 with demo:
     gr.TabbedInterface([mf_transcribe, yt_transcribe], ["Transcribe Audio", "Transcribe YouTube"])
 
-demo.launch(enable_queue=True)
+demo.launch(enable_queue=True, share=True)