# app.py
# =============
# This is a complete app.py file for an automatic Speech Recognition (ASR) using the openai/whisper-large-v3-turbo model.
# The app is built using Gradio and Hugging Face Transformers, and it runs on the CPU to avoid video memory usage.

import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
import gradio as gr

# Set device to CPU
device = "cpu"
torch_dtype = torch.float32

# Load the model and processor
model_id = "openai/whisper-large-v3-turbo"

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
)
model.to(device)

processor = AutoProcessor.from_pretrained(model_id)

# Create the ASR pipeline
pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    torch_dtype=torch_dtype,
    device=device,
    chunk_length_s=30,  # Process audio in 30-second chunks
    return_timestamps=True  # Enable timestamp prediction for long-form generation
)

# Define the transcription function
def transcribe_audio(audio_file, language):
    """
    Transcribe the given audio file using the Whisper model.
    
    Parameters:
    audio_file (str): Path to the audio file.
    language (str): Language code for transcription.
    
    Returns:
    str: Transcribed text.
    """
    generate_kwargs = {"language": language}
    result = pipe(audio_file, generate_kwargs=generate_kwargs)
    return result["text"]

# Define the Gradio interface
iface = gr.Interface(
    fn=transcribe_audio,
    inputs=[
        gr.Audio(label="Upload Audio", type="filepath"),
        gr.Dropdown(
            label="Select Language",
            choices=["en", "ru", "es", "fr", "de", "zh", "ja", "ko", "pt", "it"],
            value="en",
            info="Select the language for transcription."
        )
    ],
    outputs=gr.Textbox(label="Transcription"),
    title="Whisper ASR Demo",
    description="Upload an audio file and select the language to get the transcribed text using the openai/whisper-large-v3-turbo model.",
)

# Launch the Gradio app
if __name__ == "__main__":
    iface.launch()