lars
c1c4957e2f
Stage 2's autoregressive decoder still predicted multiplicity the v0.2 way: a one-shot n_sec_head classifier over conditioning alone, run before any secondary token existed, with the AR loop then always executing k_max slots and discarding the tail. This adds a real per-slot EOS mechanism instead: - Stage2Autoregressive gains a stop_head (build_stop_head=True) that predicts P(n_sec == k | prefix) at each slot, mutually exclusive with n_sec_head (n_sec.mode = "stop_token" builds no n_sec_head at all). - sample_secondaries_ar accepts n_sec_pred=None to drive generation off the stop head instead of a pre-resolved count: each row stops the first slot its stop logit fires (stage2_model.n_sec.stop_sampling = "greedy" — the default, threshold at 0 — or "sample", a Bernoulli draw), and the whole batch loop breaks once every row has stopped, so cost scales with the realized n_sec instead of a fixed k_max. Passing n_sec_pred explicitly (the scheduled-sampling self-sample path) is unchanged. - resolve_n_sec returns None for a stop-token decoder instead of raising; rollout.py/cli.py/validate.py now derive the realized count from sample_stage2's returned sec_valid (sec_valid.sum(-1)) after sampling, rather than resolving it up front — a no-op reordering under every other n_sec.mode, where sec_valid was already built from n_sec_pred. - Training: _stop_target_and_mask (giant/training/stage2_inputs.py) builds the per-slot target/mask (one slot wider than the existing token-content sec_mask, since the stop slot itself needs supervision) and StageTrainer._stop_loss trains it with masked BCE, gated on stop_head exactly like _n_sec_loss gates on n_sec_head. Wired into both the flow/ddpm trainer and the WGAN trainer (whose skip_g_step now also checks stop_head), weighted by the existing stage2_model.n_sec.lambda — the stop head replaces n_sec_head under this mode, so no new weight key. - validate_config now accepts stop_token (requires decoder="autoregressive" and n_sec.owner="stage2") instead of always rejecting it. Decisions made during planning: stop_sampling defaults to "greedy" for deterministic rollouts; the stop head reuses stage2_model.heads.n_sec's HeadConfig shape and stage2_model.n_sec.lambda's weight rather than adding new config keys, since the two heads never coexist. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
512 lines
22 KiB
Python
512 lines
22 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from giant.constants import CONT_SLOT_DIM, X_DIM
|
|
from giant.model.network import DdpmObjective, Stage2Autoregressive, build_objective, stage2_trunk_sec_dim
|
|
from giant.model.schedule import CosineSchedule
|
|
|
|
|
|
def _predict_n_sec_if_owned(
|
|
model: torch.nn.Module, cond_cont: torch.Tensor, cond_cat: torch.Tensor
|
|
) -> torch.Tensor | None:
|
|
"""Stage-1 `n_sec_head` is only present on a migrated v0.2 checkpoint
|
|
(fresh runs move it to stage 2 — see `Stage1Model`'s docstring). `None`
|
|
here means "ask stage 2 instead", which every caller (`giant/rollout.py`,
|
|
`giant/cli.py`) must do for a fresh checkpoint."""
|
|
if getattr(model, "n_sec_head", None) is None:
|
|
return None
|
|
logits = model.predict_n_sec(cond_cont, cond_cat)
|
|
return logits.argmax(dim=-1)
|
|
|
|
|
|
@torch.no_grad()
|
|
def sample_flow(
|
|
model: torch.nn.Module,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
steps: int = 10,
|
|
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
|
"""Euler integration of the Stage-1 vector field from t=0 to t=1.
|
|
|
|
Returns (primary_sample, n_sec_pred):
|
|
primary_sample: (B, X_DIM) — normalised 9D primary post-step output
|
|
n_sec_pred: (B,) int64 — predicted secondary count, or `None` if
|
|
`model` has no `n_sec_head` (a fresh v0.3.0 Stage1Model — see
|
|
`_predict_n_sec_if_owned`).
|
|
"""
|
|
model.eval()
|
|
B = cond_cont.size(0)
|
|
device = cond_cont.device
|
|
x = torch.randn(B, X_DIM, device=device)
|
|
dt = 1.0 / steps
|
|
for i in range(steps):
|
|
t = torch.full((B,), i * dt, device=device)
|
|
v = model(x, cond_cont, cond_cat, t=t)
|
|
x = x + v * dt
|
|
return x, _predict_n_sec_if_owned(model, cond_cont, cond_cat)
|
|
|
|
|
|
@torch.no_grad()
|
|
def sample_ddpm(
|
|
model: torch.nn.Module,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
schedule,
|
|
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
|
"""Full DDPM ancestral sampling (T reverse steps). Returns (sample, n_sec_pred)
|
|
— see `sample_flow`'s docstring for the `n_sec_pred` `None` case."""
|
|
model.eval()
|
|
B = cond_cont.size(0)
|
|
device = cond_cont.device
|
|
x = torch.randn(B, X_DIM, device=device)
|
|
T = schedule.T
|
|
for i in reversed(range(T)):
|
|
t_norm = torch.full((B,), i / T, device=device)
|
|
eps_pred = model(x, cond_cont, cond_cat, t=t_norm)
|
|
beta = schedule.betas[i]
|
|
alpha = schedule.alphas[i]
|
|
alpha_bar = schedule.alpha_bars[i]
|
|
z = torch.randn_like(x) if i > 0 else torch.zeros_like(x)
|
|
x = (1.0 / alpha.sqrt()) * (x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred) + beta.sqrt() * z
|
|
return x, _predict_n_sec_if_owned(model, cond_cont, cond_cat)
|
|
|
|
|
|
@torch.no_grad()
|
|
def sample_ddim(
|
|
model: torch.nn.Module,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
schedule,
|
|
steps: int = 50,
|
|
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
|
"""DDIM deterministic sampling (Song et al. 2020). Returns (sample, n_sec_pred)
|
|
— see `sample_flow`'s docstring for the `n_sec_pred` `None` case."""
|
|
model.eval()
|
|
B = cond_cont.size(0)
|
|
device = cond_cont.device
|
|
T = schedule.T
|
|
timesteps = torch.linspace(T - 1, 0, steps, dtype=torch.long, device=device)
|
|
x = torch.randn(B, X_DIM, device=device)
|
|
for step_idx, ts in enumerate(timesteps):
|
|
t_idx = int(ts.item())
|
|
t_norm = torch.full((B,), t_idx / T, device=device)
|
|
eps_pred = model(x, cond_cont, cond_cat, t=t_norm)
|
|
ab_t = schedule.alpha_bars[t_idx]
|
|
if step_idx + 1 < len(timesteps):
|
|
ab_prev = schedule.alpha_bars[int(timesteps[step_idx + 1].item())]
|
|
else:
|
|
ab_prev = torch.ones(1, device=device)
|
|
x0_pred = (x - (1.0 - ab_t).sqrt() * eps_pred) / ab_t.sqrt()
|
|
x = ab_prev.sqrt() * x0_pred + (1.0 - ab_prev).sqrt() * eps_pred
|
|
return x, _predict_n_sec_if_owned(model, cond_cont, cond_cat)
|
|
|
|
|
|
@torch.no_grad()
|
|
def sample_wgan(
|
|
generator: torch.nn.Module,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
|
"""Single-pass Stage-1 WGAN generator sample. Returns (sample, n_sec_pred)
|
|
— see `sample_flow`'s docstring for the `n_sec_pred` `None` case."""
|
|
generator.eval()
|
|
B = cond_cont.size(0)
|
|
z = torch.randn(B, generator.noise_dim, device=cond_cont.device)
|
|
x = generator(z, cond_cont, cond_cat)
|
|
return x, _predict_n_sec_if_owned(generator, cond_cont, cond_cat)
|
|
|
|
|
|
def _stage2_flat_width(sec_decoder: torch.nn.Module) -> int:
|
|
"""The width of `sec_decoder`'s own trunk in/out vector — folded
|
|
(continuous + type) under `particle_type.target = "physical"` or
|
|
`generator = "wgan"`, continuous-only otherwise (the type slice then
|
|
comes from `predict_type` instead — see `stage2_trunk_sec_dim`'s
|
|
docstring)."""
|
|
return stage2_trunk_sec_dim(
|
|
sec_decoder.particle_type_cfg,
|
|
sec_decoder.generator_kind,
|
|
sec_decoder.k_max,
|
|
sec_decoder.type_dim,
|
|
)
|
|
|
|
|
|
def _type_folded(sec_decoder: torch.nn.Module) -> bool:
|
|
target = sec_decoder.particle_type_cfg.target
|
|
return target == "physical" or build_objective(sec_decoder.generator_kind).folds_type_slice
|
|
|
|
|
|
def _decode_stage2_flat(
|
|
sec_decoder: torch.nn.Module,
|
|
x: torch.Tensor,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
stage1_out: torch.Tensor,
|
|
n_sec_pred: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
"""Reshape a flat `(B, flat_width)` `Stage2OneShot` output into per-slot
|
|
tensors, generator/`particle_type.target`-agnostic: shared by
|
|
`sample_secondaries`/`sample_secondaries_wgan`, which differ only in how
|
|
`x` was produced.
|
|
|
|
Returns (sec_cont, sec_type, sec_valid):
|
|
sec_cont: (B, k_max, CONT_SLOT_DIM) — [stick_logit, local_dir]
|
|
sec_type: (B, k_max, type_dim) — under `target="physical"` this is
|
|
[log_mass, charge] (normalised iff the checkpoint's sec_phys
|
|
normalizer was applied at training time — denormalize before
|
|
treating as physical units; see
|
|
giant.data.transforms.decode_secondaries); under `"onehot"` /
|
|
`"embedding"` it is raw class logits / an embedding-space vector
|
|
— decode via giant.particles.decode_topn_class /
|
|
decode_embedding_nearest (see giant/rollout.py).
|
|
sec_valid: (B, k_max) bool — True for slots i < n_sec_pred
|
|
"""
|
|
B = x.size(0)
|
|
device = x.device
|
|
k_max = sec_decoder.k_max
|
|
type_dim = sec_decoder.type_dim
|
|
if _type_folded(sec_decoder):
|
|
x_slots = x.view(B, k_max, CONT_SLOT_DIM + type_dim)
|
|
sec_cont = x_slots[:, :, :CONT_SLOT_DIM]
|
|
sec_type = x_slots[:, :, CONT_SLOT_DIM:]
|
|
else:
|
|
sec_cont = x.view(B, k_max, CONT_SLOT_DIM)
|
|
sec_type = sec_decoder.predict_type(cond_cont, cond_cat, stage1_out)
|
|
sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(1)
|
|
return sec_cont, sec_type, sec_valid
|
|
|
|
|
|
@torch.no_grad()
|
|
def sample_secondaries(
|
|
sec_decoder: torch.nn.Module,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
stage1_out: torch.Tensor,
|
|
n_sec_pred: torch.Tensor,
|
|
steps: int = 10,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
"""Euler integration of `Stage2OneShot`'s (flow/ddpm) vector field; return
|
|
raw slot outputs — see `_decode_stage2_flat`'s docstring for the returned
|
|
(sec_cont, sec_type, sec_valid) shapes/meaning.
|
|
|
|
n_sec_pred: (B,) int64 — number of valid secondaries per step
|
|
"""
|
|
sec_decoder.eval()
|
|
B = cond_cont.size(0)
|
|
device = cond_cont.device
|
|
flat_width = _stage2_flat_width(sec_decoder)
|
|
|
|
x = torch.randn(B, flat_width, device=device)
|
|
dt = 1.0 / steps
|
|
for i in range(steps):
|
|
t = torch.full((B,), i * dt, device=device)
|
|
v = sec_decoder(x, cond_cont, cond_cat, stage1_out, t=t)
|
|
x = x + v * dt
|
|
|
|
return _decode_stage2_flat(sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred)
|
|
|
|
|
|
@torch.no_grad()
|
|
def sample_secondaries_wgan(
|
|
sec_decoder: torch.nn.Module,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
stage1_out: torch.Tensor,
|
|
n_sec_pred: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
"""Single-pass `Stage2OneShot` WGAN generator sample; see
|
|
`_decode_stage2_flat`'s docstring for the returned (sec_cont, sec_type,
|
|
sec_valid) shapes/meaning."""
|
|
sec_decoder.eval()
|
|
B = cond_cont.size(0)
|
|
z = torch.randn(B, sec_decoder.noise_dim, device=cond_cont.device)
|
|
x = sec_decoder(z, cond_cont, cond_cat, stage1_out)
|
|
return _decode_stage2_flat(sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred)
|
|
|
|
|
|
@torch.no_grad()
|
|
def sample_secondaries_ar(
|
|
sec_decoder: torch.nn.Module,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
stage1_out: torch.Tensor,
|
|
n_sec_pred: torch.Tensor | None,
|
|
steps: int = 10,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
"""`Stage2Autoregressive` inference loop: one token at a time, in
|
|
descending-energy slot order, up to `k_max` sequential calls. Unlike
|
|
training (teacher forcing — a single parallel pass over ground-truth
|
|
tokens, see `giant.training.stage2_inputs._assemble_stage2_ar_inputs`),
|
|
there is no ground truth at inference: each token's conditioning is built
|
|
free-running, from the PREVIOUS TOKEN'S OWN just-generated output — the
|
|
train/inference gap that is the cost of markov history's
|
|
expressiveness.
|
|
|
|
A `{flow,ddpm}` token costs `steps` ODE substeps; `wgan` costs one pass —
|
|
the "K sequential forwards" cost applies per-token here, not
|
|
once, so a flow/ddpm AR run costs ~`k_max * steps` model calls per
|
|
physics step (or ~`n_sec * steps` under `n_sec_pred=None` below, once
|
|
every row in the batch has stopped).
|
|
|
|
`n_sec_pred`, if given, fixes each row's secondary count up front (as
|
|
resolved by `resolve_n_sec` — `n_sec.mode` in `("head", "truth")`, or a
|
|
stop-token decoder driven by `_assemble_stage2_ar_inputs_scheduled`'s
|
|
ground-truth `n_sec`, which must run the *full* `k_max`-length free-
|
|
running self-sample regardless of the decoder's own stop head — the
|
|
scheduled-sampling training contract does not truncate). This always
|
|
runs the full `k_max`-iteration loop, masking by the given count at the
|
|
end exactly as before.
|
|
|
|
`n_sec_pred=None` is only valid when `sec_decoder.stop_head` is set
|
|
(`n_sec.mode = "stop_token"`): before generating each slot's token, that
|
|
slot's own stop logit (`predict_stop`, evaluated on the same prefix
|
|
conditioning as the token itself — see `predict_type`'s docstring for
|
|
why this needs no extra state) decides whether generation should have
|
|
already stopped, per `sec_decoder.stop_sampling` ("greedy": threshold at
|
|
0; "sample": a Bernoulli draw at `sigmoid(logit)`). A row's own
|
|
`n_sec_pred` is the first slot index where this fires; once every row in
|
|
the batch has fired, the loop breaks before spending a model call on the
|
|
next slot's token — the average-case cost win the docstring above
|
|
describes. A row that never fires within `k_max` is capped there
|
|
(`K_MAX` stays a safety cap, not a modeling ceiling).
|
|
|
|
Under `history="attention"` the history encoding is computed once per
|
|
slot via `Stage2Autoregressive.history_step` (a KV-cache append)
|
|
rather than re-derived by every model call inside that slot — so an ODE
|
|
loop's `steps` substeps, and the separate `predict_type` call when the
|
|
type slice isn't folded into the trunk output, all reuse the SAME `hist`
|
|
tensor for a given `k`. Recomputing per call instead would be merely
|
|
wasteful under markov (its per-call cost is already O(1)) but wrong under
|
|
attention: `AttentionHistory.step` mutates the cache by appending, so
|
|
calling it more than once per slot would double-count that slot's own
|
|
(not-yet-existing) predecessor.
|
|
|
|
The free-running history feature stays UNSNAPPED (mirrors the
|
|
established "no snapping" precedent for `particle_type.target =
|
|
"physical"` secondaries feeding their own future conditioning):
|
|
`"physical"` carries the raw (log_mass, charge) forward as-is;
|
|
`"embedding"` carries the raw predicted vector as-is; `"onehot"` is the
|
|
one exception — its history slot must be a probability-simplex-shaped
|
|
vector (that's what `MarkovHistory`/`AttentionHistory` were trained on,
|
|
`_type_repr`'s `F.one_hot` ground truth), so it's the hard one-hot of
|
|
`argmax(logits)`, not the raw logits themselves. Discretizing further,
|
|
into a concrete PDG code, only ever happens once — at secondary-spawn
|
|
time in `giant/rollout.py` — never inside this loop.
|
|
|
|
Returns (sec_cont, sec_type, sec_valid) — same shapes/meaning as
|
|
`sample_secondaries`/`sample_secondaries_wgan`'s (see
|
|
`_decode_stage2_flat`'s docstring); `sec_type` is raw per-slot output in
|
|
all three `particle_type.target` cases (never one-hot-collapsed), so the
|
|
caller decodes it exactly the same way regardless of which decoder
|
|
produced it.
|
|
"""
|
|
sec_decoder.eval()
|
|
B = cond_cont.size(0)
|
|
device = cond_cont.device
|
|
k_max = sec_decoder.k_max
|
|
type_dim = sec_decoder.type_dim
|
|
objective = build_objective(sec_decoder.generator_kind)
|
|
target = sec_decoder.particle_type_cfg.target
|
|
type_folded = _type_folded(sec_decoder)
|
|
token_dim = CONT_SLOT_DIM + type_dim if type_folded else CONT_SLOT_DIM
|
|
|
|
sec_cont = torch.zeros(B, k_max, CONT_SLOT_DIM, device=device)
|
|
sec_type = torch.zeros(B, k_max, type_dim, device=device)
|
|
|
|
# Running per-token state, threaded from one slot to the next.
|
|
prev_repr = torch.zeros(B, CONT_SLOT_DIM + type_dim, device=device)
|
|
remaining = torch.ones(B, device=device)
|
|
history_cache = sec_decoder.init_history_cache()
|
|
|
|
use_stop_token = n_sec_pred is None
|
|
if use_stop_token:
|
|
assert getattr(sec_decoder, "stop_head", None) is not None, (
|
|
"sample_secondaries_ar called with n_sec_pred=None on a decoder "
|
|
"with no stop_head — only valid under stage2_model.n_sec.mode = "
|
|
"'stop_token'"
|
|
)
|
|
finished = torch.zeros(B, dtype=torch.bool, device=device)
|
|
derived_n_sec = torch.full((B,), k_max, dtype=torch.long, device=device)
|
|
|
|
for k in range(k_max):
|
|
has_prev = torch.full((B, 1), k >= 1, dtype=torch.bool, device=device)
|
|
history_feat = prev_repr.unsqueeze(1) # (B, 1, CONT_SLOT_DIM + type_dim)
|
|
remaining_frac = remaining.unsqueeze(1) # (B, 1)
|
|
slot_idx = torch.full((B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
|
|
hist, history_cache = sec_decoder.history_step(history_feat, has_prev, history_cache)
|
|
|
|
if use_stop_token:
|
|
stop_logit = sec_decoder.predict_stop(
|
|
cond_cont,
|
|
cond_cat,
|
|
stage1_out,
|
|
history_feat,
|
|
has_prev,
|
|
remaining_frac,
|
|
slot_idx,
|
|
hist=hist,
|
|
).squeeze(1)
|
|
if sec_decoder.stop_sampling == "sample":
|
|
stop_now = torch.rand(B, device=device) < torch.sigmoid(stop_logit)
|
|
else:
|
|
stop_now = stop_logit >= 0.0
|
|
derived_n_sec[stop_now & ~finished] = k
|
|
finished = finished | stop_now
|
|
if finished.all():
|
|
break
|
|
|
|
if objective.is_adversarial:
|
|
z = torch.randn(B, 1, sec_decoder.noise_dim, device=device)
|
|
token = sec_decoder(
|
|
z,
|
|
cond_cont,
|
|
cond_cat,
|
|
stage1_out,
|
|
history_feat,
|
|
has_prev,
|
|
remaining_frac,
|
|
slot_idx,
|
|
hist=hist,
|
|
)
|
|
else:
|
|
x = torch.randn(B, 1, token_dim, device=device)
|
|
dt = 1.0 / steps
|
|
for i in range(steps):
|
|
t = torch.full((B, 1), i * dt, device=device)
|
|
v = sec_decoder(
|
|
x,
|
|
cond_cont,
|
|
cond_cat,
|
|
stage1_out,
|
|
history_feat,
|
|
has_prev,
|
|
remaining_frac,
|
|
slot_idx,
|
|
t=t,
|
|
hist=hist,
|
|
)
|
|
x = x + v * dt
|
|
token = x
|
|
|
|
token = token.squeeze(1) # (B, token_dim)
|
|
cont_k = token[:, :CONT_SLOT_DIM]
|
|
if type_folded:
|
|
type_k = token[:, CONT_SLOT_DIM:]
|
|
else:
|
|
type_k = sec_decoder.predict_type(
|
|
cond_cont,
|
|
cond_cat,
|
|
stage1_out,
|
|
history_feat,
|
|
has_prev,
|
|
remaining_frac,
|
|
slot_idx,
|
|
hist=hist,
|
|
).squeeze(1)
|
|
|
|
sec_cont[:, k] = cont_k
|
|
sec_type[:, k] = type_k
|
|
|
|
if target == "onehot":
|
|
type_for_history = F.one_hot(type_k.argmax(dim=-1), num_classes=type_dim).float()
|
|
else:
|
|
type_for_history = type_k
|
|
|
|
stick_fraction = torch.sigmoid(cont_k[:, 0])
|
|
prev_repr = torch.cat([stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1)
|
|
remaining = torch.clamp(remaining * (1.0 - stick_fraction), min=0.0)
|
|
|
|
resolved_n_sec = derived_n_sec if use_stop_token else n_sec_pred
|
|
sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < resolved_n_sec.unsqueeze(1)
|
|
return sec_cont, sec_type, sec_valid
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Per-stage dispatch — shared by giant/rollout.py and giant/cli.py's
|
|
# `predict` command, since both need "given a stage model, produce a
|
|
# sample" without hand-picking the sampler themselves (each stage's
|
|
# generative objective is independent, read off the model's own
|
|
# `generator_kind`, not a caller-supplied `mode` string).
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def sample_stage1(
|
|
stage1_model: torch.nn.Module,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
steps: int,
|
|
ddpm_steps: int = 1000,
|
|
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
|
"""Dispatches on `stage1_model.generator_kind`."""
|
|
objective = build_objective(stage1_model.generator_kind)
|
|
if objective.is_adversarial:
|
|
return sample_wgan(stage1_model, cond_cont, cond_cat)
|
|
if isinstance(objective, DdpmObjective):
|
|
schedule = CosineSchedule(T=ddpm_steps).to(cond_cont.device)
|
|
return sample_ddpm(stage1_model, cond_cont, cond_cat, schedule)
|
|
return sample_flow(stage1_model, cond_cont, cond_cat, steps=steps)
|
|
|
|
|
|
def sample_stage2(
|
|
sec_decoder: torch.nn.Module,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
stage1_out: torch.Tensor,
|
|
n_sec_pred: torch.Tensor | None,
|
|
steps: int,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
"""Dispatches on `decoder` (one-shot vs autoregressive — the class
|
|
itself, via `isinstance`) and `sec_decoder.generator_kind` (flow/ddpm/
|
|
wgan). DDPM secondaries aren't supported — no `Stage2*` class was ever
|
|
built with `generator="ddpm"` in practice and `flow_matching_loss_secondary*`
|
|
is the only stage-2 training path that exists for the non-adversarial
|
|
case, so there's nothing to dispatch to here.
|
|
|
|
`n_sec_pred=None` (from `resolve_n_sec` on a stop-token decoder) is only
|
|
meaningful for the autoregressive path — see `sample_secondaries_ar`'s
|
|
docstring; the one-shot samplers have no per-token stop mechanism to
|
|
derive a count from, so `n_sec_pred` must already be resolved for them.
|
|
"""
|
|
if isinstance(sec_decoder, Stage2Autoregressive):
|
|
return sample_secondaries_ar(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
|
|
assert n_sec_pred is not None, "one-shot stage-2 decoders need a resolved n_sec_pred"
|
|
if build_objective(sec_decoder.generator_kind).is_adversarial:
|
|
return sample_secondaries_wgan(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred)
|
|
return sample_secondaries(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
|
|
|
|
|
|
def resolve_n_sec(
|
|
stage1_model: torch.nn.Module,
|
|
sec_decoder: torch.nn.Module,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
stage1_out: torch.Tensor,
|
|
n_sec_pred: torch.Tensor | None,
|
|
) -> torch.Tensor | None:
|
|
"""`n_sec_pred` is already populated when `stage1_model` owns a legacy
|
|
`n_sec_head` (a migrated v0.2 checkpoint — see `Stage1Model`'s
|
|
docstring); otherwise ask stage 2, which owns it by default.
|
|
|
|
Returns `None` when `sec_decoder` owns a `stop_head` (`n_sec.mode =
|
|
"stop_token"`) instead of an `n_sec_head` — there is nothing to resolve
|
|
up front in that case, since the count only exists once
|
|
`sample_secondaries_ar` has actually generated (or stopped generating)
|
|
tokens; the caller passes this `None` straight through to `sample_stage2`
|
|
and reads the real count back off its returned `sec_valid`
|
|
(`sec_valid.sum(-1)`) afterwards.
|
|
|
|
Raises if neither stage owns any n_sec mechanism at all — the only way
|
|
that happens is `stage2_model.n_sec.mode = "truth"`, which is not a valid
|
|
rollout-/predict-capable checkpoint."""
|
|
if n_sec_pred is not None:
|
|
return n_sec_pred
|
|
if getattr(sec_decoder, "stop_head", None) is not None:
|
|
return None
|
|
if getattr(sec_decoder, "n_sec_head", None) is None:
|
|
raise RuntimeError(
|
|
"checkpoint has no n_sec_head/stop_head on either stage — needs "
|
|
"stage2_model.n_sec.mode = 'head' (the default) or 'stop_token'; "
|
|
"'truth' is standalone-evaluation-only"
|
|
)
|
|
logits = sec_decoder.predict_n_sec(cond_cont, cond_cat, stage1_out)
|
|
return logits.argmax(dim=-1)
|