Spaces:

build-small-hackathon
/

AmazingDigitalPetDentures

Paused

Sync from GitHub via hub-sync

ab2b3e1 verified 5 months ago

4.39 kB

"""In-process model layer for the HTML Toy Maker — runs Nemotron on HF ZeroGPU.
The Gradio app calls `generate(messages)`; that function is wrapped in `@spaces.GPU`,
so the GPU is only attached while it runs (ZeroGPU). The big GGUF DOWNLOAD does not need
the GPU, so we kick it off in a background thread at import time (overlapping Space boot).
The model is then constructed INSIDE the GPU function the first time and cached in a module
global — so the first request only pays for model load + generation, not the download.
Constraints: NVIDIA Nemotron (RTX 5080 prize) on the llama.cpp runtime (Llama Champion),
Q8_0 quant, n_ctx = 65536. Set a HF_TOKEN Space secret for faster, rate-limited-free pulls.
⚠️ If `Llama(...)` raises an "unknown architecture" / GGML assert on load, the prebuilt
llama-cpp-python wheel is older than Nemotron-3-Nano's hybrid Mamba-2 support — bump the
wheel in requirements.txt or build from source.
"""

import os

# Must be set BEFORE huggingface_hub is imported, so the fast Rust downloader is used.
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")

import glob
import threading

import spaces
from llama_cpp import Llama

# Nemotron-3-Nano-30B-A3B, Q8_0 (~33.6 GB). Q8_0 may be sharded into *-00001-of-0000N.gguf;
# we download all matching shards and point llama.cpp at the first (it auto-loads the rest).
MODEL_REPO = "unsloth/Nemotron-3-Nano-30B-A3B-GGUF"
QUANT_GLOB = "*Q8_0*.gguf"
# Full trained context (n_ctx_train = 1,048,576). Nemotron-3-Nano is hybrid Mamba-2: only its
# few attention layers grow KV with context (Mamba layers are constant-size), so long context
# is cheap here. Q8_0 (~34 GB) leaves ~36 GB free on ZeroGPU's ~70 GB H200 for the KV cache.
# (llama.cpp allocates KV up front from this — if it ever OOMs at load, dial this back.)
N_CTX = 1048576

# Unsloth's recommended Nemotron sampling.
DEFAULTS = dict(temperature=0.6, top_p=0.95, min_p=0.01)

_llm: Llama | None = None
_model_path: str | None = None
_download_error: Exception | None = None


def _download() -> None:
    """Fetch the GGUF shards (CPU only). Runs in a background thread started at import."""
    global _model_path, _download_error
    try:
        from huggingface_hub import snapshot_download

        local_dir = snapshot_download(MODEL_REPO, allow_patterns=[QUANT_GLOB])
        shards = sorted(glob.glob(os.path.join(local_dir, "**", QUANT_GLOB), recursive=True))
        if not shards:
            raise RuntimeError(f"No {QUANT_GLOB} files found in {MODEL_REPO}")
        _model_path = shards[0]  # first shard; llama.cpp loads the rest of a split GGUF itself
    except Exception as exc:  # surfaced when _load() is first called
        _download_error = exc


# Start downloading immediately at import (overlaps Space startup; does NOT block it, so the
# port binds fast and ZeroGPU registers the GPU function right away).
_download_thread = threading.Thread(target=_download, name="gguf-download", daemon=True)
_download_thread.start()


def _load() -> Llama:
    """Construct the llama.cpp model (GPU attached), cached for reuse across requests."""
    global _llm
    if _llm is not None:
        return _llm

    _download_thread.join()  # wait for the background download (instant if already done)
    if _model_path is None:
        raise RuntimeError(f"GGUF download failed: {_download_error}")

    _llm = Llama(
        model_path=_model_path,
        n_gpu_layers=-1,        # offload everything to the GPU
        n_ctx=N_CTX,
        flash_attn=True,
        verbose=False,
    )
    return _llm


@spaces.GPU(duration=300)
def generate(messages: list[dict], max_tokens: int = 8192, **sampling) -> dict:
    """Run one chat completion. Returns {"content", "reasoning"}.
    Nemotron emits its chain-of-thought via <think>/</think> SPECIAL tokens (ids 12/13),
    which are dropped on detokenize — so reasoning usually arrives mixed into `content`.
    Newer llama.cpp may surface it separately as `reasoning_content`; we pass both back and
    let app.py split them for display.
    """
    llm = _load()
    params = {**DEFAULTS, **sampling}
    out = llm.create_chat_completion(messages=messages, max_tokens=max_tokens, **params)
    msg = out["choices"][0]["message"]
    return {
        "content": (msg.get("content") or "").strip(),
        "reasoning": (msg.get("reasoning_content") or "").strip(),
    }

	"""In-process model layer for the HTML Toy Maker — runs Nemotron on HF ZeroGPU.

	The Gradio app calls `generate(messages)`; that function is wrapped in `@spaces.GPU`,
	so the GPU is only attached while it runs (ZeroGPU). The big GGUF DOWNLOAD does not need
	the GPU, so we kick it off in a background thread at import time (overlapping Space boot).
	The model is then constructed INSIDE the GPU function the first time and cached in a module
	global — so the first request only pays for model load + generation, not the download.

	Constraints: NVIDIA Nemotron (RTX 5080 prize) on the llama.cpp runtime (Llama Champion),
	Q8_0 quant, n_ctx = 65536. Set a HF_TOKEN Space secret for faster, rate-limited-free pulls.

	⚠️ If `Llama(...)` raises an "unknown architecture" / GGML assert on load, the prebuilt
	llama-cpp-python wheel is older than Nemotron-3-Nano's hybrid Mamba-2 support — bump the
	wheel in requirements.txt or build from source.
	"""

	import os

	# Must be set BEFORE huggingface_hub is imported, so the fast Rust downloader is used.
	os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")

	import glob
	import threading

	import spaces
	from llama_cpp import Llama

	# Nemotron-3-Nano-30B-A3B, Q8_0 (~33.6 GB). Q8_0 may be sharded into *-00001-of-0000N.gguf;
	# we download all matching shards and point llama.cpp at the first (it auto-loads the rest).
	MODEL_REPO = "unsloth/Nemotron-3-Nano-30B-A3B-GGUF"
	QUANT_GLOB = "Q8_0.gguf"
	# Full trained context (n_ctx_train = 1,048,576). Nemotron-3-Nano is hybrid Mamba-2: only its
	# few attention layers grow KV with context (Mamba layers are constant-size), so long context
	# is cheap here. Q8_0 (~34 GB) leaves ~36 GB free on ZeroGPU's ~70 GB H200 for the KV cache.
	# (llama.cpp allocates KV up front from this — if it ever OOMs at load, dial this back.)
	N_CTX = 1048576

	# Unsloth's recommended Nemotron sampling.
	DEFAULTS = dict(temperature=0.6, top_p=0.95, min_p=0.01)

	_llm: Llama \| None = None
	_model_path: str \| None = None
	_download_error: Exception \| None = None


	def _download() -> None:
	"""Fetch the GGUF shards (CPU only). Runs in a background thread started at import."""
	global _model_path, _download_error
	try:
	from huggingface_hub import snapshot_download

	local_dir = snapshot_download(MODEL_REPO, allow_patterns=[QUANT_GLOB])
	shards = sorted(glob.glob(os.path.join(local_dir, "**", QUANT_GLOB), recursive=True))
	if not shards:
	raise RuntimeError(f"No {QUANT_GLOB} files found in {MODEL_REPO}")
	_model_path = shards[0] # first shard; llama.cpp loads the rest of a split GGUF itself
	except Exception as exc: # surfaced when _load() is first called
	_download_error = exc


	# Start downloading immediately at import (overlaps Space startup; does NOT block it, so the
	# port binds fast and ZeroGPU registers the GPU function right away).
	_download_thread = threading.Thread(target=_download, name="gguf-download", daemon=True)
	_download_thread.start()


	def _load() -> Llama:
	"""Construct the llama.cpp model (GPU attached), cached for reuse across requests."""
	global _llm
	if _llm is not None:
	return _llm

	_download_thread.join() # wait for the background download (instant if already done)
	if _model_path is None:
	raise RuntimeError(f"GGUF download failed: {_download_error}")

	_llm = Llama(
	model_path=_model_path,
	n_gpu_layers=-1, # offload everything to the GPU
	n_ctx=N_CTX,
	flash_attn=True,
	verbose=False,
	)
	return _llm


	@spaces.GPU(duration=300)
	def generate(messages: list[dict], max_tokens: int = 8192, **sampling) -> dict:
	"""Run one chat completion. Returns {"content", "reasoning"}.

	Nemotron emits its chain-of-thought via <think>/</think> SPECIAL tokens (ids 12/13),
	which are dropped on detokenize — so reasoning usually arrives mixed into `content`.
	Newer llama.cpp may surface it separately as `reasoning_content`; we pass both back and
	let app.py split them for display.
	"""
	llm = _load()
	params = {DEFAULTS, sampling}
	out = llm.create_chat_completion(messages=messages, max_tokens=max_tokens, **params)
	msg = out["choices"][0]["message"]
	return {
	"content": (msg.get("content") or "").strip(),
	"reasoning": (msg.get("reasoning_content") or "").strip(),
	}