Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 import torch
-from typing import List, Union
+from typing import List
 from huggingface_hub import login
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
@@ -20,22 +20,20 @@ model = AutoModelForCausalLM.from_pretrained(
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
-def build_prompt(history, new_message):
-    prompt = ""
-    if history:
-        prompt += "\n".join(history) + "\n"
-    prompt += f"User: {new_message}\nAI:"
-    return prompt
+def build_prompt(prompt: str, histories: List[str], new_message: str) -> str:
+    prompt_text = prompt.strip() + "\n" if prompt else ""
+    if histories:
+        prompt_text += "\n".join(histories) + "\n"
+    prompt_text += f"User: {new_message}\nAI:"
+    return prompt_text
 
-def chat(history: Union[str, List[str]], new_message: str) -> str:
-    if isinstance(history, str):
-        import ast
-        try:
-            history = ast.literal_eval(history)
-        except:
-            history = [history]
-    prompt = build_prompt(history, new_message)
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+def chat(
+    prompt: str,
+    histories: List[str],
+    new_message: str
+) -> str:
+    prompt_text = build_prompt(prompt, histories, new_message)
+    input_ids = tokenizer(prompt_text, return_tensors="pt").input_ids.to(device)
     with torch.no_grad():
         output = model.generate(
             input_ids,
@@ -54,15 +52,20 @@ def chat(history: Union[str, List[str]], new_message: str) -> str:
 
 with gr.Blocks() as demo:
     gr.Markdown("# MindVR Therapy Chatbot\n\nDùng thử UI hoặc gọi API!")
-    with gr.Row():
-        history = gr.Textbox(lines=8, label="History (JSON list, ví dụ: [\"User: Xin chào\"] )")
-        new_message = gr.Textbox(label="New message")
+    prompt_box = gr.Textbox(lines=2, label="Prompt (System Prompt, chỉ dẫn context cho AI, có thể bỏ trống)")
+    histories_box = gr.Textbox(lines=8, label="Histories (mỗi dòng là một message, ví dụ: User: Xin chào)")
+    new_message_box = gr.Textbox(label="New message")
     output = gr.Textbox(label="AI Response")
-    def _chat_ui(history, new_message):
-        return chat(history, new_message)
+
+    def _chat_ui(prompt, histories, new_message):
+        # histories nhập từ UI là multiline string -> chuyển thành list
+        histories_list = [line.strip() for line in histories.split('\n') if line.strip()]
+        return chat(prompt, histories_list, new_message)
+
     btn = gr.Button("Gửi")
-    btn.click(_chat_ui, inputs=[history, new_message], outputs=output)
+    btn.click(_chat_ui, inputs=[prompt_box, histories_box, new_message_box], outputs=output)
 
+    # API chuẩn RESTful với prompt, histories, new_message
     gr.api(chat, api_name="chat_ai")
 
 demo.launch()