ui_interface.py

import gradio as gr
from config import MODEL_MAP, APP_TITLE, APP_DESCRIPTION, DEFAULT_MODEL_CHOICE, DEFAULT_MAX_LENGTH, DEFAULT_NUM_BEAMS

def create_interface(translation_service):
    """Create and return the Gradio interface"""
    
    with gr.Blocks(title=APP_TITLE, theme="soft") as demo:
        gr.Markdown(APP_DESCRIPTION)
        
        # Tạo một Gradio State để lưu trữ và cập nhật lịch sử.
        history_state = gr.State(value=translation_service.get_history_display())

        # Sử dụng gr.Tab để tổ chức giao diện
        with gr.Tab("Dịch Văn bản/File"):
            with gr.Row():
                with gr.Column(scale=1):
                    with gr.Group():
                        input_text = gr.Textbox(
                            lines=5,
                            label="Input",
                            placeholder="Nhập văn bản cần dịch vào đây..."
                        )
                        translate_btn = gr.Button("Dịch Văn bản", variant="primary")
                    
                    with gr.Group():
                        file_input = gr.File(label="Tải lên File (TXT/CSV)")
                        translate_file_btn = gr.Button("Dịch File", variant="primary")
                    
                    # Các lựa chọn mô hình, độ dài, beam search
                    model_choice = gr.Dropdown(
                        choices=list(MODEL_MAP.keys()),
                        label="Mô hình dịch",
                        value=DEFAULT_MODEL_CHOICE, # Mặc định là mô hình đã được fine-tune
                        interactive=True # Cho phép người dùng chọn model
                    )
                    
                    with gr.Row():
                        max_length = gr.Slider(
                            64, 1024, 
                            value=DEFAULT_MAX_LENGTH, 
                            step=10, 
                            label="Độ dài tối đa"
                        )
                        num_beams = gr.Slider(
                            1, 10, 
                            value=DEFAULT_NUM_BEAMS, 
                            step=1, 
                            label="Số lượng Beam (Num Beams)"
                        )
                
                with gr.Column(scale=1):
                    output_text = gr.Textbox(
                        lines=5,
                        label="Output",
                        placeholder="Kết quả dịch sẽ hiển thị ở đây...",
                        interactive=False
                    )
                    # copy_btn = gr.Button("Sao chép vào ô Input", variant="secondary") # Đổi tên nút để rõ nghĩa hơn
                    # gr.Markdown("*(Để sao chép văn bản đã dịch vào ô nhập liệu)*")

        with gr.Tab("Lịch sử Dịch"):
            gr.Markdown("## Lịch sử Các bản dịch")
            history_box = gr.Dataframe(
                headers=["Input", "Output", "Model", "Max Length", "Num Beams", "Timestamp"], # Cập nhật headers
                label="Lịch sử Dịch",
                value=history_state.value, # Sử dụng giá trị từ history_state
                interactive=False,
                wrap=True,
            )
            with gr.Row():
                refresh_history_btn = gr.Button("Làm mới Lịch sử", variant="secondary")
                clear_history_btn = gr.Button("Xóa Lịch sử", variant="stop")

        # --- TAB MỚI CHO FINE-TUNING ---
        with gr.Tab("Fine-tune Mô hình"):
            gr.Markdown("## Huấn luyện lại (Fine-tune) Mô hình Dịch (Chỉ dành cho mô hình en-vi và vi-en)")
            gr.Markdown(
                "Quá trình này sẽ sử dụng dataset đã tải để cải thiện chất lượng dịch của mô hình. "
                "Nó có thể mất một khoảng thời gian đáng kể tùy thuộc vào kích thước dataset và cấu hình GPU của bạn. "
                "Vui lòng kiểm tra console để xem tiến trình và log chi tiết."
            )
            fine_tune_btn = gr.Button("BẮT ĐẦU FINE-TUNE", variant="primary")
            fine_tune_output_status = gr.Markdown("Status: Sẵn sàng để fine-tune. Quá trình fine-tune có thể làm UI tạm thời bị treo.", visible=True) # Hiển thị trạng thái fine-tuning

        # --- EVENT HANDLERS ---

        # Event handler cho dịch văn bản
        translate_btn.click(
            fn=translation_service.translate_text,
            inputs=[input_text, model_choice, max_length, num_beams],
            outputs=[output_text, history_box],
        )
        
        # Event handler cho dịch file
        translate_file_btn.click(
            fn=translation_service.translate_from_file,
            inputs=[file_input, model_choice, max_length, num_beams],
            outputs=[output_text, history_box]
        )
        
        # Event handler cho nút copy
        # copy_btn.click(
        #     fn=lambda text: text,
        #     inputs=[output_text],
        #     outputs=[input_text]
        # )
        
        # Event handler cho nút xóa lịch sử
        clear_history_btn.click(
            fn=translation_service.clear_history,
            outputs=[history_box, output_text]
        )

        # Event handler cho nút làm mới lịch sử
        refresh_history_btn.click(
            fn=translation_service.get_history_display,
            inputs=[],
            outputs=[history_box]
        )

        # Event handler cho nút Fine-tune
        fine_tune_btn.click(
            fn=lambda selected_model: translation_service.engine.fine_tune_model(selected_model),
            inputs=[model_choice],
            outputs=[fine_tune_output_status],
            show_progress='full'
        )

        # Cập nhật lịch sử ngay khi giao diện tải lên
        demo.load(
            fn=translation_service.get_history_display,
            inputs=[],
            outputs=[history_box]
        )
    
    return demo