import gradio as gr
from transformers import pipeline

# Crear pipelines de traducción específicos para cada combinación de idiomas
translation_pipelines = {
    "G": pipeline("translation_en_to_de", model="Helsinki-NLP/opus-mt-en-de"),
    "F": pipeline("translation_en_to_fr", model="Helsinki-NLP/opus-mt-en-fr"),
    "S": pipeline("translation_en_to_es", model="Helsinki-NLP/opus-mt-en-es"),
    "I": pipeline("translation_en_to_it", model="Helsinki-NLP/opus-mt-en-it")
}

# Mapear los valores de entrada de Gradio a las claves del diccionario
language_map = {
    "German": "G",
    "French": "F",
    "Spanish": "S",
    "Italian": "I"
}

def translate(text, target_language):
    if not target_language:
        return "Please select a target language."

    # Mapeamos el idioma a la clave correspondiente en el diccionario
    target_language_code = language_map.get(target_language, None)

    if not target_language_code:
        raise ValueError(f"Language '{target_language}' not supported")

    # Usar el pipeline correspondiente
    pipe = translation_pipelines[target_language_code]

    # Realizar la traducción
    translation = pipe(text)
    translated_text = translation[0]['translation_text']
    return translated_text

with gr.Blocks() as demo:
    inp = gr.Textbox(label="What do you wanna translate?")
    language = gr.Radio(["German", "French", "Spanish", "Italian"], label="Select target language")
    output = gr.Textbox(label="Translation")

    translate_btn = gr.Button("Translate")
    translate_btn.click(translate, inputs=[inp, language], outputs=output)

demo.launch(share=True)