Update app.py

app.py

@@ -1,47 +1,60 @@
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+import gradio as gr
 
-# Load model và tokenizer (chọn model mày muốn ở đây)
+# 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)
 
-# Hàm xử lý chat
-def chat_with_model(history, user_input):
-    # Lịch sử chat dưới dạng [(user, bot), ...]
-    messages = []
-    for user_msg, bot_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        messages.append({"role": "assistant", "content": bot_msg})
-    messages.append({"role": "user", "content": user_input})
-
-    # Tokenize và sinh
-    input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)
-    output = model.generate(input_ids, max_new_tokens=1024, streamer=streamer, do_sample=True, temperature=0.7)
-    decoded_output = tokenizer.decode(output[0][input_ids.shape[1]:], skip_special_tokens=True)
-
-    # Trả về output và update history
-    return decoded_output, history + [(user_input, decoded_output)]
-
-# Gradio UI
-with gr.Blocks(title="Chat với Qwen3") as demo:
-    gr.Markdown("# Chat với AI - Qwen3 Model")
-
-    chatbot = gr.Chatbot(height=400)
-    user_input = gr.Textbox(placeholder="Nhập tin nhắn...", label="Mày nói gì?")
-    clear = gr.Button("Xoá sạch")
-    
-    state = gr.State([])
+# 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."""
 
-    def respond(message, history):
-        response, updated_history = chat_with_model(history, message)
-        return updated_history, updated_history
+# Bộ nhớ hội thoại
+chat_history = []
 
-    user_input.submit(respond, [user_input, state], [chatbot, state])
-    clear.click(lambda: ([], []), None, [chatbot, state])
+# 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
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()