update app.py (context+longform+warmup+4bit fallback)

app.py

@@ -1,86 +1,104 @@
 
-import os, time, torch, gradio as gr
-os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")  # 加速首次下載
-
+import os, torch, gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, BitsAndBytesConfig
 
-TITLE    = os.getenv("SPACE_TITLE", "LanguageBridge — Multimodal Chatbot (Mistral-7B)")
+TITLE    = "LanguageBridge — Multimodal Chatbot (Mistral-7B)"
 MODEL_ID = os.getenv("MODEL_ID", "aciang/mistral7b-tk-sft-20251019-merged")
 
 SYSTEM_PROMPT = (
-    "你是『語言橋』助教。回答原則：條列、準確、可重現步驟；不足處要誠實說明。"
+    "你是『語言橋』學習助教。規則："
+    "1) 嚴謹、分段、先重點後細節；"
+    "2) 若為數學/規則題：先列步驟，再給最終答案；"
+    "3) 若資訊不足，請明確指出缺口，勿捏造；"
+    "4) 優先以繁體中文回答。"
 )
 
-_tok, _llm = None, None
 def load_llm():
-    global _tok, _llm
-    if _llm is not None: 
-        return _tok, _llm
-    # 4-bit（失敗則自動回退）
+    # 4-bit 後援；失敗則 fp16/CPU fallback
+    kwargs = dict(trust_remote_code=False)
     bnb = BitsAndBytesConfig(
         load_in_4bit=True, bnb_4bit_quant_type="nf4",
         bnb_4bit_use_double_quant=True,
         bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float16
     )
-    kwargs = dict(device_map="auto", trust_remote_code=False, quantization_config=bnb)
     try:
-        _llm = AutoModelForCausalLM.from_pretrained(MODEL_ID, **kwargs)
+        model = AutoModelForCausalLM.from_pretrained(MODEL_ID, quantization_config=bnb, device_map="auto", **kwargs)
     except Exception as e:
-        print("[4-bit failed] fallback:", e)
-        _llm = AutoModelForCausalLM.from_pretrained(
+        print("[4-bit failed] -> fp16/CPU fallback:", e)
+        model = AutoModelForCausalLM.from_pretrained(
             MODEL_ID,
-            torch_dtype=(torch.float16 if torch.cuda.is_available() else torch.float32),
-            device_map=("auto" if torch.cuda.is_available() else None),
-            trust_remote_code=False
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto" if torch.cuda.is_available() else None,
+            **kwargs
         )
-    _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"
-    if torch.cuda.is_available(): torch.backends.cuda.matmul.allow_tf32 = True
-    _llm.config.use_cache = True
-    return _tok, _llm
-
-def format_prompt(user_text:str)->str:
-    return f"{SYSTEM_PROMPT}\n\n使用者：{user_text}\n助教："
+    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"
+    if torch.cuda.is_available():
+        torch.backends.cuda.matmul.allow_tf32 = True
+    model.config.use_cache = True
+    return tok, model
 
-@torch.inference_mode()
-def generate(user_text, mx=256, temp=0.6, top_p=0.95):
-    global _tok, _llm
-    if _llm is None:
-        yield "（正在載入模型，首次需要數十秒到數分鐘，請稍候…）"
-        _tok, _llm = load_llm()
-        yield "（模型載入完成，開始回應…）"
+tokenizer, llm = load_llm(); llm.eval()
 
-    prompt = format_prompt(user_text)
-    inputs = _tok(prompt, return_tensors="pt").to(_llm.device)
+def build_prompt(context, question, longform):
+    head = SYSTEM_PROMPT
+    if context.strip():
+        head += f"\n\n[上下文]\n{context.strip()}"
+    ask  = f"\n\n[問題]\n{question.strip()}\n"
+    tail = "\n請以條列步驟與小結回覆；若可計算，先算再答。"
+    if longform:
+        tail += "\n（長文模式）請分段、標題化、最後給出『摘要重點』。"
+    return head + ask + tail
 
-    streamer = TextIteratorStreamer(_tok, skip_prompt=True, skip_special_tokens=True)
-    gen = dict(**inputs, streamer=streamer, max_new_tokens=int(mx),
-               temperature=float(temp), top_p=float(top_p),
-               do_sample=True, eos_token_id=_tok.eos_token_id, pad_token_id=_tok.pad_token_id)
+@torch.inference_mode()
+def stream_answer(context, question, longform, mx, temp, top_p, rep):
+    prompt = build_prompt(context, question, longform)
+    inputs = tokenizer(prompt, return_tensors="pt").to(llm.device)
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    gen = dict(
+        **inputs, streamer=streamer, max_new_tokens=int(mx),
+        temperature=float(temp), top_p=float(top_p),
+        repetition_penalty=float(rep),
+        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
-    t = threading.Thread(target=_llm.generate, kwargs=gen); t.start()
-
-    buf = ""
+    t = threading.Thread(target=llm.generate, kwargs=gen); t.start()
+    buf=""
     for tok in streamer:
         buf += tok
         yield buf
 
-with gr.Blocks(title=TITLE, fill_height=True) as demo:
-    gr.Markdown(f"## {TITLE}\n模型：`{MODEL_ID}`（延遲載入）")
-    chat_in = gr.Textbox(label="你的問題 / 指令", placeholder="輸入文字…", lines=4)
+def warmup():
+    try:
+        _ = list(stream_answer("", "簡述本系統的用途。", False, 96, 0.2, 0.9, 1.05))[-1]
+        print("[warmup] done")
+    except Exception as e:
+        print("[warmup] skip:", e)
+
+with gr.Blocks(title=TITLE, theme="soft") as demo:
+    gr.Markdown(f"## {TITLE}\n模型：`{MODEL_ID}`｜已啟用：上下文欄位、長文模式、流式輸出、暖機")
+    with gr.Row():
+        ctx = gr.Textbox(label="上下文（長文，可空白）", placeholder="選填的背景內容/段落/資料摘錄", lines=6)
+    with gr.Row():
+        q = gr.Textbox(label="問題 / 指令", placeholder="請清楚描述你的問題", lines=3)
+    with gr.Row():
+        longf = gr.Checkbox(label="長文模式（章節化 + 摘要）", value=True)
     with gr.Row():
-        mx   = gr.Slider(64, 1024, value=256, step=32, label="max_new_tokens")
-        temp = gr.Slider(0.1, 1.0, value=0.6, step=0.05, label="temperature")
-        top  = gr.Slider(0.5, 1.0, value=0.95, step=0.01, label="top_p")
-    go   = gr.Button("送出 🚀", variant="primary")
-    out  = gr.Textbox(label="輸出（流式）", lines=18)
-    clr  = gr.Button("清除")
+        mx   = gr.Slider(128, 1024, value=512, step=32, label="max_new_tokens")
+        temp = gr.Slider(0.0, 0.8,  value=0.2, step=0.05, label="temperature")
+        top  = gr.Slider(0.6, 1.0,  value=0.9, step=0.01, label="top_p")
+        rep  = gr.Slider(1.0, 1.3,  value=1.05, step=0.01, label="repetition_penalty")
+    go = gr.Button("送出 🚀", variant="primary")
+    out = gr.Textbox(label="輸出（流式）", lines=14)
+    clr = gr.Button("清除")
 
-    go.click(generate, inputs=[chat_in, mx, temp, top], outputs=out)
-    clr.click(lambda: "", outputs=out)
+    go.click(stream_answer, inputs=[ctx,q,longf,mx,temp,top,rep], outputs=out)
+    clr.click(lambda:"", outputs=out)
 
-    demo.queue(max_size=32, api_open=False)
+    demo.queue(concurrency_count=4, max_size=32, api_open=False)
+    warmup()
 
 if __name__ == "__main__":
-    demo.launch(share=False, show_error=True)
+    demo.launch(share=False, server_name="0.0.0.0", server_port=7860, show_error=True)