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 @@
+import librosa
+import soundfile as sf
+import onnx_asr
+import os
+import sys
+#import numpy as np
+from tqdm import tqdm
+
+print("Loading ASR model...")
+providers = [
+    "ROCMExecutionProvider",
+    #"CUDAExecutionProvider",
+    "CPUExecutionProvider",
+]
+model = onnx_asr.load_model(
+    "nemo-parakeet-tdt-0.6b-v2",
+    providers=providers,
+    #quantization="int8"
+).with_timestamps()
+
+# parakeet needs 16khz mono wav as input
+input_file = "input.wav"
+chunk_duration_seconds = 20
+sample_rate = 16000
+
+print(f"Loading and resampling {input_file}...")
+try:
+    audio, sr = librosa.load(input_file, sr=sample_rate, mono=True)
+except Exception as e:
+    print(f"Error loading audio file: {e}")
+    exit()
+
+# process audio in chunks
+chunk_size_samples = int(chunk_duration_seconds * sample_rate)
+full_transcript = []
+temp_dir = "temp_chunks"
+os.makedirs(temp_dir, exist_ok=True)
+
+print("Starting transcription process in chunks...")
+num_chunks = (len(audio) + chunk_size_samples - 1) // chunk_size_samples
+
+for i in tqdm(range(num_chunks), desc="Transcribing"):
+    start_sample = i * chunk_size_samples
+    end_sample = start_sample + chunk_size_samples
+    chunk = audio[start_sample:end_sample]
+    
+    temp_chunk_file = os.path.join(temp_dir, f"chunk_{i}.wav")
+    sf.write(temp_chunk_file, chunk, sample_rate)
+    
+    try:
+        transcript = model.recognize(temp_chunk_file)
+        if transcript:
+            full_transcript.append(transcript)
+    except Exception as e:
+        print(f"  Error processing chunk {i + 1}: {e}")
+    finally:
+        os.remove(temp_chunk_file)
+
+os.rmdir(temp_dir)
+
+print("\n" + "="*30)
+print("      FINAL TRANSCRIPT")
+print("="*30)
+print(full_transcript)
+
+# onnxruntime has some bug where it doesn't exit properly, aborting instead
+# so hard exit instead (sys.exit does graceful exit)
+os._exit(0)