Update app.py

app.py

@@ -1,25 +1,43 @@
 import os
+import torch
 from huggingface_hub import login
-login(token=os.environ["HF_TOKEN"])  # Dùng biến môi trường để lấy token
-
-import torch 
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
- 
-# Load model
-model_id = "MindVR/JohnTran_Fine-tune"  # ⚠️ Đảm bảo đây là bản mới fine-tune không dùng 4bit
-tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.environ["HF_TOKEN"])
+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)
 
+model_id = "MindVR/JohnTran_Fine-tune"
+tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     device_map="auto",
     low_cpu_mem_usage=True,
-    token=os.environ["HF_TOKEN"]
+    token=HF_TOKEN
 )
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+
+# ---- Chat Function ----
+def build_prompt(history, new_message):
+    prompt = ""
+    if history:
+        prompt += "\n".join(history) + "\n"
+    prompt += f"User: {new_message}\nAI:"
+    return prompt
 
-# Hàm xử lý yêu cầu
-def chat(prompt):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+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)
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
     with torch.no_grad():
         output = model.generate(
@@ -27,18 +45,83 @@ def chat(prompt):
             max_new_tokens=200,
             do_sample=True,
             top_p=0.95,
-            temperature=0.7
+            temperature=0.7,
+            pad_token_id=tokenizer.eos_token_id
         )
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    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
 
-# Giao diện Gradio
-demo = gr.Interface(
-    fn=chat,
-    inputs=gr.Textbox(label="Nhập câu hỏi"),
-    outputs=gr.Textbox(label="Phản hồi từ AI"),
-    title="MindVR Therapy Chatbot",
-    allow_flagging="never"
-)
+# ---- 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)
 
-demo.launch()
+if __name__ == "__main__":
+    main()