from transformers import pipeline
import torch
import gradio as gr
import os
from dotenv import load_dotenv
from huggingface_hub import login
from utils import parse_args


args = parse_args()

model = args['model']
model_name = "clinical-assistance/" + model

# Load secrets
load_dotenv()
token = os.getenv("TOKEN")
assert token is not None, 'Hugging Face token has not been specified. Please specify your token in a .env file.'

# Login to Hugging Face
login(token=token)

# Init model
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
generate_kwargs = {"language":"<|es|>", "task": "transcribe"}
pipe = pipeline(model=model_name, generate_kwargs=generate_kwargs, device = device)

def model_transcribe(audio):
    '''
    Transcribes audio to text using the Whisper model.
    '''
    text = pipe(audio)["text"]
    return text

# Define the Gradio app with the correct syntax for loading the function and handling inputs/outputs
with gr.Blocks(theme=gr.themes.Base()) as demo:
    
    with gr.Row():
        gr.Markdown(
            """
            # Transcripción Automática de Consultas Médicas 👨🏻‍⚕️🚀
            Bienvenido a la herramienta de transcripción de consultas médicas! Esta herramienta está diseñada para transcribir grabaciones de consultas médicas usando `Whisper`.
            ## Cómo usar esta herramienta?
            Usar esta herramienta es fácil! Sólo debes seguir los siguientes pasos:
            1. Sube la grabación de una consulta médica. Opcionalmente puedes grabar audio y subir la grabación.
            2. Realizar click en "Transcribir" para comenzar el proceso de transcripción.
            3. Visualizar el audio transcrito en el panel de la derecha.
            
            Eso es todo! Estás listo para transcribir grabaciones de consultas médicas de manera automática. Que lo disfrutes!
            """
        )
    
    with gr.Row():
        audio_input = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Audio de Entrada")
        text_output = gr.Textbox(label="Texto de Salida")
        
    with gr.Row():
        button = gr.Button("Transcribir")
        
    # Correct method to link function with inputs and outputs
    button.click(fn=model_transcribe, inputs=audio_input, outputs=text_output)

# Use demo.launch() to launch the interface with optional debug mode
demo.launch(debug=True, share = True)