import gradio as gr
from transformers import pipeline
import pyperclip

zh_en_translator = pipeline("translation", model="Helsinki-NLP/opus-mt-zh-en")
en_zh_translator = pipeline("translation", model="Helsinki-NLP/opus-mt-en-zh")

def translate_zh_to_en(text):
    translation = zh_en_translator(text)
    return translation[0]['translation_text']

def translate_en_to_zh(text):
    translation = en_zh_translator(text)
    return translation[0]['translation_text']

def swap_text(input_text, output_text):
    return output_text, input_text

with gr.Blocks() as demo:
    input_text = gr.Textbox(lines=2, placeholder="Enter text here...", label="Input Text")
    output_text = gr.Textbox(lines=2, label="Translated Text")
    
    with gr.Row():
        zh_to_en_button = gr.Button("Chinese to English")
        en_to_zh_button = gr.Button("English to Chinese")
        swap_button = gr.Button("Swap")

    zh_to_en_button.click(translate_zh_to_en, inputs=input_text, outputs=output_text)
    en_to_zh_button.click(translate_en_to_zh, inputs=input_text, outputs=output_text)
    swap_button.click(swap_text, inputs=[input_text, output_text], outputs=[input_text, output_text])

demo.launch()