Update app.py

app.py

@@ -14,23 +14,44 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 
 @spaces.GPU(duration=60)  # cấp GPU tạm cho 1 phút
-def chat_with_model(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    output = model.generate(
-        **inputs,
-        max_new_tokens=256,
+def chat_with_model(prompt, system_prompt):
+    
+    # 1. Tạo danh sách messages theo chuẩn ChatML
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": prompt}
+    ]
+    
+    # 2. Dùng apply_chat_template để tokenizer tự động thêm các thẻ <|im_start|>
+    inputs = tokenizer.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=True, # Thêm thẻ <|im_start|>assistant
+        return_tensors="pt"
+    ).to(model.device)
+
+    # 3. Generate
+    output_tokens = model.generate(
+        inputs,
+        max_new_tokens=1024,
         do_sample=True,
         temperature=0.7,
         top_p=0.9
     )
-    text = tokenizer.decode(output[0], skip_special_tokens=True)
-    return text
+    
+    # 4. Decode (chỉ decode phần nội dung mới, bỏ qua phần prompt)
+    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"),
+    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 (ZeroGPU)"
+    title="Xyntrai-Mistral 2.5 7B Chat"
 )
 
 if __name__ == "__main__":