import gradio as gr
from transformers import pipeline

# Load the translation pipeline
translator = pipeline("translation_en_to_zh", model="Helsinki-NLP/opus-mt-en-zh")

def translate(text, direction):
    if direction == "English to Chinese":
        translation = translator(text, max_length=400)[0]['translation_text']
    else:
        # Load the reverse translation pipeline for Chinese to English
        translator_reverse = pipeline("translation_zh_to_en", model="Helsinki-NLP/opus-mt-zh-en")
        translation = translator_reverse(text, max_length=400)[0]['translation_text']
    return translation

# Create the Gradio interface
iface = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter text here..."),
        gr.Radio(["English to Chinese", "Chinese to English"], label="Translation Direction")
    ],
    outputs="text",
    title="English-Chinese Translator",
    description="Translate text between English and Chinese using transformers."
)

# Launch the interface
iface.launch()