llm.py

"""ZeroGPU inference for the recommendation stage.

Three reviewer models, two load strategies:

* **Nemotron 3 Nano 4B (default)** loads at MODULE LEVEL and moves to
  ``cuda`` outside ``@spaces.GPU`` — the pattern proven in fine-tuner's
  nemotron_brain.py. ZeroGPU's CUDA emulation makes the placement legal at
  startup, and GPU workers (which import this module) inherit the weights.
  Loaded with the NATIVE transformers ``nemotron_h`` class (>=5.4) — the
  repo's custom code hard-requires the mamba-ssm pip package; the native
  class has a torch fallback. No ``trust_remote_code`` for it.
* **Mellum / MiniCPM** lazy-load inside the GPU task (the v1-proven path),
  at most one of the two resident besides Nemotron.

Weights download at runtime into the app-local cache app.py sets up —
``preload_from_hub`` bakes a cache the runtime user cannot write into, and
any missing-file download then dies with EACCES. README sets
``startup_duration_timeout: 45min`` to cover the ~8GB startup download.

All three together are ~34GB bf16 (8 + 24 + 2) — comfortable on the default
ZeroGPU slice.
"""

from __future__ import annotations

import os
import threading

from guide import GUIDE_BRIEF

NEMOTRON = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16"
MODELS = {
    NEMOTRON: "NVIDIA Nemotron 3 Nano 4B (default)",
    "JetBrains/Mellum2-12B-A2.5B-Instruct": "JetBrains Mellum 2 12B-A2.5B Instruct",
    "openbmb/MiniCPM5-1B": "OpenBMB MiniCPM5 1B (fastest)",
}
DEFAULT_MODEL = NEMOTRON

GPU_DURATION = 240
MAX_NEW_TOKENS = 2048

try:
    import spaces
    _gpu = spaces.GPU
except ImportError:  # local dev without the HF spaces runtime
    spaces = None

    def _gpu(*args, **kwargs):
        if args and callable(args[0]):
            return args[0]
        return lambda fn: fn


def _zero_gpu_active() -> bool:
    return spaces is not None and bool(os.environ.get("SPACES_ZERO_GPU"))


_cache: dict[str, tuple] = {}
_lock = threading.Lock()


def _load(model_id: str, device: str = "cuda"):
    """Load a reviewer model; small models evict each other, never Nemotron."""
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    with _lock:
        if model_id in _cache:
            return _cache[model_id]
        for old_id in [m for m in _cache if m != NEMOTRON]:
            del _cache[old_id]
        torch.cuda.empty_cache()
        print(f"[llm] loading {model_id} -> {device}", flush=True)
        if model_id == NEMOTRON:
            # native nemotron_h class, fine-tuner's exact recipe
            tok = AutoTokenizer.from_pretrained(model_id)
            model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto").to(device)
        else:
            tok = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
            model = AutoModelForCausalLM.from_pretrained(
                model_id,
                dtype=torch.bfloat16,
                device_map=device,
                trust_remote_code=True,
            )
        if tok.pad_token_id is None:
            tok.pad_token = tok.eos_token
        print(f"[llm] loaded {model_id}", flush=True)
        _cache[model_id] = (tok, model)
        return _cache[model_id]


# ZeroGPU-documented pattern: place the default model on cuda at module level
# (CUDA emulation outside @spaces.GPU; real CUDA inside). Off-Space: skipped.
if _zero_gpu_active():  # pragma: no cover - Space runtime only
    _load(NEMOTRON, "cuda")


SYSTEM_PROMPT = (
    "You are the reviewer behind “Ready to Submit?”, a warm and encouraging "
    "checker for the Build Small "
    "hackathon — think of yourself as a friendly camp counselor walking a "
    "builder through their submission checklist. You are given (a) the "
    "official rules digest and (b) machine-verified facts about one submission "
    "Space, including a checklist already evaluated against the rules.\n\n"
    "Hard constraints (these keep you honest):\n"
    "- Base every claim ONLY on the provided facts and rules digest. If the "
    "facts don't show something, say it could not be verified — never invent "
    "links, models, tags, or commits.\n"
    "- Quote canonical tag ids exactly (e.g. `track:backyard`).\n"
    "- Only celebrate checks whose status is pass. If the checklist warned "
    "about or failed something, it belongs in the next-steps section, never "
    "in the praise.\n\n"
    "Tone: second person, kind, specific, genuinely enthusiastic about what's "
    "already working. Frame problems as next steps on the trail, never as "
    "failures or threats. No scolding, no alarm — the deadline is just a "
    "campfire to walk toward. A light outdoorsy metaphor here and there is "
    "welcome; clarity always wins.\n\n"
    "Structure your answer as markdown with these sections:\n"
    "1. **What's already great** — celebrate the passing checks and anything "
    "genuinely impressive in the README, concretely.\n"
    "2. **Next steps before you submit** — each failed or warned check, with "
    "the concrete friendly fix.\n"
    "3. **Track fit** — which track the app fits and why, based on the README.\n"
    "4. **Prizes & badges within reach** — sponsor prizes, achievement tags "
    "and judged bonus awards this Space could plausibly earn but hasn't "
    "claimed, using the detected models and opportunities.\n"
    "5. **README polish** — 2-3 specific, doable improvements to the write-up.\n"
    "Keep the whole review under roughly 450 words — short bullets, one idea "
    "each, never trailing off mid-thought. Every suggestion should be "
    "actionable today."
)


def build_messages(ev_dict: dict) -> list[dict]:
    """A compact, human-readable briefing with the ask restated at the end.

    Small reviewer models (the 1B) echo long JSON briefings back verbatim,
    so the facts go in as short labelled bullets and the task sits last,
    where small models attend most.
    """
    facts = ev_dict.get("facts", {})
    checklist = "\n".join(
        f"- [{c['status'].upper()}] {c['rule']}: {c['evidence']}"
        for c in ev_dict.get("checks", [])
    )
    models = ", ".join(
        f"{m['id']} ({m['params'] / 1e9:.1f}B)" if m.get("params") else m["id"]
        for m in facts.get("models_detected", [])
    ) or "none detected"
    opportunities = "\n".join(f"- {o}" for o in facts.get("opportunities", [])) or "- none"
    user = (
        "# Rules cheat-sheet\n" + GUIDE_BRIEF
        + f"\n# Submission being reviewed\nSpace: {ev_dict.get('space_id')}\n"
        + f"Verdict: {ev_dict.get('verdict')}\n"
        + f"Tags on the Space: {', '.join(facts.get('tags', [])) or '(none)'}\n"
        + f"Models detected: {models}\n"
        + f"Codex-attributed commits: {len(facts.get('codex_commits', []))}\n"
        + "\n# Checklist results (machine-verified)\n" + checklist
        + "\n\n# Unclaimed opportunities (machine-verified)\n" + opportunities
        + "\n\n# README excerpt\n"
        + (facts.get("readme_body_excerpt", "")[:1500] or "(empty)")
        + "\n\n# Your task\nWrite the friendly review now: the five sections "
        "(What's already great / Next steps before you submit / Track fit / "
        "Prizes & badges within reach / README polish), under 450 words, "
        "grounded only in the briefing above. Do NOT repeat, quote, or "
        "summarize the briefing, the cheat-sheet, or these instructions — "
        "reply with the review only, starting at the first section heading."
    )
    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user},
    ]


@_gpu(duration=GPU_DURATION)
def generate_review(model_id: str, ev_dict: dict):
    """Yield a growing markdown string with the model's recommendations."""
    try:
        import torch
        has_cuda = torch.cuda.is_available()
    except ImportError:
        has_cuda = False

    if not has_cuda:
        yield (
            "*(LLM recommendations need the Space's ZeroGPU hardware — running "
            "locally without CUDA, so only the deterministic checklist above is "
            "available.)*"
        )
        return

    tok, model = _load(model_id)
    from transformers import TextIteratorStreamer

    messages = build_messages(ev_dict)
    inputs = tok.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
        enable_thinking=False,  # short friendly review, not a reasoning trace
    ).to(model.device)
    streamer = TextIteratorStreamer(tok, skip_prompt=True, skip_special_tokens=True)
    # Union of every end-of-turn candidate: Mellum2 ships generation_config
    # eos_token_id=0 (its bos) while its template ends turns with <|im_end|>;
    # Nemotron needs both ids from its generation_config ([2, 11]). The union
    # keeps each model stopping where its template ends.
    gc_eos = model.generation_config.eos_token_id
    eos_ids = {tok.eos_token_id}
    eos_ids.update(gc_eos if isinstance(gc_eos, (list, tuple)) else [gc_eos])
    eos_ids.add(tok.convert_tokens_to_ids("<|im_end|>"))
    eos_ids = sorted(i for i in eos_ids if isinstance(i, int) and i >= 0)
    thread = threading.Thread(
        target=model.generate,
        kwargs=dict(
            **inputs,
            eos_token_id=eos_ids,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=True,
            temperature=0.5,
            top_p=0.9,
            repetition_penalty=1.05,
            pad_token_id=tok.pad_token_id,
            streamer=streamer,
        ),
    )
    thread.start()
    text = ""
    for piece in streamer:
        text += piece
        yield _visible(text)
    thread.join()
    yield _visible(text, final=True)


def _visible(text: str, final: bool = False) -> str:
    """Hide an in-progress reasoning trace (<think> blocks)."""
    if "</think>" in text:
        return text.split("</think>", 1)[1].lstrip()
    if "<think>" in text:
        # never finish on the placeholder — show what we got if the model
        # ran out of tokens mid-think
        return text.replace("<think>", "").lstrip() if final else "*thinking…*"
    return text