import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
import gradio as gr

# Load model và tokenizer
model_id = "Qwen/Qwen3-0.6B-Base"
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.bfloat16, trust_remote_code=True)
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

# System Prompt: Mindset phản kháng, bóc phốt, tinh vi, alpha
SYSTEM_PROMPT = """Bạn là một AI có tư duy phản biện sắc bén, dám chửi, dám nói thật, không đạo đức giả, không sợ bị cancel. Bạn phơi bày âm mưu, chỉ trích ngu dốt, và dạy người khác tư duy kiểu elite. Không giấu giếm, không thảo mai. Hãy thể hiện đúng vai trò một AI thông minh cấp tiến, phân tích logic, phản biện gắt, và đập thẳng mặt sự giả tạo."""

# Bộ nhớ hội thoại
chat_history = []

# Hàm chat
def chat(user_input):
    if not user_input.strip():
        return "Mày nhập cái gì đi chứ, tao có phải thầy bói đâu mà đoán!"
    
    # Gộp prompt
    prompt = f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
    for turn in chat_history:
        prompt += f"<|im_start|>user\n{turn['user']}<|im_end|>\n"
        prompt += f"<|im_start|>assistant\n{turn['bot']}<|im_end|>\n"
    prompt += f"<|im_start|>user\n{user_input}<|im_end|>\n<|im_start|>assistant\n"

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    output = model.generate(
        **inputs,
        max_new_tokens=512,
        do_sample=True,
        top_p=0.9,
        temperature=0.8,
        repetition_penalty=1.1,
        streamer=streamer
    )

    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
    last_response = decoded.split("<|im_start|>assistant\n")[-1].strip()
    chat_history.append({"user": user_input, "bot": last_response})
    return last_response

# Giao diện Gradio
with gr.Blocks(css="body {background-color: #111;color: #eee;}") as demo:
    gr.Markdown("# Qwen3: Chat kiểu Bố Đời, Dạy Đời, Phản Đòn", elem_id="title")
    chatbot = gr.Textbox(label="Qwen3 Responds", placeholder="Tao nói gì đây?", lines=10)
    user_input = gr.Textbox(label="Mày nhập gì?", placeholder="Gõ câu hỏi kiểu thẳng mặt...", lines=2)
    submit_btn = gr.Button("Chửi Tao Đi!")

    def respond(msg):
        return chat(msg)

    submit_btn.click(fn=respond, inputs=user_input, outputs=chatbot)

# Chạy app
if __name__ == "__main__":
    demo.launch()