Update app.py

app.py

@@ -1,58 +1,46 @@
-from flask import Flask, request, jsonify, render_template
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from peft import PeftModel
 import torch
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-app = Flask(__name__)
-import os
-from huggingface_hub import login
-
-# Lấy token từ biến môi trường và đăng nhập HF Hub
-access_token = os.environ.get("HUGGING_FACE_HUB_TOKEN")
-if access_token is None:
-    raise ValueError("Bạn chưa đặt biến môi trường HUGGING_FACE_HUB_TOKEN")
-login(token=access_token)
-# Cấu hình model
-BASE_MODEL = "openchat/openchat-3.5-0106"
-ADAPTER_PATH = "./chatbot-gpt35-peft"
-
+MODEL_REPO = "memengoc/newchat"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load tokenizer và base model
-tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
-base_model = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL,
-    device_map="auto",
-    torch_dtype=torch.float16,
-    trust_remote_code=True,
-)
-# Load adapter PEFT
-model = PeftModel.from_pretrained(base_model, ADAPTER_PATH, device_map="auto", is_local=True)
-model = model.to(device)
+# Load model và tokenizer từ Hugging Face
+tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
+model = AutoModelForCausalLM.from_pretrained(MODEL_REPO).to(device)
 model.eval()
 
-@app.route("/")
-def home():
-    return render_template("index.html")  # Tạo file index.html ở thư mục templates
-
-@app.route("/chat", methods=["POST"])
-def chat():
-    user_input = request.json.get("message", "")
-    prompt = f"User: {user_input}\nAI:"
-    
+# Hàm sinh phản hồi
+def generate_response(prompt):
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
     outputs = model.generate(
         **inputs,
-        max_new_tokens=200,
+        max_new_tokens=300,
         do_sample=True,
         temperature=0.7,
         top_p=0.9,
-        pad_token_id=tokenizer.eos_token_id
+        pad_token_id=tokenizer.eos_token_id,
     )
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    # Loại bỏ phần prompt cũ nếu cần
-    response = response[len(prompt):].strip()
-    return jsonify({"response": response})
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# Giao diện Gradio
+def chat(user_input, history):
+    if history is None:
+        history = ""
+    prompt = f"{history}\nỨng viên: {user_input}\nNhà tuyển dụng:"
+    response = generate_response(prompt)
+    new_history = f"{prompt} {response}"
+    return response, new_history
+
+with gr.Blocks() as demo:
+    gr.Markdown("## 🤖 Chatbot Phỏng Vấn AI (OpenChat 3.5 đã Fine-tuned)")
+
+    chat_history = gr.Textbox(visible=False)
+    user_input = gr.Textbox(label="Câu trả lời của ứng viên")
+    response_output = gr.Textbox(label="Phản hồi nhà tuyển dụng")
+
+    send_btn = gr.Button("Gửi")
+
+    send_btn.click(chat, inputs=[user_input, chat_history], outputs=[response_output, chat_history])
 
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=5000, debug=True)
+demo.launch()