Update app.py

app.py

@@ -9,7 +9,7 @@ from langchain_community.vectorstores import FAISS
 from langchain_huggingface import HuggingFaceEmbeddings
 from docx import Document as DocxDocument
 from transformers import pipeline
-from huggingface_hub import login
+from huggingface_hub import login, snapshot_download
 import gradio as gr
 
 # -------------------------------
@@ -20,16 +20,49 @@ if HF_TOKEN:
     login(token=HF_TOKEN)
     print("✅ 已使用 HUGGINGFACEHUB_API_TOKEN 登入 Hugging Face")
 else:
-    print("⚠️ 沒有 HUGGINGFACEHUB_API_TOKEN，Gemma-7B 等 gated 模型可能無法使用")
+    print("⚠️ 沒有 HUGGINGFACEHUB_API_TOKEN，Gemma-7B 可能無法下載")
 
+# -------------------------------
+# 3. 模型清單
+# -------------------------------
+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"
+}
+
+# -------------------------------
+# 4. 預先下載模型到本地 ./models/
+# -------------------------------
+LOCAL_MODEL_DIRS = {}
+for name, repo in MODEL_MAP.items():
+    if repo is None:  # Auto 跳過
+        continue
+    try:
+        local_dir = f"./models/{repo.split('/')[-1]}"
+        if not os.path.exists(local_dir):
+            print(f"⬇️ 正在下載模型 {repo} ...")
+            snapshot_download(
+                repo_id=repo,
+                token=HF_TOKEN,
+                local_dir=local_dir
+            )
+        else:
+            print(f"✅ 已存在模型 {repo} -> {local_dir}")
+        LOCAL_MODEL_DIRS[name] = local_dir
+    except Exception as e:
+        print(f"⚠️ 模型 {repo} 無法下載: {e}")
+
+# -------------------------------
+# 5. 建立或載入向量資料庫
+# -------------------------------
 TXT_FOLDER = "./out_texts"
 DB_PATH = "./faiss_db"
 os.makedirs(DB_PATH, exist_ok=True)
 os.makedirs(TXT_FOLDER, exist_ok=True)
 
-# -------------------------------
-# 3. 建立或載入向量資料庫
-# -------------------------------
 EMBEDDINGS_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
 embeddings_model = HuggingFaceEmbeddings(model_name=EMBEDDINGS_MODEL_NAME)
 
@@ -51,28 +84,21 @@ else:
 retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 5})
 
 # -------------------------------
-# 4. 本地推論模型設定
+# 6. 本地 pipeline
 # -------------------------------
-MODEL_MAP = {
-    "Auto": None,
-    "Gemma-2B": "google/gemma-2b",
-    "Gemma-7B": "google/gemma-7b",   # gated，需要 HF_TOKEN
-    "BTLM-3B-8K": "cerebras/btlm-3b-8k",
-    "Mistral-7B": "mistralai/Mistral-7B-v0.1"
-}
-
 _loaded_pipelines = {}
 
 def get_pipeline(model_name):
     if model_name not in _loaded_pipelines:
-        print(f"🔄 正在載入模型 {model_name} ...")
-        model_id = MODEL_MAP[model_name]
+        local_path = LOCAL_MODEL_DIRS.get(model_name)
+        if not local_path:
+            raise ValueError(f"❌ 模型 {model_name} 尚未下載")
+        print(f"🔄 正在載入本地模型 {model_name} from {local_path}")
         generator = pipeline(
             "text-generation",
-            model=model_id,
-            tokenizer=model_id,
-            device_map="auto",
-            token=HF_TOKEN  # gated 模型會用這個
+            model=local_path,
+            tokenizer=local_path,
+            device_map="auto"
         )
         _loaded_pipelines[model_name] = generator
     return _loaded_pipelines[model_name]
@@ -86,23 +112,21 @@ def call_local_inference(model_name, prompt, max_new_tokens=512):
         return f"（生成失敗：{e}）"
 
 # -------------------------------
-# 5. 自動選模型邏輯
+# 7. Auto 模式邏輯
 # -------------------------------
 def pick_model_auto(segments):
-    """根據段落數自動挑選模型"""
     if segments <= 3:
         return "Gemma-2B"
     elif segments <= 6:
         return "BTLM-3B-8K"
     else:
-        return "Mistral-7B"  # 避免 gpt-oss-20B 太大跑不動
+        return "Mistral-7B"
 
 def generate_article_progress(query, model_name, segments=5):
     docx_file = "/tmp/generated_article.docx"
     doc = DocxDocument()
     doc.add_heading(query, level=1)
 
-    # 自動挑模型
     if model_name == "Auto":
         selected_model = pick_model_auto(int(segments))
     else:
@@ -124,11 +148,11 @@ def generate_article_progress(query, model_name, segments=5):
     yield "\n\n".join(all_text), docx_file, f"本次使用模型：{selected_model}"
 
 # -------------------------------
-# 6. Gradio 介面
+# 8. Gradio 介面
 # -------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("# 佛教經論 RAG 系統 (Gemma / BTLM / Mistral)")
-    gr.Markdown("支援 Auto 模式，並顯示實際使用的模型名稱。")
+    gr.Markdown("# 佛教經論 RAG 系統 (本地模型)")
+    gr.Markdown("支援 Gemma / BTLM / Mistral，Auto 模式會自動選擇模型。")
 
     query_input = gr.Textbox(lines=2, placeholder="請輸入文章主題", label="文章主題")
     model_dropdown = gr.Dropdown(
@@ -149,7 +173,7 @@ with gr.Blocks() as demo:
     )
 
 # -------------------------------
-# 7. 啟動 Gradio
+# 9. 啟動 Gradio
 # -------------------------------
 if __name__ == "__main__":
     demo.launch()