Update app.py

app.py

@@ -10,14 +10,9 @@ from huggingface_hub import login, snapshot_download
 import gradio as gr
 
 # -------------------------------
-# 1. 模型清單（中文 & 英文）
+# 1. 模型清單（只用中文 GPT2）
 # -------------------------------
-MODEL_MAP = {
-    "Auto": None,
-    "BTLM-3B-8K": "cerebras/btlm-3b-8k-base",
-    "GPT2-Chinese": "uer/gpt2-chinese-cluecorpusmedium",  # 中文 GPT2
-    "BART-Base": "facebook/bart-base"
-}
+MODEL_NAME = "uer/gpt2-chinese-cluecorpusmedium"
 
 HF_TOKEN = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
 if HF_TOKEN:
@@ -25,57 +20,35 @@ if HF_TOKEN:
     print("✅ 已使用 HUGGINGFACEHUB_API_TOKEN 登入 Hugging Face")
 
 # -------------------------------
-# 2. 預先下載模型
+# 2. 下載模型
 # -------------------------------
-LOCAL_MODEL_DIRS = {}
-for name, repo in MODEL_MAP.items():
-    if repo is None:
-        continue
-    try:
-        local_dir = f"./models/{repo.split('/')[-1]}"
-        if not os.path.exists(local_dir):
-            print(f"⬇️ 正在下載模型 {repo} ...")
-            snapshot_download(repo_id=repo, token=HF_TOKEN, local_dir=local_dir)
-        LOCAL_MODEL_DIRS[name] = local_dir
-    except Exception as e:
-        print(f"⚠️ 模型 {repo} 無法下載: {e}")
+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)
 
 # -------------------------------
 # 3. pipeline 載入
 # -------------------------------
-_loaded_pipelines = {}
-
-def get_pipeline(model_name):
-    if model_name not in _loaded_pipelines:
-        local_path = LOCAL_MODEL_DIRS.get(model_name)
-        print(f"🔄 正在載入模型 {model_name} from {local_path}")
-
-        if model_name == "BTLM-3B-8K":
-            tokenizer = AutoTokenizer.from_pretrained(local_path, trust_remote_code=True)
-            model = AutoModelForCausalLM.from_pretrained(local_path, trust_remote_code=True)
-        else:
-            tokenizer = AutoTokenizer.from_pretrained(local_path)
-            model = AutoModelForCausalLM.from_pretrained(local_path)
-
-        # 修正 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
-        )
-        _loaded_pipelines[model_name] = generator
-    return _loaded_pipelines[model_name]
-
-def call_local_inference(model_name, prompt, max_new_tokens=256):
+print(f"🔄 載入中文模型 {MODEL_NAME}")
+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:
-        generator = get_pipeline(model_name)
-
-        # ✅ 強制中文模式：補上提示
-        if "中文" not in prompt and "Chinese" not in prompt:
+        # 強制補充中文提示
+        if "中文" not in prompt:
             prompt += "\n（請用中文回答）"
 
         outputs = generator(
@@ -83,55 +56,41 @@ def call_local_inference(model_name, prompt, max_new_tokens=256):
             max_new_tokens=max_new_tokens,
             do_sample=True,
             temperature=0.7,
-            pad_token_id=generator.tokenizer.pad_token_id
+            pad_token_id=tokenizer.pad_token_id
         )
         return outputs[0]["generated_text"]
     except Exception as e:
         return f"（生成失敗：{e}）"
 
 # -------------------------------
-# 4. Auto 模式
+# 4. 文章生成
 # -------------------------------
-def pick_model_auto(segments):
-    if segments <= 3:
-        return "GPT2-Chinese"  # 短文 → 中文 GPT2
-    elif segments <= 6:
-        return "BTLM-3B-8K"
-    else:
-        return "BART-Base"
-
-def generate_article_progress(query, model_name, segments=5):
+def generate_article_progress(query, segments=5):
     docx_file = "/tmp/generated_article.docx"
     doc = DocxDocument()
     doc.add_heading(query, level=1)
 
-    selected_model = pick_model_auto(segments) if model_name == "Auto" else model_name
-    print(f"👉 使用模型: {selected_model}")
-
     all_text = []
     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)
+        paragraph = call_local_inference(prompt)
         all_text.append(paragraph)
         doc.add_paragraph(paragraph)
-
-        yield "\n\n".join(all_text), None, f"本次使用模型：{selected_model}"
+        yield "\n\n".join(all_text), None, f"本次使用模型：{MODEL_NAME}"
 
     doc.save(docx_file)
-    yield "\n\n".join(all_text), docx_file, f"本次使用模型：{selected_model}"
+    yield "\n\n".join(all_text), docx_file, f"本次使用模型：{MODEL_NAME}"
 
 # -------------------------------
 # 5. Gradio 介面
 # -------------------------------
 with gr.Blocks() as demo:
     gr.Markdown("# 📺 電視弘法視頻生成文章RAG系統")
-    gr.Markdown("支援 GPT2-Chinese / BTLM-3B / BART-Base，Auto 模式會自動選擇，並強制中文輸出。")
+    gr.Markdown("固定使用 **GPT2-Chinese (uer/gpt2-chinese-cluecorpusmedium)** 生成中文文章。")
 
     query_input = gr.Textbox(lines=2, placeholder="請輸入文章主題", label="文章主題")
-    model_dropdown = gr.Dropdown(choices=list(MODEL_MAP.keys()), value="Auto", 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")
@@ -140,7 +99,7 @@ with gr.Blocks() as demo:
     btn = gr.Button("生成文章")
     btn.click(
         generate_article_progress,
-        inputs=[query_input, model_dropdown, segments_input],
+        inputs=[query_input, segments_input],
         outputs=[output_text, output_file, model_used_text]
     )