reference/code/TAD/eval/eval_llada.py

import os
# os.environ["ALL_PROXY"] = "socks5h://127.0.0.1:13659"
os.environ['CURL_CA_BUNDLE'] = ''
os.environ['REQUESTS_CA_BUNDLE'] = ''
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
import accelerate
import torch
import random
import numpy as np
import torch.nn.functional as F
from datasets import Dataset
from lm_eval.__main__ import cli_evaluate
from lm_eval.api.model import LM
from lm_eval.api.registry import register_model
from tqdm import tqdm
from transformers import AutoTokenizer, AutoConfig, AutoModel
from model.modeling_llada import LLaDAModelLM
import json
import time
from accelerate import (
    Accelerator,
    InitProcessGroupKwargs,
)
from datetime import timedelta


def set_seed(seed):
    torch.manual_seed(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


def add_gumbel_noise(logits, temperature):
    if temperature == 0:
        return logits
    logits = logits.to(torch.float64)
    noise = torch.rand_like(logits, dtype=torch.float64)
    gumbel_noise = (-torch.log(noise)) ** temperature
    return logits.exp() / gumbel_noise


def get_num_transfer_tokens(mask_index, steps):
    mask_num = mask_index.sum(dim=1, keepdim=True)
    base = mask_num // steps
    remainder = mask_num % steps
    num_transfer_tokens = torch.zeros(mask_num.size(0), steps, device=mask_index.device, dtype=torch.int64) + base
    for i in range(mask_num.size(0)):
        num_transfer_tokens[i, :remainder[i]] += 1
    return num_transfer_tokens


def get_transfer_index_entropy(
    logits,
    temperature,
    remasking,
    mask_index,
    x,
    entropy_threshold=None,
    num_transfer_tokens=None,
):
    logits_with_noise = add_gumbel_noise(logits, temperature=temperature)
    x0 = torch.argmax(logits_with_noise, dim=-1)
    p = F.softmax(logits.to(torch.float64), dim=-1)
    if remasking == "low_confidence":
        entropy = -torch.sum(p * torch.log(p + 1e-12), dim=-1)
    elif remasking == "random":
        entropy = torch.rand((x0.shape[0], x0.shape[1]), device=x0.device)
    else:
        raise NotImplementedError(remasking)
    x0 = torch.where(mask_index, x0, x)
    entropy_for_selection = torch.where(mask_index, entropy, torch.inf)
    transfer_index = torch.zeros_like(x0, dtype=torch.bool, device=x0.device)
    if entropy_threshold is not None:
        transfer_index = entropy_for_selection < entropy_threshold
        for j in range(entropy_for_selection.shape[0]):
            if mask_index[j].sum() > 0 and transfer_index[j].sum() == 0:
                min_index = torch.argmin(entropy_for_selection[j])
                transfer_index[j, min_index] = True
    else:
        for j in range(entropy_for_selection.shape[0]):
            _, select_index = torch.topk(entropy_for_selection[j], k=num_transfer_tokens[j], largest=False)
            transfer_index[j, select_index] = True
    return x0, transfer_index


def _check_early_stop(x, prompt_length, eos_token_id, mask_id):
    """
    Check for EOS token in the generation region.
    If found, set all tokens after the first EOS to EOS (not mask) and
    return (True, first_eos_absolute_position).
    Otherwise return (False, None).
    """
    if eos_token_id is None:
        return False, None
    gen_region = x[:, prompt_length:]
    eos_mask = (gen_region == eos_token_id) & (gen_region != mask_id)
    if not eos_mask.any():
        return False, None
    pos = torch.arange(gen_region.shape[1], device=x.device).unsqueeze(0)
    first_eos_rel = torch.where(eos_mask, pos, gen_region.shape[1]).amin(dim=1)
    first_eos_abs = prompt_length + first_eos_rel[0].item()
    # Set everything after first EOS to EOS (not mask)
    x[:, first_eos_abs + 1:] = eos_token_id
    return True, first_eos_abs


@torch.no_grad()
def generate(
    model,
    prompt,
    steps=128,
    gen_length=128,
    block_length=128,
    temperature=0.0,
    remasking="low_confidence",
    mask_id=126336,
    threshold=None,
    eos_token_id=None,
):
    x = torch.full((1, prompt.shape[1] + gen_length), mask_id, dtype=torch.long).to(model.device)
    x[:, :prompt.shape[1]] = prompt.clone()
    prompt_length = prompt.shape[1]
    assert gen_length % block_length == 0
    num_blocks = gen_length // block_length
    assert steps % num_blocks == 0
    steps = steps // num_blocks
    nfe = 0
    for num_block in range(num_blocks):
        # Early stop: skip remaining blocks if EOS already found
        if eos_token_id is not None:
            has_eos, _ = _check_early_stop(x, prompt_length, eos_token_id, mask_id)
            if has_eos:
                break
        block_mask_index = (
            x[:, prompt.shape[1] + num_block * block_length : prompt.shape[1] + (num_block + 1) * block_length] == mask_id
        )
        num_transfer_tokens = get_num_transfer_tokens(block_mask_index, steps)
        i = 0
        while True:
            nfe += 1
            mask_index = x == mask_id
            logits = model(x).logits
            mask_index[:, prompt.shape[1] + (num_block + 1) * block_length :] = 0
            x0, transfer_index = get_transfer_index_entropy(
                logits,
                temperature,
                remasking,
                mask_index,
                x,
                entropy_threshold=threshold,
                num_transfer_tokens=num_transfer_tokens[:, i] if threshold is None else None,
            )
            x[transfer_index] = x0[transfer_index]
            i += 1
            # Early stop: check after each decode step
            if eos_token_id is not None:
                has_eos, _ = _check_early_stop(x, prompt_length, eos_token_id, mask_id)
                if has_eos:
                    break
            if (
                x[:, prompt.shape[1] + num_block * block_length : prompt.shape[1] + (num_block + 1) * block_length] == mask_id
            ).sum() == 0:
                break
    return x, nfe


@torch.no_grad()
def generate_with_prefix_cache(
    model,
    prompt,
    steps=128,
    gen_length=128,
    block_length=128,
    temperature=0.0,
    remasking="low_confidence",
    mask_id=126336,
    threshold=None,
    eos_token_id=None,
):
    x = torch.full((1, prompt.shape[1] + gen_length), mask_id, dtype=torch.long).to(model.device)
    x[:, :prompt.shape[1]] = prompt.clone()
    prompt_length = prompt.shape[1]
    assert gen_length % block_length == 0
    num_blocks = gen_length // block_length
    assert steps % num_blocks == 0
    steps = steps // num_blocks
    nfe = 0
    for num_block in range(num_blocks):
        # Early stop: skip remaining blocks if EOS already found
        if eos_token_id is not None:
            has_eos, _ = _check_early_stop(x, prompt_length, eos_token_id, mask_id)
            if has_eos:
                break
        current_block_start = prompt.shape[1] + num_block * block_length
        current_block_end = current_block_start + block_length
        block_mask_index = x[:, current_block_start:current_block_end] == mask_id
        num_transfer_tokens = get_num_transfer_tokens(block_mask_index, steps)
        output = model(x, use_cache=True)
        past_key_values = output.past_key_values
        mask_index = x == mask_id
        mask_index[:, current_block_end:] = 0
        x0, transfer_index = get_transfer_index_entropy(
            output.logits,
            temperature,
            remasking,
            mask_index,
            x,
            entropy_threshold=threshold,
            num_transfer_tokens=num_transfer_tokens[:, 0] if threshold is None else None,
        )
        x[transfer_index] = x0[transfer_index]
        new_past_key_values = []
        for i in range(len(past_key_values)):
            new_past_key_values.append(())
            for j in range(len(past_key_values[i])):
                new_past_key_values[i] += (past_key_values[i][j][:, :, :current_block_start],)
        past_key_values = new_past_key_values
        nfe += 1
        i = 1
        while True:
            nfe += 1
            mask_index = x[:, current_block_start:] == mask_id
            mask_index[:, block_length:] = 0
            logits = model(x[:, current_block_start:], past_key_values=past_key_values, use_cache=True).logits
            x0, transfer_index = get_transfer_index_entropy(
                logits,
                temperature,
                remasking,
                mask_index,
                x[:, current_block_start:],
                entropy_threshold=threshold,
                num_transfer_tokens=num_transfer_tokens[:, i] if threshold is None else None,
            )
            x[:, current_block_start:][transfer_index] = x0[transfer_index]
            # Early stop: check after each decode step
            if eos_token_id is not None:
                has_eos, _ = _check_early_stop(x, prompt_length, eos_token_id, mask_id)
                if has_eos:
                    break
            if (x[:, current_block_start:current_block_end] == mask_id).sum() == 0:
                break
            i += 1
    return x, nfe


@torch.no_grad()
def generate_with_dual_cache(
    model,
    prompt,
    steps=128,
    gen_length=128,
    block_length=128,
    temperature=0.0,
    remasking="low_confidence",
    mask_id=126336,
    threshold=None,
    eos_token_id=None,
):
    x = torch.full((1, prompt.shape[1] + gen_length), mask_id, dtype=torch.long).to(model.device)
    x[:, :prompt.shape[1]] = prompt.clone()
    prompt_length = prompt.shape[1]
    assert gen_length % block_length == 0
    num_blocks = gen_length // block_length
    assert steps % num_blocks == 0
    steps = steps // num_blocks
    nfe = 0
    for num_block in range(num_blocks):
        # Early stop: skip remaining blocks if EOS already found
        if eos_token_id is not None:
            has_eos, _ = _check_early_stop(x, prompt_length, eos_token_id, mask_id)
            if has_eos:
                break
        current_block_start = prompt.shape[1] + num_block * block_length
        current_block_end = current_block_start + block_length
        block_mask_index = x[:, current_block_start:current_block_end] == mask_id
        num_transfer_tokens = get_num_transfer_tokens(block_mask_index, steps)
        output = model(x, use_cache=True)
        past_key_values = output.past_key_values
        mask_index = x == mask_id
        mask_index[:, current_block_end:] = 0
        x0, transfer_index = get_transfer_index_entropy(
            output.logits,
            temperature,
            remasking,
            mask_index,
            x,
            entropy_threshold=threshold,
            num_transfer_tokens=num_transfer_tokens[:, 0] if threshold is None else None,
        )
        x[transfer_index] = x0[transfer_index]
        nfe += 1
        i = 1
        replace_position = torch.zeros_like(x, dtype=torch.bool)
        replace_position[:, current_block_start:current_block_end] = 1
        while True:
            nfe += 1
            mask_index = x[:, current_block_start:current_block_end] == mask_id
            logits = model(
                x[:, current_block_start:current_block_end],
                past_key_values=past_key_values,
                use_cache=True,
                replace_position=replace_position,
            ).logits
            x0, transfer_index = get_transfer_index_entropy(
                logits,
                temperature,
                remasking,
                mask_index,
                x[:, current_block_start:current_block_end],
                entropy_threshold=threshold,
                num_transfer_tokens=num_transfer_tokens[:, i] if threshold is None else None,
            )
            x[:, current_block_start:current_block_end][transfer_index] = x0[transfer_index]
            # Early stop: check after each decode step
            if eos_token_id is not None:
                has_eos, _ = _check_early_stop(x, prompt_length, eos_token_id, mask_id)
                if has_eos:
                    break
            if (x[:, current_block_start:current_block_end] == mask_id).sum() == 0:
                break
            i += 1
    return x, nfe


@torch.no_grad()
def generate_multi_block(
    model,
    prompt,
    steps=128,
    max_new_tokens=512,
    block_size=32,
    temperature=0.0,
    remasking="low_confidence",
    mask_id=126336,
    threshold=0.5,
    block_add_threshold=0.5,
    decoded_token_threshold=0.5,
    eos_token_id=None,
):
    """
    Pipelined parallel decoding without cache.

    Args:
        block_add_threshold: Add new block when last block progress >= this threshold.
                            Set to 1.0 for fully sequential processing (like generate()).
        decoded_token_threshold: Block becomes fully activated when previous block progress >= this threshold.
                                Set to 1.0 for fully sequential processing (like generate()).
        threshold: Entropy threshold for decoding (lower entropy = higher confidence).
                   Tokens with entropy > threshold are skipped.
                   Same semantics as in generate() method. Typical value: 0.5

    When block_add_threshold=1.0 and decoded_token_threshold=1.0, this method behaves
    identically to generate() with sequential block processing.
    """
    x = torch.full((1, prompt.shape[1] + max_new_tokens), mask_id, dtype=torch.long).to(
        model.device
    )
    x[:, : prompt.shape[1]] = prompt.clone()
    prompt_length = prompt.shape[1]

    # Track block states: {block_id: {start, end, mask_count, total_masks, is_complete}}
    # Initialize with prompt block
    block_states = {
        0: {
            "start": 0,
            "end": prompt.shape[1],
            "mask_count": 0,
            "total_masks": prompt.shape[1],
            "is_complete": True,
        }
    }

    # Create first generation block
    num_blocks = max_new_tokens // block_size
    next_block_id = 1
    if next_block_id <= num_blocks:
        block_start = prompt.shape[1] + (next_block_id - 1) * block_size
        block_end = min(block_start + block_size, prompt.shape[1] + max_new_tokens)
        # First block should be immediately activated since prompt (block 0) is already complete
        should_activate = 1.0 >= decoded_token_threshold  # prompt progress is always 1.0
        block_states[next_block_id] = {
            "start": block_start,
            "end": block_end,
            "mask_count": block_end - block_start,
            "total_masks": block_end - block_start,
            "is_complete": should_activate,
        }
        next_block_id += 1

    nfe = 0

    while True:
        # Check if all blocks are exhausted AND no more blocks to create
        mask_index = x == mask_id
        total_masks = mask_index[:, prompt_length:].sum()

        if total_masks == 0 and next_block_id > num_blocks:
            break

        nfe += 1

        # Early stop: check for EOS token
        if eos_token_id is not None:
            has_eos, first_eos_abs = _check_early_stop(x, prompt_length, eos_token_id, mask_id)
            if has_eos:
                # Create all remaining blocks after EOS and mark them as complete
                while next_block_id <= num_blocks:
                    block_start = prompt_length + (next_block_id - 1) * block_size
                    block_end = min(block_start + block_size, prompt_length + max_new_tokens)
                    if block_start > first_eos_abs:
                        block_states[next_block_id] = {
                            "start": block_start, "end": block_end, "mask_count": 0,
                            "total_masks": block_end - block_start, "is_complete": True,
                        }
                        next_block_id += 1
                    else:
                        break
                # Recalculate: if no masks remain, we're done
                if (x == mask_id)[:, prompt_length:].sum() == 0:
                    break

        # Update block completion states
        def update_block_activation_states():
            """Update which blocks should be fully activated based on previous block progress."""
            for bid in sorted(block_states.keys()):
                if bid > 0 and not block_states[bid]["is_complete"]:
                    prev_progress = (
                        1
                        - block_states[bid - 1]["mask_count"]
                        / block_states[bid - 1]["total_masks"]
                    )
                    if prev_progress >= decoded_token_threshold:
                        block_states[bid]["is_complete"] = True

        update_block_activation_states()

        # Add new block dynamically based on last block's progress
        if next_block_id <= num_blocks:
            last_bid = max(block_states.keys())
            if last_bid > 0:  # Not just prompt
                last_progress = (
                    1
                    - block_states[last_bid]["mask_count"]
                    / block_states[last_bid]["total_masks"]
                )
                # Create next block when:
                # 1. Last block progress >= block_add_threshold (for parallel processing), OR
                # 2. Last block is complete (mask_count == 0) for sequential processing
                should_add_block = (last_progress >= block_add_threshold) or (block_states[last_bid]["mask_count"] == 0)

                if should_add_block:
                    # Add next block
                    block_start = prompt.shape[1] + (next_block_id - 1) * block_size
                    block_end = min(block_start + block_size, prompt.shape[1] + max_new_tokens)
                    if block_end > block_start:
                        # Check how many positions in this block are actually masked
                        actual_mask_count = (x[:, block_start:block_end] == mask_id).sum().item()

                        # Determine if this block should be immediately activated
                        # Check if previous block is complete enough
                        prev_bid = next_block_id - 1
                        prev_progress = (
                            1 - block_states[prev_bid]["mask_count"] / block_states[prev_bid]["total_masks"]
                        )
                        should_activate = prev_progress >= decoded_token_threshold

                        block_states[next_block_id] = {
                            "start": block_start,
                            "end": block_end,
                            "mask_count": actual_mask_count,
                            "total_masks": block_end - block_start,
                            "is_complete": should_activate,
                        }
                        next_block_id += 1

        # Forward pass: only process up to the last complete or semi-activated block
        # Find the rightmost block that should be processed
        rightmost_active_bid = 0
        for bid in sorted(block_states.keys()):
            if block_states[bid]["is_complete"] or block_states[bid]["mask_count"] > 0:
                rightmost_active_bid = bid

        if rightmost_active_bid == 0:
            break

        active_end = block_states[rightmost_active_bid]["end"]

        # Always do forward pass on entire sequence (like generate() does)
        logits = model(x).logits

        # Mask out future blocks (positions after active_end) to prevent them from being decoded
        mask_index_for_decode = mask_index.clone()
        mask_index_for_decode[:, active_end:] = 0

        # Decode using full-sequence approach (like generate())
        x0, transfer_index = get_transfer_index_entropy(
            logits,
            temperature,
            remasking,
            mask_index_for_decode,
            x,
            entropy_threshold=threshold if threshold is not None else 999.0,
            num_transfer_tokens=None,
        )

        # For fully activated blocks, ensure at least one token is decoded (guaranteed progress)
        # Find the first fully activated block with masks
        first_fully_activated_bid = None
        for bid in sorted(block_states.keys()):
            if bid > 0 and block_states[bid]["is_complete"] and block_states[bid]["mask_count"] > 0:
                first_fully_activated_bid = bid
                break

        if first_fully_activated_bid is not None:
            # Check if any token was decoded in this block
            start, end = block_states[first_fully_activated_bid]["start"], block_states[first_fully_activated_bid]["end"]
            block_transfer = transfer_index[:, start:end]

            if not block_transfer.any():
                # Force decode the lowest entropy token in this fully activated block
                p = F.softmax(logits[:, start:end].to(torch.float64), dim=-1)
                entropy = -torch.sum(p * torch.log(p + 1e-12), dim=-1)
                block_mask = mask_index_for_decode[:, start:end]
                entropy = torch.where(block_mask, entropy, torch.inf)
                best_idx = entropy[0].argmin()
                transfer_index[0, start + best_idx] = True
                x0_resample = torch.argmax(logits[0, start + best_idx], dim=-1)
                x0[0, start + best_idx] = x0_resample

        # Apply the decoded tokens
        x[transfer_index] = x0[transfer_index]

        # Update block states based on which positions were decoded
        for bid in sorted(block_states.keys()):
            if bid > 0 and block_states[bid]["mask_count"] > 0:
                start, end = block_states[bid]["start"], block_states[bid]["end"]
                block_decoded = transfer_index[:, start:end].sum().item()
                if block_decoded > 0:
                    block_states[bid]["mask_count"] -= block_decoded

        if nfe > 10000:
            break

    return x, nfe


@register_model("llada_dist")
class LLaDAEvalHarness(LM):
    def __init__(
        self,
        model_path="/home/u2025104115/LLada_sft/SAFT-LLaDA/checkpoint/save_checkpoints/SAFT-LLaDA-5k_data",
        mask_id=126336,
        max_length=1024,
        batch_size=32,
        mc_num=128,
        is_check_greedy=True,
        steps=1024,
        gen_length=1024,
        block_length=1024,
        remasking="low_confidence",
        device="cuda",
        use_cache=False,
        threshold=None,
        save_dir=None,
        stats_dir=None,
        show_speed=False,
        dual_cache=False,
        multi_block=False,
        block_add_threshold=0.5,
        decoded_token_threshold=0.5,
        early_stop=False,
        task="null",
        **kwargs,
    ):
        super().__init__()
        accelerator_kwargs = InitProcessGroupKwargs(timeout=timedelta(weeks=52))
        accelerator = Accelerator(kwargs_handlers=[accelerator_kwargs])
        if accelerator.num_processes > 1:
            self.accelerator = accelerator
        else:
            self.accelerator = None
        model_kwargs = {}
        if self.accelerator is not None:
            model_kwargs.update({"device_map": {"": f"{self.accelerator.device}"}})
        config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
        if hasattr(config, "flash_attention"):
            config.flash_attention = True
        self.model = LLaDAModelLM.from_pretrained(
            model_path,
            trust_remote_code=True,
            torch_dtype=torch.bfloat16,
            config=config,
            **model_kwargs,
        )
        self.model.eval()
        self.device = torch.device(device)
        if self.accelerator is not None:
            self.model = self.model.to(self.accelerator.device)
            self.device = torch.device(f"{self.accelerator.device}")
            self._rank = self.accelerator.local_process_index
            self._world_size = self.accelerator.num_processes
        else:
            self.model = self.model.to(device)
            self._rank = 0
            self._world_size = 1
        self.mask_id = mask_id
        if hasattr(config, "mask_token_id") and config.mask_token_id is not None:
            self.mask_id = config.mask_token_id
        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
        self.mc_num = mc_num
        self.batch_size = int(batch_size)
        assert mc_num % self.batch_size == 0
        self.sampling_eps = 0.0
        self.max_length = max_length
        self.is_check_greedy = is_check_greedy
        self.steps = steps
        self.gen_length = gen_length
        self.block_length = block_length
        self.remasking = remasking
        self.use_cache = use_cache
        self.threshold = threshold
        # self.is_instruct = True if ("instruct" in model_path.lower() or "1.5" in model_path.lower()) else False
        self.is_instruct = True
        self.save_dir = save_dir
        self.stats_dir = stats_dir if stats_dir is not None else save_dir
        self.show_speed = show_speed
        self.dual_cache = dual_cache
        self.multi_block = multi_block
        self.block_add_threshold = block_add_threshold
        self.decoded_token_threshold = decoded_token_threshold
        self.early_stop = early_stop
        self.task = task
        self.cfg = 0

    @property
    def rank(self):
        return self._rank

    @property
    def world_size(self):
        return self._world_size

    def _forward_process(self, batch, prompt_index):
        b, l = batch.shape
        target_len = (l - prompt_index.sum()).item()
        k = torch.randint(1, target_len + 1, (), device=batch.device)
        x = torch.round(torch.linspace(float(k), k + (b - 1) * (target_len / b), steps=b, device=batch.device)).long()
        x = ((x - 1) % target_len) + 1
        assert x.min() >= 1 and x.max() <= target_len
        indices = torch.arange(target_len, device=batch.device).repeat(b, 1)
        is_mask = indices < x.unsqueeze(1)
        for i in range(b):
            is_mask[i] = is_mask[i][torch.randperm(target_len)]
        is_mask = torch.cat((torch.zeros(b, prompt_index.sum(), dtype=torch.bool, device=batch.device), is_mask), dim=1)
        noisy_batch = torch.where(is_mask, self.mask_id, batch)
        return noisy_batch, (x / target_len).unsqueeze(1).repeat(1, l)

    @torch.no_grad()
    def get_logits(self, batch, prompt_index):
        if self.cfg > 0.0:
            assert len(prompt_index) == batch.shape[1]
            prompt_index = prompt_index.unsqueeze(0).repeat(batch.shape[0], 1)
            un_batch = batch.clone()
            un_batch[prompt_index] = self.mask_id
            batch = torch.cat([batch, un_batch])
        logits = self.model(batch).logits
        if self.cfg > 0.0:
            logits, un_logits = torch.chunk(logits, 2, dim=0)
            logits = un_logits + (self.cfg + 1) * (logits - un_logits)
        return logits[:, : batch.shape[1]]

    @torch.no_grad()
    def get_loglikelihood(self, prefix, target):
        seq = torch.concatenate([prefix, target])[None, :]
        seq = seq.repeat((self.batch_size, 1)).to(self.device)
        prompt_index = torch.arange(seq.shape[1], device=self.device) < len(prefix)
        loss_acc = []
        for _ in range(self.mc_num // self.batch_size):
            perturbed_seq, p_mask = self._forward_process(seq, prompt_index)
            mask_indices = perturbed_seq == self.mask_id
            logits = self.get_logits(perturbed_seq, prompt_index)
            loss = F.cross_entropy(logits[mask_indices], seq[mask_indices], reduction="none") / p_mask[mask_indices]
            loss = loss.sum() / self.batch_size
            loss_acc.append(loss.item())
        return -sum(loss_acc) / len(loss_acc)

    @torch.no_grad()
    def suffix_greedy_prediction(self, prefix, target):
        if not self.is_check_greedy:
            return False
        seq = torch.full((1, len(prefix) + len(target)), self.mask_id, device=self.device)
        prompt_index = torch.arange(seq.shape[1], device=self.device) < len(prefix)
        prefix, target = prefix.to(self.device), target.to(self.device)
        seq[0, : len(prefix)] = prefix
        for _ in range(len(target)):
            mask_index = seq == self.mask_id
            logits = self.get_logits(seq, prompt_index)[mask_index]
            x0 = torch.argmax(logits, dim=-1)
            p = torch.softmax(logits.to(torch.float32), dim=-1)
            confidence = torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)).squeeze(dim=-1)
            _, index = torch.sort(confidence, descending=True)
            x0[index[1:]] = self.mask_id
            seq[mask_index] = x0.clone()
        correct = target == seq[0, len(prefix) :]
        correct = torch.all(correct)
        return correct

    def _encode_pair(self, context, continuation):
        n_spaces = len(context) - len(context.rstrip())
        if n_spaces > 0:
            continuation = context[-n_spaces:] + continuation
            context = context[:-n_spaces]
        whole_enc = self.tokenizer(context + continuation)["input_ids"]
        context_enc = self.tokenizer(context)["input_ids"]
        context_enc_len = len(context_enc)
        continuation_enc = whole_enc[context_enc_len:]
        return context_enc, continuation_enc

    def loglikelihood(self, requests):
        def _tokenize(e):
            prefix, target = self._encode_pair(e["prefix"], e["target"])
            return {
                "prefix_text": e["prefix"],
                "target_text": e["target"],
                "prefix": prefix,
                "target": target,
            }

        ds = [{"prefix": req.args[0], "target": req.args[1]} for req in requests]
        ds = Dataset.from_list(ds)
        ds = ds.map(_tokenize)
        ds = ds.with_format("torch")
        prompt_len = [len(x["prefix"]) + len(x["target"]) for x in ds]
        assert max(prompt_len) <= 4096
        out = []
        with torch.no_grad():
            for elem in tqdm(ds, desc="Computing likelihood..."):
                prefix = elem["prefix"]
                target = elem["target"]
                ll = self.get_loglikelihood(prefix, target)
                is_target_greedy_dec = self.suffix_greedy_prediction(prefix, target)
                out.append((ll, 1.0 if is_target_greedy_dec else 0.0))
        torch.cuda.empty_cache()
        return out

    def loglikelihood_rolling(self, requests):
        raise NotImplementedError

    def generate_until(self, requests):
        output = []
        num_tokens = 0
        num_nfe = 0
        processed_count = 0
        start_time = time.time()
        log_fh = None
        if self.save_dir is not None:
            os.makedirs(self.save_dir, exist_ok=True)
            rank = self.rank
            save_path = os.path.join(self.save_dir, f"rank_{rank}.jsonl")
            print(f"save_path: {save_path}")
            if os.path.exists(save_path):
                print(f"load from {save_path}")
                with open(save_path, "r", encoding="utf-8") as f:
                    output = [json.loads(line) for line in f]
                    processed_count = len(output)
                print(f"processed_count: {processed_count}")
        if self.stats_dir is not None:
            os.makedirs(self.stats_dir, exist_ok=True)
            stats_samples_path = os.path.join(self.stats_dir, f"rank_{self.rank}_samples.jsonl")
            log_fh = open(stats_samples_path, "a", encoding="utf-8")
        for i, req in enumerate(tqdm(requests, desc="Generating...")):
            sample_start_time = time.time()
            if i < processed_count:
                continue
            question = req.args[0]
            if self.is_instruct:
                tail = r" Please reason step by step, and put your final answer within \boxed{}."
                if self.task == "gsm8k":
                    m = [{"role": "user", "content": question + tail}]
                elif self.task == "humaneval" or self.task == "humaneval_plus":
                    start = "Write a solution to the following problem and make sure that it passes the tests:\n```python\n{{ prompt }}\n```\n "
                    question = start.replace("{{ prompt }}", question)
                    m = [{"role": "user", "content": question}]
                else:
                    m = [{"role": "user", "content": question}]
                user_input = self.tokenizer.apply_chat_template(m, add_generation_prompt=True, tokenize=False)
                input_ids = self.tokenizer(user_input)["input_ids"]
            else:
                user_input = question
                input_ids = self.tokenizer(user_input)["input_ids"]
            stop_tokens = list(req.args[1]["until"])
            if "<|eot_id|>" not in stop_tokens:
                stop_tokens.append("<|eot_id|>")
            input_ids = torch.tensor(input_ids).to(self.device).unsqueeze(0)
            # Determine eos_token_id for early stopping
            eos_token_id = self.tokenizer.eos_token_id if self.early_stop else None
            if self.multi_block:
                generated_answer, nfe = generate_multi_block(
                    self.model,
                    input_ids,
                    steps=self.steps,
                    max_new_tokens=self.gen_length,
                    block_size=self.block_length,
                    temperature=0,
                    remasking=self.remasking,
                    mask_id=self.mask_id,
                    threshold=self.threshold,
                    block_add_threshold=self.block_add_threshold,
                    decoded_token_threshold=self.decoded_token_threshold,
                    eos_token_id=eos_token_id,
                )
            elif self.use_cache:
                if self.dual_cache:
                    generated_answer, nfe = generate_with_dual_cache(
                        self.model,
                        input_ids,
                        steps=self.steps,
                        gen_length=self.gen_length,
                        block_length=self.block_length,
                        temperature=0,
                        remasking=self.remasking,
                        mask_id=self.mask_id,
                        threshold=self.threshold,
                        eos_token_id=eos_token_id,
                    )
                else:
                    generated_answer, nfe = generate_with_prefix_cache(
                        self.model,
                        input_ids,
                        steps=self.steps,
                        gen_length=self.gen_length,
                        block_length=self.block_length,
                        temperature=0,
                        remasking=self.remasking,
                        mask_id=self.mask_id,
                        threshold=self.threshold,
                        eos_token_id=eos_token_id,
                    )
            else:
                generated_answer, nfe = generate(
                    self.model,
                    input_ids,
                    steps=self.steps,
                    gen_length=self.gen_length,
                    block_length=self.block_length,
                    temperature=0,
                    remasking=self.remasking,
                    mask_id=self.mask_id,
                    threshold=self.threshold,
                    eos_token_id=eos_token_id,
                )
            if self.is_instruct and "task_id" in req.doc and str(req.doc["task_id"]).lower().startswith("humaneval"):
                generated_answer = self.tokenizer.decode(generated_answer[0][input_ids.shape[1] :], skip_special_tokens=True)
                generated_answer_ids = self.tokenizer(generated_answer)["input_ids"]
            else:
                generated_answer = self.tokenizer.decode(generated_answer[0][input_ids.shape[1] :], skip_special_tokens=False)
                print(f"stop_tokens: {stop_tokens}")
                for stop_seq in stop_tokens:
                    if stop_seq in generated_answer:
                        generated_answer = generated_answer.split(stop_seq)[0]
                generated_answer_ids = torch.tensor(self.tokenizer(generated_answer)["input_ids"])
                generated_answer = self.tokenizer.decode(generated_answer_ids, skip_special_tokens=True)
            num_tokens += len(generated_answer_ids)
            num_nfe += nfe
            output.append(generated_answer)
            if self.save_dir is not None:
                with open(save_path, "a", encoding="utf-8") as f:
                    f.write(json.dumps(generated_answer, ensure_ascii=False) + "\n")
            sample_end_time = time.time()
            if log_fh is not None:
                record = {
                    "sample_index": int(i),
                    "prompt_tokens": int(input_ids.shape[1]),
                    "generated_tokens": int(len(generated_answer_ids)),
                    "steps": int(self.steps),
                    "nfe": int(nfe),
                    "latency_seconds": float(sample_end_time - sample_start_time),
                    "timestamp": float(sample_end_time),
                }
                log_fh.write(json.dumps(record, ensure_ascii=False) + "\n")
            print("=" * 20)
            print("question: ", question)
            print("answer: ", generated_answer)
            print("=" * 20, end="\n\n")
        if log_fh is not None:
            log_fh.close()
        total_time = time.time() - start_time
        # Compute blocks per sample for logging
        blocks_per_sample = int(self.gen_length // self.block_length) if getattr(self, "block_length", 0) > 0 else 0
        if self.stats_dir is not None:
            processed_samples = int(len(output))
            avg_iters_per_sample = (float(num_nfe) / float(processed_samples)) if processed_samples > 0 else 0.0
            avg_gen_len = (float(num_tokens) / float(processed_samples)) if processed_samples > 0 else 0.0
            avg_gen_len_blocks_persample = (float(self.gen_length) / float(self.block_length)) if getattr(self, "block_length", 0) > 0 else 0.0
            final_stats = {
                "processed_samples": processed_samples,
                "total_samples": int(len(requests)),
                "total_tokens": int(num_tokens),
                "total_nfe": int(num_nfe),
                "total_time": float(total_time),
                "tokens_per_second": (float(num_tokens) / float(total_time)) if total_time > 0 else 0.0,
                "nfe_per_token": (float(num_nfe) / float(num_tokens)) if num_tokens > 0 else 0.0,
                "tokens_per_forward": (float(num_tokens)) / (float(num_nfe)),
                "blocks_per_sample": blocks_per_sample,
                "gen_length": int(self.gen_length),
                "block_length": int(self.block_length),
                "avg_iters_per_sample": avg_iters_per_sample,
                "avg_gen_len": avg_gen_len,
                "avg_gen_len_blocks_persample": avg_gen_len_blocks_persample,
                "timestamp": time.time(),
            }
            stats_path = os.path.join(self.stats_dir, f"rank_{self.rank}_final_stats.json")
            with open(stats_path, "w", encoding="utf-8") as f:
                json.dump(final_stats, f, ensure_ascii=False, indent=2)
        if self.show_speed:
            print(f"Total time taken: {total_time} seconds")
            print(f"Total NFE is {num_nfe}")
        if self.accelerator is not None and getattr(self, "world_size", 1) > 1:
            local_stats = torch.tensor(
                [
                    float(len(output)),
                    float(num_tokens),
                    float(total_time),
                    float(num_nfe),
                ],
                dtype=torch.float64,
                device=self.device,
            )
            gathered_stats = self.accelerator.gather(local_stats)
            if self.accelerator.is_local_main_process:
                gathered_stats = gathered_stats.view(self.world_size, -1)
                total_samples_all = int(gathered_stats[:, 0].sum().item())
                total_tokens_all = int(gathered_stats[:, 1].sum().item())
                sum_time_all = gathered_stats[:, 2].sum().item()
                total_nfe_all = int(gathered_stats[:, 3].sum().item())
                overall_tps = (total_tokens_all / sum_time_all) if sum_time_all > 0 else 0.0
                avg_iters_per_sample_all = (float(total_nfe_all) / float(total_samples_all)) if total_samples_all > 0 else 0.0
                avg_gen_len_all = (float(total_tokens_all) / float(total_samples_all)) if total_samples_all > 0 else 0.0
                if self.stats_dir is not None:
                    aggregated_stats = {
                        "total_processed_samples": total_samples_all,
                        "total_generated_tokens": total_tokens_all,
                        "total_wall_time": sum_time_all,
                        "overall_tokens_per_second": overall_tps,
                        "overall_nfe": total_nfe_all,
                        "overall_nfe_per_token": (float(total_nfe_all) / float(total_tokens_all)) if total_tokens_all > 0 else 0.0,
                        "overall_tokens_per_forward": (float(total_tokens_all)) / (float(total_nfe_all)),
                        "avg_iters_per_sample": avg_iters_per_sample_all,
                        "avg_gen_len": avg_gen_len_all,
                        "blocks_per_sample": blocks_per_sample,
                        "gen_length": int(self.gen_length),
                        "block_length": int(self.block_length),
                        "timestamp": time.time(),
                    }
                    all_ranks_stats_path = os.path.join(self.stats_dir, "all_ranks_final_stats.json")
                    with open(all_ranks_stats_path, "w", encoding="utf-8") as f:
                        json.dump(aggregated_stats, f, ensure_ascii=False, indent=2)
        return output


if __name__ == "__main__":
    set_seed(1234)
    cli_evaluate()