import gradio as gr
from transformers import pipeline

# Inicialización del traductor
traductor = pipeline(task='text2text-generation', model='facebook/m2m100_418M')

def traducir_texto(texto, idioma_destino):
    respuesta = traductor(texto, forced_bos_token_id=traductor.tokenizer.get_lang_id(lang=idioma_destino))
    return respuesta[0]["generated_text"]

def traducir_ingles_a_espanol(texto):
    return traducir_texto(texto, "es")

def traducir_espanol_a_ingles(texto):
    return traducir_texto(texto, "en")

def limpiar_entrada(texto):
    return ""

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=4):
            entrada_texto = gr.Textbox()
            etiqueta_respuesta = gr.Textbox()
        with gr.Column(scale=1):
            with gr.Row():
                boton_limpiar = gr.Button("Limpiar")
            with gr.Row():
                boton_ingles_espanol = gr.Button("Inglés a Español")
            with gr.Row():
                boton_espanol_ingles = gr.Button("Español a Inglés")
    

    boton_ingles_espanol.click(
        fn=traducir_ingles_a_espanol, 
        inputs=entrada_texto, 
        outputs=etiqueta_respuesta)
    
    boton_espanol_ingles.click(
        fn=traducir_espanol_a_ingles, 
        inputs=entrada_texto, 
        outputs=etiqueta_respuesta)
    
    boton_limpiar.click(
        fn=limpiar_entrada, 
        inputs=entrada_texto, 
        outputs=entrada_texto)
    

demo.launch()