import gradio as gr
from transformers import pipeline
import torch

modelo_path = "proxectonos/Nos_ASR-wav2vec2-large-xlsr-53-gl-with-lm"

asr_pipeline = pipeline(
    "automatic-speech-recognition",
    model=modelo_path,
    device=0 if torch.cuda.is_available() else -1
)

fronted_theme = "Soft"

def cargar(audio_filepath):
    if audio_filepath is None:
        return "Por favor, carga un ficheiro de audio."

    outtext = asr_pipeline(audio_filepath)

    texto_transcrito = outtext["text"]
    return texto_transcrito

with gr.Blocks(fronted_theme) as demo:
    with gr.Row():
        with gr.Column():
            gr.Markdown(
                """ ## <h1 style="text-align:center">🗣️ ASR Demo Proxecto Nós </h1> """
            )
            gr.Markdown(
                """ ## <img src="https://huggingface.co/spaces/proxectonos/README/resolve/main/title-card.png" width="100%" style="border-radius: 0.75rem;"> """
            )
        with gr.Column():
            with gr.Row():
                gr.Markdown(
                    """ <br/> <br/> <br/> <br/> Este space mostra o modelo ASR desenvolvido polo **[Proxecto Nós](https://huggingface.co/proxectonos)**. <br/> """
                )
            with gr.Row():
                input_audio = gr.Audio(label="Entrada", type="filepath")
            with gr.Row():
                output_text = gr.Textbox(label="Saída", type="text")
            with gr.Row():
                asr_button = gr.Button("Xerar", elem_id="send-btn", visible=True)
                clear_button = gr.ClearButton([input_audio, output_text], value="Limpar", elem_id="clear-btn", visible=True)

    asr_button.click(
        cargar,
        inputs=[input_audio],
        outputs=[output_text],
    )

demo.launch()