419 lines
15 KiB
Python
419 lines
15 KiB
Python
"""
|
|
Minimal EditFlow τ-leap generator for EditBase-Dream with diffusion-style visualization.
|
|
|
|
What changed vs. your original:
|
|
- tau_leap_step_minimal returns (x_next, any_edit, step_trace) preserving all intermediates.
|
|
- generate_editflow_minimal returns (final_text, trace).
|
|
- render_consecutive_trace_gif(trace, tokenizer, ...) draws a GIF where each frame shows
|
|
ONLY the current output (like the Gemini diffusion page shows progressive refinement):
|
|
* SUB tokens in the current frame are orange
|
|
* INS tokens in the current frame are blue
|
|
* KEEP tokens are black
|
|
* If any deletions happened in the step, the title shows ⌫N (red)
|
|
"""
|
|
|
|
# srun -p $PARTITION --quotatype=$QUOTATYPE --gres=gpu:1 --time=03:00:000 python examples/editflow/generate.py --model_name_or_path "models/EditFlow-Dream-Instruct-7B/tulu-3-sft-mixture/checkpoint-final" --tau 0.02 --mask_length 128 --seed 7070 --prompt "write a romantic story" --make_gif
|
|
|
|
import math
|
|
from dataclasses import dataclass
|
|
from typing import Annotated
|
|
|
|
import tyro
|
|
import torch
|
|
from transformers import AutoModel, AutoTokenizer, PreTrainedModel, PreTrainedTokenizer
|
|
|
|
from dllm.core.schedulers import BaseKappaScheduler, LinearKappaScheduler
|
|
|
|
|
|
# ------------------------------- Small utilities --------------------------------
|
|
|
|
|
|
def _bernoulli_from_rate(rate: torch.Tensor, tau: float) -> torch.Tensor:
|
|
"""First-order τ-leap Bernoulli with p ≈ rate * τ (clamped)."""
|
|
p = (rate.float() * float(tau)).clamp_(0.0, 1.0 - 1e-6)
|
|
return torch.bernoulli(p)
|
|
|
|
|
|
def _sample_from_logits(logits_row: torch.Tensor, temperature: float) -> int:
|
|
"""Sample one token id from a 1D logits row with temperature.
|
|
temperature <= 0 -> greedy (argmax).
|
|
"""
|
|
if temperature <= 0.0:
|
|
return int(torch.argmax(logits_row).item())
|
|
return int(
|
|
torch.distributions.Categorical(logits=(logits_row / temperature))
|
|
.sample()
|
|
.item()
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class GenCfg:
|
|
tau: float = 0.02 # τ step
|
|
device: str = "cuda" if torch.cuda.is_available() else "cpu"
|
|
seed: int = 1234
|
|
edit_prompt: bool = False # allow editing inside prompt region?
|
|
temperature: float = 0.7 # token sampling temperature (sub/ins)
|
|
verbose: bool = True # whether to show intermediate decoding traces
|
|
time_independent: bool = True
|
|
|
|
|
|
# -------------------------------- τ-leap one step --------------------------------
|
|
|
|
|
|
@torch.no_grad()
|
|
def tau_leap_step_minimal(
|
|
x: torch.Tensor, # [T]
|
|
model: PreTrainedModel,
|
|
tokenizer: PreTrainedTokenizer,
|
|
prompt_len: int, # number of initial prompt tokens (including BOS)
|
|
t: float,
|
|
sched: BaseKappaScheduler,
|
|
cfg: GenCfg,
|
|
prev_out: dict | None = None, # <-- pass prior step's model outputs
|
|
reuse_prev: bool = False, # <-- if True, reuse prev_out instead of forward()
|
|
) -> tuple[torch.Tensor, bool, dict, dict]:
|
|
"""
|
|
Single τ-leap step with deletion/substitution conflict resolution
|
|
and right-insert policy.
|
|
|
|
Reuse semantics:
|
|
• If cfg.time_independent == True and reuse_prev == True and prev_out is not None,
|
|
we reuse `prev_out` tensors instead of calling model() again.
|
|
• Otherwise we run a fresh forward().
|
|
|
|
Viz-only convention:
|
|
• Any local annotated as _Ann[*, "viz-only"] is used only for human-visible
|
|
tracing / debugging (console logs, GIFs) and does not affect generation.
|
|
• Such variables are also prefixed with '_' for quick visual scanning.
|
|
|
|
Returns:
|
|
x_next, any_edit, _step_trace, out_for_next (the freshly used model outputs)
|
|
"""
|
|
device = x.device
|
|
T = x.numel()
|
|
|
|
# Decide whether to reuse the previous forward results
|
|
use_reuse = bool(cfg.time_independent and reuse_prev and (prev_out is not None))
|
|
if use_reuse:
|
|
out = prev_out
|
|
else:
|
|
attn = torch.ones(1, T, dtype=torch.long, device=device)
|
|
t_tensor = torch.full((1, 1), float(t), device=device)
|
|
out = model(input_ids=x.unsqueeze(0), attention_mask=attn, t=t_tensor)
|
|
|
|
del_rate_h = out["del_rate_hat"] # [1, T]
|
|
sub_rate_h = out["sub_rate_hat"] # [1, T]
|
|
ins_rate_h = out["ins_rate_hat"] # [1, T]
|
|
sub_logits = out["sub_logits"] # [1, T, V]
|
|
ins_logits = out["ins_logits"] # [1, T, V]
|
|
|
|
# Scale normalized rates to true rates
|
|
tt = torch.tensor([[t]], device=device)
|
|
w = sched.weight(tt)
|
|
del_rate = del_rate_h * w
|
|
sub_rate = sub_rate_h * w
|
|
ins_rate = ins_rate_h * w
|
|
|
|
# Clamp prompt_len within current T (robustness)
|
|
prompt_len_clamped = int(max(1, min(prompt_len, T)))
|
|
|
|
if not cfg.edit_prompt:
|
|
# Protect the entire prompt span from del/sub
|
|
del_rate[:, :prompt_len_clamped] = 0.0
|
|
sub_rate[:, :prompt_len_clamped] = 0.0
|
|
# Disallow insertions inside the prompt EXCEPT at the last prompt token
|
|
if prompt_len_clamped >= 2:
|
|
ins_rate[:, : prompt_len_clamped - 1] = 0.0
|
|
|
|
# Combined "edit" (delete or substitute) event
|
|
comb_rate = (del_rate + sub_rate).squeeze(0) # [T]
|
|
comb_fire = _bernoulli_from_rate(comb_rate, cfg.tau).bool() # [T]
|
|
|
|
# If an edit fires at i, choose deletion with prob λ_del/(λ_del+λ_sub)
|
|
p_del = (del_rate.squeeze(0) / (comb_rate + 1e-8)).clamp(0, 1) # [T]
|
|
choose_del = (torch.rand_like(p_del) < p_del) & comb_fire # [T]
|
|
choose_sub = comb_fire & (~choose_del) # [T]
|
|
|
|
# Insertions (right of token i)
|
|
ins_fire = _bernoulli_from_rate(ins_rate.squeeze(0), cfg.tau).bool() # [T]
|
|
|
|
# Token draws (algorithmic, not viz-only)
|
|
sub_samples: list[int | None] = [
|
|
(
|
|
_sample_from_logits(sub_logits[0, i], cfg.temperature)
|
|
if choose_sub[i]
|
|
else None
|
|
)
|
|
for i in range(T)
|
|
]
|
|
ins_samples: list[int | None] = [
|
|
_sample_from_logits(ins_logits[0, i], cfg.temperature) if ins_fire[i] else None
|
|
for i in range(T)
|
|
]
|
|
|
|
# Build new sequence left→right (apply insertions to the RIGHT)
|
|
new_ids: list[int] = []
|
|
|
|
# --- viz-only per-position labels (for trace/GIF) ---
|
|
_before_ops: Annotated[list[str], "viz-only"] = (
|
|
[]
|
|
) # per 'before' position: DEL/SUB/KEEP
|
|
_after_ops: Annotated[list[str], "viz-only"] = (
|
|
[]
|
|
) # per 'after' token aligned to new_ids: INS/SUB/KEEP
|
|
|
|
for i in range(T):
|
|
if choose_del[i]:
|
|
_before_ops.append("DEL")
|
|
# deletion -> no token appended
|
|
elif choose_sub[i]:
|
|
_before_ops.append("SUB")
|
|
new_tok = sub_samples[i]
|
|
new_ids.append(int(new_tok))
|
|
_after_ops.append("SUB")
|
|
else:
|
|
_before_ops.append("KEEP")
|
|
new_ids.append(int(x[i].item()))
|
|
_after_ops.append("KEEP")
|
|
|
|
if ins_samples[i] is not None:
|
|
new_ids.append(int(ins_samples[i]))
|
|
_after_ops.append("INS")
|
|
|
|
x_next = torch.tensor(new_ids, dtype=torch.long, device=device)
|
|
any_edit = bool(comb_fire.any().item() or ins_fire.any().item())
|
|
# Provide the exact outputs we used this step for the caller to pass forward
|
|
out_for_next = out
|
|
|
|
# --- (vis) used only for verbose console trace ---
|
|
if cfg.verbose and (comb_fire.any() or ins_fire.any()):
|
|
|
|
def _tok_str(tok_id: int) -> str: # viz-only helper
|
|
try:
|
|
s = tokenizer.decode([int(tok_id)])
|
|
return s if s.strip() else f"<{int(tok_id)}>"
|
|
except Exception:
|
|
return f"<{int(tok_id)}>"
|
|
|
|
_ops_strs: Annotated[list[str], "viz-only"] = []
|
|
for i in range(T):
|
|
if choose_del[i]:
|
|
_ops_strs.append(f"DEL@{i}:{_tok_str(int(x[i]))}")
|
|
elif choose_sub[i]:
|
|
_ops_strs.append(
|
|
f"SUB@{i}:{_tok_str(int(x[i]))}->{_tok_str(sub_samples[i])}"
|
|
)
|
|
if ins_samples[i] is not None:
|
|
_ops_strs.append(f"INS@{i}->{i+1}:{_tok_str(ins_samples[i])}")
|
|
print("[time]", f"{t:.4f}")
|
|
print("[events]", "; ".join(_ops_strs))
|
|
print("[decode]\n", tokenizer.decode(new_ids, skip_special_tokens=False))
|
|
print()
|
|
|
|
# --- (vis) step trace payload (returned; used only for visualization downstream) ---
|
|
_step_trace: Annotated[dict, "viz-only"] = {
|
|
"t": float(t),
|
|
"x_before_ids": [int(i) for i in x.tolist()],
|
|
"x_after_ids": [int(i) for i in new_ids],
|
|
"before_ops": _before_ops, # viz-only labels
|
|
"after_ops": _after_ops, # viz-only labels
|
|
# below are algorithmic signals copied for visualization/analysis
|
|
"choose_del": [bool(v) for v in choose_del.tolist()],
|
|
"choose_sub": [bool(v) for v in choose_sub.tolist()],
|
|
"ins_fire": [bool(v) for v in ins_fire.tolist()],
|
|
"sub_samples": [int(s) if s is not None else None for s in sub_samples],
|
|
"ins_samples": [int(s) if s is not None else None for s in ins_samples],
|
|
"prompt_len": prompt_len_clamped,
|
|
"used_reuse": bool(use_reuse),
|
|
}
|
|
|
|
return x_next, any_edit, _step_trace, out_for_next
|
|
|
|
|
|
# -------------------------------- top-level generate -------------------------------
|
|
|
|
|
|
@torch.no_grad()
|
|
def generate_editflow_minimal(
|
|
model: PreTrainedModel,
|
|
tokenizer: PreTrainedTokenizer,
|
|
args,
|
|
cfg: GenCfg,
|
|
) -> tuple[str, dict]:
|
|
"""
|
|
Returns:
|
|
final_text, trace
|
|
|
|
Notes on annotations:
|
|
• Any local annotated with Annotated[..., "viz-only"] is only used to build
|
|
the decode trace for visualization (e.g., GIF rendering) and has no effect
|
|
on the actual generation. Such variables are also prefixed with '_' to make
|
|
this visually obvious in code.
|
|
"""
|
|
torch.manual_seed(cfg.seed)
|
|
|
|
# If prompt is None, start from BOS alone; otherwise ALWAYS prefix BOS
|
|
bos = getattr(tokenizer, "bos_token_id", None)
|
|
if bos is None:
|
|
raise ValueError("Tokenizer must have a BOS token for this sampler.")
|
|
|
|
prompt = args.prompt
|
|
if prompt is None:
|
|
ids = [bos] # BOS alone
|
|
else:
|
|
ids = tokenizer.apply_chat_template(
|
|
[{"role": "user", "content": prompt}],
|
|
tokenize=True,
|
|
add_generation_prompt=True,
|
|
)
|
|
# ids = tokenizer.encode(prompt, add_special_tokens=False)
|
|
# ids = [bos] + enc["input_ids"] # ALWAYS prefix BOS
|
|
|
|
prompt_len = len(ids)
|
|
|
|
if args.mask_length:
|
|
if getattr(tokenizer, "mask_token_id", None) is None:
|
|
raise ValueError(
|
|
"Tokenizer must define mask_token_id when --mask_length > 0."
|
|
)
|
|
ids = ids + [tokenizer.mask_token_id] * args.mask_length
|
|
|
|
x = torch.tensor(ids, dtype=torch.long, device=model.device)
|
|
|
|
sched = LinearKappaScheduler()
|
|
tau = cfg.tau
|
|
steps = math.ceil(1.0 / max(tau, 1e-9))
|
|
|
|
_trace: Annotated[dict, "viz-only: full decode trace for GIF/inspection"] = {
|
|
"steps": [],
|
|
"init": {
|
|
"t": 0.0,
|
|
"x_ids": [int(i) for i in x.tolist()],
|
|
"prompt_len": int(prompt_len),
|
|
},
|
|
"end_t": 0.0,
|
|
}
|
|
|
|
# Local-only reuse: if previous iteration had no edits, reuse its forward.
|
|
prev_out: dict | None = None
|
|
prev_had_edits = True # first iteration must run a forward
|
|
|
|
t = 0.0
|
|
for _ in range(steps):
|
|
# We can reuse prev_out only if the model is declared time-independent
|
|
# and the previous step had NO edits (sequence unchanged).
|
|
reuse_prev = (
|
|
cfg.time_independent and not prev_had_edits and (prev_out is not None)
|
|
)
|
|
|
|
x, edited, _step_trace, prev_out = tau_leap_step_minimal(
|
|
x=x,
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
prompt_len=prompt_len,
|
|
t=t,
|
|
sched=sched,
|
|
cfg=cfg,
|
|
prev_out=prev_out,
|
|
reuse_prev=reuse_prev,
|
|
)
|
|
|
|
_step_trace: Annotated[dict, "viz-only: per-step intermediates for trace"]
|
|
_trace["steps"].append(_step_trace)
|
|
|
|
prev_had_edits = edited
|
|
|
|
t = min(1.0, t + tau)
|
|
if t >= 1.0 - args.time_epsilon:
|
|
break
|
|
|
|
_trace["end_t"] = float(t)
|
|
|
|
final_text = tokenizer.decode(x.tolist(), skip_special_tokens=False)
|
|
print("[final]")
|
|
return final_text, _trace
|
|
|
|
|
|
# ---------------------------------------- CLI -------------------------------------
|
|
|
|
|
|
def main():
|
|
@dataclass
|
|
class ScriptArgs:
|
|
# Required (no default)
|
|
model_name_or_path: Annotated[str, "Path or hub id for the model"]
|
|
time_independent: Annotated[
|
|
bool, "Whether model is conditioned on time step"
|
|
] = True
|
|
|
|
prompt: Annotated[str | None, "Text prompt. If None, start from BOS alone."] = (
|
|
None
|
|
)
|
|
# Boolean flag: tyro exposes --edit-prompt / --no-edit-prompt automatically for bools
|
|
edit_prompt: Annotated[
|
|
bool,
|
|
"Allow delete/substitute and insertions in the prompt region (BOS+prompt).",
|
|
] = False
|
|
|
|
# Generation-related args
|
|
tau: Annotated[float, "τ-leap size"] = 0.01
|
|
time_epsilon: Annotated[
|
|
float, "Match this with the `time_epsilon` arg used in your EditFlowTrainer"
|
|
] = 1e-3
|
|
mask_length: Annotated[
|
|
int,
|
|
"Number of <mask> tokens appended after the prompt.\n"
|
|
"EditFlow will iteratively substitute, insert, or delete masks to form the output.",
|
|
] = 128
|
|
temperature: Annotated[float, "Token sampling temperature; 0 for greedy."] = 0.7
|
|
|
|
seed: Annotated[int, "Random seed"] = 1234
|
|
verbose: Annotated[bool, "Whether to show intermediate decoding traces"] = True
|
|
|
|
# Visualization
|
|
make_gif: Annotated[bool, "Render a decoding trace GIF after generation."] = (
|
|
False
|
|
)
|
|
gif_path: Annotated[
|
|
str | None, "Output GIF path (default: decode_trace.gif)"
|
|
] = None
|
|
frame_ms: Annotated[int, "Per-frame duration in ms"] = 120
|
|
|
|
args = tyro.cli(ScriptArgs)
|
|
|
|
cfg = GenCfg(
|
|
tau=args.tau,
|
|
seed=args.seed,
|
|
edit_prompt=args.edit_prompt,
|
|
temperature=args.temperature,
|
|
verbose=args.verbose,
|
|
time_independent=args.time_independent,
|
|
)
|
|
|
|
model = AutoModel.from_pretrained(
|
|
args.model_name_or_path,
|
|
dtype=torch.bfloat16,
|
|
device_map="auto",
|
|
).eval()
|
|
tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
|
|
|
|
final_text, trace = generate_editflow_minimal(model, tokenizer, args, cfg)
|
|
print(final_text)
|
|
|
|
if args.make_gif:
|
|
from examples.editflow.viz import render_consecutive_trace_gif
|
|
|
|
out = args.gif_path or "decode_trace.gif"
|
|
path = render_consecutive_trace_gif(
|
|
trace,
|
|
tokenizer,
|
|
out_path=out,
|
|
frame_ms=args.frame_ms,
|
|
)
|
|
print(f"[gif saved] {path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|