add timestamp option

run_demo_openai.py

@@ -4,22 +4,14 @@ import warnings
 import gradio as gr
 import pytube as pt
 import torch
+import whisper
 from huggingface_hub import hf_hub_download, model_info
 from transformers.utils.logging import disable_progress_bar
-import whisper
-
 
 warnings.filterwarnings("ignore")
 disable_progress_bar()
 
-logging.basicConfig(
-    format="%(asctime)s [%(levelname)s] [%(name)s] %(message)s",
-    datefmt="%Y-%m-%dT%H:%M:%SZ",
-)
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.DEBUG)
-
-MODEL_NAME = "bofenghuang/whisper-large-v2-cv11-french"
+DEFAULT_MODEL_NAME = "bofenghuang/whisper-large-v2-cv11-french"
 CHECKPOINT_FILENAME = "checkpoint_openai.pt"
 
 GEN_KWARGS = {
@@ -36,17 +28,63 @@ GEN_KWARGS = {
     # "no_speech_threshold": None,
 }
 
+logging.basicConfig(
+    format="%(asctime)s [%(levelname)s] [%(name)s] %(message)s",
+    datefmt="%Y-%m-%dT%H:%M:%SZ",
+)
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
 # device = 0 if torch.cuda.is_available() else "cpu"
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+logger.info(f"Model will be loaded on device `{device}`")
+
+cached_models = {}
+
+
+def print_cuda_memory_info():
+    used_mem, tot_mem = torch.cuda.mem_get_info()
+    logger.info(
+        f"CUDA memory info - Free: {used_mem / 1024 ** 3:.2f} Gb, used: {(tot_mem - used_mem) / 1024 ** 3:.2f} Gb, total: {tot_mem / 1024 ** 3:.2f} Gb"
+    )
+
+
+def print_memory_info():
+    # todo
+    if device == "cpu":
+        pass
+    else:
+        print_cuda_memory_info()
+
 
+def maybe_load_cached_pipeline(model_name):
+    model = cached_models.get(model_name)
+    if model is None:
+        downloaded_model_path = hf_hub_download(repo_id=model_name, filename=CHECKPOINT_FILENAME)
 
-downloaded_model_path = hf_hub_download(repo_id=MODEL_NAME, filename=CHECKPOINT_FILENAME)
+        model = whisper.load_model(downloaded_model_path, device=device)
+        logger.info(f"`{model_name}` has been loaded on device `{device}`")
 
-model = whisper.load_model(downloaded_model_path, device=device)
-logger.info(f"Model has been loaded on device `{device}`")
+        print_memory_info()
 
+        cached_models[model_name] = model
+    return model
 
-def transcribe(microphone, file_upload):
+
+def infer(model, filename, with_timestamps):
+    if with_timestamps:
+        model_outputs = model.transcribe(filename, **GEN_KWARGS)
+        return "\n\n".join(
+            [
+                f'Segment {segment["id"]+1} from {segment["start"]:.2f}s to {segment["end"]:.2f}s:\n{segment["text"].strip()}'
+                for segment in model_outputs["segments"]
+            ]
+        )
+    else:
+        return model.transcribe(filename, without_timestamps=True, **GEN_KWARGS)["text"]
+
+
+def transcribe(microphone, file_upload, with_timestamps, model_name=DEFAULT_MODEL_NAME):
     warn_output = ""
     if (microphone is not None) and (file_upload is not None):
         warn_output = (
@@ -59,9 +97,11 @@ def transcribe(microphone, file_upload):
 
     file = microphone if microphone is not None else file_upload
 
-    text = model.transcribe(file, **GEN_KWARGS)["text"]
+    model = maybe_load_cached_pipeline(model_name)
+    # text = model.transcribe(file, **GEN_KWARGS)["text"]
+    text = infer(model, file, with_timestamps)
 
-    logger.info(f"Transcription: {text}")
+    logger.info(f"Transcription by `{model_name}`: {text}")
 
     return warn_output + text
 
@@ -74,19 +114,24 @@ def _return_yt_html_embed(yt_url):
     return HTML_str
 
 
-def yt_transcribe(yt_url):
+def yt_transcribe(yt_url, with_timestamps, model_name=DEFAULT_MODEL_NAME):
     yt = pt.YouTube(yt_url)
     html_embed_str = _return_yt_html_embed(yt_url)
     stream = yt.streams.filter(only_audio=True)[0]
     stream.download(filename="audio.mp3")
 
-    text = model.transcribe("audio.mp3", **GEN_KWARGS)["text"]
+    model = maybe_load_cached_pipeline(model_name)
+    # text = model.transcribe("audio.mp3", **GEN_KWARGS)["text"]
+    text = infer(model, "audio.mp3", with_timestamps)
 
-    logger.info(f'Transcription of "{yt_url}": {text}')
+    logger.info(f'Transcription by `{model_name}` of "{yt_url}": {text}')
 
     return html_embed_str, text
 
 
+# load default model
+maybe_load_cached_pipeline(DEFAULT_MODEL_NAME)
+
 demo = gr.Blocks()
 
 mf_transcribe = gr.Interface(
@@ -94,6 +139,7 @@ mf_transcribe = gr.Interface(
     inputs=[
         gr.inputs.Audio(source="microphone", type="filepath", optional=True, label="Record"),
         gr.inputs.Audio(source="upload", type="filepath", optional=True, label="Upload File"),
+        gr.Checkbox(label="With timestamps?", value=True),
     ],
     # outputs="text",
     outputs=gr.outputs.Textbox(label="Transcription"),
@@ -102,7 +148,7 @@ mf_transcribe = gr.Interface(
     title="Whisper French 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"
+        f" checkpoint [{DEFAULT_MODEL_NAME}](https://huggingface.co/{DEFAULT_MODEL_NAME}) and 🤗 Transformers to transcribe audio files"
         " of arbitrary length."
     ),
     allow_flagging="never",
@@ -110,7 +156,10 @@ mf_transcribe = gr.Interface(
 
 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"],
     outputs=[
         gr.outputs.HTML(label="YouTube Page"),
@@ -121,7 +170,7 @@ yt_transcribe = gr.Interface(
     title="Whisper French 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"
+        f" [{DEFAULT_MODEL_NAME}](https://huggingface.co/{DEFAULT_MODEL_NAME}) and 🤗 Transformers to transcribe audio files of"
         " arbitrary length."
     ),
     allow_flagging="never",

run_demo_openai_merged.py

@@ -0,0 +1,169 @@
+import logging
+import warnings
+
+import gradio as gr
+import pytube as pt
+import torch
+import whisper
+from huggingface_hub import hf_hub_download, model_info
+from transformers.utils.logging import disable_progress_bar
+
+warnings.filterwarnings("ignore")
+disable_progress_bar()
+
+DEFAULT_MODEL_NAME = "bofenghuang/whisper-large-v2-cv11-french"
+CHECKPOINT_FILENAME = "checkpoint_openai.pt"
+
+GEN_KWARGS = {
+    "task": "transcribe",
+    "language": "fr",
+    # "without_timestamps": True,
+    # decode options
+    # "beam_size": 5,
+    # "patience": 2,
+    # disable fallback
+    # "compression_ratio_threshold": None,
+    # "logprob_threshold": None,
+    # vad threshold
+    # "no_speech_threshold": None,
+}
+
+logging.basicConfig(
+    format="%(asctime)s [%(levelname)s] [%(name)s] %(message)s",
+    datefmt="%Y-%m-%dT%H:%M:%SZ",
+)
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+# device = 0 if torch.cuda.is_available() else "cpu"
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+logger.info(f"Model will be loaded on device `{device}`")
+
+cached_models = {}
+
+
+def print_cuda_memory_info():
+    used_mem, tot_mem = torch.cuda.mem_get_info()
+    logger.info(
+        f"CUDA memory info - Free: {used_mem / 1024 ** 3:.2f} Gb, used: {(tot_mem - used_mem) / 1024 ** 3:.2f} Gb, total: {tot_mem / 1024 ** 3:.2f} Gb"
+    )
+
+
+def print_memory_info():
+    # todo
+    if device == "cpu":
+        pass
+    else:
+        print_cuda_memory_info()
+
+
+def maybe_load_cached_pipeline(model_name):
+    model = cached_models.get(model_name)
+    if model is None:
+        downloaded_model_path = hf_hub_download(repo_id=model_name, filename=CHECKPOINT_FILENAME)
+
+        model = whisper.load_model(downloaded_model_path, device=device)
+        logger.info(f"`{model_name}` has been loaded on device `{device}`")
+
+        print_memory_info()
+
+        cached_models[model_name] = model
+    return model
+
+
+def infer(model, filename, with_timestamps):
+    if with_timestamps:
+        model_outputs = model.transcribe(filename, **GEN_KWARGS)
+        return "\n\n".join(
+            [
+                f'Segment {segment["id"]+1} from {segment["start"]:.2f}s to {segment["end"]:.2f}s:\n{segment["text"].strip()}'
+                for segment in model_outputs["segments"]
+            ]
+        )
+    else:
+        return model.transcribe(filename, without_timestamps=True, **GEN_KWARGS)["text"]
+
+
+def download_from_youtube(yt_url, downloaded_filename="audio.wav"):
+    yt = pt.YouTube(yt_url)
+    stream = yt.streams.filter(only_audio=True)[0]
+    # stream.download(filename="audio.mp3")
+    stream.download(filename=downloaded_filename)
+    return downloaded_filename
+
+
+def transcribe(microphone, file_upload, yt_url, with_timestamps, model_name=DEFAULT_MODEL_NAME):
+    warn_output = ""
+    if (microphone is not None) and (file_upload is not None) and (yt_url is not None):
+        warn_output = (
+            "WARNING: You've uploaded an audio file, used the microphone, and pasted a YouTube URL. "
+            "The recorded file from the microphone will be used, the uploaded audio and the YouTube URL will be discarded.\n"
+        )
+
+    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"
+        )
+
+    if (microphone is not None) and (yt_url is not None):
+        warn_output = (
+            "WARNING: You've used the microphone and pasted a YouTube URL. "
+            "The recorded file from the microphone will be used and the YouTube URL will be discarded.\n"
+        )
+
+    if (file_upload is not None) and (yt_url is not None):
+        warn_output = (
+            "WARNING: You've uploaded an audio file and pasted a YouTube URL. "
+            "The uploaded audio will be used and the YouTube URL will be discarded.\n"
+        )
+
+    elif (microphone is None) and (file_upload is None) or (yt_url is None):
+        return "ERROR: You have to either use the microphone, upload an audio file or paste a YouTube URL"
+
+    if microphone is not None:
+        file = microphone
+        logging_prefix = f"Transcription by `{model_name}` of microphone:"
+    elif file_upload is not None:
+        file = file_upload
+        logging_prefix = f"Transcription by `{model_name}` of uploaded file:"
+    else:
+        file = download_from_youtube(yt_url)
+        logging_prefix = f'Transcription by `{model_name}` of "{yt_url}":'
+
+    model = maybe_load_cached_pipeline(model_name)
+    # text = model.transcribe(file, **GEN_KWARGS)["text"]
+    text = infer(model, file, with_timestamps)
+
+    logger.info(logging_prefix + "\n" + text)
+
+    return warn_output + text
+
+
+# load default model
+maybe_load_cached_pipeline(DEFAULT_MODEL_NAME)
+
+demo = gr.Interface(
+    fn=transcribe,
+    inputs=[
+        gr.inputs.Audio(source="microphone", type="filepath", label="Record", optional=True),
+        gr.inputs.Audio(source="upload", type="filepath", label="Upload File", optional=True),
+        gr.inputs.Textbox(lines=1, placeholder="Paste the URL to a YouTube video here", label="YouTube URL", optional=True),
+        gr.Checkbox(label="With timestamps?", value=True),
+    ],
+    # outputs="text",
+    outputs=gr.outputs.Textbox(label="Transcription"),
+    layout="horizontal",
+    theme="huggingface",
+    title="Whisper French 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 [{DEFAULT_MODEL_NAME}](https://huggingface.co/{DEFAULT_MODEL_NAME}) and 🤗 Transformers to transcribe audio files"
+        " of arbitrary length."
+    ),
+    allow_flagging="never",
+)
+
+
+# demo.launch(server_name="0.0.0.0", debug=True, share=True)
+demo.launch(enable_queue=True)