Hotfix: 4-bit fallback + hf_transfer + stable cache

app.py

@@ -1,132 +1,99 @@
 
-import os, time, torch, gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, BitsAndBytesConfig
-
-TITLE    = os.getenv("SPACE_TITLE", "LanguageBridge — Multimodal Chatbot (Mistral-7B)")
-MODEL_ID = os.getenv("MODEL_ID", "aciang/mistral7b-tk-sft-20251019-merged")
-
-SYSTEM_PROMPT = (
-  "你是語言橋助教。原則：1) 先條列必要重點；2) 再給最終結論；3) 嚴禁瞎掰，不足就說明。"
-)
-
-def load_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", quantization_config=bnb, trust_remote_code=False)
+import os, time, threading, torch, gradio as gr
+from huggingface_hub import snapshot_download
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextIteratorStreamer
+
+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"
+
+# ---- 加速下載 + 固定快取（/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)
+
+# ---- 嘗試先把 tokenizer 拉到本地（秒載 UI）----
+def _ensure_tokenizer(model_id):
     try:
-        model = AutoModelForCausalLM.from_pretrained(MODEL_ID, **kwargs)
+        snapshot_download(model_id, allow_patterns=["tokenizer.*","*tokenizer*","special_tokens_map.json"], local_dir=None)
     except Exception as e:
-        print("[4-bit failed] → fallback:", e)
-        kwargs.pop("quantization_config", None)
-        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
+        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
         )
-    tok = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
-    if tok.pad_token is None:
-        tok.pad_token = tok.eos_token
+        kw = dict(device_map="auto", quantization_config=bnb, trust_remote_code=False)
+    else:
+        kw = dict(device_map="auto", trust_remote_code=False)
+
+    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"
-    if torch.cuda.is_available():
-        torch.backends.cuda.matmul.allow_tf32 = True
-    model.config.use_cache = True
-    return tok, model
+    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)
 
-tokenizer, llm = load_llm(); llm.eval()
+boot_th = threading.Thread(target=boot); boot_th.start()
 
-def build_prompt(user_text:str)->str:
-    return f"{SYSTEM_PROMPT}\n\n使用者：{user_text}\n助教："
+SYSTEM = (
+    "你是語言橋助教。回覆重點：1) 條列步驟 2) 簡潔正確 3) 不確定就說明不足並提出假設。"
+)
 
 @torch.inference_mode()
-def stream_answer(history, text, mx=256, temp=0.2, top_p=0.95):
-    # history 為 list[{"role":"user|assistant","content":"..."}]
-    user_text = text or ""
-    prompt = build_prompt(user_text)
+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)
-
-    from transformers import TextIteratorStreamer
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    gen_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
-    t = threading.Thread(target=llm.generate, kwargs=gen_kwargs); t.start()
-    partial = ""
-    for piece in streamer:
-        partial += piece
-        yield partial
-
-def warmup():
-    try:
-        _ = list(stream_answer([], "π 約為多少？", mx=32))[-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}`｜已修正訊息格式；預設短答低延遲（流式 + 暖機）")
-
-    chat  = gr.Chatbot(label="Chatbot", type="messages", height=420, show_copy_button=True)
-    text  = gr.Textbox(label="你的問題 / 指令", placeholder="請輸入文字…", lines=3)
-    mx    = gr.Slider(64, 1024, value=256, step=32, label="max_new_tokens")
-    temp  = gr.Slider(0.0, 1.0, value=0.2,  step=0.05, label="temperature")
-    top   = gr.Slider(0.6, 1.0, value=0.95, step=0.01, label="top_p")
-    go    = gr.Button("送出 🚀", variant="primary")
-    clr   = gr.Button("清除")
-
-    def respond(history, text, mx, temp, top):
-        history = history or []
-        if text and text.strip():
-            history.append({"role":"user","content":text})
-            # 流式生成
-            stream = stream_answer(history, text, mx, temp, top)
-            out = ""
-            for chunk in stream:
-                out = chunk
-                yield history + [{"role":"assistant","content":out}], ""
-            history.append({"role":"assistant","content":out})
-        yield history, ""
-
-    go.click(respond, inputs=[chat, text, mx, temp, top], outputs=[chat, text])
-    clr.click(lambda: ([], ""), outputs=[chat, text])
-
-    # 不要使用舊的 queue 參數（避免 Runtime error）
-    demo.queue()
-    warmup()
+    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()
+    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)
 
 if __name__ == "__main__":
-    demo.launch(share=False, server_name="0.0.0.0", server_port=7860, show_error=True)
-
-
-    def _coerce_messages(history):
-        """確保是 [{'role','content'}] 形式。"""
-        fixed = []
-        for r, m in (history or []):
-            if isinstance(m, dict) and 'role' in m and 'content' in m:
-                fixed.append((r, m))
-            elif isinstance(m, str):
-                fixed.append((r, {"role": r, "content": m}))
-            else:
-                # 最保守兜底
-                fixed.append((r, {"role": r, "content": str(m)}))
-        return fixed
-
-    def respond(history, text, image, audio, mx, tp, top):
-        history = _coerce_messages(history)
-        history.append(("user", {"content": text}))
-        try:
-            ans = generate_reply(history, image, audio, mx, tp, top)
-        except Exception as e:
-            ans = f"(推理失敗：{e})"
-        history.append(("assistant", {"content": ans}))
-        return history, ""
-
+    demo.launch(share=False, show_error=True)