from __future__ import annotations import argparse import json import math import re import sys from dataclasses import asdict, dataclass from pathlib import Path import torch import torch.nn.functional as F from transformers import AutoModelForCausalLM, AutoTokenizer REPO_ROOT = Path(__file__).resolve().parents[1] SCRIPTS_DIR = Path(__file__).resolve().parent for path in (REPO_ROOT, SCRIPTS_DIR): if str(path) not in sys.path: sys.path.insert(0, str(path)) from eval import build_model_from_ckpt from flowtext_lab.decode import model_time_for_step, sample_noise_simplex, state_for_model from flowtext_lab.genppl import filter_generated_texts, summarize_token_diversity from flowtext_lab.tokenization import BpeTextTokenizer from standard_genppl_entropy_latest_decode import ( DecodeSetting, current_anchor, dirichlet_path_mean, dirichlet_resample, make_decode_time_grid, make_started_state, parse_final_from, parse_temps, penalty_window_scale, sample_final_ids, schedule_power_from_progress, score_with_loaded, sequence_frequency_penalty, strip_common_special, ) SPECIAL_RE = re.compile(r"\s+") def make_decode_attention_mask( *, batch_size: int, prefix_len: int, block_len: int, device: torch.device, ) -> torch.Tensor: """Mask for a truncated [clean prefix, noisy current block] decode pass.""" seq_len = int(prefix_len) + int(block_len) pos = torch.cat( [ torch.arange(prefix_len, device=device), torch.arange(prefix_len, prefix_len + block_len, device=device), ], dim=0, ) stream = torch.cat( [ torch.zeros(prefix_len, device=device, dtype=torch.long), torch.ones(block_len, device=device, dtype=torch.long), ], dim=0, ) block = pos // int(block_len) q_stream = stream[:, None] k_stream = stream[None, :] q_block = block[:, None] k_block = block[None, :] clean_query = q_stream == 0 clean_key = k_stream == 0 noisy_query = q_stream == 1 noisy_key = k_stream == 1 current_block = prefix_len // int(block_len) clean_rule = clean_query & clean_key & (k_block <= q_block) noisy_rule = noisy_query & ( (clean_key & (k_block < current_block)) | (noisy_key & (k_block == current_block)) ) keep = clean_rule | noisy_rule eye = torch.eye(seq_len, device=device, dtype=torch.bool) return (keep | eye).unsqueeze(0).expand(batch_size, -1, -1) def make_decode_position_ids(prefix_len: int, block_len: int, device: torch.device) -> torch.Tensor: return torch.cat( [ torch.arange(prefix_len, device=device, dtype=torch.long), torch.arange(prefix_len, prefix_len + block_len, device=device, dtype=torch.long), ], dim=0, ) @dataclass class BlockArDecodeConfig: max_len: int block_len: int steps: int model_t_mode: str noise_init: str dirichlet_concentration: float concentration_min: float concentration_max: float noise_sigma: float target_prob: float decode_rule: str support_power: float semantic_power: float anchor_mode: str decode_freq_penalty_alpha: float decode_freq_penalty_beta: float decode_freq_penalty_floor: float decode_freq_penalty_start: float decode_freq_penalty_end: float decode_freq_penalty_power: float start_t: float start_init: str final_sample_mode: str final_sample_temp: float final_top_k: int final_top_p: float final_freq_penalty_alpha: float final_freq_penalty_beta: float final_freq_penalty_floor: float eps: float @torch.no_grad() def decode_blockar_samples( model, tokenizer: BpeTextTokenizer, setting: DecodeSetting, cfg: BlockArDecodeConfig, *, n_samples: int, batch_size: int, decode_time_grid: list[float], device: torch.device, ) -> tuple[list[list[int]], list[str]]: if cfg.max_len % cfg.block_len != 0: raise ValueError(f"max_len={cfg.max_len} must be divisible by block_len={cfg.block_len}") all_ids: list[list[int]] = [] all_texts: list[str] = [] vocab_size = tokenizer.vocab_size remaining = int(n_samples) while remaining > 0: bs = min(int(batch_size), remaining) clean_ids = torch.zeros((bs, cfg.max_len), dtype=torch.long, device=device) for start in range(0, cfg.max_len, cfg.block_len): end = start + cfg.block_len probs = make_started_state( batch_size=bs, max_len=cfg.block_len, vocab_size=vocab_size, device=device, eps=cfg.eps, start_t=cfg.start_t, start_init=cfg.start_init, noise_init=cfg.noise_init, target_prob=cfg.target_prob, noise_sigma=cfg.noise_sigma, dirichlet_concentration=cfg.dirichlet_concentration, concentration_min=cfg.concentration_min, concentration_max=cfg.concentration_max, ) initial_probs = probs.clone() last_endpoint = probs attn = make_decode_attention_mask( batch_size=bs, prefix_len=start, block_len=cfg.block_len, device=device, ) position_ids = make_decode_position_ids(start, cfg.block_len, device) for step in range(cfg.steps): progress = float(decode_time_grid[step]) next_progress = float(decode_time_grid[step + 1]) if cfg.model_t_mode == "flow": t = torch.full((bs,), progress, device=device, dtype=torch.float32) elif cfg.model_t_mode == "post": t = torch.full((bs,), next_progress, device=device, dtype=torch.float32) else: t = model_time_for_step(cfg.model_t_mode, step, cfg.steps, bs, device, dtype=torch.float32) if start > 0: prefix_probs = F.one_hot(clean_ids[:, :start], vocab_size).to(dtype=probs.dtype) packed_probs = torch.cat([prefix_probs, probs], dim=1) else: packed_probs = probs logits = model( state_for_model(model, packed_probs, cfg.eps), t, attn, position_ids=position_ids, ).float()[:, -cfg.block_len :, :] decode_penalty_scale = penalty_window_scale( next_progress, cfg.decode_freq_penalty_start, cfg.decode_freq_penalty_end, cfg.decode_freq_penalty_power, ) if decode_penalty_scale > 0.0 and ( cfg.decode_freq_penalty_alpha > 0.0 or cfg.decode_freq_penalty_beta > 0.0 ): logits = logits - decode_penalty_scale * sequence_frequency_penalty( probs, alpha=cfg.decode_freq_penalty_alpha, beta=cfg.decode_freq_penalty_beta, floor=cfg.decode_freq_penalty_floor, eps=cfg.eps, ).unsqueeze(-2) endpoint = F.softmax(logits / float(setting.endpoint_temp), dim=-1) last_endpoint = endpoint support_t = schedule_power_from_progress(next_progress, cfg.support_power) if cfg.decode_rule == "flowmap": gamma = min((next_progress - progress) / max(1.0 - progress, cfg.eps), 1.0) probs = probs + gamma * (endpoint - probs) probs = probs.clamp_min(cfg.eps) probs = probs / probs.sum(dim=-1, keepdim=True).clamp_min(cfg.eps) elif cfg.decode_rule == "dirichlet_mean": probs = dirichlet_path_mean(endpoint, support_t, cfg.eps) elif cfg.decode_rule == "dirichlet_resample": mean = dirichlet_path_mean(endpoint, support_t, cfg.eps) probs = dirichlet_resample(mean, support_t, cfg.concentration_min, cfg.concentration_max, cfg.eps) elif cfg.decode_rule in {"dual_line_mean", "dual_line_resample"}: semantic_t = schedule_power_from_progress(next_progress, cfg.semantic_power) anchor = current_anchor(probs, cfg.anchor_mode, cfg.eps) forward_endpoint = (1.0 - semantic_t) * anchor + semantic_t * endpoint forward_endpoint = forward_endpoint.clamp_min(cfg.eps) forward_endpoint = forward_endpoint / forward_endpoint.sum(dim=-1, keepdim=True).clamp_min(cfg.eps) mean = dirichlet_path_mean(forward_endpoint, support_t, cfg.eps) if cfg.decode_rule == "dual_line_mean": probs = mean else: probs = dirichlet_resample(mean, support_t, cfg.concentration_min, cfg.concentration_max, cfg.eps) else: raise ValueError(f"Unknown decode_rule: {cfg.decode_rule}") if setting.final_from == "endpoint": final_probs = last_endpoint elif setting.final_from == "blend": final_probs = 0.5 * probs + 0.5 * last_endpoint else: final_probs = probs block_ids = sample_final_ids( final_probs, mode=cfg.final_sample_mode, temperature=cfg.final_sample_temp, top_k=cfg.final_top_k, top_p=cfg.final_top_p, freq_penalty_alpha=cfg.final_freq_penalty_alpha, freq_penalty_beta=cfg.final_freq_penalty_beta, freq_penalty_floor=cfg.final_freq_penalty_floor, eps=cfg.eps, ) clean_ids[:, start:end] = block_ids ids = clean_ids.detach().cpu().tolist() all_ids.extend(ids) all_texts.extend(tokenizer.decode(row, stop_at_eos=False, skip_special_tokens=False) for row in ids) remaining -= bs print( f"[blockar-decode] temp={setting.endpoint_temp:.2f} final={setting.final_from} " f"rule={cfg.decode_rule} anchor={cfg.anchor_mode} steps={cfg.steps} " f"generated {n_samples - remaining}/{n_samples}", flush=True, ) return all_ids, all_texts def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--checkpoint", required=True) parser.add_argument("--tokenizer_path", required=True) parser.add_argument("--scorer", required=True) parser.add_argument("--output_dir", required=True) parser.add_argument("--max_len", type=int, default=1024) parser.add_argument("--block_len", type=int, default=128) parser.add_argument("--n_samples", type=int, default=16) parser.add_argument("--decode_batch", type=int, default=2) parser.add_argument("--score_batch", type=int, default=8) parser.add_argument("--score_max_length", type=int, default=256) parser.add_argument("--steps", type=int, default=128) parser.add_argument("--model_t_mode", default="flow") parser.add_argument("--decode_time_schedule", choices=["linear", "sampled_path", "lognsr_gumbel"], default="linear") parser.add_argument("--decode_s_min_frac", type=float, default=0.0) parser.add_argument("--decode_s_max_frac", type=float, default=0.25) parser.add_argument("--decode_time_gumbel_loc", type=float, default=2.2) parser.add_argument("--decode_time_gumbel_scale", type=float, default=0.8) parser.add_argument("--no_decode_force_final_t", action="store_true") parser.add_argument( "--decode_rule", choices=["flowmap", "dirichlet_mean", "dirichlet_resample", "dual_line_mean", "dual_line_resample"], default="dual_line_resample", ) parser.add_argument("--support_power", type=float, default=1.0) parser.add_argument("--semantic_power", type=float, default=1.0) parser.add_argument("--anchor_mode", choices=["onehot", "state", "sqrt_state"], default="state") parser.add_argument("--decode_freq_penalty_alpha", type=float, default=0.0) parser.add_argument("--decode_freq_penalty_beta", type=float, default=0.0) parser.add_argument("--decode_freq_penalty_floor", type=float, default=0.0) parser.add_argument("--decode_freq_penalty_start", type=float, default=0.0) parser.add_argument("--decode_freq_penalty_end", type=float, default=1.0) parser.add_argument("--decode_freq_penalty_power", type=float, default=1.0) parser.add_argument("--start_t", type=float, default=0.0) parser.add_argument( "--start_init", choices=["noise", "uniform_mean", "uniform_dirichlet", "random_anchor_mean", "random_anchor_dirichlet"], default="noise", ) parser.add_argument("--noise_init", choices=["uniform", "logistic_normal", "dirichlet"], default="dirichlet") parser.add_argument("--noise_sigma", type=float, default=-1.0) parser.add_argument("--dirichlet_concentration", type=float, default=-1.0) parser.add_argument("--concentration_min", type=float, default=1.0) parser.add_argument("--concentration_max", type=float, default=1024.0) parser.add_argument("--target_prob", type=float, default=-1.0) parser.add_argument("--endpoint_temps", default="1.45") parser.add_argument("--final_from", default="state") parser.add_argument("--final_sample_mode", choices=["argmax", "sample", "topk", "topp"], default="argmax") parser.add_argument("--final_sample_temp", type=float, default=1.0) parser.add_argument("--final_top_k", type=int, default=64) parser.add_argument("--final_top_p", type=float, default=0.95) parser.add_argument("--final_freq_penalty_alpha", type=float, default=0.0) parser.add_argument("--final_freq_penalty_beta", type=float, default=0.0) parser.add_argument("--final_freq_penalty_floor", type=float, default=0.0) parser.add_argument("--eps", type=float, default=1e-8) parser.add_argument("--seed", type=int, default=20260522) parser.add_argument("--save_samples", type=int, default=16) args = parser.parse_args() torch.manual_seed(args.seed) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") tokenizer = BpeTextTokenizer.from_file(args.tokenizer_path) ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False) ckpt_args = ckpt.get("args", {}) step = ckpt.get("step") target_prob = float(args.target_prob if args.target_prob >= 0 else ckpt_args.get("target_prob", 1.0)) dirichlet_concentration = float( args.dirichlet_concentration if args.dirichlet_concentration > 0 else ckpt_args.get("dirichlet_concentration_min", 1.0) ) print(f"[ckpt] {args.checkpoint} step={step}", flush=True) print( f"[blockar-base] n={args.n_samples} max_len={args.max_len} block_len={args.block_len} " f"steps={args.steps} batch={args.decode_batch}", flush=True, ) decode_time_grid = make_decode_time_grid( steps=args.steps, start_t=args.start_t, schedule=args.decode_time_schedule, s_min_frac=args.decode_s_min_frac, s_max_frac=args.decode_s_max_frac, gumbel_loc=args.decode_time_gumbel_loc, gumbel_scale=args.decode_time_gumbel_scale, seed=args.seed, force_final_t=not args.no_decode_force_final_t, ) model = build_model_from_ckpt(ckpt, tokenizer.vocab_size, args.max_len * 2, device).eval() cfg = BlockArDecodeConfig( max_len=args.max_len, block_len=args.block_len, steps=args.steps, model_t_mode=args.model_t_mode, noise_init=args.noise_init, dirichlet_concentration=dirichlet_concentration, concentration_min=args.concentration_min, concentration_max=args.concentration_max, noise_sigma=args.noise_sigma, target_prob=target_prob, decode_rule=args.decode_rule, support_power=args.support_power, semantic_power=args.semantic_power, anchor_mode=args.anchor_mode, decode_freq_penalty_alpha=args.decode_freq_penalty_alpha, decode_freq_penalty_beta=args.decode_freq_penalty_beta, decode_freq_penalty_floor=args.decode_freq_penalty_floor, decode_freq_penalty_start=args.decode_freq_penalty_start, decode_freq_penalty_end=args.decode_freq_penalty_end, decode_freq_penalty_power=args.decode_freq_penalty_power, start_t=args.start_t, start_init=args.start_init, final_sample_mode=args.final_sample_mode, final_sample_temp=args.final_sample_temp, final_top_k=args.final_top_k, final_top_p=args.final_top_p, final_freq_penalty_alpha=args.final_freq_penalty_alpha, final_freq_penalty_beta=args.final_freq_penalty_beta, final_freq_penalty_floor=args.final_freq_penalty_floor, eps=args.eps, ) out_dir = Path(args.output_dir) out_dir.mkdir(parents=True, exist_ok=True) summary_path = out_dir / "summary.jsonl" samples_path = out_dir / "samples.txt" decoded_cache = [] with summary_path.open("w", encoding="utf-8") as sf, samples_path.open("w", encoding="utf-8") as tf: for setting in [ DecodeSetting(temp, final) for temp in parse_temps(args.endpoint_temps) for final in parse_final_from(args.final_from) ]: torch.manual_seed(args.seed) ids, raw_texts = decode_blockar_samples( model, tokenizer, setting, cfg, n_samples=args.n_samples, batch_size=args.decode_batch, decode_time_grid=decode_time_grid, device=device, ) stripped = [strip_common_special(t) for t in raw_texts] decoded_cache.append((setting, ids, raw_texts, stripped)) sf.write(json.dumps({"type": "decode_done", "step": step, "setting": asdict(setting)}, ensure_ascii=False) + "\n") for i in range(min(args.save_samples, len(raw_texts))): tf.write(f"===== temp={setting.endpoint_temp} final={setting.final_from} sample={i} =====\n") tf.write(stripped[i] + "\n\n") sf.flush() tf.flush() del model if torch.cuda.is_available(): torch.cuda.empty_cache() scorer_tok = AutoTokenizer.from_pretrained(args.scorer) if scorer_tok.pad_token_id is None: scorer_tok.pad_token = scorer_tok.eos_token scorer_tok.pad_token_id = scorer_tok.eos_token_id scorer = AutoModelForCausalLM.from_pretrained(args.scorer).to(device).eval() if getattr(scorer.config, "pad_token_id", None) is None: scorer.config.pad_token_id = scorer_tok.pad_token_id with summary_path.open("a", encoding="utf-8") as sf: for setting, ids, raw_texts, stripped_texts in decoded_cache: kept_raw, _ = filter_generated_texts(raw_texts, min_chars=1, normalize_whitespace=False, drop_empty=True) kept_stripped, _ = filter_generated_texts(stripped_texts, min_chars=1, normalize_whitespace=True, drop_empty=True) raw_ppl = score_with_loaded( kept_raw, scorer, scorer_tok, batch_size=args.score_batch, max_length=args.score_max_length, device=device, ) stripped_ppl = score_with_loaded( kept_stripped, scorer, scorer_tok, batch_size=args.score_batch, max_length=args.score_max_length, device=device, ) diversity = summarize_token_diversity(ids).__dict__ summary = { "type": "summary", "checkpoint": args.checkpoint, "step": step, "blockar_decode": True, "decode": { **asdict(cfg), "endpoint_temp": setting.endpoint_temp, "final_from": setting.final_from, "decode_time_schedule": args.decode_time_schedule, "decode_time_grid": decode_time_grid, "n_samples": args.n_samples, "seed": args.seed, }, "raw_genppl": raw_ppl, "stripped_genppl": stripped_ppl, "diversity": diversity, } sf.write(json.dumps(summary, ensure_ascii=False) + "\n") sf.flush() print("[summary]", json.dumps(summary, ensure_ascii=False), flush=True) print(f"[done] {out_dir}", flush=True) if __name__ == "__main__": main()