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):
    lines = input_text.split('\n')  # Tách từng dòng
    translated_lines = [pipe(line, max_length=512)[0]['translation_text'] if line.strip() else '' for line in lines]
    return '\n'.join(translated_lines)  # Gộp lại với xuống dòng

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()