import gradio as gr
from transformers import pipeline, MarianMTModel, MarianTokenizer

asr = pipeline("automatic-speech-recognition", model="openai/whisper-small")
mt_model_name = "Helsinki-NLP/opus-mt-yo-en"
tokenizer = MarianTokenizer.from_pretrained(mt_model_name)
model = MarianMTModel.from_pretrained(mt_model_name)

def translate_speech(audio):
    transcription = asr(audio)["text"]
    inputs = tokenizer(transcription, return_tensors="pt", padding=True)
    translated = model.generate(**inputs)
    translation = tokenizer.decode(translated[0], skip_special_tokens=True)
    return transcription, translation

iface = gr.Interface(fn=translate_speech,
                     inputs=gr.Audio(source="microphone", type="filepath"),
                     outputs=["text", "text"],
                     title="Yoruba to English Speech Translator")
iface.launch()