app.py and req.txt
Sridhar Sampath committed
May 21, 2025 at 07:54 UTC
30363981882fb18aa6b9ab6c60eef7e801dc3451
3 files changed
+569
README.md
new
+42
@@ -0,0 +1,42 @@
1
+# Parakeet ASR Demo
2
+
3
+Speech recognition using NVIDIA's Parakeet TDT model.
4
+
5
+## Features
6
+
7
+- 🎙️ Speech-to-text transcription using NVIDIA Parakeet TDT model
8
+- 📊 Real-time transcription with progress tracking
9
+- 📝 Support for multiple audio formats (WAV, FLAC)
10
+- 📈 Transcription history with export options
11
+- 🎯 Optimized for both short and long audio files
12
+- 💻 GPU acceleration support with fallback to CPU
13
+
14
+## Setup
15
+
16
+1. Install dependencies:
17
+```bash
18
+pip install -r requirements.txt
19
+```
20
+
21
+2. Run the application:
22
+```bash
23
+streamlit run app.py
24
+```
25
+
26
+## Requirements
27
+
28
+- Python 3.8+
29
+- NVIDIA GPU with CUDA support (strongly recommended for optimal performance)
30
+- FFmpeg (for audio processing)
31
+
32
+## Usage
33
+
34
+1. Upload an audio file or record directly in the browser
35
+2. Wait for the model to process and transcribe
36
+3. View and export transcription results
37
+
38
+## Notes
39
+
40
+- NVIDIA GPU with CUDA support is strongly recommended for optimal performance
41
+- Long audio files (>8 minutes) will automatically use optimized settings
42
+- Maximum recommended audio duration is 30 minutes
\ No newline at end of file
app.py
new
+498
@@ -0,0 +1,498 @@
1
+import streamlit as st
2
+import torch
3
+import os
4
+import tempfile
5
+import torchaudio
6
+from nemo.collections.asr.models import ASRModel
7
+from pydub import AudioSegment
8
+import numpy as np
9
+import csv
10
+import datetime
11
+import pandas as pd
12
+import time
13
+import gc
14
+
15
+# Set page config with a modern theme
16
+st.set_page_config(
17
+ page_title="Parakeet ASR Demo",
18
+ page_icon="🎙️",
19
+ layout="wide",
20
+ initial_sidebar_state="collapsed"
21
+)
22
+
23
+# Custom CSS for modern UI
24
+st.markdown("""
25
+<style>
26
+ .main {
27
+ padding: 1rem 2rem;
28
+ }
29
+ .stButton>button {
30
+ width: 100%;
31
+ background-color: #4CAF50;
32
+ color: white;
33
+ padding: 0.5rem 1rem;
34
+ border: none;
35
+ border-radius: 4px;
36
+ font-size: 1rem;
37
+ font-weight: 500;
38
+ }
39
+ .stButton>button:hover {
40
+ background-color: #45a049;
41
+ }
42
+ .stDataFrame {
43
+ border-radius: 8px;
44
+ box-shadow: 0 2px 4px rgba(0,0,0,0.1);
45
+ }
46
+ h1 {
47
+ color: #1E88E5;
48
+ font-size: 2.5rem !important;
49
+ font-weight: 700 !important;
50
+ margin: -1rem 0 1rem 0 !important;
51
+ }
52
+ h3 {
53
+ color: #1E88E5;
54
+ font-size: 1.5rem !important;
55
+ font-weight: 600 !important;
56
+ margin-top: 1.5rem !important;
57
+ }
58
+ .stProgress > div > div > div > div {
59
+ background-color: #4CAF50;
60
+ }
61
+ .info-box {
62
+ background-color: rgba(30, 136, 229, 0.1);
63
+ border-left: 5px solid #1E88E5;
64
+ padding: 0.5rem;
65
+ border-radius: 4px;
66
+ margin-bottom: 0.5rem;
67
+ font-size: 0.9rem;
68
+ }
69
+ .success-box {
70
+ background-color: rgba(76, 175, 80, 0.1);
71
+ border-left: 5px solid #4CAF50;
72
+ padding: 0.5rem;
73
+ border-radius: 4px;
74
+ margin-bottom: 0.5rem;
75
+ font-size: 0.9rem;
76
+ line-height: 1.2;
77
+ }
78
+ .warning-box {
79
+ background-color: rgba(255, 152, 0, 0.1);
80
+ border-left: 5px solid #FF9800;
81
+ padding: 0.5rem;
82
+ border-radius: 4px;
83
+ margin-bottom: 0.5rem;
84
+ font-size: 0.9rem;
85
+ }
86
+ .error-box {
87
+ background-color: rgba(244, 67, 54, 0.1);
88
+ border-left: 5px solid #F44336;
89
+ padding: 0.5rem;
90
+ border-radius: 4px;
91
+ margin-bottom: 0.5rem;
92
+ font-size: 0.9rem;
93
+ }
94
+</style>
95
+""", unsafe_allow_html=True)
96
+
97
+# Constants
98
+MODEL_NAME = "nvidia/parakeet-tdt-0.6b-v2"
99
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
100
+SUPPORTED_FORMATS = ['wav', 'flac']
101
+MAX_RECOMMENDED_DURATION = 30 * 60 # 30 minutes in seconds
102
+LONG_AUDIO_THRESHOLD = 480 # 8 minutes in seconds
103
+
104
+# Initialize session state variables
105
+if 'transcription_history' not in st.session_state:
106
+ st.session_state.transcription_history = []
107
+
108
+def custom_info(text):
109
+ st.markdown(f'<div class="info-box">{text}</div>', unsafe_allow_html=True)
110
+
111
+def custom_success(text):
112
+ st.markdown(f'<div class="success-box">{text}</div>', unsafe_allow_html=True)
113
+
114
+def custom_warning(text):
115
+ st.markdown(f'<div class="warning-box">{text}</div>', unsafe_allow_html=True)
116
+
117
+def custom_error(text):
118
+ st.markdown(f'<div class="error-box">{text}</div>', unsafe_allow_html=True)
119
+
120
+@st.cache_resource
121
+def load_model():
122
+ """Load the ASR model with proper error handling and visualization."""
123
+ try:
124
+ with st.spinner("Loading Parakeet TDT model... This may take a minute."):
125
+ model = ASRModel.from_pretrained(model_name=MODEL_NAME)
126
+ model.eval()
127
+
128
+ # Move to appropriate device
129
+ model = model.to(DEVICE)
130
+
131
+ # Use mixed precision for better performance when using GPU
132
+ if DEVICE == "cuda":
133
+ model = model.to(torch.bfloat16)
134
+
135
+ custom_success(f"Model loaded successfully on {DEVICE.upper()}!")
136
+ return model
137
+ except Exception as e:
138
+ custom_error(f"Error loading model: {str(e)}")
139
+ if "CUDA" in str(e) or "GPU" in str(e):
140
+ custom_info("GPU issues detected. Try running with CPU by setting device to 'cpu'.")
141
+ return None
142
+
143
+def process_audio(audio_path):
144
+ """Process audio file for transcription with progress reporting."""
145
+ try:
146
+ # Load audio file
147
+ audio = AudioSegment.from_file(audio_path)
148
+ duration_sec = audio.duration_seconds
149
+
150
+ if duration_sec > MAX_RECOMMENDED_DURATION:
151
+ custom_warning(f"Audio is very long ({duration_sec/60:.1f} minutes). Transcription may take a while and could encounter memory issues.")
152
+
153
+ # Progress reporting
154
+ progress_bar = st.progress(0)
155
+
156
+
157
+ # Resample to 16kHz if needed
158
+ if audio.frame_rate != 16000:
159
+ progress_bar.progress(0.2)
160
+ audio = audio.set_frame_rate(16000)
161
+
162
+ progress_bar.progress(0.5)
163
+
164
+ # Convert to mono if stereo
165
+ if audio.channels > 1:
166
+ audio = audio.set_channels(1)
167
+
168
+ progress_bar.progress(0.8)
169
+
170
+ # Save processed audio
171
+ temp_dir = tempfile.gettempdir()
172
+ processed_path = os.path.join(temp_dir, "processed_audio.wav")
173
+ audio.export(processed_path, format="wav")
174
+
175
+ progress_bar.progress(1.0)
176
+ custom_success("Audio processed successfully!")
177
+
178
+ return processed_path, duration_sec
179
+ except Exception as e:
180
+ custom_error(f"Error processing audio: {str(e)}")
181
+
182
+ # Provide more specific error messages based on common issues
183
+ if "No such file" in str(e):
184
+ custom_info("The audio file could not be found. Please upload it again.")
185
+ elif "Unsupported format" in str(e) or "unknown format" in str(e):
186
+ custom_info(f"File format not supported. Please upload one of these formats: {', '.join(SUPPORTED_FORMATS)}")
187
+ elif "memory" in str(e).lower():
188
+ custom_info("Memory error occurred. Try with a shorter audio file or restart the application.")
189
+
190
+ return None, None
191
+
192
+def format_time(seconds):
193
+ """Convert seconds to HH:MM:SS format."""
194
+ return str(datetime.timedelta(seconds=seconds)).split('.')[0]
195
+
196
+def transcribe_audio(audio_path, show_progress=True):
197
+ """Transcribe audio file using the model with detailed progress reporting."""
198
+ start_time = time.time()
199
+ try:
200
+ model = load_model()
201
+ if model is None:
202
+ return None
203
+
204
+ processed_path, duration_sec = process_audio(audio_path)
205
+ if processed_path is None:
206
+ return None
207
+
208
+ # Apply long audio settings if needed
209
+ long_audio_settings_applied = False
210
+ if duration_sec > LONG_AUDIO_THRESHOLD:
211
+ try:
212
+ custom_info(f"Audio longer than {LONG_AUDIO_THRESHOLD/60:.1f} minutes. Applying optimized settings for long transcription.")
213
+ model.change_attention_model("rel_pos_local_attn", [256, 256])
214
+ model.change_subsampling_conv_chunking_factor(1)
215
+ long_audio_settings_applied = True
216
+ except Exception as e:
217
+ custom_warning(f"Could not apply long audio settings: {str(e)}")
218
+
219
+ try:
220
+ if show_progress:
221
+ with st.spinner("Transcribing audio... This may take a while for longer files."):
222
+ progress_bar = st.progress(0)
223
+
224
+ # Create a progress indicator that updates based on estimated time
225
+ # This is an estimation since we can't track actual progress
226
+ estimated_total_time = duration_sec * 0.5 # Rough estimate: processing takes ~50% of audio duration
227
+
228
+ def update_progress():
229
+ for i in range(1, 101):
230
+ elapsed = time.time() - start_time
231
+ if elapsed >= estimated_total_time:
232
+ break
233
+ progress = min(elapsed / estimated_total_time, 0.95) # Max at 95% until completion
234
+ progress_bar.progress(progress)
235
+ time.sleep(estimated_total_time / 100)
236
+
237
+ # Start progress update in separate thread
238
+ import threading
239
+ progress_thread = threading.Thread(target=update_progress)
240
+ progress_thread.daemon = True
241
+ progress_thread.start()
242
+
243
+ # Actual transcription
244
+ output = model.transcribe([processed_path], timestamps=True)
245
+
246
+ # Complete progress
247
+ progress_bar.progress(1.0)
248
+ else:
249
+ output = model.transcribe([processed_path], timestamps=True)
250
+
251
+ if not output or not isinstance(output, list) or not output[0] or not hasattr(output[0], 'timestamp'):
252
+ custom_error("Transcription failed or produced unexpected output format.")
253
+ return None
254
+
255
+ segment_timestamps = output[0].timestamp['segment']
256
+
257
+ # Generate CSV content with better headers
258
+ csv_data = [["From (s)", "To (s)", "From (time)", "To (time)", "Duration", "Transcription"]]
259
+
260
+ for ts in segment_timestamps:
261
+ start_s = ts['start']
262
+ end_s = ts['end']
263
+ start_formatted = format_time(start_s)
264
+ end_formatted = format_time(end_s)
265
+ duration = end_s - start_s
266
+
267
+ csv_data.append([
268
+ f"{start_s:.2f}",
269
+ f"{end_s:.2f}",
270
+ start_formatted,
271
+ end_formatted,
272
+ f"{duration:.2f}",
273
+ ts['segment']
274
+ ])
275
+
276
+ processing_time = time.time() - start_time
277
+ custom_success(f"Transcription completed in {processing_time:.1f} seconds!")
278
+
279
+ # Save to history
280
+ filename = os.path.basename(audio_path)
281
+ st.session_state.transcription_history.append({
282
+ 'filename': filename,
283
+ 'duration': duration_sec,
284
+ 'timestamp': datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
285
+ 'csv_data': csv_data
286
+ })
287
+
288
+ return csv_data
289
+
290
+ finally:
291
+ # Cleanup
292
+ if long_audio_settings_applied:
293
+ try:
294
+ model.change_attention_model("rel_pos")
295
+ model.change_subsampling_conv_chunking_factor(-1)
296
+ except Exception as e:
297
+ custom_warning(f"Issue reverting model settings: {str(e)}")
298
+
299
+ if os.path.exists(processed_path):
300
+ os.remove(processed_path)
301
+
302
+ # Force garbage collection
303
+ torch.cuda.empty_cache()
304
+ gc.collect()
305
+
306
+ except Exception as e:
307
+ custom_error(f"Error during transcription: {str(e)}")
308
+
309
+ # Offer more helpful advice based on error
310
+ if "CUDA out of memory" in str(e):
311
+ custom_info("GPU ran out of memory. Try processing a shorter audio file, or restart the application.")
312
+ elif "timeout" in str(e).lower():
313
+ custom_info("The operation timed out. This could be due to the file size or server load.")
314
+
315
+ return None
316
+
317
+def export_to_formats(csv_data):
318
+ """Create exportable data in multiple formats"""
319
+ # For CSV
320
+ csv_string = "\n".join([",".join([f'"{cell}"' if ',' in cell else cell for cell in row]) for row in csv_data])
321
+
322
+ # For plain text (just the transcription)
323
+ text_string = "\n\n".join([row[5] for row in csv_data[1:]])
324
+
325
+ # For SRT (subtitle format)
326
+ srt_string = ""
327
+ for i, row in enumerate(csv_data[1:], 1):
328
+ start_s = float(row[0])
329
+ end_s = float(row[1])
330
+
331
+ # Convert to SRT time format (HH:MM:SS,mmm)
332
+ start_srt = f"{int(start_s//3600):02d}:{int((start_s%3600)//60):02d}:{int(start_s%60):02d},{int((start_s%1)*1000):03d}"
333
+ end_srt = f"{int(end_s//3600):02d}:{int((end_s%3600)//60):02d}:{int(end_s%60):02d},{int((end_s%1)*1000):03d}"
334
+
335
+ srt_string += f"{i}\n{start_srt} --> {end_srt}\n{row[5]}\n\n"
336
+
337
+ return csv_string, text_string, srt_string
338
+
339
+# Main UI
340
+st.title("🎙️ Speech Transcription with Parakeet TDT")
341
+
342
+# Sidebar for navigation
343
+with st.sidebar:
344
+ st.title("Navigation")
345
+ page = st.radio("Go to", ["Transcribe", "About"])
346
+
347
+# Main content based on selected page
348
+if page == "Transcribe":
349
+ # Description
350
+ st.markdown("""
351
+ This demo showcases [parakeet-tdt-0.6b-v2](https://huggingface.co/nvidia/parakeet-tdt-0.6b-v2), a 600-million-parameter model designed for high-quality English speech recognition.
352
+
353
+ **Key Features:**
354
+ - Automatic punctuation and capitalization
355
+ - Accurate word-level timestamps
356
+ - Efficiently transcribes long audio segments
357
+ - Robust performance on spoken numbers and song lyrics transcription
358
+ """)
359
+
360
+ # Create two columns for better layout
361
+ col1, col2 = st.columns([2, 1])
362
+
363
+ with col1:
364
+ # File uploader with expanded format support
365
+ uploaded_file = st.file_uploader(f"Upload an audio file (WAV or FLAC)",
366
+ type=SUPPORTED_FORMATS)
367
+
368
+ with col2:
369
+ if uploaded_file:
370
+ st.markdown("### Audio Preview")
371
+ st.audio(uploaded_file)
372
+
373
+ # Display basic audio info
374
+ file_info = f"**File:** {uploaded_file.name}<br>"
375
+ file_info += f"**Size:** {uploaded_file.size / (1024*1024):.2f} MB<br>"
376
+ st.markdown(f'<div class="info-box">{file_info}</div>', unsafe_allow_html=True)
377
+
378
+ if uploaded_file:
379
+ # Save uploaded file
380
+ temp_dir = tempfile.gettempdir()
381
+ audio_path = os.path.join(temp_dir, uploaded_file.name)
382
+ with open(audio_path, "wb") as f:
383
+ f.write(uploaded_file.getbuffer())
384
+
385
+ if st.button("🎯 Transcribe Audio", type="primary"):
386
+ with st.spinner("Processing and transcribing audio..."):
387
+ csv_data = transcribe_audio(audio_path)
388
+
389
+ if csv_data:
390
+ # Display transcription results
391
+ st.markdown("### 📝 Transcription Results")
392
+
393
+ # Create a dataframe for display
394
+ df = pd.DataFrame(csv_data[1:], columns=csv_data[0])
395
+
396
+ # Display as table with modern styling
397
+ st.dataframe(
398
+ df,
399
+ column_config={
400
+ "From (s)": st.column_config.NumberColumn(format="%.2f", width="small"),
401
+ "To (s)": st.column_config.NumberColumn(format="%.2f", width="small"),
402
+ "From (time)": st.column_config.TextColumn(width="small"),
403
+ "To (time)": st.column_config.TextColumn(width="small"),
404
+ "Duration": st.column_config.NumberColumn(format="%.2f", width="small"),
405
+ "Transcription": st.column_config.TextColumn(width="large")
406
+ },
407
+ hide_index=True,
408
+ use_container_width=True
409
+ )
410
+
411
+ # Create export strings for different formats
412
+ csv_string, text_string, srt_string = export_to_formats(csv_data)
413
+
414
+ # Show export options
415
+ st.markdown("### 📥 Export Options")
416
+ col1, col2, col3 = st.columns(3)
417
+
418
+ with col1:
419
+ st.download_button(
420
+ "📄 Download as CSV",
421
+ data=csv_string,
422
+ file_name=f"{os.path.splitext(uploaded_file.name)[0]}_transcript.csv",
423
+ mime="text/csv",
424
+ use_container_width=True
425
+ )
426
+
427
+ with col2:
428
+ st.download_button(
429
+ "📝 Download as Text",
430
+ data=text_string,
431
+ file_name=f"{os.path.splitext(uploaded_file.name)[0]}_transcript.txt",
432
+ mime="text/plain",
433
+ use_container_width=True
434
+ )
435
+
436
+ with col3:
437
+ st.download_button(
438
+ "🎬 Download as SRT",
439
+ data=srt_string,
440
+ file_name=f"{os.path.splitext(uploaded_file.name)[0]}_subtitle.srt",
441
+ mime="text/plain",
442
+ use_container_width=True
443
+ )
444
+
445
+ # Word count analysis
446
+ total_words = sum(len(row[5].split()) for row in csv_data[1:])
447
+ total_duration = float(csv_data[-1][1]) - float(csv_data[1][0])
448
+ words_per_minute = (total_words / total_duration) * 60 if total_duration > 0 else 0
449
+
450
+ st.markdown("### 📊 Analysis")
451
+ col1, col2, col3 = st.columns(3)
452
+ col1.metric("Total Words", f"{total_words}")
453
+ col2.metric("Speech Duration", f"{format_time(total_duration)}")
454
+ col3.metric("Words per Minute", f"{words_per_minute:.1f}")
455
+
456
+ # Cleanup
457
+ if os.path.exists(audio_path):
458
+ os.remove(audio_path)
459
+
460
+elif page == "About":
461
+ st.title("About this Application")
462
+
463
+ st.markdown("""
464
+ ## Parakeet Speech Recognition
465
+
466
+ This application uses NVIDIA's Parakeet-TDT, a powerful speech recognition model designed for accurate transcription with timestamps and punctuation.
467
+
468
+ ### Model Details
469
+
470
+ - **Model**: [nvidia/parakeet-tdt-0.6b-v2](https://huggingface.co/nvidia/parakeet-tdt-0.6b-v2)
471
+ - **Parameters**: 600 million
472
+ - **Features**: Automatic punctuation, capitalization, and word-level timestamps
473
+ - **Language**: English
474
+
475
+
476
+ ### About Speech Recognition Technology
477
+
478
+ Modern speech recognition systems like Parakeet use advanced neural networks trained on thousands of hours of speech data. These systems analyze audio waveforms to predict the most likely sequence of words being spoken, taking into account language patterns and context.
479
+ """)
480
+
481
+ # System information
482
+ st.subheader("System Information")
483
+ col1, col2 = st.columns(2)
484
+
485
+ with col1:
486
+ st.markdown(f"**System Platform:** {os.name.upper()}")
487
+ st.markdown(f"**PyTorch Version:** {torch.__version__}")
488
+
489
+ with col2:
490
+ if torch.cuda.is_available():
491
+ st.markdown(f"**CUDA Available:** Yes (Version {torch.version.cuda})")
492
+ st.markdown(f"**GPU:** {torch.cuda.get_device_name(0)}")
493
+ else:
494
+ st.markdown("**CUDA Available:** No (Using CPU)")
495
+
496
+# Footer
497
+st.markdown("---")
498
+st.markdown("Made with NVIDIA's Parakeet-TDT model")
\ No newline at end of file
requirements.txt
new
+29
@@ -0,0 +1,29 @@
1
+streamlit>=1.32.0
2
+torch>=2.0.0
3
+torchaudio>=2.0.0
4
+numpy>=1.20.0
5
+pandas>=1.3.0
6
+matplotlib>=3.4.0
7
+git+https://github.com/NVIDIA/NeMo.git@main
8
+soundfile>=0.12.1
9
+transformers>=4.51.3
10
+sentencepiece>=0.2.0
11
+jiwer>=3.0.3
12
+pyannote.core>=3.1.1
13
+pyannote.metrics>=3.1.1
14
+librosa>=0.9.0
15
+omegaconf>=2.1.1
16
+pyparsing>=3.0.7
17
+pyyaml>=6.0
18
+scipy>=1.7.0
19
+tensorboard>=2.8.0
20
+unidecode>=1.3.4
21
+webdataset>=0.2.20
22
+wget>=3.2
23
+hydra-core>=1.1.0
24
+sacrebleu>=2.0.0
25
+datasets>=2.18.0
26
+editdistance==0.8.1
27
+ipython>=8.0.0
28
+pydub>=0.25.1
29
+ffmpeg-python>=0.2.0
\ No newline at end of file