app.py

import os
import tempfile
import sys
import subprocess
import gradio as gr
import numpy as np
import soundfile as sf
import librosa
import torch
import torch.cuda
import gc
import json
import datetime
from pathlib import Path

# Check if required packages are installed, if not install them
try:
    from espnet2.bin.s2t_inference import Speech2Text
    import torchaudio
    # Try importing espnet_model_zoo specifically
    try:
        import espnet_model_zoo
        print("All packages already installed.")
    except ModuleNotFoundError:
        print("Installing espnet_model_zoo. This may take a few minutes...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-U", "espnet_model_zoo"])
        import espnet_model_zoo
        print("espnet_model_zoo installed successfully.")
except ModuleNotFoundError as e:
    missing_module = str(e).split("'")[1]
    print(f"Installing missing module: {missing_module}")
    
    if missing_module == "espnet2":
        print("Installing ESPnet. This may take a few minutes...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "espnet"])
    elif missing_module == "torchaudio":
        print("Installing torchaudio. This may take a few minutes...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "torchaudio"])
    
    # Try importing again
    try:
        from espnet2.bin.s2t_inference import Speech2Text
        import torchaudio
        # Also check for espnet_model_zoo
        try:
            import espnet_model_zoo
        except ModuleNotFoundError:
            print("Installing espnet_model_zoo. This may take a few minutes...")
            subprocess.check_call([sys.executable, "-m", "pip", "install", "-U", "espnet_model_zoo"])
            import espnet_model_zoo
        print("All required packages installed successfully.")
    except ModuleNotFoundError as e:
        print(f"Failed to install {str(e).split('No module named ')[1]}. Please install manually.")
        raise

# Initialize the model with language option
def load_model():
    # Force garbage collection
    gc.collect()
    torch.cuda.empty_cache()
    
    # Set memory-efficient options
    torch.cuda.set_per_process_memory_fraction(0.95)  # Use 95% of available memory
    
    # Check if CUDA is available
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")
    
    # For memory efficiency, you could try loading with 8-bit quantization
    # This requires the bitsandbytes library
    # pip install bitsandbytes
    
    model = Speech2Text.from_pretrained(
        "espnet/owls_4B_180K",
        task_sym="<asr>",
        beam_size=1,
        device=device
    )
    return model

# Load the model at startup with English as default
print("Loading multilingual model...")
model = load_model()
print("Model loaded successfully!")

def transcribe_audio(audio_file, language):
    """Process the audio file and return the transcription"""
    if audio_file is None:
        return "Please upload an audio file or record audio."
    
    # If audio is a tuple (from microphone recording)
    if isinstance(audio_file, tuple):
        sr, audio_data = audio_file
        # Create a temporary file to save the audio
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_audio:
            temp_path = temp_audio.name
            sf.write(temp_path, audio_data, sr)
            audio_file = temp_path
    
    # Load and resample the audio file to 16kHz
    speech, _ = librosa.load(audio_file, sr=16000)
    
    # Update the language symbol if needed
    model.beam_search.hyps = None
    model.beam_search.pre_beam_score_key = None

    if language != None:
        model.lang_sym = language
    
    # Perform ASR
    text, *_ = model(speech)[0]
    
    # Clean up temporary file if created
    if isinstance(audio_file, str) and audio_file.startswith(tempfile.gettempdir()):
        os.unlink(audio_file)
    
    return text

# Function to handle English transcription
def transcribe_english(audio_file):
    return transcribe_audio(audio_file, "<eng>")

# Function to handle Chinese transcription
def transcribe_chinese(audio_file):
    return transcribe_audio(audio_file, "<zho>")

# Function to handle Japanese transcription
def transcribe_japanese(audio_file):
    return transcribe_audio(audio_file, "<jpn>")

# Function to handle Korean transcription
def transcribe_korean(audio_file):
    return transcribe_audio(audio_file, "<kor>")

# Function to handle Thai transcription
def transcribe_thai(audio_file):
    return transcribe_audio(audio_file, "<tha>")

# Function to handle Italian transcription
def transcribe_italian(audio_file):
    return transcribe_audio(audio_file, "<ita>")

# Function to handle German transcription
def transcribe_german(audio_file):
    return transcribe_audio(audio_file, "<deu>")

# Create a function to save feedback
def save_feedback(transcription, rating, language, audio_path=None):
    """Save user feedback to a JSON file"""
    # Create feedback directory if it doesn't exist
    feedback_dir = Path("feedback_data")
    feedback_dir.mkdir(exist_ok=True)
    
    # Create a unique filename based on timestamp
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    feedback_file = feedback_dir / f"feedback_{timestamp}.json"
    
    # Prepare feedback data
    feedback_data = {
        "timestamp": timestamp,
        "language": language,
        "transcription": transcription,
        "rating": rating,
        "audio_path": str(audio_path) if audio_path else None
    }
    
    # Save to JSON file
    with open(feedback_file, "w", encoding="utf-8") as f:
        json.dump(feedback_data, f, ensure_ascii=False, indent=2)
    
    return "Thank you for your feedback!"

# Create the Gradio interface with tabs
demo = gr.Blocks(title="NVIDIA Research Multilingual Demo")

with demo:
    gr.Markdown("# NVIDIA Research Multilingual Demo")
    gr.Markdown("Upload or record audio to transcribe up to 150 human languages using the NVIDIA Research (NVR) 9B model. Audio will be automatically resampled to 16kHz.")
    
    with gr.Tabs():
        with gr.TabItem("Microphone Recording"):
            language_mic = gr.Radio(
                ["English", "English-Mandarin", "Japanese", "Korean", "Thai", "Italian", "German"], 
                label="Select Language", 
                value="English"
            )
            
            with gr.Row():
                with gr.Column():
                    mic_input = gr.Audio(sources=["microphone"], type="filepath", label="Record Audio")
                    mic_button = gr.Button("Transcribe Recording")
                with gr.Column():
                    mic_output = gr.Textbox(label="Transcription")
            
            # Add feedback components
            with gr.Row():
                mic_rating = gr.Slider(minimum=1, maximum=5, step=1, value=3, 
                                      label="Rate the transcription quality (1=worst, 5=best)")
                mic_feedback_btn = gr.Button("Submit Feedback")
            mic_feedback_msg = gr.Textbox(label="Feedback Status", visible=True)
            
            def transcribe_mic(audio, lang):
                lang_map = {
                    "English": "<eng>", 
                    "Chinese": "<zho>", 
                    "Japanese": "<jpn>", 
                    "Korean": "<kor>",
                    "Thai": "<tha>",
                    "Italian": "<ita>",
                    "German": "<deu>"
                }
                return transcribe_audio(audio, lang_map.get(lang, "<eng>"))
            
            mic_button.click(fn=transcribe_mic, inputs=[mic_input, language_mic], outputs=mic_output)
            
            # Add feedback submission function
            def submit_mic_feedback(transcription, rating, language):
                lang_name = language  # Already a string like "English"
                return save_feedback(transcription, rating, lang_name)
            
            mic_feedback_btn.click(
                fn=submit_mic_feedback, 
                inputs=[mic_output, mic_rating, language_mic], 
                outputs=mic_feedback_msg
            )
        
        with gr.TabItem("English"):
            with gr.Row():
                with gr.Column():
                    en_input = gr.Audio(sources=["upload"], type="filepath", label="Upload Audio")
                    en_button = gr.Button("Transcribe Speech")
                with gr.Column():
                    en_output = gr.Textbox(label="Speech Transcription")
            
            # Add feedback components
            with gr.Row():
                en_rating = gr.Slider(minimum=1, maximum=5, step=1, value=3, 
                                     label="Rate the transcription quality (1=worst, 5=best)")
                en_feedback_btn = gr.Button("Submit Feedback")
            en_feedback_msg = gr.Textbox(label="Feedback Status", visible=True)
            
            # Add example if the file exists
            if os.path.exists("wav_en_sample_48k.wav"):
                gr.Examples(
                    examples=[["wav_en_sample_48k.wav"]],
                    inputs=en_input
                )
            
            en_button.click(fn=transcribe_english, inputs=en_input, outputs=en_output)
            
            # Add feedback submission
            def submit_en_feedback(transcription, rating, audio_path):
                return save_feedback(transcription, rating, "English", audio_path)
            
            en_feedback_btn.click(
                fn=submit_en_feedback, 
                inputs=[en_output, en_rating, en_input], 
                outputs=en_feedback_msg
            )
        
        with gr.TabItem("Mandarin"):
            with gr.Row():
                with gr.Column():
                    zh_input = gr.Audio(sources=["upload"], type="filepath", label="Upload Audio")
                    zh_button = gr.Button("Transcribe Speech")
                with gr.Column():
                    zh_output = gr.Textbox(label="Speech Transcription")
            
            # Add feedback components
            with gr.Row():
                zh_rating = gr.Slider(minimum=1, maximum=5, step=1, value=3, 
                                     label="Rate the transcription quality (1=worst, 5=best)")
                zh_feedback_btn = gr.Button("Submit Feedback")
            zh_feedback_msg = gr.Textbox(label="Feedback Status", visible=True)
            
            # Add example if the file exists
            if os.path.exists("wav_zh_tw_sample_16k.wav"):
                gr.Examples(
                    examples=[["wav_zh_tw_sample_16k.wav"]],
                    inputs=zh_input
                )
            
            zh_button.click(fn=transcribe_chinese, inputs=zh_input, outputs=zh_output)
            
            # Add feedback submission
            def submit_zh_feedback(transcription, rating, audio_path):
                return save_feedback(transcription, rating, "Mandarin", audio_path)
            
            zh_feedback_btn.click(
                fn=submit_zh_feedback, 
                inputs=[zh_output, zh_rating, zh_input], 
                outputs=zh_feedback_msg
            )
        
        with gr.TabItem("Japanese"):
            with gr.Row():
                with gr.Column():
                    jp_input = gr.Audio(sources=["upload"], type="filepath", label="Upload Audio")
                    jp_button = gr.Button("Transcribe Speech")
                with gr.Column():
                    jp_output = gr.Textbox(label="Speech Transcription")
            
            # Add feedback components
            with gr.Row():
                jp_rating = gr.Slider(minimum=1, maximum=5, step=1, value=3, 
                                     label="Rate the transcription quality (1=worst, 5=best)")
                jp_feedback_btn = gr.Button("Submit Feedback")
            jp_feedback_msg = gr.Textbox(label="Feedback Status", visible=True)
            
            # Add example if the file exists
            if os.path.exists("wav_jp_sample_48k.wav"):
                gr.Examples(
                    examples=[["wav_jp_sample_48k.wav"]],
                    inputs=jp_input
                )
            
            jp_button.click(fn=transcribe_japanese, inputs=jp_input, outputs=jp_output)
            
            # Add feedback submission
            def submit_jp_feedback(transcription, rating, audio_path):
                return save_feedback(transcription, rating, "Japanese", audio_path)
            
            jp_feedback_btn.click(
                fn=submit_jp_feedback, 
                inputs=[jp_output, jp_rating, jp_input], 
                outputs=jp_feedback_msg
            )
        
        with gr.TabItem("Korean"):
            with gr.Row():
                with gr.Column():
                    kr_input = gr.Audio(sources=["upload"], type="filepath", label="Upload Audio")
                    kr_button = gr.Button("Transcribe Speech")
                with gr.Column():
                    kr_output = gr.Textbox(label="Speech Transcription")
            
            # Add feedback components
            with gr.Row():
                kr_rating = gr.Slider(minimum=1, maximum=5, step=1, value=3, 
                                     label="Rate the transcription quality (1=worst, 5=best)")
                kr_feedback_btn = gr.Button("Submit Feedback")
            kr_feedback_msg = gr.Textbox(label="Feedback Status", visible=True)
            
            # Add example if the file exists
            if os.path.exists("wav_kr_sample_48k.wav"):
                gr.Examples(
                    examples=[["wav_kr_sample_48k.wav"]],
                    inputs=kr_input
                )
            
            kr_button.click(fn=transcribe_korean, inputs=kr_input, outputs=kr_output)
            
            # Add feedback submission
            def submit_kr_feedback(transcription, rating, audio_path):
                return save_feedback(transcription, rating, "Korean", audio_path)
            
            kr_feedback_btn.click(
                fn=submit_kr_feedback, 
                inputs=[kr_output, kr_rating, kr_input], 
                outputs=kr_feedback_msg
            )
        
        with gr.TabItem("Thai"):
            with gr.Row():
                with gr.Column():
                    th_input = gr.Audio(sources=["upload"], type="filepath", label="Upload Audio")
                    th_button = gr.Button("Transcribe Speech")
                with gr.Column():
                    th_output = gr.Textbox(label="Speech Transcription")
            
            # Add feedback components
            with gr.Row():
                th_rating = gr.Slider(minimum=1, maximum=5, step=1, value=3, 
                                     label="Rate the transcription quality (1=worst, 5=best)")
                th_feedback_btn = gr.Button("Submit Feedback")
            th_feedback_msg = gr.Textbox(label="Feedback Status", visible=True)
            
            # Add example if the file exists
            if os.path.exists("wav_thai_sample.wav"):
                gr.Examples(
                    examples=[["wav_thai_sample.wav"]],
                    inputs=th_input
                )
            
            th_button.click(fn=transcribe_thai, inputs=th_input, outputs=th_output)
            
            # Add feedback submission
            def submit_th_feedback(transcription, rating, audio_path):
                return save_feedback(transcription, rating, "Thai", audio_path)
            
            th_feedback_btn.click(
                fn=submit_th_feedback, 
                inputs=[th_output, th_rating, th_input], 
                outputs=th_feedback_msg
            )
        
        with gr.TabItem("Italian"):
            with gr.Row():
                with gr.Column():
                    it_input = gr.Audio(sources=["upload"], type="filepath", label="Upload Audio")
                    it_button = gr.Button("Transcribe Speech")
                with gr.Column():
                    it_output = gr.Textbox(label="Speech Transcription")
            
            # Add feedback components
            with gr.Row():
                it_rating = gr.Slider(minimum=1, maximum=5, step=1, value=3, 
                                     label="Rate the transcription quality (1=worst, 5=best)")
                it_feedback_btn = gr.Button("Submit Feedback")
            it_feedback_msg = gr.Textbox(label="Feedback Status", visible=True)
            
            # Add example if the file exists
            if os.path.exists("wav_it_sample.wav"):
                gr.Examples(
                    examples=[["wav_it_sample.wav"]],
                    inputs=it_input
                )
            
            it_button.click(fn=transcribe_italian, inputs=it_input, outputs=it_output)
            
            # Add feedback submission
            def submit_it_feedback(transcription, rating, audio_path):
                return save_feedback(transcription, rating, "Italian", audio_path)
            
            it_feedback_btn.click(
                fn=submit_it_feedback, 
                inputs=[it_output, it_rating, it_input], 
                outputs=it_feedback_msg
            )

        with gr.TabItem("German"):
            with gr.Row():
                with gr.Column():
                    de_input = gr.Audio(sources=["upload"], type="filepath", label="Upload Audio")
                    de_button = gr.Button("Transcribe Speech")
                with gr.Column():
                    de_output = gr.Textbox(label="Speech Transcription")
            
            # Add feedback components
            with gr.Row():
                de_rating = gr.Slider(minimum=1, maximum=5, step=1, value=3, 
                                     label="Rate the transcription quality (1=worst, 5=best)")
                de_feedback_btn = gr.Button("Submit Feedback")
            de_feedback_msg = gr.Textbox(label="Feedback Status", visible=True)
            
            # Add example if the file exists
            if os.path.exists("wav_de_sample.wav"):
                gr.Examples(
                    examples=[["wav_de_sample.wav"]],
                    inputs=de_input
                )
            
            de_button.click(fn=transcribe_german, inputs=de_input, outputs=de_output)
            
            # Add feedback submission
            def submit_de_feedback(transcription, rating, audio_path):
                return save_feedback(transcription, rating, "German", audio_path)
            
            de_feedback_btn.click(
                fn=submit_de_feedback, 
                inputs=[de_output, de_rating, de_input], 
                outputs=de_feedback_msg
            )

# Launch the app with Hugging Face Spaces compatible settings
if __name__ == "__main__":
    demo.launch(share=False)