import gradio as gr
from transformers import pipeline

# Tải mô hình đã upload từ Hugging Face
model_name = "TDN-M/tknv1"  # Thay thế bằng tên người dùng và tên repository của bạn
pipe = pipeline("text-generation", model=model_name, max_new_tokens=512)

def chatbot(message, history):
    # Tạo tin nhắn từ người dùng
    messages = [{"role": "user", "content": message}]
    
    # Sử dụng mô hình để tạo phản hồi
    response = pipe(messages)
    
    # Trả về phản hồi từ mô hình
    return response[0]['generated_text']

# Tạo ứng dụng Gradio với giao diện chatbot hiện đại
iface = gr.ChatInterface(
    fn=chatbot,
    title="Chatbot sử dụng tiếng Việt",
    description="Một chatbot sử dụng mô hình đã được huấn luyện bằng tiếng Việt.",
    theme="default",  # Chọn theme cho giao diện
    css=".gradio-container {background-color: #fef0f0;}"  # Thêm CSS tùy chỉnh
)

# Chạy ứng dụng
iface.launch(share=True)  # Thêm share=True để tạo public link