support nvdia Parakeet TDT ASR

app.py

@@ -44,14 +44,18 @@ def configure_page():
         </style>
     """, unsafe_allow_html=True)
 
-def handle_file_processing(upload_path):
+def handle_file_processing(upload_path, asr_model="whisper"):
     """
     Execute the complete processing pipeline:
     1. Speech-to-Text (STT)
     2. Machine Translation
     3. Text-to-Speech (TTS)
+    
+    Args:
+        upload_path: Path to the uploaded audio file
+        asr_model: ASR model to use (whisper or parakeet)
     """
-    logger.info(f"Starting processing for: {upload_path}")
+    logger.info(f"Starting processing for: {upload_path} using {asr_model} model")
     progress_bar = st.progress(0)
     status_text = st.empty()
     
@@ -59,8 +63,8 @@ def handle_file_processing(upload_path):
         # STT Phase
         logger.info("Beginning STT processing")
         status_text.markdown("🔍 **Performing Speech Recognition...**")
-        with st.spinner("Initializing Whisper model..."):
-            english_text = transcribe_audio(upload_path)
+        with st.spinner(f"Initializing {asr_model.capitalize()} model..."):
+            english_text = transcribe_audio(upload_path, model_name=asr_model)
         progress_bar.progress(30)
         logger.info(f"STT completed. Text length: {len(english_text)} characters")
 
@@ -172,6 +176,14 @@ def main():
         format_func=lambda x: x
     )
     speed = st.sidebar.slider("Speech Speed", 0.5, 2.0, 1.0, 0.1)
+    
+    # Model selection
+    asr_model = st.selectbox(
+        "Select Speech Recognition Model",
+        options=["whisper", "parakeet"],
+        index=0,
+        help="Choose the ASR model for speech recognition"
+    )
 
     uploaded_file = st.file_uploader(
         "Select Audio File (MP3/WAV)",
@@ -185,7 +197,7 @@ def main():
         with open(upload_path, "wb") as f:
             f.write(uploaded_file.getbuffer())
         
-        results = handle_file_processing(upload_path)
+        results = handle_file_processing(upload_path, asr_model=asr_model)
         if results:
             render_results(*results)
 

requirements.txt

@@ -12,4 +12,6 @@ accelerate>=1.2.0
 soundfile>=0.13.0
 kokoro>=0.7.9
 ordered-set>=4.1.0
-phonemizer-fork>=3.3.2
+phonemizer-fork>=3.3.2
+# NeMo Toolkit with ASR support
+nemo_toolkit[asr]

utils/stt.py

@@ -1,51 +1,77 @@
 """
-Speech Recognition Module using Whisper Large-v3
+Speech Recognition Module
+Supports multiple ASR models including Whisper and Parakeet
 Handles audio preprocessing and transcription
 """
 
 import logging
 import numpy as np
+import os
+from abc import ABC, abstractmethod
+
 logger = logging.getLogger(__name__)
 
 import torch
 from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor
 from pydub import AudioSegment
-import soundfile as sf  # Add this import
+import soundfile as sf
 
-def transcribe_audio(audio_path):
-    """
-    Convert audio file to text using Whisper ASR model
-    Args:
-        audio_path: Path to input audio file
-    Returns:
-        Transcribed English text
-    """
-    logger.info(f"Starting transcription for: {audio_path}")
+class ASRModel(ABC):
+    """Base class for ASR models"""
     
-    try:
-        # Audio conversion
+    @abstractmethod
+    def load_model(self):
+        """Load the ASR model"""
+        pass
+    
+    @abstractmethod
+    def transcribe(self, audio_path):
+        """Transcribe audio to text"""
+        pass
+    
+    def preprocess_audio(self, audio_path):
+        """Convert audio to required format"""
         logger.info("Converting audio format")
         audio = AudioSegment.from_file(audio_path)
         processed_audio = audio.set_frame_rate(16000).set_channels(1)
-        wav_path = audio_path.replace(".mp3", ".wav")
+        wav_path = audio_path.replace(".mp3", ".wav") if audio_path.endswith(".mp3") else audio_path
+        if not wav_path.endswith(".wav"):
+            wav_path = f"{os.path.splitext(wav_path)[0]}.wav"
         processed_audio.export(wav_path, format="wav")
         logger.info(f"Audio converted to: {wav_path}")
+        return wav_path
 
-        # Model initialization
-        logger.info("Loading Whisper model")
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        logger.info(f"Using device: {device}")
 
-        model = AutoModelForSpeechSeq2Seq.from_pretrained(
+class WhisperModel(ASRModel):
+    """Whisper ASR model implementation"""
+    
+    def __init__(self):
+        self.model = None
+        self.processor = None
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        
+    def load_model(self):
+        """Load Whisper model"""
+        logger.info("Loading Whisper model")
+        logger.info(f"Using device: {self.device}")
+        
+        self.model = AutoModelForSpeechSeq2Seq.from_pretrained(
             "openai/whisper-large-v3",
             torch_dtype=torch.float32,
             low_cpu_mem_usage=True,
             use_safetensors=True
-        ).to(device)
+        ).to(self.device)
+        
+        self.processor = AutoProcessor.from_pretrained("openai/whisper-large-v3")
+        logger.info("Whisper model loaded successfully")
+    
+    def transcribe(self, audio_path):
+        """Transcribe audio using Whisper"""
+        if self.model is None or self.processor is None:
+            self.load_model()
+            
+        wav_path = self.preprocess_audio(audio_path)
         
-        processor = AutoProcessor.from_pretrained("openai/whisper-large-v3")
-        logger.info("Model loaded successfully")
-
         # Processing
         logger.info("Processing audio input")
         logger.debug("Loading audio data")
@@ -53,20 +79,20 @@ def transcribe_audio(audio_path):
         audio_data = audio_data.astype(np.float32)
         
         # Increase chunk length and stride for longer transcriptions
-        inputs = processor(
+        inputs = self.processor(
             audio_data,
             sampling_rate=16000,
             return_tensors="pt",
             # Increase chunk length to handle longer segments
-            chunk_length_s=60,  # Increased from 30
-            stride_length_s=10  # Increased from 5
-        ).to(device)
+            chunk_length_s=60,
+            stride_length_s=10
+        ).to(self.device)
 
         # Transcription
         logger.info("Generating transcription")
         with torch.no_grad():
             # Add max_length parameter to allow for longer outputs
-            outputs = model.generate(
+            outputs = self.model.generate(
                 **inputs, 
                 language="en", 
                 task="transcribe",
@@ -74,11 +100,84 @@ def transcribe_audio(audio_path):
                 no_repeat_ngram_size=3  # Prevent repetition in output
             )
         
-        result = processor.batch_decode(outputs, skip_special_tokens=True)[0]
-        logger.info(f"transcription: %s" % result)
+        result = self.processor.batch_decode(outputs, skip_special_tokens=True)[0]
+        logger.info(f"Transcription completed successfully")
+        return result
+
+
+class ParakeetModel(ASRModel):
+    """Parakeet ASR model implementation"""
+    
+    def __init__(self):
+        self.model = None
+        
+    def load_model(self):
+        """Load Parakeet model"""
+        try:
+            import nemo.collections.asr as nemo_asr
+            logger.info("Loading Parakeet model")
+            self.model = nemo_asr.models.ASRModel.from_pretrained(model_name="nvidia/parakeet-tdt-0.6b-v2")
+            logger.info("Parakeet model loaded successfully")
+        except ImportError:
+            logger.error("Failed to import nemo_toolkit. Please install with: pip install -U 'nemo_toolkit[asr]'")
+            raise
+    
+    def transcribe(self, audio_path):
+        """Transcribe audio using Parakeet"""
+        if self.model is None:
+            self.load_model()
+            
+        wav_path = self.preprocess_audio(audio_path)
+        
+        # Transcription
+        logger.info("Generating transcription with Parakeet")
+        output = self.model.transcribe([wav_path])
+        result = output[0].text
         logger.info(f"Transcription completed successfully")
         return result
 
+
+class ASRFactory:
+    """Factory for creating ASR model instances"""
+    
+    @staticmethod
+    def get_model(model_name="whisper"):
+        """
+        Get ASR model by name
+        Args:
+            model_name: Name of the model to use (whisper or parakeet)
+        Returns:
+            ASR model instance
+        """
+        if model_name.lower() == "whisper":
+            return WhisperModel()
+        elif model_name.lower() == "parakeet":
+            return ParakeetModel()
+        else:
+            logger.warning(f"Unknown model: {model_name}, falling back to Whisper")
+            return WhisperModel()
+
+
+def transcribe_audio(audio_path, model_name="whisper"):
+    """
+    Convert audio file to text using specified ASR model
+    Args:
+        audio_path: Path to input audio file
+        model_name: Name of the ASR model to use (whisper or parakeet)
+    Returns:
+        Transcribed English text
+    """
+    logger.info(f"Starting transcription for: {audio_path} using {model_name} model")
+    
+    try:
+        # Get the appropriate model
+        asr_model = ASRFactory.get_model(model_name)
+        
+        # Transcribe audio
+        result = asr_model.transcribe(audio_path)
+        logger.info(f"transcription: %s" % result)
+        return result
+
     except Exception as e:
         logger.error(f"Transcription failed: {str(e)}", exc_info=True)
         raise