import torch
import spaces
import gradio as gr
from transformers import pipeline

MODEL_NAME = "tarob0ba/whisper-small-eo-v0.2"
BATCH_SIZE = 1

device = 0 if torch.cuda.is_available() else "cpu"

pipe = pipeline(
    task="automatic-speech-recognition",
    model=MODEL_NAME,
    device=device,
)

@spaces.GPU
def transcribe(inputs):
    if not inputs:
        raise gr.Error("No audio file submitted! Please upload or record an audio file before submitting your request.")
    # Perform ASR with timestamps, forcing the task to "transcribe"
    result = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": "transcribe"}, return_timestamps=True)
    return result["text"]

with gr.Blocks() as demo:
    gr.Markdown(f"# Whisper Small (Esperanto) ASR\n\nDemo with model [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME})")
    
    audio_input = gr.Audio(sources=["microphone", "upload"], type="filepath", label="Audio Input")
    transcribe_button = gr.Button("Transcribe")
    transcript_output = gr.Textbox(label="Transcript")
    
    transcribe_button.click(
        fn=transcribe, 
        inputs=audio_input, 
        outputs=transcript_output
    )

demo.launch()