Update ai_service.py

ai_service.py

@@ -1,10 +1,12 @@
 # ai_service.py
+import torch
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 from config import (
-    HUGGING_FACE_TOKEN, LLM_MODEL, LLM_MAX_NEW_TOKENS,
+    LLM_MODEL, LLM_MAX_NEW_TOKENS,
     LLM_TOP_K, LLM_TEMPERATURE
 )
 
-# 用於延遲載入語言模型的字典，現在只儲存 pipeline 物件
+# 用於延遲載入語言模型的字典
 _LLM = {"loaded": False, "ok": False, "err": None, "model": None}
 
 def _ensure_llm():
@@ -13,52 +15,38 @@ def _ensure_llm():
         return _LLM["ok"], _LLM["err"]
     _LLM["loaded"] = True
 
-    # 檢查 HF Token 是否存在
-    if not HUGGING_FACE_TOKEN:
-        _LLM["err"] = "HUGGING_FACE_TOKEN secret not set in deployment environment."
-        _LLM["ok"] = False
-        return False, _LLM["err"]
-
     try:
-        import torch
-        from transformers import pipeline
-        from huggingface_hub import login
+        # [修改] 採用更穩健的載入方式
+        # 1. 決定裝置 (GPU or CPU)
+        device = "cuda" if torch.cuda.is_available() else "cpu"
 
-        # 使用 Token 登入
-        login(token=HUGGING_FACE_TOKEN)
+        # 2. 分別載入 tokenizer 和 model
+        tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
+        model = AutoModelForCausalLM.from_pretrained(LLM_MODEL).to(device)
 
-        # 建立 text-generation pipeline
-        # device_map="auto" 會自動使用 GPU (如果可用)
+        # 3. 將載入好的 tokenizer 和 model 傳入 pipeline
         pipe = pipeline(
             "text-generation",
-            model=LLM_MODEL,
-            torch_dtype=torch.bfloat16,
-            device_map="auto",
+            model=model,
+            tokenizer=tokenizer,
+            device=device
         )
+        
         _LLM.update({"ok": True, "model": pipe})
         return True, None
     except Exception as e:
-        # 如果 Token 無效或未同意模型條款，會在此處拋出錯誤
         _LLM["err"] = f"{e}"
         _LLM["ok"] = False
         return False, _LLM["err"]
 
 def generate_ai_text(user_prompt: str) -> str:
-    """使用已載入的 Gemma 模型生成文字回應。"""
+    """使用已載入的 AI 模型生成文字回應。"""
     ok, err = _ensure_llm()
     if not ok:
-        return (
-            "🤖 AI 模型無法使用。\n"
-            "可能原因：\n"
-            "1. 未在 Hugging Face Spaces 設定名為 HUGGING_FACE_TOKEN 的 Secret。\n"
-            "2. 尚未在 Hugging Face 網站上同意 gemma-2b-it 模型的使用條款。\n"
-            f"\n詳細錯誤：{err}"
-        )
+        return f"🤖 AI 模型無法使用。\n詳細錯誤：{err}"
 
     pipe = _LLM["model"]
-    
-    # 建立符合 Gemma 指令微調模型的 Prompt 格式
-    prompt = f"<start_of_turn>user\n你是一個多功能的台灣在地LINE助理，請用繁體中文簡潔有力地回答問題。{user_prompt}<end_of_turn>\n<start_of_turn>model\n"
+    prompt = user_prompt
 
     try:
         outputs = pipe(
@@ -67,11 +55,14 @@ def generate_ai_text(user_prompt: str) -> str:
             do_sample=True,
             temperature=LLM_TEMPERATURE,
             top_k=LLM_TOP_K,
-            top_p=0.95,
         )
         # 從 pipeline 的輸出中解析出模型生成的部分
-        full_text = outputs[0]["generated_text"]
-        response = full_text.split("<start_of_turn>model\n")[-1]
+        response = outputs[0]["generated_text"]
+        
+        # 移除原始 prompt 以獲得乾淨的回應
+        if prompt in response:
+            response = response.split(prompt, 1)[-1]
+            
         return response.strip() or "（AI 沒有產生任何內容）"
     except Exception as e:
         return f"AI 產生內容時發生錯誤：{e}"