#!/usr/bin/env python3 """chat.py: ChatML chat with MetaDiffusion-600M checkpoints. Left-to-right block commit (semi-autoregressive): the leftmost masked positions are filled first, so <|im_end|> cannot win the race at position 0 (which produced empty responses on the 150M architecture). Usage: Interactive: python chat.py --model-path checkpoints/step_30000.pt One-shot: python chat.py --model-path checkpoints/step_30000.pt \ --prompt "What is 2+2?" --watch """ import argparse import json import math import sys from pathlib import Path import torch import torch.nn.functional as F from transformers import AutoTokenizer sys.path.insert(0, str(Path(__file__).resolve().parent)) from model import MetaDiffusionConfig, MetaDiffusionLM # noqa: E402 IM_START, IM_END = "<|im_start|>", "<|im_end|>" RAINBOW_TOKENS = [f"<|r{i}|>" for i in range(1, 8)] def build_config(config_dict): valid = {k: v for k, v in config_dict.items() if k in MetaDiffusionConfig.__dataclass_fields__} return MetaDiffusionConfig(**valid) def load_model(model_path, device): path = Path(model_path) if path.is_dir(): with open(path / "config.json") as f: config = build_config(json.load(f)) model = MetaDiffusionLM(config).to(device) from safetensors.torch import load_file sd = load_file(path / "model.safetensors") sd = {k[len("model."):] if k.startswith("model.") else k: v for k, v in sd.items()} model.load_state_dict(sd, strict=True) elif path.suffix == ".safetensors": from safetensors.torch import load_file sd = load_file(model_path) emb = sd["model.embed_tokens.weight"] config = MetaDiffusionConfig(hidden_size=emb.shape[1], mask_vocab_size=emb.shape[0]) model = MetaDiffusionLM(config).to(device) sd = {k[len("model."):] if k.startswith("model.") else k: v for k, v in sd.items()} model.load_state_dict(sd, strict=True) else: ckpt = torch.load(model_path, map_location=device, weights_only=False) config = build_config(ckpt["config"]) model = MetaDiffusionLM(config).to(device) sd = {k.replace("_orig_mod.", "", 1) if k.startswith("_orig_mod.") else k: v for k, v in ckpt["model_state_dict"].items()} model.load_state_dict(sd, strict=True) model.eval() print(f" Loaded {sum(p.numel() for p in model.parameters())/1e6:.1f}M params, " f"vocab={config.mask_vocab_size}") return model def ensure_special_tokens(tokenizer): """Add [MASK] + rainbow if missing (source tokenizer case).""" added = [] if tokenizer.convert_tokens_to_ids("[MASK]") == tokenizer.unk_token_id: added.append("[MASK]") missing = [t for t in RAINBOW_TOKENS if tokenizer.convert_tokens_to_ids(t) == tokenizer.unk_token_id] if missing: added.extend(missing) if added: tokenizer.add_special_tokens({"additional_special_tokens": added}) return tokenizer def format_messages(messages): parts = [] for m in messages: parts.append(f"{IM_START}{m['role']}\n{m['content']}{IM_END}") return "\n".join(parts) def cumulative_unmask_frac(i, N): return 0.5 * (1 - math.cos(math.pi * i / N)) @torch.no_grad() def generate_response(model, tokenizer, prompt_ids, gen_len, num_steps, temperature, repetition_penalty, device, watch=False, stop_on_end=True, cfg_scale=0.0, ban_ids=None, top_p=0.0, min_p=0.0, refine=False, refine_frac=0.3, refine_steps=16, im_end_bias=0.0, im_end_bias_t=0.3, smart_remask=False, smart_remask_thresh=0.5, smart_remask_iters=2): mask_id = model.config.mask_token_id im_end_id = tokenizer.convert_tokens_to_ids(IM_END) eos_id = tokenizer.eos_token_id rainbow_ids = [tokenizer.convert_tokens_to_ids(t) for t in RAINBOW_TOKENS] prompt_len = prompt_ids.shape[1] total_len = prompt_len + gen_len x = torch.full((1, total_len), mask_id, device=device, dtype=torch.long) x[0, :prompt_len] = prompt_ids # commit-confidence map for --smart-remask (top-1 prob at commit time); # 1.0 for prompt/uncommitted so only real commits can fall below the bar conf = (torch.ones((1, total_len), dtype=torch.float32, device=device) if smart_remask else None) terminated = False for i in range(num_steps): frac_now = cumulative_unmask_frac(i, num_steps) frac_next = cumulative_unmask_frac(i + 1, num_steps) n_masked = (x == mask_id).sum().item() if i == num_steps - 1: n_unmask = n_masked else: n_total = int((frac_next - frac_now) * gen_len + 0.5) n_unmask = max(n_total, 1) if n_masked > 0 else 0 if n_unmask == 0: break t = 1.0 - frac_now t_val = torch.full((1,), t, device=device) logits = model(x, t_val).float() # fp32 sampling path: stable softmax # Mid-run curriculum probes extrapolate t beyond what the model has # seen (e.g. t=0.99 at step 8.5K when the ramp max is ~0.48). The # timestep embedding can then blow up to NaN/inf inside the bf16 # forward. Sanitize once here so CFG, softmax and multinomial never # see a poisoned distribution. logits = torch.nan_to_num(logits, nan=0.0, posinf=50.0, neginf=-50.0) if cfg_scale > 0: # classifier-free guidance: unconditional branch sees the prompt # region masked too; logits = cond + s*(cond - uncond). # The all-mask input is off the training manifold, so its logits # can be extreme; extrapolating them in bf16 overflows to inf and # poisons softmax/multinomial. Compute in fp32 and clamp. uncond_x = torch.full_like(x, mask_id) uncond_logits = model(uncond_x, t_val).float() uncond_logits = torch.nan_to_num(uncond_logits, nan=0.0, posinf=50.0, neginf=-50.0) logits = (logits + cfg_scale * (logits - uncond_logits)).clamp(-50.0, 50.0) # padding placeholders (rainbow) and [MASK] are never legitimate output logits[:, :, mask_id] = -1e9 logits[:, :, rainbow_ids] = -1e9 if ban_ids: # partial-byte vocab entries that cannot decode to valid UTF-8: # the literal "�" characters; never legitimate output either logits[:, :, ban_ids] = -1e9 if im_end_bias != 0.0 and t < im_end_bias_t: # pragmatic terminator nudge at the end of denoising: the model's # continuation knowledge runs out before its terminator probability # rises, so make <|im_end|> competitive in the frontier distribution logits[:, :, im_end_id] = logits[:, :, im_end_id] + im_end_bias if repetition_penalty != 1.0: committed = x[0, prompt_len:] committed = committed[committed != mask_id] if committed.numel() > 0: for tok in committed.unique(): ti = tok.item() logits[0, :, ti] = torch.where( logits[0, :, ti] < 0, logits[0, :, ti] * repetition_penalty, logits[0, :, ti] / repetition_penalty, ) mask_positions = x == mask_id sampled, probs = sample_masked(logits, mask_positions, temperature, top_p, min_p) p_max = probs.max(dim=-1).values if conf is not None else None mask_flat = mask_positions.nonzero(as_tuple=False) if n_unmask < mask_positions.sum(): # Left-to-right commit: fill the leftmost masked positions first fill_positions = mask_flat[:n_unmask] for idx, tok in zip(fill_positions, sampled[:n_unmask]): x[idx[0], idx[1]] = tok if conf is not None: conf[fill_positions[:, 0], fill_positions[:, 1]] = p_max[:n_unmask] else: x[mask_positions] = sampled if conf is not None: conf[mask_positions] = p_max if watch: remaining = (x == mask_id).sum().item() live = [t for t in x[0, prompt_len:].tolist() if t != mask_id] partial = tokenizer.decode(cut_response(live, tokenizer), skip_special_tokens=True).strip()[:70] line = f"step {i+1:3d}/{num_steps} | t={t:.3f} | masks={remaining:3d} | {partial}" if sys.stdout.isatty(): sys.stdout.write("\r" + line[:99].ljust(99)) sys.stdout.flush() elif i % max(1, num_steps // 8) == 0: print(line) if stop_on_end and ((x[0, prompt_len:] == im_end_id).any() or (x[0, prompt_len:] == eos_id).any()): terminated = True break if smart_remask and conf is not None: # confidence-gated remasking: re-mask only the low-confidence commits # (the junk-prone tokens) and re-denoise them with the head fixed x = smart_remask_pass(model, x, prompt_len, gen_len, conf, im_end_id, eos_id, mask_id, rainbow_ids, ban_ids, refine_steps, smart_remask_thresh, smart_remask_iters, temperature, repetition_penalty, device, top_p, min_p, im_end_bias, im_end_bias_t) elif refine and not terminated: x = refine_tail(model, x, prompt_len, gen_len, im_end_id, eos_id, mask_id, rainbow_ids, ban_ids, refine_steps, refine_frac, temperature, repetition_penalty, device, top_p, min_p) if watch and sys.stdout.isatty(): sys.stdout.write("\n") return x def sample_masked(logits, mask_positions, temperature, top_p=0.0, min_p=0.0): """Truncation-sampled tokens for the masked positions. top-p nucleus (Holtzman 2020) or Min-P (Nguyen 2024) truncate the unreliable tail of the distribution, which is exactly where junk and rare tokens live when the model runs out of budget. Min-P scales the cutoff by the top token's probability (pbase 0.05-0.1 recommended; use ONE of them). Also guards degenerate rows so multinomial never sees inf/nan/negatives. Returns (sampled, probs).""" probs = F.softmax(logits[mask_positions] / max(temperature, 1e-8), dim=-1) probs = torch.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0) if top_p > 0.0: sorted_probs, indices = probs.sort(dim=-1, descending=True) drop = (sorted_probs.cumsum(dim=-1) - sorted_probs) > top_p sorted_probs = sorted_probs.masked_fill(drop, 0.0) sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True).clamp(min=1e-12) probs = torch.zeros_like(probs).scatter_(-1, indices, sorted_probs) elif min_p > 0.0: threshold = min_p * probs.max(dim=-1, keepdim=True).values probs = probs.masked_fill(probs < threshold, 0.0) probs = probs / probs.sum(dim=-1, keepdim=True).clamp(min=1e-12) # degenerate rows (all-zero after truncation/nan handling) fall back to # uniform so multinomial never sees an invalid distribution zero_rows = probs.sum(dim=-1, keepdim=True) <= 0 if zero_rows.any(): probs = probs + zero_rows.to(probs.dtype) probs = probs / probs.sum(dim=-1, keepdim=True).clamp(min=1e-12) return torch.multinomial(probs, 1).squeeze(-1), probs def refine_tail(model, x, prompt_len, gen_len, im_end_id, eos_id, mask_id, rainbow_ids, ban_ids, steps, frac, temperature, repetition_penalty, device, top_p=0.0, min_p=0.0, im_end_bias=0.0, im_end_bias_t=0.3): """PURE/TOLERATOR-style post-hoc refinement: when a response never committed <|im_end|>, re-mask the tail (keeping the head fixed) and re-denoise it with a short chain. Training-free; converts leftover compute into coherence instead of letting committed junk stick.""" cut = prompt_len + int(gen_len * (1.0 - frac)) x[0, cut:] = mask_id tail_len = gen_len - (cut - prompt_len) for i in range(steps): n_masked = (x[0, cut:] == mask_id).sum().item() if n_masked == 0: break if i == steps - 1: n_unmask = n_masked else: n_unmask = max(int((cumulative_unmask_frac(i + 1, steps) - cumulative_unmask_frac(i, steps)) * tail_len + 0.5), 1) t_val = torch.full((1,), 1.0 - cumulative_unmask_frac(i, steps), device=device) logits = model(x, t_val).float() logits = torch.nan_to_num(logits, nan=0.0, posinf=50.0, neginf=-50.0) logits[:, :, mask_id] = -1e9 logits[:, :, rainbow_ids] = -1e9 if ban_ids: logits[:, :, ban_ids] = -1e9 t_now = 1.0 - cumulative_unmask_frac(i, steps) if im_end_bias != 0.0 and t_now < im_end_bias_t: logits[:, :, im_end_id] = logits[:, :, im_end_id] + im_end_bias if repetition_penalty != 1.0: committed = x[0, prompt_len:] committed = committed[committed != mask_id] if committed.numel() > 0: for tok in committed.unique(): ti = tok.item() logits[0, :, ti] = torch.where( logits[0, :, ti] < 0, logits[0, :, ti] * repetition_penalty, logits[0, :, ti] / repetition_penalty) mask_positions = x == mask_id sampled, _ = sample_masked(logits, mask_positions, temperature, top_p, min_p) mask_flat = mask_positions.nonzero(as_tuple=False) if n_unmask < mask_positions.sum(): for idx, tok in zip(mask_flat[:n_unmask], sampled[:n_unmask]): x[idx[0], idx[1]] = tok else: x[mask_positions] = sampled if (x[0, prompt_len:] == im_end_id).any() or (x[0, prompt_len:] == eos_id).any(): break return x def smart_remask_pass(model, x, prompt_len, gen_len, conf, im_end_id, eos_id, mask_id, rainbow_ids, ban_ids, steps, thresh, max_iters, temperature, repetition_penalty, device, top_p=0.0, min_p=0.0, im_end_bias=0.0, im_end_bias_t=0.3): """Confidence-gated re-denoising (PURE-style smart remasking). The blind --refine tail remask wastes budget on tokens the model already committed with high confidence. Here, re-mask exactly the tokens whose top-1 commit probability fell below `thresh` (the junk-prone ones, often the budget-tail fillers) and re-denoise them with the head fixed. Runs even when im_end committed: it also cleans low-confidence junk sitting before the terminator. Repeats up to max_iters rounds and stops early once the terminator commits or nothing is below the bar.""" lo = prompt_len hi = prompt_len + gen_len for _ in range(max_iters): resp = x[0, lo:hi] term = (resp == im_end_id) | (resp == eos_id) if term.any(): # never touch the terminator or anything past it hi = lo + term.nonzero(as_tuple=True)[0][0].item() if hi <= lo: break low = (conf[0, lo:hi] < thresh).nonzero(as_tuple=True)[0] if low.numel() == 0: break n_remask = low.numel() x[0, lo + low] = mask_id conf[0, lo + low] = 1.0 # re-commits below the bar get caught again for i in range(steps): n_masked = (x[0, lo:hi] == mask_id).sum().item() if n_masked == 0: break if i == steps - 1: n_unmask = n_masked else: n_unmask = max(int((cumulative_unmask_frac(i + 1, steps) - cumulative_unmask_frac(i, steps)) * n_remask + 0.5), 1) n_unmask = min(n_unmask, n_masked) t_now = 1.0 - cumulative_unmask_frac(i, steps) t_val = torch.full((1,), t_now, device=device) logits = model(x, t_val).float() logits = torch.nan_to_num(logits, nan=0.0, posinf=50.0, neginf=-50.0) logits[:, :, mask_id] = -1e9 logits[:, :, rainbow_ids] = -1e9 if ban_ids: logits[:, :, ban_ids] = -1e9 if im_end_bias != 0.0 and t_now < im_end_bias_t: logits[:, :, im_end_id] = logits[:, :, im_end_id] + im_end_bias if repetition_penalty != 1.0: committed = x[0, prompt_len:] committed = committed[committed != mask_id] if committed.numel() > 0: for tok in committed.unique(): ti = tok.item() logits[0, :, ti] = torch.where( logits[0, :, ti] < 0, logits[0, :, ti] * repetition_penalty, logits[0, :, ti] / repetition_penalty) mask_positions = x == mask_id sampled, probs = sample_masked(logits, mask_positions, temperature, top_p, min_p) p_max = probs.max(dim=-1).values mask_flat = mask_positions.nonzero(as_tuple=False) n_fill = min(n_unmask, mask_flat.shape[0]) if n_fill: idxs = mask_flat[:n_fill] x[idxs[:, 0], idxs[:, 1]] = sampled[:n_fill] conf[idxs[:, 0], idxs[:, 1]] = p_max[:n_fill] if (x[0, lo:hi] == im_end_id).any() or \ (x[0, lo:hi] == eos_id).any(): break if (x[0, lo:hi] == im_end_id).any() or (x[0, lo:hi] == eos_id).any(): break return x def cut_response(tokens, tokenizer): """Cut at <|im_end|> / eos; drop rainbow and pad tokens.""" im_end_id = tokenizer.convert_tokens_to_ids(IM_END) eos_id = tokenizer.eos_token_id rainbow_ids = {tokenizer.convert_tokens_to_ids(t) for t in RAINBOW_TOKENS} out = [] for t in tokens: if t == im_end_id or t == eos_id: break if t in rainbow_ids or t == tokenizer.pad_token_id: continue out.append(t) return out def invalid_utf8_ids(tokenizer): """Ids whose decode is *only* U+FFFD. Byte-fallback tokens that merely contain a replacement char when decoded alone stay; those are how Qwen builds rare unicode.""" ban = [] for i in range(len(tokenizer)): s = tokenizer.decode([i], skip_special_tokens=True) if s and all(c == "\uFFFD" for c in s): ban.append(i) return ban def trim_messages(messages, tokenizer, max_context, max_new_tokens): """Drop oldest non-system turns until prompt + gen budget fits.""" budget = max(32, max_context - max_new_tokens) kept = list(messages) while kept: prompt = format_messages(kept) + f"\n{IM_START}assistant\n" n = len(tokenizer.encode(prompt, add_special_tokens=False)) if n <= budget: return kept drop_at = next((i for i, m in enumerate(kept) if m["role"] != "system"), None) if drop_at is None: return kept del kept[drop_at] return kept def run_turn(model, tokenizer, messages, args, device): max_ctx = getattr(args, "max_context", 4096) cap = min(getattr(model.config, "max_position_embeddings", 40960), max_ctx) messages = trim_messages(messages, tokenizer, cap, args.max_new_tokens) prompt = format_messages(messages) + f"\n{IM_START}assistant\n" prompt_ids = torch.tensor([tokenizer.encode(prompt, add_special_tokens=False)], device=device) for attempt in range(3): x = generate_response(model, tokenizer, prompt_ids, args.max_new_tokens, args.num_steps, args.temperature * (1 + 0.15 * attempt), args.repetition_penalty, device, watch=args.watch, cfg_scale=args.cfg_scale, ban_ids=getattr(args, "bad_token_ids", None), top_p=args.top_p, min_p=args.min_p, refine=args.refine, refine_frac=args.refine_frac, refine_steps=args.refine_steps, im_end_bias=args.im_end_bias, im_end_bias_t=args.im_end_bias_t, smart_remask=args.smart_remask, smart_remask_thresh=args.smart_remask_thresh, smart_remask_iters=args.smart_remask_iters) text = tokenizer.decode(cut_response(x[0, prompt_ids.shape[1]:].tolist(), tokenizer), skip_special_tokens=True).strip() if text: return text return "(empty response)" def main(): p = argparse.ArgumentParser(description="MetaDiffusion-600M chat") p.add_argument("--model-path", required=True) p.add_argument("--tokenizer", default=None, help="Tokenizer dir (needed for .pt checkpoints)") p.add_argument("--prompt", default=None) p.add_argument("--system", default="You are a helpful assistant.") p.add_argument("--max-new-tokens", type=int, default=96) p.add_argument("--num-steps", type=int, default=128) p.add_argument("--temperature", type=float, default=0.7) p.add_argument("--repetition-penalty", type=float, default=1.5) p.add_argument("--cfg-scale", type=float, default=0.0, help="Classifier-free guidance scale (0 = off; try 0.5-1.2). " "Unconditional branch masks the prompt too.") p.add_argument("--top-p", type=float, default=0.0, help="Nucleus sampling: keep tokens covering this mass (0=off; " "use EITHER --top-p or --min-p, not both)") p.add_argument("--min-p", type=float, default=0.1, help="Min-P truncation: keep tokens >= min_p x top-token prob " "(0=off; 0.05-0.1 recommended). Truncates the junk tail.") p.add_argument("--refine", action="store_true", help="Post-hoc tail refinement: if im_end never commits, re-mask " "the tail and re-denoise it (PURE/TOLERATOR-style)") p.add_argument("--refine-frac", type=float, default=0.3, help="Fraction of the response tail to re-denoise (--refine)") p.add_argument("--refine-steps", type=int, default=16, help="Denoising steps for the refinement pass") p.add_argument("--smart-remask", action="store_true", help="Confidence-gated remasking (PURE-style): re-mask only " "the tokens committed with low top-1 probability and " "re-denoise them with the head fixed. Runs even when " "im_end committed (cleans pre-terminator junk); takes " "precedence over --refine.") p.add_argument("--smart-remask-thresh", type=float, default=0.5, help="Commit-confidence bar (top-1 token prob at commit " "time); tokens below it are re-masked (--smart-remask)") p.add_argument("--smart-remask-iters", type=int, default=2, help="Max refinement rounds; stops early when the " "terminator commits or nothing is below the bar") p.add_argument("--im-end-bias", type=float, default=0.0, help="Logit bonus on <|im_end|> when t < --im-end-bias-t " "(pragmatic terminator nudge; try 1.5-3.0)") p.add_argument("--im-end-bias-t", type=float, default=0.3, help="t threshold below which --im-end-bias applies") p.add_argument("--device", default="cuda") p.add_argument("--max-context", type=int, default=4096, help="Trim multi-turn history so prompt+gen fits this many tokens") p.add_argument("--watch", action="store_true") args = p.parse_args() device = torch.device(args.device if torch.cuda.is_available() else "cpu") print(f"[*] Loading model from {args.model_path}") model = load_model(args.model_path, device) tok_path = args.tokenizer if tok_path is None: model_path = Path(args.model_path) if model_path.is_dir(): cand = model_path / "tokenizer" if not cand.exists() and (model_path / "tokenizer.json").exists(): cand = model_path tok_path = str(cand) if not tok_path or not Path(tok_path).exists(): raise SystemExit("No tokenizer found; pass --tokenizer (data/tokenizer)") tokenizer = ensure_special_tokens(AutoTokenizer.from_pretrained(str(tok_path))) print(f"[*] Tokenizer: {tok_path} (vocab {len(tokenizer)})") # hard-ban vocab entries that cannot decode to valid UTF-8 (partial-byte # tokens): they surface as "�" garbage and are never legitimate output args.bad_token_ids = invalid_utf8_ids(tokenizer) if args.bad_token_ids: print(f"[*] Banning {len(args.bad_token_ids)} standalone-U+FFFD tokens") if args.prompt: text = run_turn(model, tokenizer, [{"role": "user", "content": args.prompt}], args, device) print(f"\nUser: {args.prompt}\nAssistant: {text}\n") return print("\nMetaDiffusion-600M chat. Type 'exit' to leave.\n") messages = [{"role": "system", "content": args.system}] while True: try: user_input = input("You: ").strip() except (EOFError, KeyboardInterrupt): print() break if user_input.lower() in ("exit", "quit"): break if not user_input: continue messages.append({"role": "user", "content": user_input}) text = run_turn(model, tokenizer, messages, args, device) print(f"Assistant: {text}\n") messages.append({"role": "assistant", "content": text}) if __name__ == "__main__": main()