Update app.py

app.py

@@ -1,7 +1,4 @@
 # app.py
-# -------------------------------
-# 1. 套件載入
-# -------------------------------
 import os, glob
 from langchain.docstore.document import Document
 from langchain.text_splitter import RecursiveCharacterTextSplitter
@@ -13,28 +10,21 @@ from huggingface_hub import login, snapshot_download
 import gradio as gr
 
 # -------------------------------
-# 2. 環境變數與登入
-# -------------------------------
-HF_TOKEN = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
-if HF_TOKEN:
-    login(token=HF_TOKEN)
-    print("✅ 已使用 HUGGINGFACEHUB_API_TOKEN 登入 Hugging Face")
-else:
-    print("⚠️ 沒有 HUGGINGFACEHUB_API_TOKEN，部分 gated 模型可能無法下載")
-
-# -------------------------------
-# 3. 模型清單（CPU 免費可跑）
+# 1. 模型清單（全部公開）
 # -------------------------------
 MODEL_MAP = {
     "Auto": None,
-    "Gemma-2B": "google/gemma-2b",             # gated，需要 Access repository
-    "BTLM-3B-8K": "cerebras/btlm-3b-8k-base",  # 正確 repo
-    "DistilGPT2": "distilgpt2",                # 小模型
-    "BART-Base": "facebook/bart-base"          # 小模型
+    "BTLM-3B-8K": "cerebras/btlm-3b-8k-base",  # 3B 模型，公開
+    "DistilGPT2": "distilgpt2",                 # 小模型
+    "BART-Base": "facebook/bart-base"           # 小模型
 }
 
+HF_TOKEN = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
+if HF_TOKEN:
+    login(token=HF_TOKEN)
+
 # -------------------------------
-# 4. 預先下載模型到本地 ./models/
+# 2. 預先下載模型到 ./models/
 # -------------------------------
 LOCAL_MODEL_DIRS = {}
 for name, repo in MODEL_MAP.items():
@@ -44,11 +34,7 @@ for name, repo in MODEL_MAP.items():
         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
-            )
+            snapshot_download(repo_id=repo, token=HF_TOKEN, local_dir=local_dir)
         else:
             print(f"✅ 已存在模型 {repo} -> {local_dir}")
         LOCAL_MODEL_DIRS[name] = local_dir
@@ -56,63 +42,13 @@ for name, repo in MODEL_MAP.items():
         print(f"⚠️ 模型 {repo} 無法下載: {e}")
 
 # -------------------------------
-# 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} 可用")
-        except Exception as e:
-            print(f"❌ 模型 {name} 無法載入: {e}")
-
-test_models()
-
-# -------------------------------
-# 6. 建立或載入向量資料庫
-# -------------------------------
-TXT_FOLDER = "./out_texts"
-DB_PATH = "./faiss_db"
-os.makedirs(DB_PATH, exist_ok=True)
-os.makedirs(TXT_FOLDER, exist_ok=True)
-
-EMBEDDINGS_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
-embeddings_model = HuggingFaceEmbeddings(model_name=EMBEDDINGS_MODEL_NAME)
-
-if os.path.exists(os.path.join(DB_PATH, "index.faiss")):
-    print("載入現有向量資料庫...")
-    db = FAISS.load_local(DB_PATH, embeddings_model, allow_dangerous_deserialization=True)
-else:
-    print("沒有資料庫，開始建立新向量資料庫...")
-    txt_files = glob.glob(f"{TXT_FOLDER}/*.txt")
-    docs = []
-    for filepath in txt_files:
-        with open(filepath, "r", encoding="utf-8") as f:
-            docs.append(Document(page_content=f.read(), metadata={"source": os.path.basename(filepath)}))
-    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
-    split_docs = splitter.split_documents(docs)
-    db = FAISS.from_documents(split_docs, embeddings_model)
-    db.save_local(DB_PATH)
-
-retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 5})
-
-# -------------------------------
-# 7. 本地 pipeline
+# 3. pipeline 載入
 # -------------------------------
 _loaded_pipelines = {}
 
 def get_pipeline(model_name):
     if model_name not in _loaded_pipelines:
         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=local_path,
@@ -131,56 +67,46 @@ def call_local_inference(model_name, prompt, max_new_tokens=512):
         return f"（生成失敗：{e}）"
 
 # -------------------------------
-# 8. Auto 模式邏輯
+# 4. Auto 模式邏輯
 # -------------------------------
 def pick_model_auto(segments):
     if segments <= 3:
-        return "DistilGPT2"   # 短文用最小模型，快
+        return "DistilGPT2"
     elif segments <= 6:
-        return "Gemma-2B"     # 中篇用 Gemma-2B
-    elif segments <= 8:
-        return "BTLM-3B-8K"   # 長文用 BTLM
+        return "BTLM-3B-8K"
     else:
-        return "BART-Base"    # 超長用 Bart-base
+        return "BART-Base"
 
 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:
-        selected_model = model_name
+    selected_model = pick_model_auto(segments) if model_name == "Auto" else model_name
     print(f"👉 使用模型: {selected_model}")
 
     all_text = []
     prompt = f"請依據下列主題生成段落：{query}\n\n每段約150-200字。"
 
-    for i in range(int(segments)):
+    for i in range(segments):
         paragraph = call_local_inference(selected_model, prompt)
         all_text.append(paragraph)
         doc.add_paragraph(paragraph)
         prompt = f"請接續上一段生成下一段：\n{paragraph}\n\n下一段："
-
         yield "\n\n".join(all_text), None, f"本次使用模型：{selected_model}"
 
     doc.save(docx_file)
     yield "\n\n".join(all_text), docx_file, f"本次使用模型：{selected_model}"
 
 # -------------------------------
-# 9. Gradio 介面
+# 5. Gradio 介面
 # -------------------------------
 with gr.Blocks() as demo:
     gr.Markdown("# 佛教經論 RAG 系統 (CPU 免費版)")
-    gr.Markdown("支援 Gemma-2B / BTLM-3B / DistilGPT2 / BART-Base，Auto 模式會自動選擇。")
+    gr.Markdown("支援 DistilGPT2 / BTLM-3B / BART-Base，Auto 模式會自動選擇。")
 
     query_input = gr.Textbox(lines=2, placeholder="請輸入文章主題", label="文章主題")
-    model_dropdown = gr.Dropdown(
-        choices=list(MODEL_MAP.keys()),
-        value="Auto",
-        label="選擇生成模型"
-    )
+    model_dropdown = gr.Dropdown(choices=list(MODEL_MAP.keys()), value="Auto", label="選擇生成模型")
     segments_input = gr.Slider(minimum=1, maximum=10, step=1, value=5, label="段落數")
     output_text = gr.Textbox(label="生成文章")
     output_file = gr.File(label="下載 DOCX")
@@ -193,8 +119,5 @@ with gr.Blocks() as demo:
         outputs=[output_text, output_file, model_used_text]
     )
 
-# -------------------------------
-# 10. 啟動 Gradio
-# -------------------------------
 if __name__ == "__main__":
     demo.launch()