adapt asr.py for multiple file inputs

1 files changed, 30 insertions, 32 deletions

diff --git a/asr.py b/asr.py
index 54b83eb..1f4bf86 100644
--- a/asr.py
+++ b/asr.py
@@ -6,6 +6,23 @@ import json
 import numpy as np
 from tqdm import tqdm
 
+print("Loading ASR model...")
+providers = [
+    "ROCMExecutionProvider",
+    "CPUExecutionProvider",
+]
+try:
+    model = onnx_asr.load_model(
+        "nemo-parakeet-tdt-0.6b-v3",
+        providers=providers,
+    ).with_timestamps()
+except Exception as e:
+    print(f"Failed to load the ASR model. Ensure ONNX Runtime for ROCm is installed correctly. Error: {e}")
+    sys.exit(1)
+
+chunk_duration_seconds = 20
+sample_rate = 16000
+
 # converts TimestampedResult to JSON-serializable dictionary
 def result_to_dict(result, offset_seconds):
     token_map = [
@@ -24,37 +41,18 @@ def result_to_dict(result, offset_seconds):
     }
 
 def main(input_file, output_file):
-    print("Loading ASR model...")
-    providers = [
-        "ROCMExecutionProvider",
-        "CPUExecutionProvider",
-    ]
-    try:
-        model = onnx_asr.load_model(
-            "nemo-parakeet-tdt-0.6b-v3",
-            providers=providers,
-        ).with_timestamps()
-    except Exception as e:
-        print(f"Failed to load the ASR model. Ensure ONNX Runtime for ROCm is installed correctly. Error: {e}")
-        sys.exit(1)
-
-    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}")
+        tqdm.write(f"Error loading audio file: {e}")
         sys.exit(1)
 
     chunk_size_samples = int(chunk_duration_seconds * sample_rate)
     full_transcript_data = []
 
-    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"):
+    for i in tqdm(range(num_chunks), desc="Transcribing", position=1, leave=False):
         start_sample = i * chunk_size_samples
         end_sample = start_sample + chunk_size_samples
         chunk = audio[start_sample:end_sample]
@@ -73,21 +71,21 @@ def main(input_file, output_file):
                 for result in results_list:
                     full_transcript_data.append(result_to_dict(result, chunk_offset_seconds))
         except Exception as e:
-            print(f"  Error processing chunk {i + 1}: {e}")
+            tqdm.write(f"  Error processing chunk {i + 1}: {e}")
 
-    print(f"\nTranscription complete. Saving results to {output_file}...")
     with open(output_file, "w") as f:
         json.dump(full_transcript_data, f, indent=2)
 
-    print("Done.")
-    # os._exit(0) workaround for ONNX Runtime bug
-    os._exit(0)
-
 if __name__ == "__main__":
-    if len(sys.argv) < 3:
-        print("Usage: python asr.py <input_audio_file> <output_json_file>")
+    if len(sys.argv) < 2:
+        print("Usage: python asr.py input_audio_file1 [input_audio_file2 [input_audio_file3 [...]]]")
         sys.exit(1)
     
-    input_path = sys.argv[1]
-    output_path = sys.argv[2]
-    main(input_path, output_path)
+    all_inputs = sys.argv[1:]
+    inputs = [f for f in all_inputs if not os.path.exists(os.path.splitext(f)[0] + ".json")]
+
+    for i in tqdm(inputs, desc="Processing files", position=0):
+        main(i, os.path.splitext(i)[0] + ".json")
+
+# os._exit(0) workaround for ONNX Runtime bug
+os._exit(0)