import gradio, requests, os


API_URL = os.environ["API_URL"]
HEADERS = {"Authorization": os.environ["TOKEN"]}


def predict(audio, lang="pa"):
    try:
        response = requests.post(
            API_URL,
            headers=HEADERS,
            files=[("file", ("audio.wav", open(audio, "rb"), "audio/wave"))],
            data={
                "lang": lang,
                "include_translation": True,
            },
        )

        data = response.json()
        return data["transcription"], data["translation"]
    except Exception:
        return "Model was sleeping and a request has been sent to wake it up. Please try again after 3 mins.", ""


gradio.Interface(
    fn=predict,
    inputs=[
        gradio.Audio(
            type="filepath",
            format="wav",
            label="Speak in Panjabi or English",
            sources=["microphone"],
            max_length=60,
        ),
        gradio.Radio(
            choices={("English", "en"), ("Panjabi", "pa")},
            value="pa",
            label="Select your speech language",
        ),
    ],
    outputs=[gradio.Label(label="Transcription"), gradio.Label(label="Translation")],
    title="Faster whisper model trained to recognise Panjabi-English speech (WER 12.9)",
).launch()