import gradio as gr
from transformers import pipeline
import numpy as np


traductor = pipeline("translation", model="Helsinki-NLP/opus-mt-tc-big-en-es")
emitidorAudio = pipeline("text-to-speech", model="suno/bark-small")
callback = gr.CSVLogger()


def traducir(frase):
    traduccion = traductor(frase) 
    return traduccion[0]["translation_text"]

def emitirAudio(texto_traducido):

    audio = emitidorAudio(texto_traducido)
    
    audio_array = np.array(audio["audio"][0], dtype=np.float32)
    sampling_rate = audio["sampling_rate"]
    
    return (sampling_rate,audio_array)

with gr.Blocks() as demo:
    gr.Markdown("ESTA ES NUESTRA APLICACIÓN")
    with gr.Row():
        with gr.Column(scale=1):
            entrada = gr.Textbox(placeholder="Frase a traducir", label="Frase")
        with gr.Column(scale=2):
            traduccion = gr.Textbox(label="Frase en Español")
            btn = gr.Button("Emitir Audio")
        with gr.Column(scale=1):
            salida = gr.Audio(label="Audio en Español")
    with gr.Row():
        btnFlag = gr.Button("Flag")
  
    entrada.change(fn=traducir, inputs=entrada, outputs=traduccion)
  
    btn.click(fn=emitirAudio, inputs=traduccion, outputs=salida)
    
    # This needs to be called at some point prior to the first call to callback.flag()
    callback.setup([entrada, traduccion, salida], "flagged_data_points")
    
      # We can choose which components to flag -- in this case, we'll flag all of them
    btnFlag.click(lambda *args: callback.flag(list(args)), [entrada, traduccion, salida], None, preprocess=False)

demo.launch()