import gradio as gr

from transformers import pipeline

to_german = pipeline("translation", model="t5-base")
to_spanish = pipeline("translation", model="Helsinki-NLP/opus-mt-en-es")


def translate(text, translator):
    return translator(text)[0]["translation_text"]


example_text = ["I went to the supermarket yesterday.",
                "Helen is a good swimmer."]


with gr.Blocks() as demo:
    with gr.Tabs():
        with gr.TabItem("English to German"):
            with gr.Row():
                with gr.Column():
                    english = gr.Textbox(label="English Text")
                    translate_to_german = gr.Button(value="Translate To German")
                with gr.Column():
                    german = gr.Textbox(label="German Text")
            translate_to_german.click(lambda text: translate(text, to_german), inputs=english, outputs=german)
            gr.Examples(examples=example_text,
                        inputs=[english])
        with gr.TabItem("English to Spanish"):
            with gr.Row():
                with gr.Column():
                    english_2 = gr.Textbox(label="English Text")
                    translate_to_spanish = gr.Button(value="Translate To Spanish")
                with gr.Column():
                    spanish = gr.Textbox(label="Spanish Text")
            translate_to_spanish.click(lambda text: translate(text, to_spanish), inputs=english_2, outputs=spanish)
            gr.Examples(examples=example_text,
                        inputs=[english_2])


if __name__ == "__main__":
    demo.launch()