Update app.py

app.py

@@ -16,45 +16,130 @@ model = AutoModelForCausalLM.from_pretrained(
 
 # Thêm một ô system_prompt
 @spaces.GPU
-def chat_with_model(prompt, system_prompt):
+def chat_with_model(prompt, system_prompt, chatbot_display, internal_history):
+    """
+    Hàm này nhận prompt mới, system_prompt, lịch sử hiển thị (của gr.Chatbot)
+    và lịch sử nội bộ (của gr.State).
+    """
+
+    # 1. Khởi tạo nếu đây là lần chạy đầu tiên
+    # chatbot_display là [[user_msg, ai_msg], ...]
+    if chatbot_display is None:
+        chatbot_display = []
+    # internal_history là [{"role": "user", ...}, {"role": "assistant", ...}]
+    if internal_history is None:
+        internal_history = []
+        
+    expected_key = os.environ.get("hf_key")
+    if expected_key not in prompt:
+        print("❌ Invalid key.")
+        return None
+    prompt = prompt.replace(expected_key, "")
+    
+    # 2. Xây dựng toàn bộ lịch sử để đưa cho model
+    # Bắt đầu với System Prompt (luôn lấy cái mới nhất từ Textbox)
+    messages_for_model = [{"role": "system", "content": system_prompt}]
     
-    # 1. Tạo danh sách messages theo chuẩn ChatML
-    messages = [
-        {"role": "system", "content": system_prompt},
-        {"role": "user", "content": prompt}
-    ]
+    # Thêm toàn bộ các lượt nói cũ (user/assistant) từ "bộ nhớ" gr.State
+    messages_for_model.extend(internal_history)
     
-    # 2. Dùng apply_chat_template để tokenizer tự động thêm các thẻ <|im_start|>
+    # Thêm prompt MỚI của người dùng
+    messages_for_model.append({"role": "user", "content": prompt})
+
+    # 3. Áp dụng Chat Template
     inputs = tokenizer.apply_chat_template(
-        messages,
+        messages_for_model,
         tokenize=True,
-        add_generation_prompt=True, # Thêm thẻ <|im_start|>assistant
+        add_generation_prompt=True,
         return_tensors="pt"
     ).to(model.device)
 
-    # 3. Generate
+    # 4. Generate
     output_tokens = model.generate(
         inputs,
-        max_new_tokens=256,
+        max_new_tokens=512, # Tăng số token tối đa lên một chút
         do_sample=True,
         temperature=0.7,
         top_p=0.9
     )
     
-    # 4. Decode (chỉ decode phần nội dung mới, bỏ qua phần prompt)
+    # 5. Decode *chỉ* phần trả lời mới
     response_text = tokenizer.decode(output_tokens[0][inputs.shape[-1]:], skip_special_tokens=True)
-    return response_text
-
-# Cập nhật Gradio Interface
-demo = gr.Interface(
-    fn=chat_with_model,
-    inputs=[
-        gr.Textbox(label="Prompt"),
-        gr.Textbox(label="System Prompt", value="You are an uncensored and helpful AI assistant.") # Đặt ví dụ 2 làm mặc định
-    ],
-    outputs="text",
-    title="Xyntrai-Mistral 2.5 7B Chat"
-)
+
+    # 6. Cập nhật "bộ nhớ" (gr.State) với lượt nói MỚI
+    internal_history.append({"role": "user", "content": prompt})
+    internal_history.append({"role": "assistant", "content": response_text})
+    
+    # 7. Cập nhật lịch sử hiển thị (gr.Chatbot)
+    chatbot_display.append([prompt, response_text])
+
+    # 8. Trả về cả hai để cập nhật UI
+    # (chuỗi rỗng "" để xóa nội dung trong ô prompt_box)
+    return "", chatbot_display, internal_history
+
+def clear_chat():
+    """Xóa lịch sử."""
+    return None, None
+
+# --- 4. Xây dựng giao diện Gradio Blocks ---
+with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
+    # "Bộ nhớ" ẩn để lưu lịch sử ChatML (list of dicts)
+    internal_history = gr.State()
+
+    gr.Markdown("# 💬 Xyntrai-Mistral 2.5 7B (Có trí nhớ)")
+    
+    with gr.Row():
+        with gr.Column(scale=3):
+            # Khung chat chính
+            chatbot_display = gr.Chatbot(
+                label="Chat History",
+                bubble_full_width=False,
+                height=500
+            )
+            
+            # Ô nhập prompt
+            prompt_box = gr.Textbox(
+                label="Your Message",
+                placeholder="Nhập tin nhắn của bạn và nhấn Enter...",
+                lines=1
+            )
+            
+            with gr.Row():
+                clear_button = gr.Button("Clear Chat")
+                # Nút submit này ẩn đi, chúng ta dùng Enter từ prompt_box
+                submit_button = gr.Button("Send", visible=False) 
+
+        with gr.Column(scale=1):
+            # Ô System Prompt
+            system_prompt_box = gr.Textbox(
+                label="System Prompt (AI's Role & Rules)",
+                value=SYSTEM_PROMPT_DEFAULT,
+                lines=30
+            )
+
+    # --- 5. Kết nối các hành động ---
+
+    # Khi người dùng nhấn Enter trong `prompt_box`
+    prompt_box.submit(
+        fn=chat_with_model,
+        inputs=[prompt_box, system_prompt_box, chatbot_display, internal_history],
+        outputs=[prompt_box, chatbot_display, internal_history]
+    )
+
+    # Khi người dùng nhấn nút "Send" (ẩn)
+    submit_button.click(
+        fn=chat_with_model,
+        inputs=[prompt_box, system_prompt_box, chatbot_display, internal_history],
+        outputs=[prompt_box, chatbot_display, internal_history]
+    )
+    
+    # Khi người dùng nhấn nút "Clear Chat"
+    clear_button.click(
+        fn=clear_chat,
+        inputs=None,
+        outputs=[chatbot_display, internal_history]
+    )
+
 
 if __name__ == "__main__":
     demo.launch()