Update app.py

app.py

@@ -1,27 +1,51 @@
-from flask import Flask, request, jsonify, send_from_directory
+from flask import Flask, render_template, request, jsonify, session
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-from transformers import pipeline
+from datetime import datetime
 
-# Khởi tạo ứng dụng Flask
 app = Flask(__name__)
+app.secret_key = "supersecretkey"  # Dùng để quản lý session
 
-# Khởi tạo mô hình
-model_id = "Arrcttacsrks/Llama3.2-3B-Instruct"
-pipe = pipeline("text-generation", model=model_id, torch_dtype=torch.bfloat16)
+# Tải mô hình và tokenizer từ Hugging Face
+model_path = "phamhai/Llama-3.2-3B-Instruct-Frog"
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = AutoModelForCausalLM.from_pretrained(model_path)
 
 @app.route('/')
-def home():
-    return send_from_directory('.', 'index.html')
-
-@app.route('/style.css')
-def style():
-    return send_from_directory('.', 'style.css')
+def index():
+    session.setdefault('history', [])  # Tạo session nếu chưa có
+    return render_template('index.html', history=session['history'])
 
 @app.route('/chat', methods=['POST'])
 def chat():
-    user_input = request.json.get('message')
-    response = pipe(user_input, max_new_tokens=256)
-    return jsonify({'response': response[0]['generated_text']})
+    user_input = request.json.get("message")
+    if not user_input:
+        return jsonify({"response": "Xin hãy nhập tin nhắn!"})
+
+    # Tokenize input và sinh output
+    messages = [
+        {"role": "system", "content": "Bạn là trợ lý của tôi, tên là Vivi."},
+        {"role": "user", "content": user_input}
+    ]
+    tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
+    outputs = model.generate(tokenized_chat, max_new_tokens=128)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # Lưu lại vào lịch sử chat
+    chat_entry = {
+        "user": user_input,
+        "bot": response,
+        "timestamp": datetime.now().strftime("%H:%M:%S")  # Thêm thời gian
+    }
+    session['history'].append(chat_entry)
+    session.modified = True  # Cập nhật session
+
+    return jsonify({"response": response, "timestamp": chat_entry["timestamp"]})
+
+@app.route('/clear', methods=['POST'])
+def clear():
+    session['history'] = []  # Xóa lịch sử chat
+    return jsonify({"status": "ok"})
 
-if __name__ == '__main__':
-    app.run()
+if __name__ == "__main__":
+    app.run(debug=True)