S2 hotfix: app.py

app.py

@@ -1,99 +1,86 @@
 
-import os, time, threading, torch, gradio as gr
-from huggingface_hub import snapshot_download
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextIteratorStreamer
+import os, time, torch, gradio as gr
+os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")  # 加速首次下載
 
-SPACE_TITLE    = "LanguageBridge — Multimodal Chatbot (Mistral-7B)"
-PRIMARY_MODEL  = "aciang/mistral7b-tk-sft-20251019-merged"
-FALLBACK_MODEL = "unsloth/mistral-7b-instruct-v0.2-bnb-4bit"
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, BitsAndBytesConfig
 
-# ---- 加速下載 + 固定快取（/data 在 Spaces 會持久化）----
-os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER","1")
-os.environ.setdefault("HF_HOME","/data/.cache/hf")         # 持久化 cache
-os.environ.setdefault("TRANSFORMERS_CACHE","/data/.cache/hf/transformers")
-os.makedirs(os.environ["HF_HOME"], exist_ok=True)
+TITLE    = os.getenv("SPACE_TITLE", "LanguageBridge — Multimodal Chatbot (Mistral-7B)")
+MODEL_ID = os.getenv("MODEL_ID", "aciang/mistral7b-tk-sft-20251019-merged")
 
-# ---- 嘗試先把 tokenizer 拉到本地（秒載 UI）----
-def _ensure_tokenizer(model_id):
+SYSTEM_PROMPT = (
+    "你是『語言橋』助教。回答原則：條列、準確、可重現步驟；不足處要誠實說明。"
+)
+
+_tok, _llm = None, None
+def load_llm():
+    global _tok, _llm
+    if _llm is not None: 
+        return _tok, _llm
+    # 4-bit（失敗則自動回退）
+    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:
-        snapshot_download(model_id, allow_patterns=["tokenizer.*","*tokenizer*","special_tokens_map.json"], local_dir=None)
+        _llm = AutoModelForCausalLM.from_pretrained(MODEL_ID, **kwargs)
     except Exception as e:
-        print("[tok prefetch] skip:", e)
-
-# ---- 模型載入（含 4-bit 後援）----
-def load_llm(prefer_primary=True):
-    model_id = PRIMARY_MODEL if prefer_primary else FALLBACK_MODEL
-    use_4bit  = (model_id != PRIMARY_MODEL)
-
-    if use_4bit:
-        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
+        print("[4-bit failed] fallback:", e)
+        _llm = 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
         )
-        kw = dict(device_map="auto", quantization_config=bnb, trust_remote_code=False)
-    else:
-        kw = dict(device_map="auto", trust_remote_code=False)
+    _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
 
-    print(f"[load] try model = {model_id} | 4bit={use_4bit}")
-    t0 = time.time()
-    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"
-    mdl = AutoModelForCausalLM.from_pretrained(model_id, **kw)
-    mdl.eval()
-    print(f"[load] ok in {time.time()-t0:.1f}s")
-    return tok, mdl, model_id
-
-# ---- 啟動邏輯：先開一條背景線程載入 PRIMARY；若超時改載入 FALLBACK ----
-tokenizer = None
-llm = None
-active_model = None
-
-def boot():
-    global tokenizer, llm, active_model
-    _ensure_tokenizer(PRIMARY_MODEL)
-    deadline = time.time() + 14*60   # 14 分鐘內載不完就切換（留 16 分緩衝 < 30 分鐘）
-    try:
-        tokenizer, llm, active_model = load_llm(prefer_primary=True)
-    except Exception as e:
-        print("[boot] primary failed early:", e)
-    if llm is None or time.time() > deadline:
-        print("[boot] switching to FALLBACK for fast availability...")
-        tokenizer, llm, active_model = load_llm(prefer_primary=False)
+def format_prompt(user_text:str)->str:
+    return f"{SYSTEM_PROMPT}\n\n使用者：{user_text}\n助教："
 
-boot_th = threading.Thread(target=boot); boot_th.start()
+@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 "（模型載入完成，開始回應…）"
 
-SYSTEM = (
-    "你是語言橋助教。回覆重點：1) 條列步驟 2) 簡潔正確 3) 不確定就說明不足並提出假設。"
-)
+    prompt = format_prompt(user_text)
+    inputs = _tok(prompt, return_tensors="pt").to(_llm.device)
 
-@torch.inference_mode()
-def stream_answer(q, mx=256, temp=0.6, top_p=0.95):
-    boot_th.join()  # 確保載入完成
-    prompt = f"{SYSTEM}\\n\\n使用者：{q}\\n助教："
-    inputs = tokenizer(prompt, return_tensors="pt").to(llm.device)
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    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 if float(temp)>0 else False,
-               eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.pad_token_id)
-    t = threading.Thread(target=llm.generate, kwargs=gen); t.start()
+               do_sample=True, eos_token_id=_tok.eos_token_id, pad_token_id=_tok.pad_token_id)
+    import threading
+    t = threading.Thread(target=_llm.generate, kwargs=gen); t.start()
+
     buf = ""
     for tok in streamer:
         buf += tok
         yield buf
 
-with gr.Blocks(title=SPACE_TITLE, fill_height=True) as demo:
-    gr.Markdown(f"### {SPACE_TITLE}\\n目前模型：`{active_model or 'loading…'}`\\n（首次啟動若超時將自動切到 4-bit 權重）")
-    q   = gr.Textbox(label="你的問題 / 指令")
-    mx  = gr.Slider(64, 1024, value=512, step=32, label="max_new_tokens")
-    tp  = gr.Slider(0.0, 1.2, 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=12)
-    go.click(stream_answer, inputs=[q, mx, tp, top], outputs=out)
-    demo.queue(api_open=False)
+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)
+    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("清除")
+
+    go.click(generate, inputs=[chat_in, mx, temp, top], outputs=out)
+    clr.click(lambda: "", outputs=out)
+
+    demo.queue(max_size=32, api_open=False)
 
 if __name__ == "__main__":
     demo.launch(share=False, show_error=True)