Update app.py

app.py

@@ -1,16 +1,16 @@
 # app.py
-import os, torch
+import os
 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 transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from huggingface_hub import login, snapshot_download
 import gradio as gr
 
 # -------------------------------
-# 0. 載入向量資料庫
+# 0. 向量資料庫載入
 # -------------------------------
 EMBEDDINGS_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
 embeddings_model = HuggingFaceEmbeddings(model_name=EMBEDDINGS_MODEL_NAME)
@@ -25,48 +25,25 @@ else:
 retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 5})
 
 # -------------------------------
-# 1. 模型設定（中文 GPT2 + fallback）
+# 1. 中文模型（Randeng-T5）
 # -------------------------------
-PRIMARY_MODEL = "uer/gpt2-chinese-cluecorpusmedium"
-FALLBACK_MODEL = "uer/gpt2-chinese-cluecorpussmall"
+MODEL_NAME = "IDEA-CCNL/Randeng-T5-784M-Summary-Chinese"
 
 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}")
+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)
 
-# -------------------------------
-# 2. pipeline 載入
-# -------------------------------
 tokenizer = AutoTokenizer.from_pretrained(LOCAL_MODEL_DIR)
-model = AutoModelForCausalLM.from_pretrained(LOCAL_MODEL_DIR)
-
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
+model = AutoModelForSeq2SeqLM.from_pretrained(LOCAL_MODEL_DIR)
 
 generator = pipeline(
-    "text-generation",
+    "text2text-generation",
     model=model,
     tokenizer=tokenizer,
     device=-1  # CPU
@@ -74,21 +51,18 @@ 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,
-            do_sample=True,
-            temperature=0.7,
-            pad_token_id=tokenizer.pad_token_id
+            do_sample=False,   # 用摘要模型 → 不建議隨機取樣
+            temperature=0.7
         )
         return outputs[0]["generated_text"]
     except Exception as e:
         return f"（生成失敗：{e}）"
 
 # -------------------------------
-# 3. 僅基於 RAG 的文章生成
+# 2. 基於 RAG 的文章生成
 # -------------------------------
 def generate_article_rag_only(query, segments=3):
     docx_file = "/tmp/generated_article.docx"
@@ -103,8 +77,8 @@ def generate_article_rag_only(query, segments=3):
     context_texts = [d.page_content for d in retrieved_docs]
     full_context = "\n".join(context_texts)
 
-    # 切分成小片段，避免模型超載
-    splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=50)
+    # 切分 context，避免太長
+    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
     chunks = splitter.split_text(full_context)
 
     for i, chunk in enumerate(chunks[:segments]):
@@ -113,7 +87,7 @@ def generate_article_rag_only(query, segments=3):
             f"以下是唯一可用的參考內容：\n{chunk}\n\n"
             f"請基於這些內容，寫一段約150-200字的中文文章，"
             f"主題：{query}。\n"
-            f"⚠️ 僅能使用參考內容，不可加入其他知識。"
+            f"⚠️ 僅能使用參考內容，不可加入外部知識。"
         )
         paragraph = call_local_inference(prompt)
         all_text.append(paragraph)
@@ -129,11 +103,11 @@ def generate_article_rag_only(query, segments=3):
     yield "\n\n".join(all_text), docx_file, f"本次使用模型：{MODEL_NAME}", full_context, final_progress
 
 # -------------------------------
-# 4. Gradio 介面
+# 3. Gradio 介面
 # -------------------------------
 with gr.Blocks() as demo:
     gr.Markdown("# 📺 電視弘法視頻生成文章RAG系統")
-    gr.Markdown("只基於 faiss_db 內容生成文章，不加入外部知識。")
+    gr.Markdown("只基於 faiss_db 內容生成中文文章。")
 
     query_input = gr.Textbox(lines=2, placeholder="請輸入文章主題", label="文章主題")
     segments_input = gr.Slider(minimum=1, maximum=10, step=1, value=3, label="段落數")