import gradio as gr
import time
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import torch

# Load mô hình
model_checkpoint = "hHoai/model_spelll"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, revision="checkpoint-10026")

model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)

def spellcheck(text):
    start_time = time.time()
    
    # Mã hóa đầu vào
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    
    # Sinh văn bản đã sửa lỗi
    with torch.no_grad():
        outputs = model.generate(**inputs)
    corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    end_time = time.time()
    processing_time = f"{end_time - start_time:.4f} giây"
    
    # Giả lập phát hiện lỗi chính tả (có thể thay bằng logic so sánh input-output)
    has_errors = text.strip() != corrected_text.strip()
    
    if not has_errors:
        return "Không có lỗi chính tả", corrected_text, "", "", processing_time
    
    # Hiện tại chưa có logic tách từ sai và gợi ý, cần thêm post-processing nếu cần
    return "Có lỗi chính tả", corrected_text, "Đang cập nhật", "Đang cập nhật", processing_time

with gr.Blocks() as demo:
    gr.Markdown("# Demo Kiểm tra lỗi chính tả tiếng Việt")

    input_text = gr.Textbox(label="Nhập văn bản")
    output_status = gr.Textbox(label="Trạng thái", interactive=False)
    output_corrected = gr.Textbox(label="Văn bản sau khi sửa", interactive=False)

    with gr.Row():
        output_errors = gr.Textbox(label="Từ sai", lines=4, interactive=False)
        output_suggestions = gr.Textbox(label="Gợi ý sửa", lines=4, interactive=False)

    output_time = gr.Textbox(label="Thời gian xử lý", interactive=False)

    submit_btn = gr.Button("Kiểm tra lỗi")

    submit_btn.click(spellcheck, inputs=input_text, outputs=[output_status, output_corrected, output_errors, output_suggestions, output_time])

    gr.Markdown("### Ví dụ: Nhập văn bản có lỗi chính tả để kiểm tra")

if __name__ == "__main__":
    demo.launch()