sanchit-gandhi HF staff commited on
Commit
7091430
1 Parent(s): bd5a509

create app.py

Browse files
Files changed (2) hide show
  1. app.py +150 -0
  2. 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