Spaces:
Running
on
Zero
Running
on
Zero
Commit
•
7091430
1
Parent(s):
bd5a509
create app.py
Browse files- app.py +150 -0
- requirements.txt +5 -0
app.py
ADDED
@@ -0,0 +1,150 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
|
2 |
+
from transformers.utils import is_flash_attn_2_available
|
3 |
+
import torch
|
4 |
+
import gradio as gr
|
5 |
+
import matplotlib.pyplot as plt
|
6 |
+
import time
|
7 |
+
import os
|
8 |
+
|
9 |
+
BATCH_SIZE = 16
|
10 |
+
TOKEN = os.environ.get("HF_TOKEN", None)
|
11 |
+
|
12 |
+
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
13 |
+
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
|
14 |
+
use_flash_attention_2 = is_flash_attn_2_available()
|
15 |
+
|
16 |
+
model = AutoModelForSpeechSeq2Seq.from_pretrained(
|
17 |
+
"openai/whisper-large-v2", torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True, use_flash_attention_2=use_flash_attention_2
|
18 |
+
)
|
19 |
+
distilled_model = AutoModelForSpeechSeq2Seq.from_pretrained(
|
20 |
+
"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
|
21 |
+
)
|
22 |
+
|
23 |
+
if not use_flash_attention_2:
|
24 |
+
model = model.bettertransformer()
|
25 |
+
distilled_model = distilled_model.bettertransformer()
|
26 |
+
|
27 |
+
processor = AutoProcessor.from_pretrained("openai/whisper-tiny.en")
|
28 |
+
|
29 |
+
model.to(device)
|
30 |
+
distilled_model.to(device)
|
31 |
+
|
32 |
+
pipe = pipeline(
|
33 |
+
"automatic-speech-recognition",
|
34 |
+
model=model,
|
35 |
+
tokenizer=processor.tokenizer,
|
36 |
+
feature_extractor=processor.feature_extractor,
|
37 |
+
max_new_tokens=128,
|
38 |
+
chunk_length_s=30,
|
39 |
+
torch_dtype=torch_dtype,
|
40 |
+
device=device,
|
41 |
+
language="en",
|
42 |
+
task="transcribe",
|
43 |
+
)
|
44 |
+
pipe_forward = pipe._forward
|
45 |
+
|
46 |
+
distil_pipe = pipeline(
|
47 |
+
"automatic-speech-recognition",
|
48 |
+
model=distilled_model,
|
49 |
+
tokenizer=processor.tokenizer,
|
50 |
+
feature_extractor=processor.feature_extractor,
|
51 |
+
max_new_tokens=128,
|
52 |
+
chunk_length_s=15,
|
53 |
+
torch_dtype=torch_dtype,
|
54 |
+
device=device,
|
55 |
+
language="en",
|
56 |
+
task="transcribe",
|
57 |
+
)
|
58 |
+
distil_pipe_forward = distil_pipe._forward
|
59 |
+
|
60 |
+
def transcribe(inputs):
|
61 |
+
if inputs is None:
|
62 |
+
raise gr.Error("No audio file submitted! Please record or upload an audio file before submitting your request.")
|
63 |
+
|
64 |
+
def _forward_distil_time(*args, **kwargs):
|
65 |
+
global distil_runtime
|
66 |
+
start_time = time.time()
|
67 |
+
result = distil_pipe_forward(*args, **kwargs)
|
68 |
+
distil_runtime = time.time() - start_time
|
69 |
+
return result
|
70 |
+
|
71 |
+
distil_pipe._forward = _forward_distil_time
|
72 |
+
distil_text = distil_pipe(inputs, batch_size=BATCH_SIZE)["text"]
|
73 |
+
yield distil_text, distil_runtime, None, None, None
|
74 |
+
|
75 |
+
def _forward_time(*args, **kwargs):
|
76 |
+
global runtime
|
77 |
+
start_time = time.time()
|
78 |
+
result = pipe_forward(*args, **kwargs)
|
79 |
+
runtime = time.time() - start_time
|
80 |
+
return result
|
81 |
+
|
82 |
+
pipe._forward = _forward_time
|
83 |
+
text = pipe(inputs, batch_size=BATCH_SIZE)["text"]
|
84 |
+
|
85 |
+
relative_latency = runtime / distil_runtime
|
86 |
+
|
87 |
+
# Create figure and axis
|
88 |
+
fig, ax = plt.subplots(figsize=(5, 5))
|
89 |
+
|
90 |
+
# Define bar width and positions
|
91 |
+
bar_width = 0.1
|
92 |
+
positions = [0, 0.1] # Adjusted positions to bring bars closer
|
93 |
+
|
94 |
+
# Plot data
|
95 |
+
ax.bar(positions[0], distil_runtime, bar_width, edgecolor='black')
|
96 |
+
ax.bar(positions[1], runtime, bar_width, edgecolor='black')
|
97 |
+
|
98 |
+
# Set title, labels, and xticks
|
99 |
+
ax.set_ylabel('Transcription time (s)')
|
100 |
+
ax.set_xticks(positions)
|
101 |
+
ax.set_xticklabels(['Distil-Whisper', 'Whisper'])
|
102 |
+
|
103 |
+
# Gridlines and other styling
|
104 |
+
ax.grid(which='major', axis='y', linestyle='--', linewidth=0.5)
|
105 |
+
|
106 |
+
# Use tight layout to avoid overlaps
|
107 |
+
plt.tight_layout()
|
108 |
+
|
109 |
+
yield distil_text, distil_runtime, text, runtime, plt
|
110 |
+
|
111 |
+
if __name__ == "__main__":
|
112 |
+
with gr.Blocks() as demo:
|
113 |
+
gr.HTML(
|
114 |
+
"""
|
115 |
+
<div style="text-align: center; max-width: 700px; margin: 0 auto;">
|
116 |
+
<div
|
117 |
+
style="
|
118 |
+
display: inline-flex; align-items: center; gap: 0.8rem; font-size: 1.75rem;
|
119 |
+
"
|
120 |
+
>
|
121 |
+
<h1 style="font-weight: 900; margin-bottom: 7px; line-height: normal;">
|
122 |
+
Distil-Whisper VS Whisper
|
123 |
+
</h1>
|
124 |
+
</div>
|
125 |
+
</div>
|
126 |
+
"""
|
127 |
+
)
|
128 |
+
gr.HTML(
|
129 |
+
f"""
|
130 |
+
This demo evaluates the <a href="https://huggingface.co/distil-whisper/distil-large-v2"> Distil-Whisper </a> model
|
131 |
+
against the <a href="https://huggingface.co/openai/whisper-large-v2"> Whisper </a> model.
|
132 |
+
"""
|
133 |
+
)
|
134 |
+
audio = gr.components.Audio(source="upload", type="filepath", label="Audio file")
|
135 |
+
button = gr.Button("Transcribe")
|
136 |
+
plot = gr.components.Plot()
|
137 |
+
with gr.Row():
|
138 |
+
distil_runtime = gr.components.Textbox(label="Distil-Whisper Transcription Time (s)")
|
139 |
+
runtime = gr.components.Textbox(label="Whisper Transcription Time (s)")
|
140 |
+
with gr.Row():
|
141 |
+
distil_transcription = gr.components.Textbox(label="Distil-Whisper Transcription").style(show_copy_button=True)
|
142 |
+
transcription = gr.components.Textbox(label="Whisper Transcription").style(show_copy_button=True)
|
143 |
+
|
144 |
+
button.click(
|
145 |
+
fn=transcribe,
|
146 |
+
inputs=audio,
|
147 |
+
outputs=[distil_transcription, distil_runtime, transcription, runtime, plot],
|
148 |
+
)
|
149 |
+
|
150 |
+
demo.queue().launch()
|
requirements.txt
ADDED
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
--extra-index-url https://download.pytorch.org/whl/cu113
|
2 |
+
torch
|
3 |
+
pip install git+https://github.com/huggingface/transformers
|
4 |
+
accelerate
|
5 |
+
optimum
|