import gradio as gr
from transformers import pipeline

pipe = pipeline("translation", model="t5-base", tokenizer_kwargs={"model_max_length": 1024})

def translate(text, target_language):
    if text is None or text == "":
        return ""
    return pipe(text, target_language=target_language)[0]["translation_text"]

def translate_callback(english, language_select):
    if english is None or english == "":
        return ""
    target_language = "de" if language_select == "German" else "fr" if language_select == "French" else "lo"
    translated_text = translate(english, target_language)
    return translated_text

def clear_callback(english):
    return ""

def main():
    english_textbox = gr.inputs.Textbox(label="English text")
    language_select = gr.inputs.Dropdown(label="Translate to", choices=["German", "French", "Lao"])
    translated_textbox = gr.outputs.Textbox(label="Translated Text")

    translate_interface = gr.Interface(
        fn=translate_callback,
        inputs=[english_textbox, language_select],
        outputs=translated_textbox,
        title="Text Translation",
        description="Translate English text to German, French, or Lao."
    )

    clear_interface = gr.Interface(
        fn=clear_callback,
        inputs=english_textbox,
        outputs=translated_textbox,
        title="Text Translation",
        description="Translate English text to German, French, or Lao."
    )

    iface = gr.Interface(
        fn=translate_callback,
        inputs=[english_textbox, language_select],
        outputs=translated_textbox,
        title="Text Translation",
        description="Translate English text to German, French, or Lao."
    )

    iface.launch()

if __name__ == "__main__":
    main()