Update app.py

app.py

@@ -1,57 +1,62 @@
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, Response
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import requests
 
+# 初始化 Flask
 app = Flask(__name__)
 
-# 模型名稱
+# Telegram Token 與 API URL
+TELEGRAM_TOKEN = "7967078631:AAH9viY8zWZ6mi7krxw1RSz5eycrI9Lce8Q"
+TELEGRAM_URL = f"https://api.telegram.org/bot{TELEGRAM_TOKEN}/sendMessage"
+
+# 載入中文輕量語言模型
 model_name = "souljoy/gpt2-small-chinese-cluecorpussmall"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Telegram 設定
-TELEGRAM_TOKEN = "7967078631:AAH9viY8zWZ6mi7krxw1RSz5eycrI9Lce8Q"
-TELEGRAM_URL = f"https://api.telegram.org/bot{TELEGRAM_TOKEN}/sendMessage"
-
-# 基本對話生成函數
-def generate_reply(text):
-    inputs = tokenizer(text, return_tensors="pt")
+# 建立主回應函數
+def generate_reply(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
     outputs = model.generate(
         **inputs,
         max_new_tokens=40,
         do_sample=True,
         top_k=50,
         top_p=0.95,
-        temperature=0.8
+        temperature=0.8,
+        pad_token_id=tokenizer.eos_token_id
     )
     reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return reply
+    return reply.strip()
 
-# 根目錄
+# ✅ Space首頁顯示
 @app.route("/", methods=["GET"])
 def index():
-    return "✅ Flask server is running."
+    return Response("✅ Flask server is running and ready to receive Telegram messages!", mimetype="text/plain")
 
-# Telegram webhook 端點
+# ✅ Telegram webhook endpoint
 @app.route("/telegram", methods=["POST"])
 def telegram_webhook():
     data = request.get_json()
-    if not data:
-        return jsonify({"error": "no data"}), 400
+    if not data or "message" not in data:
+        return jsonify({"status": "ignored"}), 200
+
+    message = data["message"]
+    chat_id = message["chat"]["id"]
+    user_text = message.get("text", "")
 
-    message = data.get("message", {})
-    chat_id = message.get("chat", {}).get("id")
-    text = message.get("text", "")
+    if not user_text:
+        return jsonify({"status": "no text"}), 200
 
-    if text:
-        reply = generate_reply(text)
-        requests.post(TELEGRAM_URL, json={
-            "chat_id": chat_id,
-            "text": reply
-        })
+    try:
+        reply = generate_reply(user_text)
+    except Exception as e:
+        reply = "⚠️ 系統錯誤，請稍後再試。"
 
-    return jsonify({"ok": True})
+    requests.post(TELEGRAM_URL, json={
+        "chat_id": chat_id,
+        "text": reply
+    })
 
-if __name__ == "__main__":
-    app.run()
+    return jsonify({"status": "ok"}), 200