from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline from transformers.utils import is_flash_attn_2_available import torch import gradio as gr import matplotlib.pyplot as plt import time import os BATCH_SIZE = 16 # TODO: remove token before release and update ckpt path TOKEN = os.environ.get("HF_TOKEN", None) device = "cuda:0" if torch.cuda.is_available() else "cpu" torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32 use_flash_attention_2 = is_flash_attn_2_available() model = AutoModelForSpeechSeq2Seq.from_pretrained( "openai/whisper-large-v2", torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True, use_flash_attention_2=use_flash_attention_2 ) distilled_model = AutoModelForSpeechSeq2Seq.from_pretrained( "sanchit-gandhi/distil-large-v2-private", torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True, use_flash_attention_2=use_flash_attention_2, token=TOKEN ) if not use_flash_attention_2: model = model.to_bettertransformer() distilled_model = distilled_model.to_bettertransformer() processor = AutoProcessor.from_pretrained("openai/whisper-tiny.en") model.to(device) distilled_model.to(device) pipe = pipeline( "automatic-speech-recognition", model=model, tokenizer=processor.tokenizer, feature_extractor=processor.feature_extractor, max_new_tokens=128, chunk_length_s=30, torch_dtype=torch_dtype, device=device, ) pipe_forward = pipe._forward distil_pipe = pipeline( "automatic-speech-recognition", model=distilled_model, tokenizer=processor.tokenizer, feature_extractor=processor.feature_extractor, max_new_tokens=128, chunk_length_s=15, torch_dtype=torch_dtype, device=device, ) distil_pipe_forward = distil_pipe._forward def transcribe(inputs): if inputs is None: raise gr.Error("No audio file submitted! Please record or upload an audio file before submitting your request.") def _forward_distil_time(*args, **kwargs): global distil_runtime start_time = time.time() result = distil_pipe_forward(*args, **kwargs) distil_runtime = time.time() - start_time return result distil_pipe._forward = _forward_distil_time distil_text = distil_pipe(inputs, batch_size=BATCH_SIZE)["text"] yield distil_text, distil_runtime, None, None, None def _forward_time(*args, **kwargs): global runtime start_time = time.time() result = pipe_forward(*args, **kwargs) runtime = time.time() - start_time return result pipe._forward = _forward_time text = pipe(inputs, batch_size=BATCH_SIZE, language="en", task="transcribe")["text"] # Create figure and axis fig, ax = plt.subplots(figsize=(5, 5)) # Define bar width and positions bar_width = 0.1 positions = [0, 0.1] # Adjusted positions to bring bars closer # Plot data ax.bar(positions[0], distil_runtime, bar_width, edgecolor='black') ax.bar(positions[1], runtime, bar_width, edgecolor='black') # Set title, labels, and xticks ax.set_ylabel('Transcription time (s)') ax.set_xticks(positions) ax.set_xticklabels(['Distil-Whisper', 'Whisper']) # Gridlines and other styling ax.grid(which='major', axis='y', linestyle='--', linewidth=0.5) # Use tight layout to avoid overlaps plt.tight_layout() yield distil_text, distil_runtime, text, runtime, plt if __name__ == "__main__": with gr.Blocks() as demo: gr.HTML( """

Distil-Whisper VS Whisper

""" ) gr.HTML( f""" This demo evaluates the Distil-Whisper model against the Whisper model. """ ) audio = gr.components.Audio(type="filepath", label="Audio input") button = gr.Button("Transcribe") plot = gr.components.Plot() with gr.Row(): distil_runtime = gr.components.Textbox(label="Distil-Whisper Transcription Time (s)") runtime = gr.components.Textbox(label="Whisper Transcription Time (s)") with gr.Row(): distil_transcription = gr.components.Textbox(label="Distil-Whisper Transcription", show_copy_button=True) transcription = gr.components.Textbox(label="Whisper Transcription", show_copy_button=True) button.click( fn=transcribe, inputs=audio, outputs=[distil_transcription, distil_runtime, transcription, runtime, plot], ) demo.queue().launch()