LTA_openwebtext_dualt/scripts/blockar_twostream_genppl_entropy_decode.py

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()