import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM

# 1. Định nghĩa đường dẫn model
# Dùng dấu chấm "." để báo cho code biết model nằm ngay tại thư mục này
model_path = "." 

print("Đang tải model từ thư mục hiện tại...")

# 2. Load Tokenizer và Model thủ công để kiểm soát lỗi tốt hơn
try:
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    
    # Tạo pipeline từ model đã load
    pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
except Exception as e:
    print(f"Lỗi khi load model: {e}")
    # Fallback: Nếu không load được local, thử load từ ID gốc (phòng hờ)
    pipe = pipeline("text2text-generation", model="Whelxi/bartpho-teencode")

# 3. Định nghĩa hàm xử lý
def dich_teencode(text):
    if not text:
        return ""
    # max_length=128 là đủ cho các câu teencode thông thường
    result = pipe(text, max_length=128)
    return result[0]['generated_text']

# 4. Tạo giao diện
demo = gr.Interface(
    fn=dich_teencode,
    inputs=gr.Textbox(label="Nhập Teencode", placeholder="Ví dụ: k hum nay di hok khong?", lines=2),
    outputs=gr.Textbox(label="Kết quả Tiếng Việt"),
    title="Whelxi - Teencode Converter",
    description="Chuyển đổi Teencode sang Tiếng Việt chuẩn sử dụng mô hình BARTpho-syllable đã được fine-tune.",
    examples=[["k hum nay di hok khong?"], ["hqa m lam j?"]]
)

# 5. Chạy ứng dụng
if __name__ == "__main__":
    demo.launch()