Update app.py

app.py

@@ -10,27 +10,10 @@ from huggingface_hub import login, snapshot_download
 import gradio as gr
 
 # -------------------------------
-# 0. 向量資料庫載入
-# -------------------------------
-EMBEDDINGS_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
-embeddings_model = HuggingFaceEmbeddings(model_name=EMBEDDINGS_MODEL_NAME)
-
-DB_PATH = "./faiss_db"
-if os.path.exists(DB_PATH):
-    print("✅ 載入現有向量資料庫...")
-    db = FAISS.load_local(DB_PATH, embeddings_model, allow_dangerous_deserialization=True)
-else:
-    raise ValueError("❌ 沒找到 faiss_db，請先建立向量資料庫")
-
-retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 5})
-
-# -------------------------------
-# 1. 中文模型（T5 Pegasus）
+# 1. 模型設定（專門中文 T5 Pegasus）
 # -------------------------------
 MODEL_NAME = "imxly/t5-pegasus-small"
 
-
-
 HF_TOKEN = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
 if HF_TOKEN:
     login(token=HF_TOKEN)
@@ -41,7 +24,15 @@ 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)
 
-tokenizer = AutoTokenizer.from_pretrained(LOCAL_MODEL_DIR)
+print(f"👉 最終使用模型：{MODEL_NAME}")
+
+# -------------------------------
+# 2. pipeline 載入
+# -------------------------------
+tokenizer = AutoTokenizer.from_pretrained(
+    LOCAL_MODEL_DIR,
+    use_fast=False  # ✅ 避免 tiktoken 錯誤
+)
 model = AutoModelForSeq2SeqLM.from_pretrained(LOCAL_MODEL_DIR)
 
 generator = pipeline(
@@ -56,7 +47,7 @@ def call_local_inference(prompt, max_new_tokens=256):
         outputs = generator(
             prompt,
             max_new_tokens=max_new_tokens,
-            do_sample=False,   # 用摘要模型 → 不建議隨機取樣
+            do_sample=True,
             temperature=0.7
         )
         return outputs[0]["generated_text"]
@@ -64,67 +55,86 @@ def call_local_inference(prompt, max_new_tokens=256):
         return f"（生成失敗：{e}）"
 
 # -------------------------------
-# 2. 基於 RAG 的文章生成
+# 3. 建立或載入向量資料庫
+# -------------------------------
+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/paraphrase-multilingual-MiniLM-L12-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("⚠️ 沒有資料庫，建立新向量資料庫...")
+    docs = []
+    txt_files = [f for f in os.listdir(TXT_FOLDER) if f.endswith(".txt")]
+    for filename in txt_files:
+        with open(os.path.join(TXT_FOLDER, filename), "r", encoding="utf-8") as f:
+            docs.append(Document(page_content=f.read(), metadata={"source": filename}))
+    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+    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})
+
+# -------------------------------
+# 4. 文章生成（結合 RAG）
 # -------------------------------
-def generate_article_rag_only(query, segments=3):
+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)
 
     all_text = []
 
-    # 🔍 RAG 檢索
+    # 🔍 使用 RAG 檢索
     retrieved_docs = retriever.get_relevant_documents(query)
     context_texts = [d.page_content for d in retrieved_docs]
-    full_context = "\n".join(context_texts)
+    context = "\n".join([f"{i+1}. {txt}" for i, txt in enumerate(context_texts[:3])])
 
-    # 切分 context，避免太長
-    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-    chunks = splitter.split_text(full_context)
-
-    for i, chunk in enumerate(chunks[:segments]):
-        progress_text = f"⏳ 正在生成第 {i+1}/{segments} 段..."
+    for i in range(segments):
         prompt = (
-            f"以下是唯一可用的參考內容：\n{chunk}\n\n"
-            f"請基於這些內容，寫一段約150-200字的中文文章，"
-            f"主題：{query}。\n"
-            f"⚠️ 僅能使用參考內容，不可加入外部知識。"
+            f"以下是佛教經論的相關段落：\n{context}\n\n"
+            f"請依據上面內容，寫一段約150-200字的中文文章，"
+            f"主題：{query}。\n第{i+1}段："
         )
+
         paragraph = call_local_inference(prompt)
         all_text.append(paragraph)
+        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}"
 
-        yield "\n\n".join(all_text), None, f"本次使用模型：{MODEL_NAME}", full_context, progress_text
-
-    final_progress = f"✅ 已完成全部 {segments} 段生成！"
-    yield "\n\n".join(all_text), docx_file, f"本次使用模型：{MODEL_NAME}", full_context, final_progress
+    doc.save(docx_file)
+    yield "\n\n".join(all_text), docx_file, f"本次使用模型：{MODEL_NAME}"
 
 # -------------------------------
-# 3. Gradio 介面
+# 5. Gradio 介面
 # -------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("# 📺 電視弘法視頻生成文章RAG系統")
-    gr.Markdown("只基於 faiss_db 內容生成中文文章。")
+    gr.Markdown("# 📺 電視弘法視頻生成文章 RAG 系統")
+    gr.Markdown("使用 FAISS + 中文 T5 模型，根據資料庫生成中文文章。")
 
     query_input = gr.Textbox(lines=2, placeholder="請輸入文章主題", label="文章主題")
-    segments_input = gr.Slider(minimum=1, maximum=10, step=1, value=3, 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)
+    output_model = gr.Textbox(label="使用的模型")
 
     btn = gr.Button("生成文章")
     btn.click(
-        generate_article_rag_only,
+        generate_article_progress,
         inputs=[query_input, segments_input],
-        outputs=[output_text, output_file, model_used_text, context_text, progress_text]
+        outputs=[output_text, output_file, output_model]
     )
 
+# -------------------------------
+# 6. 啟動
+# -------------------------------
 if __name__ == "__main__":
     demo.launch()