Update app.py

app.py

@@ -1,3 +1,4 @@
+# app.py
 import os, torch
 from langchain.docstore.document import Document
 from langchain.text_splitter import RecursiveCharacterTextSplitter
@@ -9,20 +10,34 @@ from huggingface_hub import login, snapshot_download
 import gradio as gr
 
 # -------------------------------
-# 1. 模型設定（中文 T5）
+# 1. 模型設定（中文 T5 / Pegasus）
 # -------------------------------
-MODEL_NAME = "Langboat/mengzi-t5-base"   # ✅ 換成穩定的中文 T5
+PRIMARY_MODEL = "imxly/t5-pegasus-small"   # 適合中文摘要/生成
+FALLBACK_MODEL = "uer/gpt2-chinese-cluecorpussmall"  # 若 T5 無法下載就 fallback GPT2
 
 HF_TOKEN = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
 if HF_TOKEN:
     login(token=HF_TOKEN)
     print("✅ 已使用 HUGGINGFACEHUB_API_TOKEN 登入 Hugging Face")
 
-# 嘗試下載模型
-LOCAL_MODEL_DIR = f"./models/{MODEL_NAME.split('/')[-1]}"
-if not os.path.exists(LOCAL_MODEL_DIR):
-    print(f"⬇️ 嘗試下載模型 {MODEL_NAME} ...")
-    snapshot_download(repo_id=MODEL_NAME, token=HF_TOKEN, local_dir=LOCAL_MODEL_DIR)
+def try_download_model(repo_id):
+    local_dir = f"./models/{repo_id.split('/')[-1]}"
+    if not os.path.exists(local_dir):
+        print(f"⬇️ 嘗試下載模型 {repo_id} ...")
+        try:
+            snapshot_download(repo_id=repo_id, token=HF_TOKEN, local_dir=local_dir)
+        except Exception as e:
+            print(f"⚠️ 模型 {repo_id} 無法下載: {e}")
+            return None
+    return local_dir
+
+LOCAL_MODEL_DIR = try_download_model(PRIMARY_MODEL)
+if LOCAL_MODEL_DIR is None:
+    print("⚠️ 切換到 fallback 模型：小型 GPT2-Chinese")
+    LOCAL_MODEL_DIR = try_download_model(FALLBACK_MODEL)
+    MODEL_NAME = FALLBACK_MODEL
+else:
+    MODEL_NAME = PRIMARY_MODEL
 
 print(f"👉 最終使用模型：{MODEL_NAME}")
 
@@ -31,22 +46,28 @@ print(f"👉 最終使用模型：{MODEL_NAME}")
 # -------------------------------
 tokenizer = AutoTokenizer.from_pretrained(
     LOCAL_MODEL_DIR,
-    use_fast=False   # ✅ 避免 tiktoken / fast tokenizer 問題
+    use_fast=False  # 防止 sentencepiece 問題
 )
-model = AutoModelForSeq2SeqLM.from_pretrained(LOCAL_MODEL_DIR)
+
+# 判斷 GPU (CL3) 或 CPU
+device = 0 if torch.cuda.is_available() else -1
+print(f"💻 使用裝置：{'GPU' if device == 0 else 'CPU'}")
+
+try:
+    model = AutoModelForSeq2SeqLM.from_pretrained(LOCAL_MODEL_DIR)
+except:
+    from transformers import AutoModelForCausalLM
+    model = AutoModelForCausalLM.from_pretrained(LOCAL_MODEL_DIR)
 
 generator = pipeline(
-    "text2text-generation",  # ✅ Seq2Seq 用這個
+    "text2text-generation" if "t5" in MODEL_NAME or "pegasus" in MODEL_NAME else "text-generation",
     model=model,
     tokenizer=tokenizer,
-    device=-1  # CPU
+    device=device
 )
 
 def call_local_inference(prompt, max_new_tokens=256):
     try:
-        if "中文" not in prompt:
-            prompt += "\n（請用中文回答）"
-
         outputs = generator(
             prompt,
             max_new_tokens=max_new_tokens,
@@ -58,7 +79,7 @@ def call_local_inference(prompt, max_new_tokens=256):
         return f"（生成失敗：{e}）"
 
 # -------------------------------
-# 3. RAG 部分：向量資料庫
+# 3. FAISS 向量資料庫載入
 # -------------------------------
 DB_PATH = "./faiss_db"
 EMBEDDINGS_MODEL_NAME = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
@@ -68,10 +89,10 @@ 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("⚠️ 沒有找到資料庫，請先建立 faiss_db")
-    db = None
+    print("⚠️ 找不到向量資料庫，將建立空的 DB")
+    db = FAISS.from_documents([], embeddings_model)
 
-retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 3}) if db else None
+retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 5})
 
 # -------------------------------
 # 4. 文章生成（結合 RAG）
@@ -83,12 +104,9 @@ def generate_article_progress(query, segments=3):
 
     all_text = []
 
-    # 🔍 從資料庫檢索
-    context = ""
-    if retriever:
-        retrieved_docs = retriever.get_relevant_documents(query)
-        context_texts = [d.page_content for d in retrieved_docs]
-        context = "\n".join([f"{i+1}. {txt}" for i, txt in enumerate(context_texts[:3])])
+    retrieved_docs = retriever.get_relevant_documents(query)
+    context_texts = [d.page_content for d in retrieved_docs]
+    context = "\n".join([f"{i+1}. {txt}" for i, txt in enumerate(context_texts[:3])])
 
     for i in range(segments):
         prompt = (
@@ -100,27 +118,29 @@ def generate_article_progress(query, segments=3):
         all_text.append(paragraph)
         doc.add_paragraph(paragraph)
 
-        yield "\n\n".join(all_text), None, f"本次使用模型：{MODEL_NAME}"
+        yield "\n\n".join(all_text), None, f"本次使用模型：{MODEL_NAME}，裝置：{'GPU' if device == 0 else 'CPU'}"
 
     doc.save(docx_file)
-    yield "\n\n".join(all_text), docx_file, f"本次使用模型：{MODEL_NAME}"
+    yield "\n\n".join(all_text), docx_file, f"本次使用模型：{MODEL_NAME}，裝置：{'GPU' if device == 0 else 'CPU'}"
 
 # -------------------------------
 # 5. Gradio 介面
 # -------------------------------
 with gr.Blocks() as demo:
     gr.Markdown("# 📺 電視弘法視頻生成文章 RAG 系統")
+    gr.Markdown("使用 Hugging Face 本地模型 + FAISS RAG，僅基於資料庫生成文章。")
+
     query_input = gr.Textbox(lines=2, placeholder="請輸入文章主題", label="文章主題")
     segments_input = gr.Slider(minimum=1, maximum=10, step=1, value=3, label="段落數")
     output_text = gr.Textbox(label="生成文章")
     output_file = gr.File(label="下載 DOCX")
-    model_info = gr.Textbox(label="模型資訊")
+    status_info = gr.Label(label="狀態")
 
     btn = gr.Button("生成文章")
     btn.click(
         generate_article_progress,
         inputs=[query_input, segments_input],
-        outputs=[output_text, output_file, model_info]
+        outputs=[output_text, output_file, status_info]
     )
 
 if __name__ == "__main__":