import spaces
import gradio as gr
from AinaTheme import theme

MODEL_NAME = "projecte-aina/whisper-large-v3-ca-es-synth-cs"
BATCH_SIZE = 8
device = 0 if torch.cuda.is_available() else "cpu"

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

@spaces.GPU
def transcribe(inputs):
    if inputs is None:
        raise gr.Error("Cap fitxer d'àudio introduit! Si us plau pengeu un fitxer "\
                       "o enregistreu un àudio abans d'enviar la vostra sol·licitud")
    text = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": "transcribe"}, return_timestamps=True)["text"]
    return text


description_string = "Transcripció automàtica de micròfon o de fitxers d'àudio.\n Aquest demostrador s'ha desenvolupat per"\
              " comprovar els models de reconeixement de parla per a móbils. Per ara utilitza el checkpoint "\
              f"[{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}) i la llibreria de 🤗 Transformers per a la transcripció."


def clear():
     return (
          None
     )


with gr.Blocks(theme=theme) as demo:
    gr.Markdown(description_string)
    with gr.Row():
        with gr.Column(scale=1):
            #input = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Audio")
            input = gr.Audio(sources=["upload"], type="filepath", label="Audio")

        with gr.Column(scale=1):
            output = gr.Textbox(label="Output", lines=8)
    
    with gr.Row(variant="panel"):
            clear_btn = gr.Button("Clear")
            submit_btn = gr.Button("Submit", variant="primary")


    submit_btn.click(fn=transcribe, inputs=[input], outputs=[output])
    clear_btn.click(fn=clear,inputs=[], outputs=[input], queue=False,)


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