import whisper
import gradio as gr

model = whisper.load_hf_model(repo_id="jerpint/whisper", filename="base.pt")


def transcribe(audio, translate):

    task = "translate" if translate else None
    result = model.transcribe(audio, task=task)

    return result["text"]

title = "BabelFish"

description = "Record your voice in any language, babelfish will output a transcript of what was said. Check 'Translate to english' to get an english transcription. Based on the OpenAI Whisper model"

gr.Interface(
    fn=transcribe,
    inputs=[
        gr.Audio(source="microphone", type="filepath"),
        gr.Checkbox(label="Translate to english"),
    ],
    title=title,
    description=description,
    examples=[
        ["samples/french_hello.wav", True],
        ["samples/english_hello.wav", True],
        ["samples/hebrew_hello.wav", True],
        ["samples/spanish_hello.wav", True],
    ],
    outputs="text",
).launch()