Update app.py

app.py

@@ -1,5 +1,4 @@
-# app.py
-import os
+import os, torch
 from langchain.docstore.document import Document
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import FAISS
@@ -10,15 +9,16 @@ from huggingface_hub import login, snapshot_download
 import gradio as gr
 
 # -------------------------------
-# 1. 模型設定（專門中文 T5 Pegasus）
+# 1. 模型設定（中文 T5）
 # -------------------------------
-MODEL_NAME = "imxly/t5-pegasus-small"
+MODEL_NAME = "Langboat/mengzi-t5-base"   # ✅ 換成穩定的中文 T5
 
 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} ...")
@@ -31,12 +31,12 @@ print(f"👉 最終使用模型：{MODEL_NAME}")
 # -------------------------------
 tokenizer = AutoTokenizer.from_pretrained(
     LOCAL_MODEL_DIR,
-    use_fast=False  # ✅ 避免 tiktoken 錯誤
+    use_fast=False   # ✅ 避免 tiktoken / fast tokenizer 問題
 )
 model = AutoModelForSeq2SeqLM.from_pretrained(LOCAL_MODEL_DIR)
 
 generator = pipeline(
-    "text2text-generation",
+    "text2text-generation",  # ✅ Seq2Seq 用這個
     model=model,
     tokenizer=tokenizer,
     device=-1  # CPU
@@ -44,6 +44,9 @@ generator = pipeline(
 
 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,
@@ -55,13 +58,9 @@ def call_local_inference(prompt, max_new_tokens=256):
         return f"（生成失敗：{e}）"
 
 # -------------------------------
-# 3. 建立或載入向量資料庫
+# 3. RAG 部分：向量資料庫
 # -------------------------------
-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)
 
@@ -69,41 +68,34 @@ 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})
+    print("⚠️ 沒有找到資料庫，請先建立 faiss_db")
+    db = None
+
+retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 3}) if db else None
 
 # -------------------------------
 # 4. 文章生成（結合 RAG）
 # -------------------------------
-def generate_article_progress(query, segments=5):
+def generate_article_progress(query, segments=3):
     docx_file = "/tmp/generated_article.docx"
     doc = DocxDocument()
     doc.add_heading(query, level=1)
 
     all_text = []
 
-    # 🔍 使用 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])])
+    # 🔍 從資料庫檢索
+    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])])
 
     for i in range(segments):
         prompt = (
-            f"以下是佛教經論的相關段落：\n{context}\n\n"
+            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)
@@ -118,23 +110,18 @@ def generate_article_progress(query, segments=5):
 # -------------------------------
 with gr.Blocks() as demo:
     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=5, 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")
-    output_model = gr.Textbox(label="使用的模型")
+    model_info = gr.Textbox(label="模型資訊")
 
     btn = gr.Button("生成文章")
     btn.click(
         generate_article_progress,
         inputs=[query_input, segments_input],
-        outputs=[output_text, output_file, output_model]
+        outputs=[output_text, output_file, model_info]
     )
 
-# -------------------------------
-# 6. 啟動
-# -------------------------------
 if __name__ == "__main__":
     demo.launch()