import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load model từ Hugging Face Hub
model_name = "RickyGM15/vit5-gec"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

def correct_text(input_text):
    if not input_text.strip():
        return ""
    
    inputs = tokenizer(
        input_text,
        return_tensors="pt",
        truncation=True,
        padding=True,
        max_length=128
    )

    outputs = model.generate(
        **inputs,
        max_length=128,
        num_beams=5,
        repetition_penalty=1.2,
        no_repeat_ngram_size=2,
        length_penalty=0.9,
        early_stopping=True
    )

    corrected = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return corrected

# Gradio UI
demo = gr.Interface(
    fn=correct_text,
    inputs=gr.Textbox(lines=3, placeholder="Nhập câu cần sửa lỗi chính tả..."),
    outputs="text",
    title="📝 ViT5 - Chỉnh sửa lỗi chính tả tiếng Việt",
    description="Dựa trên mô hình ViT5 huấn luyện riêng.",
    examples=[
        ["A ko bik gì hết"],
        ["Toi dang di hoc trua"],
        ["Mai e đi học về"],
    ]
)

demo.launch()