Update transcribe_audio function to include stream_name parameter and adjust usage in retrieve_audio_information

offline_pipeline.py

@@ -20,8 +20,8 @@ def retrieve_audio_information(
     enhanced_spec_path = f"/tmp/{sample_id}_enhanced_spectrogram.png"
     spec_image(original_array).save(noisy_spec_path)
     spec_image(enhanced_array).save(enhanced_spec_path)
-    enhanced_transcript = transcribe_audio(enhanced_array, DEFAULT_SR, stt_model, mode="ENHANCED")
-    noisy_transcript = transcribe_audio(original_array, DEFAULT_SR, stt_model, mode="NOISY")
+    enhanced_transcript = transcribe_audio(enhanced_array, DEFAULT_SR, stt_model, stream_name=f"{sample_id}_enhanced")
+    noisy_transcript = transcribe_audio(original_array, DEFAULT_SR, stt_model, stream_name=f"{sample_id}_noisy")
     try: 
         original_transcript = get_transcript(sample_id)
         wer_enhanced = compute_wer(original_transcript, enhanced_transcript)

utils.py

@@ -6,8 +6,7 @@ from PIL import Image
 import io
 import matplotlib.pyplot as plt
 import resampy
-from stt_streamers import DeepgramStreamer, SonioxStreamer
-from constants import DEFAULT_SR
+from constants import DEFAULT_SR, STREAMER_CLASSES
 
 def to_gradio_audio(x: np.ndarray, sr: int) -> tuple[int, np.ndarray]:
     x = np.asarray(x)
@@ -95,7 +94,7 @@ def compute_wer(reference: str, hypothesis: str) -> float:
 
 
 
-def transcribe_audio(audio_array: np.ndarray, sr: int, streamer_type: str = "Deepgram", mode: str = "RAW"):
+def transcribe_audio(audio_array: np.ndarray, sr: int, streamer_type: str = "Deepgram", stream_name: str = "RAW"):
     """
     Transcribe an audio array using the specified STT streamer.
 
@@ -103,7 +102,7 @@ def transcribe_audio(audio_array: np.ndarray, sr: int, streamer_type: str = "Dee
         audio_array (np.ndarray): Audio data array
         sr (int): Sample rate of the audio array
         streamer_type (str): "Soniox" or "Deepgram"
-        mode (str): Optional label for streamer instance ("RAW", "ENHANCED", etc.)
+        stream_name (str): Optional label for streamer instance ("RAW", "ENHANCED", etc.)
 
     Returns:
         str: Transcript text
@@ -113,20 +112,15 @@ def transcribe_audio(audio_array: np.ndarray, sr: int, streamer_type: str = "Dee
         audio_array = resampy.resample(audio_array, sr, DEFAULT_SR)
         sr = DEFAULT_SR
         
-    api_map = {
-        "Soniox": SonioxStreamer,
-        "Deepgram": DeepgramStreamer,
-    }
 
-    streamer_key = streamer_type
-    if streamer_key not in api_map:
+    if streamer_type not in STREAMER_CLASSES:
         raise ValueError(
             f"Invalid streamer_type '{streamer_type}'. "
-            f"Choose from: {', '.join(api_map.keys())}"
+            f"Choose from: {', '.join(STREAMER_CLASSES.keys())}"
         )
 
-    StreamerClass = api_map[streamer_key]
-    streamer = StreamerClass(sr, mode)
+    StreamerClass = STREAMER_CLASSES[streamer_type]
+    streamer = StreamerClass(sr, stream_name)
 
     transcript = streamer.stream_array(audio_array)