Update app.py

app.py

@@ -1,3 +1,14 @@
+# app.py
+import os, torch
+from langchain.docstore.document import Document
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import FAISS
+from langchain_huggingface import HuggingFaceEmbeddings
+from docx import Document as DocxDocument
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from huggingface_hub import login, snapshot_download
+import gradio as gr
+
 # -------------------------------
 # 0. 載入向量資料庫
 # -------------------------------
@@ -6,7 +17,7 @@ embeddings_model = HuggingFaceEmbeddings(model_name=EMBEDDINGS_MODEL_NAME)
 
 DB_PATH = "./faiss_db"
 if os.path.exists(DB_PATH):
-    print("載入現有向量資料庫...")
+    print("✅ 載入現有向量資料庫...")
     db = FAISS.load_local(DB_PATH, embeddings_model, allow_dangerous_deserialization=True)
 else:
     raise ValueError("❌ 沒找到 faiss_db，請先建立向量資料庫")
@@ -14,17 +25,81 @@ else:
 retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 5})
 
 # -------------------------------
-# 文章生成（RAG + 檢索片段 + 進度提示 + 即時寫入DOCX）
+# 1. 模型設定（中文 GPT2 + fallback）
+# -------------------------------
+PRIMARY_MODEL = "uer/gpt2-chinese-cluecorpusmedium"
+FALLBACK_MODEL = "uer/gpt2-chinese-cluecorpussmall"
+
+HF_TOKEN = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
+if HF_TOKEN:
+    login(token=HF_TOKEN)
+    print("✅ 已使用 HUGGINGFACEHUB_API_TOKEN 登入 Hugging Face")
+
+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}")
+
+# -------------------------------
+# 2. pipeline 載入
+# -------------------------------
+tokenizer = AutoTokenizer.from_pretrained(LOCAL_MODEL_DIR)
+model = AutoModelForCausalLM.from_pretrained(LOCAL_MODEL_DIR)
+
+# 修正 pad_token 缺失問題
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+generator = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device=-1  # CPU
+)
+
+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,
+            do_sample=True,
+            temperature=0.7,
+            pad_token_id=tokenizer.pad_token_id
+        )
+        return outputs[0]["generated_text"]
+    except Exception as e:
+        return f"（生成失敗：{e}）"
+
+# -------------------------------
+# 3. 文章生成（RAG + 即時寫入DOCX + 檢索片段 + 進度提示）
 # -------------------------------
 def generate_article_progress(query, segments=5):
     docx_file = "/tmp/generated_article.docx"
     doc = DocxDocument()
     doc.add_heading(query, level=1)
-    doc.save(docx_file)  # 先建立空的 docx，避免後面保存出錯
+    doc.save(docx_file)  # 先建立空的檔案
 
     all_text = []
 
-    # 🔍 使用 RAG 從 FAISS 檢索相關文獻
+    # 🔍 RAG 檢索佛經段落
     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])])
@@ -39,9 +114,9 @@ def generate_article_progress(query, segments=5):
         paragraph = call_local_inference(prompt)
         all_text.append(paragraph)
 
-        # ✅ 每段生成後立即寫入 DOCX
-        doc = DocxDocument(docx_file)   # 重新打開現有的檔案
-        doc.add_paragraph(paragraph)
+        # ✅ 每段生成後即時寫入 DOCX
+        doc = DocxDocument(docx_file)
+        doc.add_paragraph(f"第{i+1}段：\n{paragraph}")
         doc.save(docx_file)
 
         yield "\n\n".join(all_text), None, f"本次使用模型：{MODEL_NAME}", context, progress_text
@@ -49,25 +124,26 @@ def generate_article_progress(query, segments=5):
     final_progress = f"✅ 已完成全部 {segments} 段生成！"
     yield "\n\n".join(all_text), docx_file, f"本次使用模型：{MODEL_NAME}", context, final_progress
 
-
 # -------------------------------
-# Gradio 介面
+# 4. Gradio 介面
 # -------------------------------
 with gr.Blocks() as demo:
     gr.Markdown("# 📺 電視弘法視頻生成文章RAG系統")
-    gr.Markdown("使用 GPT2-Chinese + FAISS RAG，生成文章。")
+    gr.Markdown("使用 GPT2-Chinese + FAISS RAG，生成佛教主題文章。")
 
     query_input = gr.Textbox(lines=2, placeholder="請輸入文章主題", 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")
     model_used_text = gr.Textbox(label="實際使用模型", interactive=False)
+    context_text = gr.Textbox(label="檢索到的佛經片段", interactive=False, lines=6)
+    progress_text = gr.Textbox(label="生成進度", interactive=False)
 
     btn = gr.Button("生成文章")
     btn.click(
         generate_article_progress,
         inputs=[query_input, segments_input],
-        outputs=[output_text, output_file, model_used_text]
+        outputs=[output_text, output_file, model_used_text, context_text, progress_text]
     )
 
 if __name__ == "__main__":