Update app.py

app.py

@@ -20,17 +20,17 @@ if HF_TOKEN:
     login(token=HF_TOKEN)
     print("✅ 已使用 HUGGINGFACEHUB_API_TOKEN 登入 Hugging Face")
 else:
-    print("⚠️ 沒有 HUGGINGFACEHUB_API_TOKEN，Gemma-7B 可能無法下載")
+    print("⚠️ 沒有 HUGGINGFACEHUB_API_TOKEN，部分 gated 模型可能無法下載")
 
 # -------------------------------
-# 3. 模型清單
+# 3. 模型清單（CPU 免費可跑）
 # -------------------------------
 MODEL_MAP = {
     "Auto": None,
     "Gemma-2B": "google/gemma-2b",
-    "Gemma-7B": "google/gemma-7b",   # gated
-    "BTLM-3B-8K": "cerebras/btlm-3b-8k",
-    "Mistral-7B": "mistralai/Mistral-7B-v0.1"
+    "BTLM-3B-8K": "tiiuae/btlm-3b-8k-base",
+    "DistilGPT2": "distilgpt2",
+    "BART-Base": "facebook/bart-base"
 }
 
 # -------------------------------
@@ -38,7 +38,7 @@ MODEL_MAP = {
 # -------------------------------
 LOCAL_MODEL_DIRS = {}
 for name, repo in MODEL_MAP.items():
-    if repo is None:  # Auto 跳過
+    if repo is None:
         continue
     try:
         local_dir = f"./models/{repo.split('/')[-1]}"
@@ -56,7 +56,26 @@ for name, repo in MODEL_MAP.items():
         print(f"⚠️ 模型 {repo} 無法下載: {e}")
 
 # -------------------------------
-# 5. 建立或載入向量資料庫
+# 5. 模型可用性檢查
+# -------------------------------
+def test_models():
+    print("\n🔍 啟動模型檢查：")
+    for name, local_dir in LOCAL_MODEL_DIRS.items():
+        try:
+            _ = pipeline(
+                "text-generation",
+                model=local_dir,
+                tokenizer=local_dir,
+                device_map="cpu"
+            )
+            print(f"✅ 模型 {name} ({local_dir}) 可用")
+        except Exception as e:
+            print(f"❌ 模型 {name} ({local_dir}) 無法載入: {e}")
+
+test_models()
+
+# -------------------------------
+# 6. 建立或載入向量資料庫
 # -------------------------------
 TXT_FOLDER = "./out_texts"
 DB_PATH = "./faiss_db"
@@ -84,7 +103,7 @@ else:
 retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 5})
 
 # -------------------------------
-# 6. 本地 pipeline
+# 7. 本地 pipeline
 # -------------------------------
 _loaded_pipelines = {}
 
@@ -98,7 +117,7 @@ def get_pipeline(model_name):
             "text-generation",
             model=local_path,
             tokenizer=local_path,
-            device_map="auto"
+            device_map="cpu"
         )
         _loaded_pipelines[model_name] = generator
     return _loaded_pipelines[model_name]
@@ -112,7 +131,7 @@ def call_local_inference(model_name, prompt, max_new_tokens=512):
         return f"（生成失敗：{e}）"
 
 # -------------------------------
-# 7. Auto 模式邏輯
+# 8. Auto 模式邏輯
 # -------------------------------
 def pick_model_auto(segments):
     if segments <= 3:
@@ -120,7 +139,7 @@ def pick_model_auto(segments):
     elif segments <= 6:
         return "BTLM-3B-8K"
     else:
-        return "Mistral-7B"
+        return "BART-Base"
 
 def generate_article_progress(query, model_name, segments=5):
     docx_file = "/tmp/generated_article.docx"
@@ -148,11 +167,11 @@ def generate_article_progress(query, model_name, segments=5):
     yield "\n\n".join(all_text), docx_file, f"本次使用模型：{selected_model}"
 
 # -------------------------------
-# 8. Gradio 介面
+# 9. Gradio 介面
 # -------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("# 佛教經論 RAG 系統 (本地模型)")
-    gr.Markdown("支援 Gemma / BTLM / Mistral，Auto 模式會自動選擇模型。")
+    gr.Markdown("# 佛教經論 RAG 系統 (CPU 免費版)")
+    gr.Markdown("支援 Gemma-2B / BTLM-3B / DistilGPT2 / BART-Base，Auto 模式會自動選擇。")
 
     query_input = gr.Textbox(lines=2, placeholder="請輸入文章主題", label="文章主題")
     model_dropdown = gr.Dropdown(
@@ -173,7 +192,7 @@ with gr.Blocks() as demo:
     )
 
 # -------------------------------
-# 9. 啟動 Gradio
+# 10. 啟動 Gradio
 # -------------------------------
 if __name__ == "__main__":
     demo.launch()