"""Writer contract for producing schema-valid radio scripts."""
from __future__ import annotations
import json
import os
import re
import sys
from dataclasses import asdict
from functools import cached_property
from typing import Protocol
from .fixtures import fixture_script
from .schema import Genre, Script, script_from_json_dict
class ScriptWriter(Protocol):
"""Interface shared by fixture, Nemotron, and fine-tuned writers."""
name: str
def write(
self, premise: str, genre: Genre = Genre.WEIRD, feedback: str | None = None
) -> Script:
"""Return a validated script for the given premise and genre.
`feedback` carries the error from a previous failed attempt so model
writers can self-correct on a retry; deterministic writers ignore it.
"""
class FixtureWriter:
"""Deterministic schema-valid writer used until model integration lands."""
name = "fixture"
def write(
self, premise: str, genre: Genre = Genre.WEIRD, feedback: str | None = None
) -> Script:
return fixture_script(premise, genre=genre)
class NemotronWriter:
"""Nemotron-backed writer using Transformers generation.
The model is lazy-loaded so importing this class never downloads weights.
Runtime configuration:
- `MIDNIGHT_NEMOTRON_MODEL`: HF model id.
- `MIDNIGHT_NEMOTRON_DEVICE_MAP`: Transformers device map, default `auto`.
- `MIDNIGHT_NEMOTRON_MAX_NEW_TOKENS`: generation cap, default `4000`.
Note: Nemotron 3 Nano is a reasoning model — it emits thinking tokens before
the JSON answer, so the cap must be generous (a GPU smoke showed ~8.6k chars
of output; 1800 tokens truncated the JSON). See modal/smoke_nemotron.py.
"""
name = "nemotron"
def __init__(
self,
model_id: str | None = None,
device_map: str | None = None,
max_new_tokens: int | None = None,
) -> None:
self.model_id = model_id or os.getenv(
"MIDNIGHT_NEMOTRON_MODEL",
"nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16",
)
self.device_map = device_map or os.getenv("MIDNIGHT_NEMOTRON_DEVICE_MAP", "auto")
self.max_new_tokens = max_new_tokens or int(
os.getenv("MIDNIGHT_NEMOTRON_MAX_NEW_TOKENS", "4000")
)
def write(
self, premise: str, genre: Genre = Genre.WEIRD, feedback: str | None = None
) -> Script:
prompt = build_writer_prompt(premise, genre, feedback=feedback)
raw = self._generate_text(prompt)
return parse_writer_json(raw)
@cached_property
def _model_bundle(self):
try:
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
except ImportError as exc: # pragma: no cover - environment guard
raise RuntimeError(
"NemotronWriter requires the writer/model dependencies. "
"Run `uv sync --extra dev --extra tts` for the current scaffold."
) from exc
tokenizer = AutoTokenizer.from_pretrained(self.model_id, trust_remote_code=True)
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
model = AutoModelForCausalLM.from_pretrained(
self.model_id,
torch_dtype=dtype,
device_map=self.device_map,
trust_remote_code=True,
)
return tokenizer, model
def _generate_text(self, prompt: str) -> str:
tokenizer, model = self._model_bundle
messages = [
{
"role": "system",
"content": "You are Midnight Static's radio-drama scriptwriter.",
},
{"role": "user", "content": prompt},
]
if hasattr(tokenizer, "apply_chat_template"):
input_ids = tokenizer.apply_chat_template(
messages,
add_generation_prompt=True,
return_tensors="pt",
)
else:
input_ids = tokenizer(prompt, return_tensors="pt").input_ids
input_ids = input_ids.to(model.device)
outputs = model.generate(
input_ids,
max_new_tokens=self.max_new_tokens,
do_sample=True,
temperature=0.75,
top_p=0.9,
pad_token_id=tokenizer.eos_token_id,
)
generated = outputs[0][input_ids.shape[-1] :]
return tokenizer.decode(generated, skip_special_tokens=True)
class LlamaCppNemotronWriter:
"""The real Nemotron 3 Nano 4B running **locally on CPU** via llama.cpp.
Uses the GGUF build (`nemotron_h` architecture) so it needs neither a GPU nor
the `mamba-ssm` CUDA kernels — which is what lets the live cpu-basic Space run
the genuine model and keep the off-the-grid claim. Reasoning is disabled
(the chat template's empty ````) for speed, and a wall-clock
deadline + token cap bound the latency; on overrun the streamed text is
returned as-is, so an incomplete script fails validation and
``generate_script`` falls back to the fixture.
Runtime configuration:
- `MIDNIGHT_GGUF_REPO` / `MIDNIGHT_GGUF_FILE`: override the GGUF source.
- `MIDNIGHT_GGUF_MAX_TOKENS`: generation cap, default `800`.
- `MIDNIGHT_GGUF_DEADLINE`: wall-clock budget in seconds, default `150`.
- `MIDNIGHT_GGUF_THREADS`: llama.cpp threads, default `0` (llama.cpp auto).
"""
name = "nemotron"
def __init__(
self,
repo_id: str | None = None,
filename: str | None = None,
max_tokens: int | None = None,
deadline_seconds: float | None = None,
) -> None:
self.repo_id = repo_id or os.getenv(
"MIDNIGHT_GGUF_REPO", "lmstudio-community/NVIDIA-Nemotron-3-Nano-4B-GGUF"
)
self.filename = filename or os.getenv(
"MIDNIGHT_GGUF_FILE", "NVIDIA-Nemotron-3-Nano-4B-Q4_K_M.gguf"
)
self.max_tokens = max_tokens or int(os.getenv("MIDNIGHT_GGUF_MAX_TOKENS", "800"))
self.deadline_seconds = deadline_seconds or float(
os.getenv("MIDNIGHT_GGUF_DEADLINE", "150")
)
@cached_property
def _llm(self):
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
threads = int(os.getenv("MIDNIGHT_GGUF_THREADS", "0")) or None
model_path = hf_hub_download(self.repo_id, self.filename)
return Llama(
model_path=model_path,
n_ctx=4096,
n_threads=threads,
n_gpu_layers=0,
verbose=False,
)
def write(
self, premise: str, genre: Genre = Genre.WEIRD, feedback: str | None = None
) -> Script:
import time
prompt = _nemotron_chatml(
system="You are Midnight Static's radio-drama scriptwriter. Reply with ONLY the JSON script, no prose.",
user=build_writer_prompt(premise, genre, feedback=feedback),
)
deadline = time.monotonic() + self.deadline_seconds
chunks: list[str] = []
for part in self._llm.create_completion(
prompt,
max_tokens=self.max_tokens,
temperature=0.7,
top_p=0.9,
stop=["<|im_end|>"],
stream=True,
):
chunks.append(part["choices"][0]["text"])
if time.monotonic() > deadline:
break
return parse_writer_json("".join(chunks))
def _nemotron_chatml(system: str, user: str) -> str:
"""Reasoning-off ChatML prompt for Nemotron 3 Nano (empty )."""
return (
f"<|im_start|>system\n{system}<|im_end|>\n"
f"<|im_start|>user\n{user}<|im_end|>\n"
f"<|im_start|>assistant\n"
)
def get_writer(name: str = "fixture") -> ScriptWriter:
if name == "fixture":
return FixtureWriter()
if name == "nemotron":
# Live default: the real model on CPU via llama.cpp (GGUF).
return LlamaCppNemotronWriter()
if name == "nemotron-hf":
# Transformers/GPU path (needs CUDA + mamba-ssm); used off the Space.
return NemotronWriter()
raise ValueError(f"unknown writer: {name}")
# Failures we treat as "the model produced an unusable script": malformed JSON,
# missing/extra keys, or cross-field validation errors from the schema.
WRITER_FAILURES = (ValueError, KeyError, TypeError, json.JSONDecodeError)
# Failures that mean the selected writer cannot run in this environment at all —
# e.g. Nemotron on the cpu-basic Space (no GPU, no mamba-ssm CUDA kernels) or a
# missing torch/transformers. Retrying won't help; we drop straight to the
# fixture writer so the station never shows an error page.
WRITER_UNAVAILABLE = (ImportError, RuntimeError, OSError)
def generate_script(
writer: ScriptWriter,
premise: str,
genre: Genre = Genre.WEIRD,
*,
retries: int = 1,
fixture_fallback: bool = True,
) -> Script:
"""Produce a script with one retry on invalid output, then a fixture fallback.
Mirrors the ARCHITECTURE.md policy: a schema/JSON failure triggers one retry
with the error fed back to the writer; a second failure improvises with the
genre fixture (woven around the user's premise) so the station never shows an
error page.
"""
last_error: Exception | None = None
feedback: str | None = None
for _ in range(retries + 1):
try:
return writer.write(premise, genre=genre, feedback=feedback)
except WRITER_FAILURES as exc:
last_error = exc
feedback = _format_feedback(exc)
print(
f"[writer] {getattr(writer, 'name', '?')} produced unusable output "
f"({type(exc).__name__}: {exc}); will retry/fallback",
file=sys.stderr,
flush=True,
)
except WRITER_UNAVAILABLE as exc:
# The writer can't load here (e.g. Nemotron on CPU). No point
# retrying — go straight to the fixture fallback below.
last_error = exc
print(
f"[writer] {getattr(writer, 'name', '?')} unavailable "
f"({type(exc).__name__}: {exc}); falling back to fixture",
file=sys.stderr,
flush=True,
)
break
if fixture_fallback:
return fixture_script(premise, genre=genre)
raise last_error # type: ignore[misc]
def _format_feedback(error: Exception) -> str:
return (
"Your previous attempt was rejected: "
f"{type(error).__name__}: {error}. "
"Return one corrected JSON object that satisfies every schema rule."
)
def script_to_json_dict(script: Script) -> dict:
script.validate()
return asdict(script)
def parse_writer_json(raw: str) -> Script:
payload = json.loads(extract_json_object(raw))
return script_from_json_dict(payload)
def extract_json_object(raw: str) -> str:
"""Extract the first top-level JSON object from model output."""
text = raw.strip()
if text.startswith("```"):
text = re.sub(r"^```(?:json)?\s*", "", text)
text = re.sub(r"\s*```$", "", text)
start = text.find("{")
if start < 0:
raise ValueError("writer output did not contain a JSON object")
depth = 0
in_string = False
escaped = False
for index, char in enumerate(text[start:], start=start):
if escaped:
escaped = False
continue
if char == "\\":
escaped = True
continue
if char == '"':
in_string = not in_string
continue
if in_string:
continue
if char == "{":
depth += 1
elif char == "}":
depth -= 1
if depth == 0:
return text[start : index + 1]
raise ValueError("writer output contained incomplete JSON")
def build_writer_prompt(premise: str, genre: Genre, feedback: str | None = None) -> str:
premise = " ".join(premise.strip().split())[:300] or "a signal with no sender"
correction = f"\nCorrection required: {feedback}\n" if feedback else ""
voices = [
"am_michael",
"am_onyx",
"af_bella",
"am_puck",
"am_fenrir",
"bm_george",
"hf_alpha",
"hm_omega",
]
deliveries = ["neutral", "slow", "urgent", "whisper", "booming", "deadpan", "agitated"]
return f"""
Write a 60-90 second vintage radio-drama script as strict JSON only.
Premise: {premise}
Genre: {genre.value}
{correction}
Rules:
- Return one JSON object and no markdown.
- Use exactly the schema keys shown below.
- `genre` and `music.genre` must be "{genre.value}".
- Use 2-4 cast members.
- Cast `voice` must be one of: {", ".join(voices)}.
- Every line `cast` must reference a cast member id.
- Every line `delivery` must be one of: {", ".join(deliveries)}.
- Use 2-3 scenes and 8-18 total dialogue lines.
- Keep content PG-13. Avoid real public figures, explicit gore, and slurs.
- For hindi_melodrama, dialogue must be romanized Hinglish only.
Schema example:
{{
"title": "The Case of the Vanishing Signal",
"logline": "A one-sentence hook.",
"genre": "{genre.value}",
"cast": [
{{"id": "announcer", "name": "The Announcer", "voice": "am_michael", "description": "Warm station voice."}},
{{"id": "caller", "name": "The Caller", "voice": "af_bella", "description": "Nervous lead."}}
],
"scenes": [
{{
"title": "Scene title",
"sfx": ["radio static", "distant thunder"],
"lines": [
{{"cast": "announcer", "text": "Line text.", "delivery": "booming"}}
]
}}
],
"music": {{
"genre": "{genre.value}",
"opening_sting": "short music prompt",
"bed": "loopable bed prompt",
"closing_sting": "short closing prompt"
}},
"estimated_seconds": 75
}}
""".strip()
SMOKE_PREMISES = [
"a phone that only receives calls from 1962",
"two rival lighthouse keepers, one light",
"static",
"the last petrol pump before the salt flats",
"a wedding where every guest recognizes the wrong bride",
]