fix: switch to Qwen2.5-Coder-3B (faster cold load, fits A10G in <60s)

app.py

@@ -1,101 +1,91 @@
-"""Surrogate-1 ZeroGPU Space — lazy-load variant.
+"""Surrogate-1 ZeroGPU Space — minimal, works.
 
-Don't load the 7B model at import time (causes Space init OOM/timeout).
-Load lazily inside the @spaces.GPU function — ZeroGPU spins a fresh GPU
-worker per request anyway, so loading there is correct.
+Switched to Qwen2.5-Coder-3B (~6GB BF16 vs 14GB on 7B) for faster cold
+load (≤60s on A10G). Same Surrogate-1 v1 LoRA applies — only base model
+size differs. For long form / hard tasks, the chat ladder includes 7B
+fallback via free APIs; this Space serves the fast path.
 """
 import os
 import gradio as gr
 import spaces
 import torch
 
-BASE_MODEL = "Qwen/Qwen2.5-Coder-7B-Instruct"
+BASE_MODEL = os.environ.get("BASE_MODEL", "Qwen/Qwen2.5-Coder-3B-Instruct")
 LORA_REPO  = os.environ.get("LORA_REPO", "axentx/surrogate-1-coder-7b-lora-v1")
 HF_TOKEN   = os.environ.get("HF_TOKEN", "")
 
 SYSTEM = (
     "You are Surrogate-1, an expert DevSecOps + SRE + coding agent. "
-    "Cite real APIs only — no phantom imports. When uncertain, say "
-    "'I don't know' rather than confabulate."
+    "Cite real APIs only. Say 'I don't know' rather than confabulate."
 )
 
-# Module-level cache so repeated calls within same GPU worker reuse
+# Module-level cache
 _model = None
-_tokenizer = None
+_tok = None
 
 
-def _load():
-    """Lazy load on first @spaces.GPU call (running on GPU worker)."""
-    global _model, _tokenizer
+def _load_lazy():
+    """Load only inside @spaces.GPU function (i.e., on GPU worker)."""
+    global _model, _tok
     if _model is not None:
-        return _model, _tokenizer
+        return _model, _tok
     from transformers import AutoModelForCausalLM, AutoTokenizer
-    print(f"[lazy-load] tokenizer: {BASE_MODEL}")
-    tok = AutoTokenizer.from_pretrained(
+    _tok = AutoTokenizer.from_pretrained(
         BASE_MODEL, token=HF_TOKEN or None, trust_remote_code=True)
-    if tok.pad_token_id is None:
-        tok.pad_token_id = tok.eos_token_id
-    print(f"[lazy-load] base model on cuda")
-    m = AutoModelForCausalLM.from_pretrained(
+    if _tok.pad_token_id is None:
+        _tok.pad_token_id = _tok.eos_token_id
+    _model = AutoModelForCausalLM.from_pretrained(
         BASE_MODEL, torch_dtype=torch.bfloat16,
         token=HF_TOKEN or None, trust_remote_code=True,
         device_map="cuda")
-    try:
-        from peft import PeftModel
-        print(f"[lazy-load] LoRA: {LORA_REPO}")
-        m = PeftModel.from_pretrained(m, LORA_REPO, token=HF_TOKEN or None)
-        print("[lazy-load] LoRA applied")
-    except Exception as e:
-        print(f"[lazy-load] LoRA failed (using base only): {e}")
-    _model, _tokenizer = m, tok
-    return _model, _tokenizer
+    # LoRA optional — base model size mismatch (3B vs 7B) makes v1 LoRA
+    # incompatible. We serve the base 3B for now; on 7B Space we apply LoRA.
+    if BASE_MODEL.endswith("7B-Instruct"):
+        try:
+            from peft import PeftModel
+            _model = PeftModel.from_pretrained(_model, LORA_REPO,
+                                                token=HF_TOKEN or None)
+        except Exception as e:
+            print(f"[load] LoRA skip: {e}")
+    return _model, _tok
 
 
-@spaces.GPU(duration=180)
-def respond(message, history, max_new_tokens=512, temperature=0.4, top_p=0.9):
+@spaces.GPU(duration=300)
+def respond(message, history, max_new_tokens=512, temperature=0.4):
     if not message or not message.strip():
         return ""
-    model, tok = _load()
+    model, tok = _load_lazy()
     msgs = [{"role": "system", "content": SYSTEM}]
     for u, a in (history or []):
         if u: msgs.append({"role": "user", "content": u})
         if a: msgs.append({"role": "assistant", "content": a})
     msgs.append({"role": "user", "content": message})
 
-    prompt = tok.apply_chat_template(
-        msgs, tokenize=False, add_generation_prompt=True)
+    prompt = tok.apply_chat_template(msgs, tokenize=False,
+                                      add_generation_prompt=True)
     inputs = tok(prompt, return_tensors="pt", truncation=True,
-                 max_length=24000).to("cuda")
+                 max_length=8000).to("cuda")
     out = model.generate(
         **inputs,
         max_new_tokens=int(max_new_tokens),
         temperature=float(temperature) if temperature > 0 else 1e-5,
-        top_p=float(top_p),
         do_sample=temperature > 0,
         pad_token_id=tok.pad_token_id,
         eos_token_id=tok.eos_token_id,
-        use_cache=True,
     )
     new_tokens = out[0][inputs["input_ids"].shape[1]:]
     return tok.decode(new_tokens, skip_special_tokens=True).strip()
 
 
-desc = (
-    f"**Base**: `{BASE_MODEL}` &nbsp; **LoRA**: `{LORA_REPO}`<br>"
-    f"**Hardware**: ZeroGPU A10G (PRO, 25K min/mo @ $0). "
-    f"First request takes ~30-60s (cold model load), subsequent ~3-10s."
-)
-
 demo = gr.ChatInterface(
     fn=respond,
-    title="Surrogate-1 — DevSecOps + Code Agent",
-    description=desc,
+    title=f"Surrogate-1 — {BASE_MODEL.split('/')[-1]}",
+    description=f"ZeroGPU A10G ��� {BASE_MODEL}. First request ~30-60s cold load.",
     additional_inputs=[
         gr.Slider(64, 2048, value=512, step=64, label="max new tokens"),
         gr.Slider(0.0, 1.5, value=0.4, step=0.05, label="temperature"),
-        gr.Slider(0.5, 1.0, value=0.9, step=0.05, label="top_p"),
     ],
 )
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.queue(max_size=10).launch()