Update app.py

app.py

@@ -3,10 +3,7 @@ import torch
 from huggingface_hub import login
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
-from fastapi import FastAPI, Request
-from pydantic import BaseModel
 
-# ---- Load Model ----
 HF_TOKEN = os.environ.get("HF_TOKEN")
 if HF_TOKEN:
     login(token=HF_TOKEN)
@@ -22,7 +19,6 @@ model = AutoModelForCausalLM.from_pretrained(
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
-# ---- Chat Function ----
 def build_prompt(history, new_message):
     prompt = ""
     if history:
@@ -30,14 +26,9 @@ def build_prompt(history, new_message):
     prompt += f"User: {new_message}\nAI:"
     return prompt
 
-def chat_gradio(message, history):
-    history_text = []
-    if history:
-        # history là dạng list các cặp [msg, response]
-        for user_msg, ai_msg in history:
-            history_text.append(f"User: {user_msg}")
-            history_text.append(f"AI: {ai_msg}")
-    prompt = build_prompt(history_text, message)
+def chat(history, new_message):
+    # history: list[str], new_message: str
+    prompt = build_prompt(history, new_message)
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
     with torch.no_grad():
         output = model.generate(
@@ -49,79 +40,22 @@ def chat_gradio(message, history):
             pad_token_id=tokenizer.eos_token_id
         )
     output_text = tokenizer.decode(output[0], skip_special_tokens=True)
-    # Lấy đoạn trả lời AI cuối cùng
     if "AI:" in output_text:
         response = output_text.split("AI:")[-1].strip()
     else:
         response = output_text.strip()
     return response
 
-# ---- Gradio Interface ----
-with gr.Blocks() as demo:
-    gr.Markdown("# MindVR Therapy Chatbot")
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox(label="Nhập câu hỏi")
-    send = gr.Button("Gửi")
-
-    def user_chat(message, history):
-        response = chat_gradio(message, history)
-        return response
-
-    send.click(
-        fn=user_chat,
-        inputs=[msg, chatbot],
-        outputs=chatbot,
-        queue=False
-    )
-    msg.submit(
-        fn=user_chat,
-        inputs=[msg, chatbot],
-        outputs=chatbot,
-        queue=False
-    )
-
-# ---- REST API Endpoint ----
-app = FastAPI()
-
-class ChatRequest(BaseModel):
-    history: list
-    new_message: str
-
-@app.post("/generate")
-async def generate(data: ChatRequest):
-    # history dạng ["User: ...", "AI: ...", ...]
-    prompt = build_prompt(data.history, data.new_message)
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
-    with torch.no_grad():
-        output = model.generate(
-            input_ids,
-            max_new_tokens=200,
-            do_sample=True,
-            top_p=0.95,
-            temperature=0.7,
-            pad_token_id=tokenizer.eos_token_id
-        )
-    output_text = tokenizer.decode(output[0], skip_special_tokens=True)
-    if "AI:" in output_text:
-        response = output_text.split("AI:")[-1].strip()
-    else:
-        response = output_text.strip()
-    return {"response": response}
-
-# ---- Export both Gradio and API ----
-import uvicorn
-
-def main():
-    import threading
-    import time
-
-    # Run FastAPI on background
-    def run_api():
-        uvicorn.run(app, host="0.0.0.0", port=7861)
-    threading.Thread(target=run_api, daemon=True).start()
-
-    # Run Gradio interface
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
+# CHỈ DÙNG Gradio Interface, input là [history, new_message]
+iface = gr.Interface(
+    fn=chat,
+    inputs=[
+        gr.inputs.Textbox(lines=8, label="History (JSON list, ví dụ: [\"User: Xin chào\"] )"),
+        gr.inputs.Textbox(label="New message")
+    ],
+    outputs="text",
+    title="MindVR Therapy Chatbot",
+    allow_flagging="never"
+)
 
-if __name__ == "__main__":
-    main()
+iface.launch()