v3.5: top-level imports + @spaces.GPU (canonical GRM-2.6-Opus pattern)

app.py

@@ -1,50 +1,87 @@
-import os, json
-import gradio as gr
-
-SAFE = {"qwen2.5-0.5b": ("Qwen/Qwen2.5-0.5B-Instruct", 1.0), "llama3.2-1b": ("meta-llama/Llama-3.2-1B-Instruct", 1.5)}
-SPECS = {**SAFE, "deepseek-r1-1.5b": ("deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B", 2.2), "gemma3-4b": ("google/gemma-3-4b-it", 6.0)}
-
-def gpu_ok():
-    try:
-        import torch
-        return torch.cuda.is_available()
-    except: return False
-
-def gen(m, p, t):
-    if not gpu_ok():
-        return "ERROR: No GPU. Personal accounts need PRO for ZeroGPU. Use MCP client for 25min/day free, or run hf jobs for pay-as-you-go GPU. https://huggingface.co/mcp"
-    import spaces, torch
-    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-    hid, vram = SPECS[m]
-    bnb = torch.classes if False else None  # dummy
-    bnb = None
-    try:
-        bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16)
-    except: pass
-    tok = AutoTokenizer.from_pretrained(hid, trust_remote_code=True)
-    if tok.pad_token is None: tok.pad_token = tok.eos_token
-    kw = {"trust_remote_code": True, "low_cpu_mem_usage": True, "torch_dtype": torch.bfloat16}
-    if bnb: kw["quantization_config"] = bnb; kw["device_map"] = "auto"
-    mod = AutoModelForCausalLM.from_pretrained(hid, **kw)
-    msgs = [{"role":"user","content":p}]
-    txt = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
-    inp = tok(txt, return_tensors="pt", truncation=True, max_length=16384).to(next(mod.parameters()).device)
-    out = mod.generate(**inp, max_new_tokens=512, do_sample=t>0.01, temperature=max(t,1e-5), top_p=0.95, top_k=40, pad_token_id=tok.pad_token_id, eos_token_id=tok.eos_token_id)
-    return tok.decode(out[0], skip_special_tokens=True)
-
-with gr.Blocks(title="NEXUS LAB v3.4") as demo:
-    gr.Markdown("# NEXUS LAB v3.4")
-    gr.Markdown("**Status:** " + ("GPU Ready" if gpu_ok() else "CPU Only — No GPU detected"))
+import os, time, json
+import spaces, torch, gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
+from threading import Thread
+
+SPECS = {
+  "llama3.2-1b": ("meta-llama/Llama-3.2-1B-Instruct", 1.5),
+  "qwen2.5-0.5b": ("Qwen/Qwen2.5-0.5B-Instruct", 1.0),
+  "deepseek-r1-1.5b": ("deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B", 2.2),
+  "gemma3-4b": ("google/gemma-3-4b-it", 6.0),
+}
+
+BNB = BitsAndBytesConfig(
+  load_in_4bit=True, bnb_4bit_quant_type="nf4",
+  bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16,
+)
+
+def load_model(spec_id):
+  hid, vram = SPECS[spec_id]
+  tok = AutoTokenizer.from_pretrained(hid, trust_remote_code=True)
+  if tok.pad_token is None: tok.pad_token = tok.eos_token
+  mod = AutoModelForCausalLM.from_pretrained(
+    hid, trust_remote_code=True, quantization_config=BNB,
+    device_map="auto", torch_dtype=torch.bfloat16, low_cpu_mem_usage=True,
+  )
+  return mod, tok
+
+_model_cache = {}
+
+def get_model(spec_id):
+  if spec_id not in _model_cache:
+    _model_cache[spec_id] = load_model(spec_id)
+  return _model_cache[spec_id]
+
+@spaces.GPU(duration=lambda *a: 180, size="large")
+def generate(spec_id, prompt, system, temp, top_p, max_tokens):
+  mod, tok = get_model(spec_id)
+  msgs = []
+  if system: msgs.append({"role":"system","content":system})
+  msgs.append({"role":"user","content":prompt})
+  txt = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+  inp = tok(txt, return_tensors="pt", truncation=True, max_length=8192).to(next(mod.parameters()).device)
+  s = TextIteratorStreamer(tok, skip_prompt=True, skip_special_tokens=True)
+  kw = dict(
+    **inp, streamer=s, max_new_tokens=max_tokens,
+    do_sample=temp>0.01, temperature=max(temp,1e-5), top_p=top_p,
+    pad_token_id=tok.pad_token_id, eos_token_id=tok.eos_token_id,
+  )
+  Thread(target=mod.generate, kwargs=kw).start()
+  raw = ""
+  for chunk in s:
+    raw += chunk
+    yield raw
+
+@spaces.GPU(duration=lambda *a: 180, size="large")
+def temp_sweep(spec_id, prompt, temp_range, top_p, max_tokens):
+  temps = [float(x.strip()) for x in temp_range.split(",") if x.strip()]
+  out = []
+  for t in temps:
+    txt = ""
+    for p in generate(spec_id, prompt, "", t, top_p, max_tokens):
+      txt = p
+    out.append(f"--- T={t:.2f} ---\n{txt.strip()}\n")
+  return "\n".join(out)
+
+with gr.Blocks(title="NEXUS LAB v3.5") as demo:
+  gr.Markdown("# NEXUS LAB v3.5 - ZeroGPU Team Edition")
+  with gr.Tab("Chat"):
     m = gr.Dropdown(choices=list(SPECS.keys()), value="llama3.2-1b", label="Model")
-    p = gr.Textbox(label="Prompt", lines=3)
+    p = gr.Textbox(label="Prompt", lines=3, value="Explain the NEXUS OS TICOS guard system.")
+    sy = gr.Textbox(label="System", value="Think step by step.")
     t = gr.Slider(0.0, 2.0, 0.7, step=0.05, label="Temperature")
-    o = gr.Textbox(label="Output", lines=15)
-    def wrapped(m, p, t):
-        if not gpu_ok(): return "ERROR: No GPU. Enable ZeroGPU (PRO), use MCP for 25min/day free: https://huggingface.co/mcp, or run hf jobs."
-        return gen(m, p, t)
-    gr.Button("Generate", variant="primary").click(wrapped, [m, p, t], o)
-    gr.Markdown("---")
-    gr.Markdown("MCP: `/gradio_api/mcp/sse` — Use MCP client for free GPU via agents.md")
-    gr.Markdown("Jobs: `hf jobs run --flavor gpu-t4-small python train.py` — Pay-as-you-go, no PRO needed")
-
+    pp = gr.Slider(0.1, 1.0, 0.95, step=0.05, label="Top-p")
+    mt = gr.Slider(64, 2048, 1024, step=64, label="Max Tokens")
+    o = gr.Textbox(label="Output", lines=20)
+    gr.Button("Generate", variant="primary").click(generate, [m,p,sy,t,pp,mt], o)
+  with gr.Tab("Temp Sweep"):
+    m2 = gr.Dropdown(choices=list(SPECS.keys()), value="llama3.2-1b", label="Model")
+    p2 = gr.Textbox(label="Prompt")
+    tr = gr.Textbox(label="Temps", value="0.0,0.3,0.6,0.9,1.2")
+    pp2 = gr.Slider(0.1, 1.0, 0.95, step=0.05, label="Top-p")
+    mt2 = gr.Slider(64, 1024, 256, step=64, label="Max Tokens")
+    o2 = gr.Textbox(label="Results", lines=20)
+    gr.Button("Run Sweep", variant="primary").click(temp_sweep, [m2,p2,tr,pp2,mt2], o2)
+  gr.Markdown("---")
+  gr.Markdown("MCP: /gradio_api/mcp/sse | ZeroGPU Team 40min/day | Bucket: specimba/nexus-vap-bucket")
 demo.launch(mcp_server=True)