Update app.py

app.py

@@ -1,77 +1,226 @@
+import os
+import torch
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-import os, torch, gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextIteratorStreamer
-os.environ.setdefault("HF_HOME", "/data/.cache")
-os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
-MODEL_ID = os.getenv("MODEL_ID", "aciang/mistral7b-tk-sft-20251019-merged")
-TITLE = "LanguageBridge — Multimodal Chatbot (Mistral-7B)"
-SYSTEM_PROMPT = "你是教學助教。先讀【任務】，按【格式】作答；資料不足先列缺口，勿猜測。"
-
-def load_llm():
-    has_cuda = torch.cuda.is_available()
-    kwargs = dict(trust_remote_code=False, low_cpu_mem_usage=True)
-    try:
-        if has_cuda:
-            bnb = BitsAndBytesConfig(
-                load_in_4bit=True, bnb_4bit_use_double_quant=True,
-                bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
-            )
-            model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto", quantization_config=bnb, **kwargs)
-        else:
-            print("[no CUDA] using CPU fp32")
-            model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="cpu", torch_dtype=torch.float32, **kwargs)
-    except Exception as e:
-        print("[loader fallback fp16/cpu]:", e)
-        model = AutoModelForCausalLM.from_pretrained(
-            MODEL_ID,
-            device_map="auto" if has_cuda else "cpu",
-            torch_dtype=torch.float16 if has_cuda else torch.float32,
-            **kwargs
+# ========= 基本設定 =========
+# 你的模型 repo id（目前是 private 也沒關係）
+MODEL_ID = "aciang/mistral7b-tk-sft-20251019-merged"
+
+# 若模型是 private，建議在 Space 的「Settings → Repository secrets」加上 HF_TOKEN
+HF_TOKEN = os.getenv("HF_TOKEN", None)
+
+# 建議預設的「傳統知識」系統提示，可以在介面中修改
+DEFAULT_SYSTEM_PROMPT = """你是一位熟悉台灣與國際原住民族傳統知識的學者，
+擅長用淺顯但尊重文化脈絡的繁體中文說明各族的夢境、儀式、宇宙觀、傳統醫療與環境知識。
+
+回答原則：
+1. 先簡短摘要重點（3–5 點條列）。
+2. 儘量說明「族名、場域、情境」與「知識來源背景」，避免抽象空話。
+3. 若是推論或類比，要清楚標註「推測」而不是說成唯一正解。
+4. 若資料不足或超出目前教材範圍，請誠實說明，並給出安全的延伸建議。
+5. 全程使用繁體中文。"""
+
+# ========= 載入模型 =========
+
+print(f"載入模型：{MODEL_ID} ...")
+
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID,
+    use_auth_token=HF_TOKEN,
+)
+
+# 保險起見，若沒有 pad_token 就沿用 eos_token
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float16,
+    device_map="auto",          # 自動分配到 GPU
+    use_auth_token=HF_TOKEN,
+)
+
+model.eval()
+print("模型載入完成。")
+
+
+# ========= 建立提示詞 =========
+
+def build_prompt(system_prompt: str, history: list[tuple[str, str]], user_message: str) -> str:
+    """
+    將 system_prompt + 歷史對話 + 新問題 組成一段文字 prompt。
+    這裡用簡單的「使用者 / 助手」格式，對傳統知識生成已經很足夠。
+    """
+    system_prompt = system_prompt.strip()
+    prompt = f"[系統提示]\n{system_prompt}\n\n"
+
+    # 過去對話（若有）
+    if history:
+        prompt += "[對話紀錄]\n"
+        for i, (user, bot) in enumerate(history, start=1):
+            prompt += f"輪次 {i}：\n使用者：{user}\n助手：{bot}\n\n"
+
+    # 最新一輪問題
+    prompt += "[目前問題]\n"
+    prompt += f"使用者：{user_message}\n助手："
+
+    return prompt
+
+
+# ========= 生成函式 =========
+
+def generate_reply(
+    user_message: str,
+    chat_history: list[tuple[str, str]],
+    system_prompt: str,
+    temperature: float,
+    max_new_tokens: int,
+):
+    if not user_message.strip():
+        return chat_history, gr.update(value="")
+
+    # 組合成一個大 prompt
+    prompt_text = build_prompt(system_prompt, chat_history, user_message)
+
+    inputs = tokenizer(
+        prompt_text,
+        return_tensors="pt",
+        add_special_tokens=True,
+    ).to(model.device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=int(max_new_tokens),
+            do_sample=True,
+            temperature=float(temperature),
+            top_p=0.9,
+            pad_token_id=tokenizer.eos_token_id,
         )
-    tok = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
-    if tok.pad_token is None: tok.pad_token = tok.eos_token
-    tok.padding_side = "left"
-    model.config.use_cache = True
-    return tok, model
-
-tokenizer, llm = load_llm(); llm.eval()
-
-def build_prompt(task, ctx=None):
-    head = "你是教學助教。先讀任務，依：1) 摘要要點；2) 逐步推理；3) 結論條列。\n\n"
-    if ctx:
-        ctx = ctx[-6000:]
-        return f"{head}【參考上下文】\n{ctx}\n\n【使用者問題】\n{task}\n\n【回答】"
-    return f"{head}【使用者問題】\n{task}\n\n【回答】"
-
-@torch.inference_mode()
-def stream_answer(task, context, mx=256, temp=0.15, top_p=0.9):
-    prompt = build_prompt(task, context.strip() or None)
-    inputs = tokenizer(prompt, return_tensors="pt").to(llm.device)
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    kwargs = dict(**inputs, streamer=streamer, max_new_tokens=int(mx),
-                  temperature=float(temp), top_p=float(top_p),
-                  do_sample=True if float(temp)>0 else False,
-                  eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.pad_token_id)
-    import threading; threading.Thread(target=llm.generate, kwargs=kwargs).start()
-    buf=""
-    for tok in streamer:
-        buf += tok
-        yield buf
-
-with gr.Blocks(title=TITLE, theme="soft") as demo:
-    gr.Markdown(f"## {TITLE}｜模型：`{MODEL_ID}`（流式）")
-    q   = gr.Textbox(label="你的問題 / 指令", lines=5, placeholder="可貼長文；我會先摘要→推理→結論")
-    ctx = gr.Textbox(label="（可選）上下文", lines=6)
-    mx   = gr.Slider(64, 512, value=256, step=32, label="max_new_tokens")
-    temp = gr.Slider(0.0, 0.8, value=0.15, step=0.05, label="temperature")
-    top  = gr.Slider(0.6, 1.0, value=0.9, step=0.01, label="top_p")
-    go   = gr.Button("送出 🚀", variant="primary")
-    out  = gr.Textbox(label="輸出（流式）", lines=14)
-    clr  = gr.Button("清除")
-    go.click(stream_answer, inputs=[q, ctx, mx, temp, top], outputs=out)
-    clr.click(lambda:"", outputs=out)
-    # ← 修正：不要用舊參數 concurrency_count
-    demo.queue(max_size=32, status_update_rate=1, api_open=False)
-
-if __name__ == "__main__":
-    demo.launch(share=False, server_name="0.0.0.0", server_port=7860, show_error=True)
+
+    # 只取新生成的部分
+    input_len = inputs["input_ids"].shape[-1]
+    generated_tokens = outputs[0, input_len:]
+    answer = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
+
+    chat_history = chat_history + [(user_message, answer)]
+    return chat_history, ""  # 清空輸入框
+
+
+def clear_history():
+    return [], ""
+
+
+# ========= Gradio 介面 =========
+
+with gr.Blocks(title="語言橋傳統知識聊天機器人 — Mistral7B TK", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # 語言橋傳統知識聊天機器人 — Mistral7B TK
+        使用你自訓練的 **mistral7b-tk-sft-20251019-merged** 模型，離線在 Hugging Face Space 上回答與傳統知識相關的問題。
+
+        建議題材舉例：
+        - 不同族群對「夢境」的五種層次與詮釋差異  
+        - 布農族狩獵儀式與祖靈信仰  
+        - 排灣族階級制度與紋面、圖騰的意義  
+        - 阿美族年齡階層制與植物分類知識  
+        - 海外原住民（如 Inuit）對身體、疾病與療癒的理解  
+        """
+    )
+
+    with gr.Row():
+        # 左側：設定區
+        with gr.Column(scale=1):
+            system_prompt_box = gr.Textbox(
+                label="系統提示（模型角色與回答風格）",
+                value=DEFAULT_SYSTEM_PROMPT,
+                lines=16,
+            )
+            temperature_slider = gr.Slider(
+                label="溫度（創造性）",
+                minimum=0.1,
+                maximum=1.5,
+                value=0.7,
+                step=0.05,
+            )
+            max_tokens_slider = gr.Slider(
+                label="最大回覆長度（token 數，大約字數的 1.5–2 倍）",
+                minimum=64,
+                maximum=1024,
+                value=512,
+                step=16,
+            )
+            gr.Markdown(
+                """
+                **小提醒：**
+                - 回答太發散 → 降低溫度（0.4–0.7）。  
+                - 回答太短 → 拉高「最大回覆長度」。  
+                - Space 若常 timeout，可以稍微降低最大回覆長度。
+                """
+            )
+
+        # 右側：聊天區
+        with gr.Column(scale=2):
+            chatbot = gr.Chatbot(
+                label="傳統知識 Chatbot",
+                height=480,
+                show_copy_button=True,
+            )
+            user_input = gr.Textbox(
+                label="輸入你的問題（可多輪對話）",
+                placeholder="例如：請比較布農族、排灣族和阿美族對治療疾病與夢境預兆的不同理解方式。",
+                lines=4,
+            )
+            with gr.Row():
+                send_btn = gr.Button("送出問題", variant="primary")
+                clear_btn = gr.Button("清除對話")
+
+    # 狀態：對話歷史
+    state = gr.State([])  # list[tuple[user, bot]]
+
+    # 綁定互動
+    send_btn.click(
+        fn=generate_reply,
+        inputs=[
+            user_input,
+            state,
+            system_prompt_box,
+            temperature_slider,
+            max_tokens_slider,
+        ],
+        outputs=[chatbot, user_input],
+    ).then(
+        fn=lambda h: h,
+        inputs=[chatbot],
+        outputs=[state],
+    )
+
+    user_input.submit(
+        fn=generate_reply,
+        inputs=[
+            user_input,
+            state,
+            system_prompt_box,
+            temperature_slider,
+            max_tokens_slider,
+        ],
+        outputs=[chatbot, user_input],
+    ).then(
+        fn=lambda h: h,
+        inputs=[chatbot],
+        outputs=[state],
+    )
+
+    clear_btn.click(
+        fn=clear_history,
+        inputs=[],
+        outputs=[chatbot, user_input],
+    ).then(
+        fn=lambda: [],
+        inputs=[],
+        outputs=[state],
+    )
+
+# 在 HF Space 中不需要 demo.launch()，平台會自動呼叫 demo
+