import gradio as gr
import librosa
from asr import transcribe, ASR_EXAMPLES, ASR_NOTE

# Define only English and Swahili languages
ASR_LANGUAGES = {
    "eng": "English",
    "swh": "Swahili",
}

# Speech-to-Text Interface
mms_transcribe = gr.Interface(
    fn=transcribe,
    inputs=[
        gr.Audio(),
        gr.Dropdown(
            [f"{k} ({v})" for k, v in ASR_LANGUAGES.items()],
            label="Language",
            value="eng English",
        ),
    ],
    outputs="text",
    examples=ASR_EXAMPLES,
    title="Speech-to-Text",
    description="Transcribe audio in either English or Swahili.",
    article=ASR_NOTE,
    allow_flagging="never",
)

# Main Gradio App
with gr.Blocks() as demo:
    gr.Markdown("<p align='center' style='font-size: 20px;'>MMS Speech-to-Text</p>")
    gr.HTML("<center>Convert speech to text in English and Swahili.</center>")

    mms_transcribe.render()

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