Update app.py

app.py

@@ -10,12 +10,12 @@ from huggingface_hub import login, snapshot_download
 import gradio as gr
 
 # -------------------------------
-# 1. 模型清單
+# 1. 模型清單（中文 & 英文）
 # -------------------------------
 MODEL_MAP = {
     "Auto": None,
     "BTLM-3B-8K": "cerebras/btlm-3b-8k-base",
-    "DistilGPT2": "distilgpt2",
+    "GPT2-Chinese": "uer/gpt2-chinese-cluecorpusmedium",  # 中文 GPT2
     "BART-Base": "facebook/bart-base"
 }
 
@@ -29,7 +29,7 @@ if HF_TOKEN:
 # -------------------------------
 LOCAL_MODEL_DIRS = {}
 for name, repo in MODEL_MAP.items():
-    if repo is None: 
+    if repo is None:
         continue
     try:
         local_dir = f"./models/{repo.split('/')[-1]}"
@@ -65,7 +65,7 @@ def get_pipeline(model_name):
             "text-generation",
             model=model,
             tokenizer=tokenizer,
-            device=-1  # CPU 強制
+            device=-1  # CPU
         )
         _loaded_pipelines[model_name] = generator
     return _loaded_pipelines[model_name]
@@ -73,6 +73,11 @@ def get_pipeline(model_name):
 def call_local_inference(model_name, prompt, max_new_tokens=256):
     try:
         generator = get_pipeline(model_name)
+
+        # ✅ 強制中文模式：補上提示
+        if "中文" not in prompt and "Chinese" not in prompt:
+            prompt += "\n（請用中文回答）"
+
         outputs = generator(
             prompt,
             max_new_tokens=max_new_tokens,
@@ -89,7 +94,7 @@ def call_local_inference(model_name, prompt, max_new_tokens=256):
 # -------------------------------
 def pick_model_auto(segments):
     if segments <= 3:
-        return "DistilGPT2"
+        return "GPT2-Chinese"  # 短文 → 中文 GPT2
     elif segments <= 6:
         return "BTLM-3B-8K"
     else:
@@ -104,13 +109,15 @@ def generate_article_progress(query, model_name, segments=5):
     print(f"👉 使用模型: {selected_model}")
 
     all_text = []
-    prompt = f"請依據下列主題生成段落：{query}\n\n每段約150-200字。"
+    base_prompt = f"請依據下列主題生成一篇中文文章，主題：{query}\n每段約150-200字。\n"
 
     for i in range(segments):
+        # ✅ 每段獨立生成
+        prompt = base_prompt + f"第{i+1}段："
         paragraph = call_local_inference(selected_model, prompt)
         all_text.append(paragraph)
         doc.add_paragraph(paragraph)
-        prompt = f"請接續上一段生成下一段：\n{paragraph}\n\n下一段："
+
         yield "\n\n".join(all_text), None, f"本次使用模型：{selected_model}"
 
     doc.save(docx_file)
@@ -120,8 +127,8 @@ def generate_article_progress(query, model_name, segments=5):
 # 5. Gradio 介面
 # -------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("# 📺 電視弘法視頻生成文章RAG系統(CPU免費版))")
-    gr.Markdown("支援 DistilGPT2 / BTLM-3B / BART-Base，Auto 模式會自動選擇。")
+    gr.Markdown("# 📺 電視弘法視頻生成文章RAG系統")
+    gr.Markdown("支援 GPT2-Chinese / BTLM-3B / BART-Base，Auto 模式會自動選擇，並強制中文輸出。")
 
     query_input = gr.Textbox(lines=2, placeholder="請輸入文章主題", label="文章主題")
     model_dropdown = gr.Dropdown(choices=list(MODEL_MAP.keys()), value="Auto", label="選擇生成模型")