Update backend/asr.py

backend/asr.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 import numpy as np
 
 from backend.utils import device
-# import nemo.collections.asr as nemo_asr
+import nemo.collections.asr as nemo_asr
 
 try:
     import torch
@@ -31,35 +31,35 @@ def _huggingface_device() -> int | str | None:
     return "cpu"
 
 
-# def _initialize_typhoon_pipeline():
-#     if torch is None or pipeline is None:
-#         return None
-#     device = 'cuda' if torch.cuda.is_available() else 'mps'
-#     print(f"Using device: {device}")
-#     print("Initializing Typhoon ASR pipeline...")
-#     asr_model = nemo_asr.models.ASRModel.from_pretrained(
-#     model_name="scb10x/typhoon-asr-realtime",
-#     map_location=device
-# )
-#     print("Typhoon ASR pipeline initialized.")
-#     return asr_model
-
-def _initialize_whisper_pipeline():
-    pipe = pipeline(
-    task="automatic-speech-recognition",
-    model="nectec/Pathumma-whisper-th-medium",
-    chunk_length_s=30,
-    device=device,
-    model_kwargs={"torch_dtype": torch.bfloat16},
-    )
-    pipe.model.config.forced_decoder_ids = pipe.tokenizer.get_decoder_prompt_ids(
-    language='th',
-    task="transcribe"   
-    )
-    return pipe
-_ASR_TYPHOON = None
-# _ASR_TYPHOON = _initialize_typhoon_pipeline()
-_ASR_WHISPER = _initialize_whisper_pipeline()
+def _initialize_typhoon_pipeline():
+    if torch is None or pipeline is None:
+        return None
+    device = 'cuda' if torch.cuda.is_available() else 'mps'
+    print(f"Using device: {device}")
+    print("Initializing Typhoon ASR pipeline...")
+    asr_model = nemo_asr.models.ASRModel.from_pretrained(
+    model_name="scb10x/typhoon-asr-realtime",
+    map_location=device
+)
+    print("Typhoon ASR pipeline initialized.")
+    return asr_model
+
+# def _initialize_whisper_pipeline():
+#     pipe = pipeline(
+#     task="automatic-speech-recognition",
+#     model="nectec/Pathumma-whisper-th-medium",
+#     chunk_length_s=30,
+#     device=device,
+#     model_kwargs={"torch_dtype": torch.bfloat16},
+#     )
+#     pipe.model.config.forced_decoder_ids = pipe.tokenizer.get_decoder_prompt_ids(
+#     language='th',
+#     task="transcribe"   
+#     )
+#     return pipe
+# _ASR_TYPHOON = None
+_ASR_TYPHOON = _initialize_typhoon_pipeline()
+# _ASR_WHISPER = _initialize_whisper_pipeline()
 
 
 def _transcribe_with_pipeline(audio_array: np.ndarray) -> str:
@@ -98,17 +98,17 @@ def transcribe_audio(audio_array: np.ndarray) -> str:
     """Transcribe user audio with the best available backend."""
     if audio_array is None or not np.any(audio_array):
         return ""
-    # if _ASR_TYPHOON:
-    #     try:
-    #         transcriptions = _ASR_PIPELINE.transcribe(audio=audio_array)
-    #     except Exception as exc:
-    #         print(f"Typhoon ASR pipeline failed: {exc}")
-    if _ASR_WHISPER:
+    if _ASR_TYPHOON:
         try:
-            transcription = _ASR_WHISPER(audio_array)["text"]
-            return transcription
+            transcriptions = _ASR_TYPHOON.transcribe(audio=audio_array)
         except Exception as exc:
             print(f"Typhoon ASR pipeline failed: {exc}")
+    # if _ASR_WHISPER:
+    #     try:
+    #         transcription = _ASR_WHISPER(audio_array)["text"]
+    #         return transcription
+    #     except Exception as exc:
+    #         print(f"Typhoon ASR pipeline failed: {exc}")
 
     try:
         return _transcribe_with_google(audio_array)