import gradio as gr
from transformers import pipeline

# Load translation pipeline
pipe = pipeline("translation", model="chi-vi/hirashiba-mt-tiny-zh-vi")

def translate_text(input_text):
    output_text = pipe(input_text, max_length=512)[0]['translation_text']
    return output_text

if __name__ == '__main__':
    with gr.Blocks() as app:
        gr.Markdown('## Chinese to Vietnamese Translation')
        
        with gr.Row():
            with gr.Column(scale=1):
                input_text = gr.Textbox(label='Input Chinese Text', lines=5, placeholder='Enter Chinese text here...')
                translate_button = gr.Button('Translate')
                output_text = gr.Textbox(label='Output Vietnamese Text', lines=5, interactive=False)

        translate_button.click(
            fn=translate_text,
            inputs=input_text,
            outputs=output_text
        )

    app.launch()