reference/code/TAD/eval/eval_dream.py

import accelerate
import torch
import random
import numpy as np
import types
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
import os
from transformers import AutoTokenizer, AutoConfig, AutoModel
import json
import time
from pathlib import Path
from accelerate import (
    Accelerator,
    InitProcessGroupKwargs,
)
from datetime import timedelta
from model.generation_utils_dream import DreamGenerationMixin


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


@register_model("dream_dist")
class DreamEvalHarness(LM):
    def __init__(
        self,
        model_path="Dream-org/Dream-v0-Instruct-7B",
        max_length=2048,
        batch_size=1,
        mc_num=128,
        diffusion_steps=128,
        max_new_tokens=128,
        temperature=0.0,
        top_p=None,
        top_k=None,
        alg="entropy",
        alg_temp=0.0,
        device="cuda",
        remasking="low_confidence",
        block_length=None,
        threshold=0.9,
        save_dir=None,
        stats_dir=None,
        show_speed=False,
        multi_block=False,
        block_add_threshold=0.5,
        decoded_token_threshold=0.5,
        early_stop=False,
        task="null",
        nll_type="mc",
        log_type="nll",
        classifier_free_guidance=0.0,
        sampling_eps=0.0,
        add_bos_token=False,
        escape_until=False,
        **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)

        self.model = AutoModel.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.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

        # Dream uses mask_token_id from tokenizer or config
        if hasattr(self.tokenizer, "mask_token_id") and self.tokenizer.mask_token_id is not None:
            self.mask_id = self.tokenizer.mask_token_id
        elif hasattr(config, "mask_token_id") and config.mask_token_id is not None:
            self.mask_id = config.mask_token_id
        else:
            # Default Dream mask token id
            self.mask_id = 151666

        self.mc_num = mc_num
        self.batch_size = int(batch_size)
        assert mc_num % self.batch_size == 0
        self.sampling_eps = float(sampling_eps)
        self.max_length = int(max_length)
        self.diffusion_steps = int(diffusion_steps)
        self.max_new_tokens = int(max_new_tokens)
        self.temperature = float(temperature)
        self.top_p = float(top_p) if top_p is not None else None
        self.top_k = int(top_k) if top_k is not None else None
        self.alg = str(alg)
        self.alg_temp = float(alg_temp) if alg_temp is not None else None
        self.remasking = remasking
        self.classifier_free_guidance = float(classifier_free_guidance)
        self.add_bos_token = add_bos_token if isinstance(add_bos_token, bool) else str(add_bos_token).lower() == "true"
        self.escape_until = escape_until if isinstance(escape_until, bool) else str(escape_until).lower() == "true"

        self.is_instruct = True if ("instruct" in model_path.lower()) else False
        self.save_dir = save_dir
        self.stats_dir = stats_dir if stats_dir is not None else save_dir
        self.show_speed = show_speed if isinstance(show_speed, bool) else str(show_speed).lower() == "true"
        self.multi_block = multi_block if isinstance(multi_block, bool) else str(multi_block).lower() == "true"
        self.block_add_threshold = float(block_add_threshold)
        self.decoded_token_threshold = float(decoded_token_threshold)
        self.early_stop = early_stop if isinstance(early_stop, bool) else str(early_stop).lower() == "true"
        self.task = task
        self.nll_type = nll_type
        self.log_type = log_type
        self.block_length = int(block_length) if block_length is not None else None
        self.threshold = float(threshold)

        # Monkey-patch CDLM's diffusion_generate and _sample onto the model
        # This replaces the HuggingFace official methods with CDLM's versions
        # that support confidence_threshold dynamic decoding and block-wise generation
        self.model.diffusion_generate = types.MethodType(
            DreamGenerationMixin.diffusion_generate, self.model
        )
        self.model._sample = types.MethodType(
            DreamGenerationMixin._sample, self.model
        )
        self.model.generate_multi_block = types.MethodType(
            DreamGenerationMixin.generate_multi_block, self.model
        )
        self.model._sample_multi_block = types.MethodType(
            DreamGenerationMixin._sample_multi_block, self.model
        )

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

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

    @property
    def tokenizer_name(self) -> str:
        return self.tokenizer.name_or_path.replace("/", "__")

    def apply_chat_template(
        self, chat_history, add_generation_prompt: bool = True
    ) -> str:
        """
        Method to apply a chat template to a list of chat history between user and model.
        """
        chat_templated = self.tokenizer.apply_chat_template(
            chat_history,
            tokenize=False,
            add_generation_prompt=add_generation_prompt,
            continue_final_message=not add_generation_prompt,
        )
        return chat_templated

    # ── loglikelihood helpers ──────────────────────────────────────────────

    def _forward_process(self, batch, prompt_index):
        """
        Apply forward diffusion noise to the target portion of the batch.
        Returns noisy_batch and the noise level p_mask for each position.
        """
        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):
        """
        Get logits from the Dream model.
        Dream uses a shifted logits convention:
            logits = cat([logits[:, :1], logits[:, :-1]], dim=1)
        so that logits[i] predicts token[i] (not token[i+1] as in AR models).
        """
        if self.classifier_free_guidance > 0.0:
            assert len(prompt_index) == batch.shape[1]
            pi = prompt_index.unsqueeze(0).repeat(batch.shape[0], 1)
            un_batch = batch.clone()
            un_batch[pi] = self.mask_id
            batch = torch.cat([batch, un_batch])

        logits = self.model(batch).logits

        # Dream logits shift: logits = cat([logits[:,:1], logits[:,:-1]], dim=1)
        logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)

        if self.classifier_free_guidance > 0.0:
            logits, un_logits = torch.chunk(logits, 2, dim=0)
            logits = un_logits + (self.classifier_free_guidance + 1) * (logits - un_logits)

        return logits[:, :batch.shape[1]]

    @torch.no_grad()
    def get_loglikelihood(self, prefix, target):
        """
        Monte Carlo estimation of log-likelihood for the target given the prefix.
        """
        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):
        """
        Check if greedy decoding from the model matches the target.
        """
        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) <= self.max_length, (
            f"Max prompt length {max(prompt_len)} exceeds max_length {self.max_length}"
        )

        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

    # ── generation ────────────────────────────────────────────────────────

    def _generate_batch(self, prompts):
        """
        Call Dream model's diffusion_generate method (CDLM version) for a batch of prompts.
        Accepts a list of prompt strings, tokenizes internally, generates, and decodes.
        Returns (responses: List[str], nfe: int).
        """
        prompt_ids = self.tokenizer(prompts, return_tensors="pt", padding=True, padding_side="left").input_ids
        prompt_ids = prompt_ids.to(device=self.device)
        attn_mask = prompt_ids.ne(self.tokenizer.pad_token_id).to(device=self.device)

        generation_result = self.model.diffusion_generate(
            prompt_ids,
            attention_mask=attn_mask,
            max_new_tokens=self.max_new_tokens,
            output_history=False,
            return_dict_in_generate=True,
            steps=self.diffusion_steps,
            temperature=self.temperature,
            top_p=self.top_p,
            top_k=self.top_k,
            alg=self.alg,
            alg_temp=self.alg_temp,
            threshold=self.threshold,
            block_length=self.block_length,
            early_stop=self.early_stop,
        )

        nfe = int(getattr(generation_result, 'nfe', self.diffusion_steps))

        # decode only the generated part (after prompt)
        responses = [
            self.tokenizer.decode(g[len(p):].tolist()).split(self.tokenizer.eos_token)[0]
            for p, g in zip(prompt_ids, generation_result.sequences)
        ]

        return responses, nfe

    def _generate_batch_multi_block(self, prompts):
        """
        Call Dream model's generate_multi_block method for a batch of prompts.
        Uses pipelined parallel decoding with multi-block strategy.
        Accepts a list of prompt strings, tokenizes internally, generates, and decodes.
        Returns (responses: List[str], nfe: int).
        """
        prompt_ids = self.tokenizer(prompts, return_tensors="pt", padding=True, padding_side="left").input_ids
        prompt_ids = prompt_ids.to(device=self.device)
        attn_mask = prompt_ids.ne(self.tokenizer.pad_token_id).to(device=self.device)

        generation_result, nfe = self.model.generate_multi_block(
            prompt_ids,
            attention_mask=attn_mask,
            max_new_tokens=self.max_new_tokens,
            output_history=False,
            return_dict_in_generate=True,
            steps=self.diffusion_steps,
            temperature=self.temperature,
            top_p=self.top_p,
            top_k=self.top_k,
            alg=self.alg,
            alg_temp=self.alg_temp,
            threshold=self.threshold,
            block_size=self.block_length if self.block_length is not None else 32,
            block_add_threshold=self.block_add_threshold,
            decoded_token_threshold=self.decoded_token_threshold,
            early_stop=self.early_stop,
        )

        nfe = int(nfe)

        # decode only the generated part (after prompt)
        responses = [
            self.tokenizer.decode(g[len(p):].tolist()).split(self.tokenizer.eos_token)[0]
            for p, g in zip(prompt_ids, generation_result.sequences)
        ]

        return responses, nfe

    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

        # ── Checkpoint resume ─────────────────────────────────────────────
        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")

        # ── Main generation loop (batch-based, aligned with dParallel) ────
        pbar = tqdm(total=len(requests), desc="Generating...")

        for batch_idx in range(0, len(requests), self.batch_size):
            if batch_idx < processed_count:
                pbar.update(min(self.batch_size, len(requests) - batch_idx))
                continue

            sample_start_time = time.time()

            batch_requests = requests[batch_idx : batch_idx + self.batch_size]
            contexts, gen_args = zip(*[req.arguments for req in batch_requests])

            # Generate — contexts are already chat-templated by the framework
            if self.multi_block:
                responses, nfe = self._generate_batch_multi_block(contexts)
            else:
                responses, nfe = self._generate_batch(contexts)
            num_nfe += nfe

            sample_end_time = time.time()

            for i, r in enumerate(responses):
                # Count tokens before truncation
                generated_answer_ids = self.tokenizer.encode(r)
                num_tokens += len(generated_answer_ids)

                # Truncate at stop tokens
                for s in gen_args[i]['until']:
                    r = r.split(s)[0]
                responses[i] = r

            output.extend(responses)

            if self.rank == 0:
                print(f"Context:\n{contexts[0]}\nResponse:\n{responses[0]}\n")
                print("nfe step:", nfe)

            # Save checkpoint
            if self.save_dir is not None:
                with open(save_path, "a", encoding="utf-8") as f:
                    for r in responses:
                        f.write(json.dumps(r, ensure_ascii=False) + "\n")

            # Log per-sample stats
            if log_fh is not None:
                for i, r in enumerate(responses):
                    record = {
                        "sample_index": int(batch_idx + i),
                        "generated_tokens": int(len(self.tokenizer.encode(r))),
                        "steps": int(self.diffusion_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")

            pbar.update(len(batch_requests))

        pbar.close()

        if log_fh is not None:
            log_fh.close()

        total_time = time.time() - start_time

        # ── Per-rank final stats ──────────────────────────────────────────
        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_latency_per_sample = (
                (float(total_time) / float(processed_samples)) if processed_samples > 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)) if num_nfe > 0 else 0.0
                ),
                "avg_iters_per_sample": avg_iters_per_sample,
                "avg_gen_len": avg_gen_len,
                "avg_latency_per_sample": avg_latency_per_sample,
                "diffusion_steps": int(self.diffusion_steps),
                "max_new_tokens": int(self.max_new_tokens),
                "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 num_tokens > 0:
                print(f"Tokens per second (TPS): {num_tokens / total_time:.2f}")
                print(f"Tokens per forward (TPF): {num_tokens / num_nfe:.4f}")
                print(f"NFE per token: {num_nfe / num_tokens:.4f}")
            if len(output) > 0:
                print(f"Avg latency per sample: {total_time / len(output):.2f}s")

        # ── Multi-GPU aggregation ─────────────────────────────────────────
        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
                )
                overall_avg_latency = (
                    (sum_time_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))
                            if total_nfe_all > 0
                            else 0.0
                        ),
                        "avg_iters_per_sample": avg_iters_per_sample_all,
                        "avg_gen_len": avg_gen_len_all,
                        "overall_avg_latency_per_sample": overall_avg_latency,
                        "diffusion_steps": int(self.diffusion_steps),
                        "max_new_tokens": int(self.max_new_tokens),
                        "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()