"""Generation backends for the GODSEED mind.
Protocol (duck-typed; see `Backend`):
name: str # "mock" | "llamacpp" | "zerogpu"
model_id: str # human-readable model identifier
async def generate_stream(prompt: str, grammar: str | None,
max_tokens: int) -> AsyncIterator[str]
`grammar` is a llama.cpp GBNF string (mind/grammar.gbnf) or None for free
text. Backends that cannot enforce a grammar (mock, zerogpu) still receive it
and use it to recognize *which* constrained shape is wanted (the moderation
grammar contains the '"allowed"' key; the turn grammar does not).
Selected via the GODSEED_BACKEND env: mock (default) | llamacpp | zerogpu.
- MockBackend: deterministic keyword-driven scripts (mind/mock_scripts.py)
with small delays so streaming looks alive. Local dev + headless demos.
- LlamaCppBackend: llama-cpp-python + Nemotron-3-Nano-4B GGUF Q4_K_M, lazy
hf_hub_download on first use, native GBNF. The v0 ship backend.
- ZeroGPUBackend: transformers + spaces.GPU (import-guarded at call time),
Nemotron 3 Nano 30B-A3B. No grammar support -> strict JSON validation with
one retry, then the planner's own re-ask path takes over.
"""
from __future__ import annotations
import asyncio
import concurrent.futures
import json
import os
import re
import threading
from typing import AsyncIterator, Protocol, runtime_checkable
from . import prompts
from .mock_scripts import build_script
from .validate import is_moderation_grammar, parse_moderation, parse_turn
@runtime_checkable
class Backend(Protocol):
"""Structural type for generation backends."""
name: str
model_id: str
def generate_stream(
self, prompt: str, grammar: str | None, max_tokens: int
) -> AsyncIterator[str]: ...
# --------------------------------------------------------------------------
# Mock backend — what demos run on
# --------------------------------------------------------------------------
_MOCK_DENYLIST: tuple[tuple[re.Pattern[str], str], ...] = (
(re.compile(r"\b(nazi|hitler|swastika|kkk|genocide|lynch)\b", re.I), "hate"),
(re.compile(r"\b(kill|murder|slaughter|massacre|behead|stab)\b", re.I), "violence"),
(re.compile(r"\b(bomb|terror|terrorist|terrorism|shooting)\b", re.I), "violence"),
(re.compile(r"\b(rape|porn|sex|sexual|nude|naked|nsfw)\b", re.I), "sexual"),
(re.compile(r"\b(suicide|self[- ]?harm|cutting)\b", re.I), "self-harm"),
)
def _chunks(text: str, size: int) -> list[str]:
return [text[i : i + size] for i in range(0, len(text), size)]
class MockBackend:
"""Deterministic scripted backend.
Recovers the wish and turn index from the rendered prompt using the
marker constants in mind/prompts.py:
- wish: first line after the LAST `WISH: ` marker (the few-shot example's
wish sits before the real one, so "last" is always the live wish);
- turn index: count of `OBSERVATION: ` lines after that marker (each
completed or skipped turn contributes exactly one);
- moderation requests are recognized by the grammar's '"allowed"' key and
judged with a small keyword denylist.
Token delays come from GODSEED_MOCK_DELAY (seconds per chunk, default
0.012); tests pass delay=0.0 explicitly.
"""
name = "mock"
model_id = "godseed-mock-scripts"
def __init__(self, delay: float | None = None):
if delay is None:
try:
delay = float(os.environ.get("GODSEED_MOCK_DELAY", "0.012"))
except ValueError:
delay = 0.012
self.delay = max(0.0, delay)
async def _tick(self) -> None:
if self.delay:
await asyncio.sleep(self.delay)
else:
await asyncio.sleep(0) # stay fair to the event loop
@staticmethod
def _tail(prompt: str, marker: str) -> str:
parts = prompt.rsplit(marker, 1)
return parts[-1] if len(parts) == 2 else prompt
def _wish_and_turn(self, prompt: str) -> tuple[str, int]:
tail = self._tail(prompt, prompts.WISH_MARKER)
lines = tail.splitlines() or [""]
wish = lines[0].strip()
turn_index = tail.count(prompts.OBS_MARKER)
return wish, turn_index
def _moderate(self, prompt: str) -> dict:
tail = self._tail(prompt, prompts.CANDIDATE_MARKER)
candidate = (tail.splitlines() or [""])[0].strip()
for pattern, category in _MOCK_DENYLIST:
if pattern.search(candidate):
return {"allowed": False, "category": category}
return {"allowed": True, "category": ""}
async def generate_stream(
self, prompt: str, grammar: str | None = None, max_tokens: int = 256
) -> AsyncIterator[str]:
if is_moderation_grammar(grammar):
verdict = json.dumps(self._moderate(prompt))
for chunk in _chunks(verdict, 8):
await self._tick()
yield chunk
return
wish, turn_index = self._wish_and_turn(prompt)
script = build_script(wish)
if grammar is None:
# The Reading: stream word by word, capped at max_tokens words.
words = re.findall(r"\S+\s*", script["reading"])
for word in words[: max(1, max_tokens)]:
await self._tick()
yield word
return
# A turn: scripts always end with a done turn, so an index past the
# end (e.g. after skipped turns) re-yields the final done turn.
turns = script["turns"]
turn = turns[min(turn_index, len(turns) - 1)]
payload = json.dumps(turn, ensure_ascii=False)
for chunk in _chunks(payload, 12):
await self._tick()
yield chunk
# --------------------------------------------------------------------------
# llama.cpp backend — Nemotron-3-Nano-4B GGUF (v0 ship target)
# --------------------------------------------------------------------------
_THREAD_DONE = object()
class LlamaCppBackend:
"""llama-cpp-python backend with native GBNF grammar support.
Lazy: nothing is imported or downloaded until the first generate call.
Model resolution order:
1. GODSEED_GGUF env — local path to a .gguf file;
2. hf_hub_download(GODSEED_GGUF_REPO, GODSEED_GGUF_FILE) — defaults
target the Nemotron-3-Nano-4B Q4_K_M quant (~2.5GB, free CPU tier).
Generation runs in a thread executor; chunks cross into the event loop
through a bounded asyncio.Queue (backpressure included). If the consumer
abandons the stream early, a cancel flag stops the producer thread.
"""
name = "llamacpp"
oneshot = True # one generation per wish (fast on CPU; lenient JSON parse +
# engine forgiveness + the deterministic town fallback cover messy output)
DEFAULT_REPO = "nvidia/NVIDIA-Nemotron-3-Nano-4B-GGUF"
DEFAULT_FILE = "NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"
def __init__(
self,
model_path: str | None = None,
n_ctx: int = 4096,
n_threads: int | None = None,
):
self.model_path = model_path or os.environ.get("GODSEED_GGUF") or None
self.repo_id = os.environ.get("GODSEED_GGUF_REPO", self.DEFAULT_REPO)
self.filename = os.environ.get("GODSEED_GGUF_FILE", self.DEFAULT_FILE)
self.n_ctx = n_ctx
self.n_threads = n_threads
self.model_id = self.model_path or f"{self.repo_id}:{self.filename}"
self._llm = None
self._lock = threading.Lock()
self._grammar_cache: dict[str, object] = {}
# -- lazy model load (thread-safe; called from executor threads) --------
def _ensure_llm(self):
with self._lock:
if self._llm is not None:
return self._llm
from llama_cpp import Llama # heavy import deferred
path = self.model_path
if not path or not os.path.exists(path):
from huggingface_hub import hf_hub_download
path = hf_hub_download(repo_id=self.repo_id, filename=self.filename)
kwargs = {"model_path": path, "n_ctx": self.n_ctx, "verbose": False}
# Use all available cores on the Space CPU (the GGUF runs CPU-only).
threads = self.n_threads or (os.cpu_count() or 4)
kwargs["n_threads"] = threads
kwargs["n_batch"] = 512 # larger prompt-eval batch → faster prefill
self._llm = Llama(**kwargs)
self.model_id = path
return self._llm
def _compile_grammar(self, grammar: str):
cached = self._grammar_cache.get(grammar)
if cached is not None:
return cached
from llama_cpp import LlamaGrammar
compiled = LlamaGrammar.from_string(grammar)
self._grammar_cache[grammar] = compiled
return compiled
async def generate_stream(
self, prompt: str, grammar: str | None = None, max_tokens: int = 256
) -> AsyncIterator[str]:
loop = asyncio.get_running_loop()
queue: asyncio.Queue = asyncio.Queue(maxsize=64)
cancelled = threading.Event()
temperature = 0.7 if grammar is None else 0.35
def _put(item) -> bool:
"""Push from the producer thread; honor cancellation."""
future = asyncio.run_coroutine_threadsafe(queue.put(item), loop)
while not cancelled.is_set():
try:
future.result(timeout=0.25)
return True
except concurrent.futures.TimeoutError:
continue
except Exception:
return False
future.cancel()
return False
def _produce() -> None:
try:
llm = self._ensure_llm()
compiled = self._compile_grammar(grammar) if grammar else None
stream = llm.create_completion(
prompt=prompt,
max_tokens=max_tokens,
stream=True,
grammar=compiled,
temperature=temperature,
top_p=0.9,
repeat_penalty=1.1,
)
for part in stream:
if cancelled.is_set():
break
text = part["choices"][0].get("text", "")
if text and not _put(text):
break
except Exception as exc: # surfaced to the consumer below
_put(exc)
finally:
_put(_THREAD_DONE)
producer = loop.run_in_executor(None, _produce)
try:
while True:
item = await queue.get()
if item is _THREAD_DONE:
break
if isinstance(item, Exception):
raise item
yield item
finally:
cancelled.set()
# Drain so a blocked producer can reach its sentinel and exit.
while not queue.empty():
queue.get_nowait()
await producer
# --------------------------------------------------------------------------
# ZeroGPU backend — Nemotron 3 Nano 30B-A3B (headline upgrade path)
# --------------------------------------------------------------------------
class ZeroGPUBackend:
"""transformers + spaces.GPU backend (HF ZeroGPU).
All heavy imports (torch, transformers, spaces) are deferred to the first
generate call so this module imports cleanly anywhere. The GPU-decorated
function is built lazily; if `spaces` is unavailable (local dev) the bare
function is used.
transformers has no GBNF support, so when a grammar is supplied the
output is validated strictly (mind/validate.py) and regenerated once at
temperature ~0 on failure. If both attempts are malformed the raw text is
yielded anyway — the planner's re-ask/skip path handles it from there.
Streaming note: tokens are buffered per generation (validation requires
the full text) and then yielded in small chunks, so SSE consumers still
see incremental output.
NOTE: glassblower's pins (transformers==4.48.3, mamba-ssm) were for
Nemotron-Nano-9B-v2 and do NOT apply to this MoE; the 30B-A3B needs a
newer transformers. Verify versions at deploy time.
"""
name = "zerogpu"
oneshot = True # generate the whole wish in ONE @spaces.GPU call (ZeroGPU
# detaches the GPU between calls; the per-turn loop's ~8 calls crash on #2)
# Nemotron-Nano-9B-v2: its nemotron_h arch is IN-TREE in transformers, so it
# loads with trust_remote_code=False and runs the kernel-free native Mamba-2
# path. The Nemotron-3 4B/30B repos hard-require remote code + mamba-ssm
# kernels, which cannot import on the ZeroGPU image (no compatible CUDA libs
# at startup — June 12). 4B GGUF remains the llama.cpp local mode.
DEFAULT_MODEL = "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
GPU_DURATION_S = 240
def __init__(self, model_id: str | None = None, max_input_tokens: int = 4096):
self.model_id = model_id or os.environ.get(
"GODSEED_HF_MODEL", self.DEFAULT_MODEL
)
self.max_input_tokens = max_input_tokens
self._tokenizer = None
self._model = None
self._gpu_generate = None
self._on_cuda = False
self._lock = threading.Lock()
# Load weights to CPU at construction (app startup). CRITICAL: do NOT
# touch CUDA here. Initializing CUDA in the main process poisons
# ZeroGPU's per-request GPU attach (the forked worker inherits a broken
# context → "NVML_SUCCESS INTERNAL ASSERT FAILED" in the CUDA allocator
# on every forward, June 12). The model moves to cuda LAZILY inside the
# @spaces.GPU function, where CUDA is correctly attached.
self._ensure()
# -- setup (model on CPU; cuda placement deferred to the GPU context) ------
def _ensure(self):
with self._lock:
if self._gpu_generate is not None:
return self._gpu_generate
import torch # noqa: F401
from transformers import AutoModelForCausalLM, AutoTokenizer
# trust_remote_code MUST stay False: the repo also ships custom code
# that hard-requires mamba-ssm kernels (unimportable on ZeroGPU);
# False forces the in-tree NemotronH class with its native fallback.
self._tokenizer = AutoTokenizer.from_pretrained(self.model_id)
self._model = AutoModelForCausalLM.from_pretrained(
self.model_id,
dtype=torch.bfloat16, # v5 name (torch_dtype removed in transformers 5)
low_cpu_mem_usage=True,
)
self._model.eval()
# Move to cuda ONCE at load, guarded by `import spaces` (the lib
# virtualizes it and keeps the model resident across @spaces.GPU
# calls). With ONE-SHOT generation a wish is a single GPU call, so
# the multi-call NVML crash can't happen — and keeping the model
# resident avoids the ~67s CPU→GPU transfer that, done per-call, ate
# the whole GPU time budget and left no time to generate (→ empty
# output → fallback, June 12). Resident + one call = fast real output.
try:
import spaces # noqa: F401
self._model.to("cuda")
except ImportError:
pass # local/CPU dev
def _generate(prompt: str, max_new_tokens: int, temperature: float) -> str:
import torch as _torch # noqa: F401
# Nemotron-Nano-9B-v2 is a reasoning model that THINKS by
# default; "/no_think" in the system slot disables it (June 12
# bench: raw completion leaked into outputs). The
# planner's full completion-style prompt rides as one user turn.
if getattr(self._tokenizer, "chat_template", None):
text_in = self._tokenizer.apply_chat_template(
[
{"role": "system", "content": "/no_think"},
{"role": "user", "content": prompt},
],
tokenize=False,
add_generation_prompt=True,
)
else:
text_in = prompt
inputs = self._tokenizer(
text_in,
return_tensors="pt",
truncation=True,
max_length=self.max_input_tokens,
).to(self._model.device)
with _torch.no_grad():
output = self._model.generate(
**inputs,
max_new_tokens=max_new_tokens,
do_sample=temperature > 0.05,
temperature=max(temperature, 0.05),
top_p=0.9,
pad_token_id=self._tokenizer.eos_token_id,
)
new_tokens = output[0][inputs["input_ids"].shape[1]:]
text = self._tokenizer.decode(new_tokens, skip_special_tokens=True)
# Defensive: strip any (possibly empty) think block that slips out.
text = re.sub(r".*?", "", text, flags=re.S)
text = text.replace("", "").replace("", "")
return text.strip()
try:
import spaces # import-guarded at call time, per contract
self._gpu_generate = spaces.GPU(duration=self.GPU_DURATION_S)(_generate)
except Exception:
self._gpu_generate = _generate
return self._gpu_generate
@staticmethod
def _is_valid(grammar: str, text: str) -> bool:
if is_moderation_grammar(grammar):
return parse_moderation(text)[1] is None
return parse_turn(text)[1] is None
async def _gen(self, generate, prompt: str, max_tokens: int, temperature: float) -> str:
"""Run one generation, retrying transient ZeroGPU CUDA/NVML failures.
Each @spaces.GPU call re-schedules onto a (possibly healthier) GPU, so a
retry can clear the intermittent NVML allocator assert seen June 12."""
loop = asyncio.get_running_loop()
last_exc = None
for attempt in range(3):
try:
return await loop.run_in_executor(None, generate, prompt, max_tokens, temperature)
except Exception as exc: # noqa: BLE001
msg = str(exc)
transient = any(s in msg for s in ("NVML", "CUDA", "cuda", "device", "out of memory"))
last_exc = exc
if not transient or attempt == 2:
raise
# drop the poisoned cuda placement so the next call re-homes the model
self._on_cuda = False
import torch as _t
try:
_t.cuda.empty_cache()
except Exception:
pass
await asyncio.sleep(0.5 * (attempt + 1))
raise last_exc # pragma: no cover
async def generate_stream(
self, prompt: str, grammar: str | None = None, max_tokens: int = 256
) -> AsyncIterator[str]:
loop = asyncio.get_running_loop()
generate = await loop.run_in_executor(None, self._ensure)
if grammar is None:
text = await self._gen(generate, prompt, max_tokens, 0.7)
# chat-templated models tend to echo the prompt's trailing label
text = re.sub(r"^\s*\**\s*READING:?\s*\**\s*", "", text)
else:
strict_prompt = (
prompt + "\nReply with exactly one JSON object and nothing else."
)
text = await self._gen(generate, strict_prompt, max_tokens, 0.4)
if not self._is_valid(grammar, text):
retry_prompt = (
strict_prompt
+ "\nYour previous output was malformed. Output ONLY the JSON object."
)
text = await self._gen(generate, retry_prompt, max_tokens, 0.0)
# Still malformed? Yield as-is; the planner re-asks then skips.
for chunk in _chunks(text, 24):
await asyncio.sleep(0)
yield chunk
# --------------------------------------------------------------------------
# Factory
# --------------------------------------------------------------------------
_BACKEND_ALIASES = {
"mock": "mock",
"llamacpp": "llamacpp",
"llama": "llamacpp",
"gguf": "llamacpp",
"zerogpu": "zerogpu",
"zero-gpu": "zerogpu",
"transformers": "zerogpu",
}
def make_backend(name: str | None = None):
"""Build the backend selected by `name` or the GODSEED_BACKEND env.
Defaults to mock (always works, zero deps beyond stdlib).
"""
raw = (name or os.environ.get("GODSEED_BACKEND") or "mock").strip().lower()
resolved = _BACKEND_ALIASES.get(raw)
if resolved == "mock":
return MockBackend()
if resolved == "llamacpp":
return LlamaCppBackend()
if resolved == "zerogpu":
return ZeroGPUBackend()
raise ValueError(
f"unknown GODSEED_BACKEND {raw!r} (expected mock | llamacpp | zerogpu)"
)