Spaces:
Runtime error
Runtime error
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( | |
""" | |
<div style="text-align: center; max-width: 700px; margin: 0 auto;"> | |
<div | |
style=" | |
display: inline-flex; align-items: center; gap: 0.8rem; font-size: 1.75rem; | |
" | |
> | |
<h1 style="font-weight: 900; margin-bottom: 7px; line-height: normal;"> | |
Distil-Whisper VS Whisper | |
</h1> | |
</div> | |
</div> | |
""" | |
) | |
gr.HTML( | |
f""" | |
This demo evaluates the <a href="https://huggingface.co/distil-whisper/distil-large-v2"> Distil-Whisper </a> model | |
against the <a href="https://huggingface.co/openai/whisper-large-v2"> Whisper </a> 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() |