Update app.py

app.py

@@ -1,40 +1,188 @@
-import os, threading
+#!/usr/bin/env python3
+# ---
+# title: ZeusMM Chat
+# emoji: 🤖
+# colorFrom: indigo
+# colorTo: purple
+# sdk: gradio
+# sdk_version: 5.0.1
+# app_file: app.py
+# pinned: false
+# ---
+
+import os
+import threading
 import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    AutoConfig,
+    TextIteratorStreamer,
+)
+
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
+# ===== Env & Model config =====
+os.environ.setdefault("ACCELERATE_DISABLE_MAPPED_DEVICE", "1")   # avoid meta-tensors on CPU
+os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")          # faster downloads in Spaces
 
 MODEL_ID = os.environ.get("MODEL_ID", "Wonder-Griffin/ZeusMM-SFT-oasst1")
-HF_TOKEN = os.environ.get("HF_TOKEN")
+HF_TOKEN = os.environ.get("HF_TOKEN")  # add as a Space secret if the model is private
+IS_GPU = torch.cuda.is_available()
+
+# Optional: pin to a specific revision to avoid surprise code updates
+MODEL_REVISION = os.environ.get("MODEL_REVISION")  # e.g., a commit SHA; leave unset to use latest
+
+
+# ===== Robust CPU loader: builds real tensors, no meta, then loads weights =====
+def load_cpu_no_meta(model_id: str, hf_token: str | None = None, revision: str | None = None):
+    cfg = AutoConfig.from_pretrained(
+        model_id,
+        trust_remote_code=True,
+        token=hf_token,
+        revision=revision,
+    )
+    model = AutoModelForCausalLM.from_config(
+        cfg,
+        trust_remote_code=True,
+        torch_dtype=torch.float32,
+    )
+    # Allocate real storage on CPU for all params/buffers
+    model.to_empty(device="cpu")
+
+    # Find and load the primary weight file
+    # (adjust filename if your repo uses something else)
+    weights_path = hf_hub_download(
+        repo_id=model_id,
+        filename="model.safetensors",
+        token=hf_token,
+        revision=revision,
+    )
+    state = load_file(weights_path)  # safetensors -> state_dict
+
+    missing, unexpected = model.load_state_dict(state, strict=False)
+    if missing or unexpected:
+        # Print to Space logs; non-fatal if they are non-critical heads/keys
+        print("Missing keys:", missing)
+        print("Unexpected keys:", unexpected)
+
+    model.eval()
+    return model
 
-# Avoid Accelerate mapped-device heuristics that can create meta tensors on CPU
-os.environ.setdefault("ACCELERATE_DISABLE_MAPPED_DEVICE", "1")
-os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
 
-# --- Load tokenizer ---
+# ===== Tokenizer (shared) =====
 tok_kwargs = {"trust_remote_code": True}
-if HF_TOKEN: tok_kwargs["token"] = HF_TOKEN
+if HF_TOKEN:
+    tok_kwargs["token"] = HF_TOKEN
+if MODEL_REVISION:
+    tok_kwargs["revision"] = MODEL_REVISION
+
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, **tok_kwargs)
 
-# --- Load model (CPU-safe / GPU-smart) ---
-IS_GPU = torch.cuda.is_available()
+
+# ===== Model (GPU uses device_map, CPU uses robust loader) =====
 if IS_GPU:
-    # GPU: allow device_map and auto dtype, but force eager attention
     mdl_kwargs = dict(
         trust_remote_code=True,
         torch_dtype="auto",
         device_map="auto",
-        attn_implementation="eager",
+        attn_implementation="eager",  # stable across kernels
     )
-    if HF_TOKEN: mdl_kwargs["token"] = HF_TOKEN
+    if HF_TOKEN:
+        mdl_kwargs["token"] = HF_TOKEN
+    if MODEL_REVISION:
+        mdl_kwargs["revision"] = MODEL_REVISION
+
     model = AutoModelForCausalLM.from_pretrained(MODEL_ID, **mdl_kwargs)
 else:
-    # CPU: NO device_map, NO low_cpu_mem_usage -> real tensors (not meta)
-    mdl_kwargs = dict(
-        trust_remote_code=True,
-        torch_dtype=torch.float32,
-        low_cpu_mem_usage=False,
-        attn_implementation="eager",
+    model = load_cpu_no_meta(MODEL_ID, HF_TOKEN, MODEL_REVISION)
+
+
+# ===== Prompt building =====
+def build_prompt(system_message: str, history: list[tuple[str, str]], user_message: str) -> str:
+    messages = []
+    if system_message:
+        messages.append({"role": "system", "content": system_message})
+
+    for u, a in (history or []):
+        if u:
+            messages.append({"role": "user", "content": u})
+        if a:
+            messages.append({"role": "assistant", "content": a})
+
+    messages.append({"role": "user", "content": user_message})
+
+    if hasattr(tokenizer, "apply_chat_template"):
+        try:
+            return tokenizer.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=True
+            )
+        except Exception:
+            pass
+
+    # Fallback (generic)
+    out = []
+    if system_message:
+        out.append(f"[SYSTEM] {system_message}\n")
+    for m in messages:
+        role = (m.get("role") or "user").upper()
+        out.append(f"[{role}] {m.get('content','')}\n")
+    out.append("[ASSISTANT] ")
+    return "".join(out)
+
+
+# ===== Generation (streaming) =====
+def respond(message, history, system_message, max_tokens, temperature, top_p):
+    prompt = build_prompt(system_message, history, message)
+    inputs = tokenizer(prompt, return_tensors="pt")
+
+    # Send inputs to the same device as the first model parameter (works for CPU/GPU)
+    first_param_device = next(model.parameters()).device
+    inputs = {k: v.to(first_param_device) for k, v in inputs.items()}
+
+    streamer = TextIteratorStreamer(
+        tokenizer,
+        skip_prompt=True,
+        skip_special_tokens=True,
     )
-    if HF_TOKEN: mdl_kwargs["token"] = HF_TOKEN
-    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, **mdl_kwargs)
-    model.to("cpu")
+
+    gen_kwargs = dict(
+        **inputs,
+        max_new_tokens=int(max_tokens),
+        temperature=float(temperature),
+        top_p=float(top_p),
+        do_sample=True,
+        streamer=streamer,
+    )
+
+    t = threading.Thread(target=model.generate, kwargs=gen_kwargs)
+    t.start()
+
+    partial = ""
+    for chunk in streamer:
+        partial += chunk
+        yield partial
+
+
+# ===== UI =====
+demo = gr.ChatInterface(
+    fn=respond,
+    additional_inputs=[
+        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+        gr.Slider(minimum=1, maximum=4096, value=512, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
+    ],
+    title="ZeusMM Chat",
+    description="Chat with your ZeusMM-SFT model with streaming responses.",
+)
+
+# Expose for Spaces
+app = demo
+
+if __name__ == "__main__":
+    # queue helps avoid cold-start timeouts and enables token streaming
+    demo.queue(max_size=32, concurrency_count=1).launch(server_name="0.0.0.0", server_port=7860)