1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
|
import librosa
import onnx_asr
import os
import sys
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 = [
[token, round(ts + offset_seconds, 2)]
for token, ts in zip(result.tokens, result.timestamps)
]
start_time = token_map[0][1] if token_map else offset_seconds
end_time = token_map[-1][1] if token_map else offset_seconds
return {
"text": result.text,
"start_time": start_time,
"end_time": end_time,
"token_map": token_map
}
def main(input_file, output_file):
try:
audio, sr = librosa.load(input_file, sr=sample_rate, mono=True)
except Exception as e:
tqdm.write(f"Error loading audio file: {e}")
sys.exit(1)
chunk_size_samples = int(chunk_duration_seconds * sample_rate)
full_transcript_data = []
num_chunks = (len(audio) + chunk_size_samples - 1) // chunk_size_samples
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]
chunk_offset_seconds = start_sample / sample_rate
try:
transcript_results = model.recognize(np.copy(chunk))
if transcript_results:
if not isinstance(transcript_results, list):
results_list = [transcript_results]
else:
results_list = transcript_results
for result in results_list:
full_transcript_data.append(result_to_dict(result, chunk_offset_seconds))
except Exception as e:
tqdm.write(f" Error processing chunk {i + 1}: {e}")
with open(output_file, "w") as f:
json.dump(full_transcript_data, f, indent=2)
if __name__ == "__main__":
if len(sys.argv) < 2:
print("Usage: python asr.py input_audio_file1 [input_audio_file2 [input_audio_file3 [...]]]")
sys.exit(1)
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)
|
