chat.py

#!/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()