podcast-sponsor-remove

Attempt at identify sponsored segments in audio transcripts and removing them.
Log | Files | Refs

asr.py (2875B)

      1 import librosa
      2 import onnx_asr
      3 import os
      4 import sys
      5 import json
      6 import numpy as np
      7 from tqdm import tqdm
      8 
      9 print("Loading ASR model...")
     10 providers = [
     11     "ROCMExecutionProvider",
     12     "CPUExecutionProvider",
     13 ]
     14 try:
     15     model = onnx_asr.load_model(
     16         "nemo-parakeet-tdt-0.6b-v3",
     17         providers=providers,
     18     ).with_timestamps()
     19 except Exception as e:
     20     print(f"Failed to load the ASR model. Ensure ONNX Runtime for ROCm is installed correctly. Error: {e}")
     21     sys.exit(1)
     22 
     23 chunk_duration_seconds = 20
     24 sample_rate = 16000
     25 
     26 # converts TimestampedResult to JSON-serializable dictionary
     27 def result_to_dict(result, offset_seconds):
     28     token_map = [
     29         [token, round(ts + offset_seconds, 2)]
     30         for token, ts in zip(result.tokens, result.timestamps)
     31     ]
     32 
     33     start_time = token_map[0][1] if token_map else offset_seconds
     34     end_time = token_map[-1][1] if token_map else offset_seconds
     35 
     36     return {
     37         "text": result.text,
     38         "start_time": start_time,
     39         "end_time": end_time,
     40         "token_map": token_map
     41     }
     42 
     43 def main(input_file, output_file):
     44     try:
     45         audio, sr = librosa.load(input_file, sr=sample_rate, mono=True)
     46     except Exception as e:
     47         tqdm.write(f"Error loading audio file: {e}")
     48         sys.exit(1)
     49 
     50     chunk_size_samples = int(chunk_duration_seconds * sample_rate)
     51     full_transcript_data = []
     52 
     53     num_chunks = (len(audio) + chunk_size_samples - 1) // chunk_size_samples
     54 
     55     for i in tqdm(range(num_chunks), desc="Transcribing", position=1, leave=False):
     56         start_sample = i * chunk_size_samples
     57         end_sample = start_sample + chunk_size_samples
     58         chunk = audio[start_sample:end_sample]
     59         
     60         chunk_offset_seconds = start_sample / sample_rate
     61         
     62         try:
     63             transcript_results = model.recognize(np.copy(chunk))
     64             
     65             if transcript_results:
     66                 if not isinstance(transcript_results, list):
     67                     results_list = [transcript_results]
     68                 else:
     69                     results_list = transcript_results
     70 
     71                 for result in results_list:
     72                     full_transcript_data.append(result_to_dict(result, chunk_offset_seconds))
     73         except Exception as e:
     74             tqdm.write(f"  Error processing chunk {i + 1}: {e}")
     75 
     76     with open(output_file, "w") as f:
     77         json.dump(full_transcript_data, f, indent=2)
     78 
     79 if __name__ == "__main__":
     80     if len(sys.argv) < 2:
     81         print("Usage: python asr.py input_audio_file1 [input_audio_file2 [input_audio_file3 [...]]]")
     82         sys.exit(1)
     83     
     84     all_inputs = sys.argv[1:]
     85     inputs = [f for f in all_inputs if not os.path.exists(os.path.splitext(f)[0] + ".json")]
     86 
     87     for i in tqdm(inputs, desc="Processing files", position=0):
     88         main(i, os.path.splitext(i)[0] + ".json")
     89 
     90 # os._exit(0) workaround for ONNX Runtime bug
     91 os._exit(0)