import gradio as gr
from transformers import T5ForConditionalGeneration, T5Tokenizer

print("Đang tải model...")
model_name = "tarudesu/ViHateT5-base-HSD"
try:
    tokenizer = T5Tokenizer.from_pretrained(model_name)
    model = T5ForConditionalGeneration.from_pretrained(model_name)
    print("Tải model thành công!")
except Exception as e:
    print(f"Lỗi khi tải model: {e}")
    model = None

def classify_vietnamese_text(text_to_check):
    if model is None or tokenizer is None:
        return "Lỗi: Model chưa được khởi tạo."
    
    input_text = f"hate-speech-detection: {text_to_check}"
    
    input_ids = tokenizer(input_text, return_tensors="pt").input_ids
    outputs = model.generate(input_ids, max_length=10)
    
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return result

demo = gr.Interface(
    fn=classify_vietnamese_text,
    inputs=gr.Textbox(lines=2, placeholder="Nhập câu Tiếng Việt cần kiểm tra ở đây..."),
    outputs="text",
    title="Vietnamese Hate Speech Detection",
    description="Một API demo cho model tarudesu/ViHateT5-base-HSD."
)

demo.launch()