import spaces
import torch
import gradio as gr
from transformers import WhisperForConditionalGeneration, WhisperProcessor

MODEL_ID = "phonsobon/Whisper-Small-Khmer-continued"

device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if device == "cuda" else torch.float32

processor = WhisperProcessor.from_pretrained(MODEL_ID)
model = WhisperForConditionalGeneration.from_pretrained(
    MODEL_ID, torch_dtype=torch_dtype
)
model.to(device)
model.eval()


@spaces.GPU
def transcribe(audio_path):
    if audio_path is None:
        return "Please provide an audio input."

    import librosa

    speech, sr = librosa.load(audio_path, sr=16000)

    inputs = processor(
        speech, sampling_rate=16000, return_tensors="pt"
    )
    input_features = inputs.input_features.to(device=device, dtype=torch_dtype)

    forced_decoder_ids = processor.get_decoder_prompt_ids(
        language="khmer", task="transcribe"
    )

    with torch.no_grad():
        predicted_ids = model.generate(
            input_features,
            forced_decoder_ids=forced_decoder_ids,
            max_new_tokens=444,
        )

    transcription = processor.batch_decode(
        predicted_ids, skip_special_tokens=True
    )[0]

    return transcription.strip()


with gr.Blocks(title="ASR - Project") as demo:
    gr.Markdown(
        """
        # ASR - Project
        Khmer speech-to-text using **phonsobon/Whisper-Small-Khmer-continued**
        (fine-tuned Whisper-Small for Khmer), running on ZeroGPU.
        """
    )

    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio(
                sources=["microphone", "upload"],
                type="filepath",
                label="Audio Input (Khmer speech)",
            )
            transcribe_btn = gr.Button("Transcribe", variant="primary")

        with gr.Column():
            output_text = gr.Textbox(
                label="Transcription",
                lines=6,
                show_copy_button=True,
            )

    transcribe_btn.click(
        fn=transcribe,
        inputs=audio_input,
        outputs=output_text,
    )

    gr.Examples(
        examples=[],
        inputs=audio_input,
        label="Examples (add sample audio files here)",
    )

if __name__ == "__main__":
    demo.queue().launch()