Update app.py

app.py

@@ -1,45 +1,17 @@
-from flask import Flask, request, jsonify
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from flask_cors import CORS
+import gradio as gr
+from transformers import pipeline
 
-app = Flask(__name__)
-CORS(app)
+# Load chatbot pipeline từ Hugging Face
+chatbot = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
 
-# Load model TinyLlama tối ưu cho CPU
-MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float32, device_map="cpu")
+# Hàm xử lý chat
+def chat_with_ai(user_message):
+    prompt = f"User: {user_message}\nAssistant:"
+    response = chatbot(prompt, max_length=200, do_sample=True, temperature=0.7)
+    return response[0]["generated_text"].split("Assistant:")[-1].strip()
 
-# Hàm xử lý chat không lưu lịch sử
-def get_answer(user_text):
-    prompt = f"User: {user_text}\nAssistant:"
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
+# Giao diện Gradio
+iface = gr.Interface(fn=chat_with_ai, inputs="text", outputs="text", title="TinyLlama Chatbot")
 
-    with torch.no_grad():  # Giảm tải tính toán
-        output = model.generate(input_ids, max_length=150, do_sample=True, temperature=0.7)
-
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    return response.replace(prompt, "").strip()
-
-@app.route("/")
-def home():
-    return "TinyLlama Chatbot API đang chạy!"
-
-@app.route("/chat", methods=["POST"])
-def chat():
-    try:
-        data = request.json
-        user_message = data.get("message", "").strip()
-
-        if not user_message:
-            return jsonify({"error": "Tin nhắn không được để trống!"}), 400
-
-        bot_reply = get_answer(user_message)
-        return jsonify({"FROM": "AI Assistant", "MESSAGE": bot_reply})
-
-    except Exception as e:
-        return jsonify({"error": str(e)}), 500
-
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860, debug=True)
+# Chạy ứng dụng
+iface.launch()