summaryrefslogtreecommitdiff
path: root/asr.py
diff options
context:
space:
mode:
authorvin <git@vineetk.net>2025-08-13 11:16:58 -0400
committervin <git@vineetk.net>2025-08-13 11:17:46 -0400
commit8c01ad5cd727b4844e26e8a050c4094b545c6938 (patch)
tree1e7cb310dc82a88fd49d3d365adb11fcc52773fc /asr.py
parent94a64876146c9ca4707c5d750dfe1e464db41735 (diff)
add parakeet asr via onnx
found out that nvidia's parakeet has onnx support. it's about 4-5x faster than openai's whisper-large-v3-turbo (via whisper.cpp) and seems more accurate too without lots of errors near the end of the transcript. I thought I lost this repo, turns out I wasn't stupid and already had it pushed.
Diffstat (limited to 'asr.py')
-rw-r--r--asr.py68
1 files changed, 68 insertions, 0 deletions
diff --git a/asr.py b/asr.py
new file mode 100644
index 0000000..b773931
--- /dev/null
+++ b/asr.py
@@ -0,0 +1,68 @@
1import librosa
2import soundfile as sf
3import onnx_asr
4import os
5import sys
6#import numpy as np
7from tqdm import tqdm
8
9print("Loading ASR model...")
10providers = [
11 "ROCMExecutionProvider",
12 #"CUDAExecutionProvider",
13 "CPUExecutionProvider",
14]
15model = onnx_asr.load_model(
16 "nemo-parakeet-tdt-0.6b-v2",
17 providers=providers,
18 #quantization="int8"
19).with_timestamps()
20
21# parakeet needs 16khz mono wav as input
22input_file = "input.wav"
23chunk_duration_seconds = 20
24sample_rate = 16000
25
26print(f"Loading and resampling {input_file}...")
27try:
28 audio, sr = librosa.load(input_file, sr=sample_rate, mono=True)
29except Exception as e:
30 print(f"Error loading audio file: {e}")
31 exit()
32
33# process audio in chunks
34chunk_size_samples = int(chunk_duration_seconds * sample_rate)
35full_transcript = []
36temp_dir = "temp_chunks"
37os.makedirs(temp_dir, exist_ok=True)
38
39print("Starting transcription process in chunks...")
40num_chunks = (len(audio) + chunk_size_samples - 1) // chunk_size_samples
41
42for i in tqdm(range(num_chunks), desc="Transcribing"):
43 start_sample = i * chunk_size_samples
44 end_sample = start_sample + chunk_size_samples
45 chunk = audio[start_sample:end_sample]
46
47 temp_chunk_file = os.path.join(temp_dir, f"chunk_{i}.wav")
48 sf.write(temp_chunk_file, chunk, sample_rate)
49
50 try:
51 transcript = model.recognize(temp_chunk_file)
52 if transcript:
53 full_transcript.append(transcript)
54 except Exception as e:
55 print(f" Error processing chunk {i + 1}: {e}")
56 finally:
57 os.remove(temp_chunk_file)
58
59os.rmdir(temp_dir)
60
61print("\n" + "="*30)
62print(" FINAL TRANSCRIPT")
63print("="*30)
64print(full_transcript)
65
66# onnxruntime has some bug where it doesn't exit properly, aborting instead
67# so hard exit instead (sys.exit does graceful exit)
68os._exit(0)