generation_utils.py

# coding=utf-8
# Copyright 2024 The Dream team, HKUNLP Group and the HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import warnings
import copy
from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple, Union

import torch
import torch.distributions as dists
from torch.nn import functional as F
from transformers import __version__
from transformers.generation.configuration_utils import (
    GenerationConfig
)
from transformers.utils import (
    ModelOutput,
    is_torchdynamo_compiling,
    logging,
)
logger = logging.get_logger(__name__)
from tqdm import tqdm


def top_p_logits(logits, top_p=None):
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
    sorted_indices_to_remove = cumulative_probs > top_p
    # Shift the indices to the right to keep the first token above the threshold
    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
    sorted_indices_to_remove[..., 0] = 0

    mask = torch.zeros_like(logits, dtype=torch.bool, device=logits.device)
    mask = mask.scatter_(-1, sorted_indices, sorted_indices_to_remove)
    logits = logits.masked_fill(mask, torch.finfo(logits.dtype).min)
    return logits

def top_k_logits(logits, top_k=None):
    top_k = min(top_k, logits.size(-1))  # Safety check
    # Remove all tokens with a probability less than the last token of the top-k
    indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
    logits = logits.masked_fill(indices_to_remove, torch.finfo(logits.dtype).min)
    return logits


def sample_tokens(logits, temperature=0.0, top_p=None, top_k=None, margin_confidence=False, neg_entropy=False, 
                  repeat_penalty=1.0, max_position_penalty=1.0, past_x=None, mask_id=None,):
    if temperature > 0:
        logits = logits / temperature

    if top_p is not None and top_p < 1:
        logits = top_p_logits(logits, top_p)
    if top_k is not None:
        logits = top_k_logits(logits, top_k)

    if repeat_penalty != 1.0:
        select_mask = torch.logical_and((past_x != 0), (past_x != mask_id))
        generated_tokens = set(past_x[select_mask].tolist())
        for token in set(generated_tokens):
            logits[:, token][logits[:, token] < 0] *= repeat_penalty
            logits[:, token][logits[:, token] >= 0] /= repeat_penalty

    if max_position_penalty != 1.0:
        token_length = logits.shape[-2]
        if token_length > 100:
            penalty_map = [i / (token_length - 100) * (max_position_penalty - 1.0) + 1.0 
                           for i in range(token_length - 100)]
            penalty_map = torch.tensor(penalty_map).unsqueeze(-1).to(logits.device).to(logits.dtype)
            penalty_map = torch.cat([torch.ones_like(logits[:100, :1]), penalty_map], dim=0)
            penalty_map = penalty_map.repeat(1, logits.shape[-1])
        
            logits[logits < 0] *= penalty_map[logits < 0]
            logits[logits >= 0] /= penalty_map[logits >= 0]

    probs = torch.softmax(logits, dim=-1)
    if temperature > 0:
        try:
            x0 = dists.Categorical(probs=probs).sample()
            confidence = torch.gather(probs, -1, x0.unsqueeze(-1)).squeeze(-1)
        except:
            confidence, x0 = probs.max(dim=-1)
    else:
        confidence, x0 = probs.max(dim=-1)

    if margin_confidence:
        sorted_probs, _ = torch.sort(probs, dim=-1, descending=True)
        top1_probs = sorted_probs[:, 0]
        top2_probs = sorted_probs[:, 1]
        confidence = top1_probs - top2_probs

    if neg_entropy:
        epsilon = 1e-10
        log_probs = torch.log(probs + epsilon)
        confidence = torch.sum(probs * log_probs, dim=-1)

    return confidence, x0


@dataclass
class DreamModelOutput(ModelOutput):
    sequences: torch.LongTensor = None
    history: Optional[Tuple[torch.FloatTensor]] = None


class DreamGenerationConfig(GenerationConfig):
    def __init__(self, **kwargs):
        self.temperature: float = kwargs.pop("temperature", 0.0)
        self.top_p: Optional[float] = kwargs.pop("top_p", None)
        self.top_k: Optional[int] = kwargs.pop("top_k", None)
        self.max_length = kwargs.pop("max_length", 20)
        self.max_new_tokens = kwargs.pop("max_new_tokens", None)
        # diffusion specific params
        self.eps: float = kwargs.pop("eps", 1e-3)
        self.steps: int = kwargs.pop("steps", 512)
        self.alg: str = kwargs.pop("alg", 'origin')
        self.alg_temp: Optional[float] = kwargs.pop("alg_temp", None)

        # Parameters that define the output variables of `generate`
        self.num_return_sequences: int = kwargs.pop("num_return_sequences", 1)
        self.return_dict_in_generate: bool = kwargs.pop("return_dict_in_generate", False)
        self.output_history: bool = kwargs.pop("output_history", False)

        # Special tokens that can be used at generation time
        self.mask_token_id = kwargs.pop("mask_token_id", None)
        self.pad_token_id = kwargs.pop("pad_token_id", None)
        self.bos_token_id = kwargs.pop("bos_token_id", None)
        self.eos_token_id = kwargs.pop("eos_token_id", None)

        # Wild card
        self.generation_kwargs = kwargs.pop("generation_kwargs", {})

        # The remaining attributes do not parametrize `.generate()`, but are informative and/or used by the hub
        # interface.
        self._from_model_config = kwargs.pop("_from_model_config", False)
        self._commit_hash = kwargs.pop("_commit_hash", None)
        self.transformers_version = kwargs.pop("transformers_version", __version__)

        # Additional attributes without default values
        if not self._from_model_config:
            # we don't want to copy values from the model config if we're initializing a `GenerationConfig` from a
            # model's default configuration file
            for key, value in kwargs.items():
                try:
                    setattr(self, key, value)
                except AttributeError as err:
                    logger.error(f"Can't set {key} with value {value} for {self}")
                    raise err

        # Validate the values of the attributes
        self.validate(is_init=True)

    def validate(self, is_init=False):
        pass

class DreamGenerationMixin:
    @staticmethod
    def _expand_inputs_for_generation(
        expand_size: int = 1,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.LongTensor] = None
    ) -> Tuple[torch.LongTensor, Dict[str, Any]]:
        """Expands tensors from [batch_size, ...] to [batch_size * expand_size, ...]"""
        # Do not call torch.repeat_interleave if expand_size is 1 because it clones
        # the input tensor and thus requires more memory although no change is applied
        if expand_size == 1:
            return input_ids, attention_mask
        if input_ids is not None:
            input_ids = input_ids.repeat_interleave(expand_size, dim=0)
        if attention_mask is not None:
            attention_mask = attention_mask.repeat_interleave(expand_size, dim=0)
        return input_ids, attention_mask

    def _validate_generated_length(self, generation_config, input_ids_length, has_default_max_length):
        """Performs validation related to the resulting generated length"""

        # Can't throw warnings/exceptions during compilation
        if is_torchdynamo_compiling():
            return

        # 1. Max length warnings related to poor parameterization
        if has_default_max_length and generation_config.max_new_tokens is None and generation_config.max_length == 20:
            # 20 is the default max_length of the generation config
            warnings.warn(
                f"Using the model-agnostic default `max_length` (={generation_config.max_length}) to control the "
                "generation length. We recommend setting `max_new_tokens` to control the maximum length of the "
                "generation.",
                UserWarning,
            )
        if input_ids_length >= generation_config.max_length:
            input_ids_string = "input_ids"
            raise ValueError(
                f"Input length of {input_ids_string} is {input_ids_length}, but `max_length` is set to"
                f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
                " increasing `max_length` or, better yet, setting `max_new_tokens`."
            )

    def _prepare_generated_length(
        self,
        generation_config,
        has_default_max_length,
        input_ids_length,
    ):
        """Prepared max and min length in generation configs to avoid clashes between similar attributes"""

        if generation_config.max_new_tokens is not None:
            if not has_default_max_length and generation_config.max_length is not None:
                logger.warning(
                    f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
                    f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
                    "Please refer to the documentation for more information. "
                    "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)"
                )
            generation_config.max_length = generation_config.max_new_tokens + input_ids_length

        elif has_default_max_length:
            if generation_config.max_length == DreamGenerationConfig().max_length:
                generation_config.max_length = generation_config.max_length + input_ids_length
                max_position_embeddings = getattr(self.config, "max_position_embeddings", None)
                if max_position_embeddings is not None:
                    generation_config.max_length = min(generation_config.max_length, max_position_embeddings)

        return generation_config

    def _prepare_generation_config(
        self, generation_config: Optional[DreamGenerationConfig], **kwargs: Dict
    ) -> DreamGenerationConfig:
        """
        Prepares the base generation config, then applies any generation configuration options from kwargs. This
        function handles retrocompatibility with respect to configuration files.
        """
        # priority: `generation_config` argument > `model.generation_config` (the default generation config)
        using_model_generation_config = False
        if generation_config is None:
            generation_config = DreamGenerationConfig.from_model_config(self.config)
            using_model_generation_config = True

        # `torch.compile` can't compile `copy.deepcopy`, arguments in `kwargs` that are part of `generation_config`
        # will mutate the object with `.update`. As such, passing these arguments through `kwargs` is disabled -- an
        # exception will be raised in `_validate_model_kwargs`
        if not is_torchdynamo_compiling():
            generation_config = copy.deepcopy(generation_config)
            _kwargs = generation_config.update(**kwargs)
            # If `generation_config` is provided, let's fallback ALL special tokens to the default values for the model
            if not using_model_generation_config:
                if generation_config.bos_token_id is None:
                    generation_config.bos_token_id = self.generation_config.bos_token_id
                if generation_config.eos_token_id is None:
                    generation_config.eos_token_id = self.generation_config.eos_token_id
                if generation_config.pad_token_id is None:
                    generation_config.pad_token_id = self.generation_config.pad_token_id
                if generation_config.mask_token_id is None:
                    generation_config.mask_token_id = self.generation_config.mask_token_id

        return generation_config

    def _prepare_special_tokens(
        self,
        generation_config: DreamGenerationConfig,
        device: Optional[Union[torch.device, str]] = None,
    ):
        """
        Prepares the special tokens for generation, overwriting the generation config with their processed versions
        converted to tensor.

        Note that `generation_config` is changed in place and stops being serializable after this method is called.
        That is no problem if called within `generate` (`generation_config` is a local copy that doesn't leave the
        function). However, if called outside `generate`, consider creating a copy of `generation_config` first.
        """

        # Convert special tokens to tensors
        def _tensor_or_none(token, device=None):
            if token is None:
                return token

            device = device if device is not None else self.device
            if isinstance(token, torch.Tensor):
                return token.to(device)
            return torch.tensor(token, device=device, dtype=torch.long)

        bos_token_tensor = _tensor_or_none(generation_config.bos_token_id, device=device)
        eos_token_tensor = _tensor_or_none(generation_config.eos_token_id, device=device)
        pad_token_tensor = _tensor_or_none(generation_config.pad_token_id, device=device)
        mask_token_tensor = _tensor_or_none(generation_config.mask_token_id, device=device)

        # We can have more than one eos token. Always treat it as a 1D tensor (when it exists).
        if eos_token_tensor is not None and eos_token_tensor.ndim == 0:
            eos_token_tensor = eos_token_tensor.unsqueeze(0)

        # Set pad token if unset (and there are conditions to do so)
        if pad_token_tensor is None and eos_token_tensor is not None:
            pad_token_tensor = eos_token_tensor[0]
            logger.warning(f"Setting `pad_token_id` to `eos_token_id`:{pad_token_tensor} for open-end generation.")

        # Update generation config with the updated special tokens tensors
        # NOTE: this must be written into a different attribute name than the one holding the original special tokens
        # (in their non-tensor form), in order to enable end-to-end compilation. See
        # https://pytorch.org/docs/stable/torch.compiler_cudagraph_trees.html#limitations
        generation_config._bos_token_tensor = bos_token_tensor
        generation_config._eos_token_tensor = eos_token_tensor
        generation_config._pad_token_tensor = pad_token_tensor
        generation_config._mask_token_tensor = mask_token_tensor

    @torch.no_grad()
    def diffusion_generate(
        self,
        inputs: Optional[torch.Tensor] = None,
        generation_config: Optional[DreamGenerationConfig] = None,
        inputs_embeds=None,
        prefix_lm=False,
        alg=None,
        block_size=-1,
        cfg=0.0,
        add_boa_token=False,
        **kwargs,
    ) -> Union[DreamModelOutput, torch.LongTensor]:
        # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call
        generation_config = self._prepare_generation_config(generation_config, **kwargs)
        generation_tokens_hook_func = kwargs.pop("generation_tokens_hook_func", lambda step, x, logits: x)
        generation_logits_hook_func = kwargs.pop("generation_logits_hook_func", lambda step, x, logits: logits)
        # breakpoint()
        # 2. Define model inputs
        if inputs is not None:
            input_ids = inputs
            device = input_ids.device
            input_ids_length = input_ids.shape[-1]
        else:
            input_ids = None
            device = inputs_embeds.device
            input_ids_length = inputs_embeds.shape[1]
        attention_mask = kwargs.pop("attention_mask", None)
        self._prepare_special_tokens(generation_config, device=device)

        # 3. Prepare `max_length`.
        
        has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
        generation_config = self._prepare_generated_length(
            generation_config=generation_config,
            has_default_max_length=has_default_max_length,
            input_ids_length=input_ids_length,
        )

        self._validate_generated_length(generation_config, input_ids_length, has_default_max_length)
        # import pdb;pdb.set_trace()
        # 4. Check input_ids
        #if not is_torchdynamo_compiling() and self.device.type != input_ids.device.type:
        if not is_torchdynamo_compiling() and self.device.type != device.type:
            warnings.warn(
                "You are calling .generate() with the `input_ids` being on a device type different"
                f" than your model's device. `input_ids` is on {input_ids.device.type}, whereas the model"
                f" is on {self.device.type}. You may experience unexpected behaviors or slower generation."
                " Please make sure that you have put `input_ids` to the"
                f" correct device by calling for example input_ids = input_ids.to('{self.device.type}') before"
                " running `.generate()`.",
                UserWarning,
            )
        # breakpoint()
        if (
            hasattr(generation_config, "pad_token_id") and
            input_ids is not None and
            torch.any(input_ids == generation_config.pad_token_id) and 
            attention_mask is None
        ):
            warnings.warn(
                "Padding was detected but no attention mask is passed here. For correct "
                "generation results, please set `attention_mask` when batch-padding inputs.",
                UserWarning,
            )
        assert generation_config.num_return_sequences == 1, \
            "Currently, we only support num_return_sequences = 1 for diffusion generation."
        input_ids, attention_mask = self._expand_inputs_for_generation(
            expand_size=generation_config.num_return_sequences,
            input_ids=input_ids,
            attention_mask=attention_mask 
        )
        result = self._sample(
            input_ids,
            attention_mask=attention_mask,
            generation_config=generation_config,
            generation_tokens_hook_func=generation_tokens_hook_func,
            generation_logits_hook_func=generation_logits_hook_func,
            inputs_embeds=inputs_embeds,
            device=device,
            prefix_lm=prefix_lm,
            alg=alg,
            block_size=block_size,
            cfg=cfg,
            add_boa_token=add_boa_token,
            **kwargs,
        )
        return result

    def _sample(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.LongTensor],
        generation_config: DreamGenerationConfig,
        generation_tokens_hook_func,
        generation_logits_hook_func,
        inputs_embeds=None,
        prefix_lm=False,
        device=None,
        step_ratio=None,
        penalty=1.2,
        alg=None,
        block_size=None,
        add_boa_token=False,
        max_position_penalty=1.0,
        repeat_penalty=1.0,
        cfg=0.0,
        **kwargs,
    ) -> Union[DreamModelOutput, torch.LongTensor]:
        output_history = True
        return_dict_in_generate = generation_config.return_dict_in_generate  
        max_length = generation_config.max_length                     
        mask_token_id = generation_config.mask_token_id               
        max_new_tokens = generation_config.max_new_tokens             
        steps = min(generation_config.steps, max_new_tokens) 
        eps = generation_config.eps                                    
        alg = generation_config.alg if alg is None else alg                       
        print("denoise algorithm: " + alg)
        alg_temp = generation_config.alg_temp                          
        temperature = generation_config.temperature                    
        top_p = generation_config.top_p                                
        top_k = generation_config.top_k              

        histories = [] if (return_dict_in_generate and output_history) else None
        all_logit = []
        generated_tokens = []
        block_size = max_new_tokens if block_size < 0 else block_size

        if input_ids is None:
            assert device is not None
            assert inputs_embeds is not None
            bsz, seq_len = inputs_embeds.shape[:2]                  
            max_length = seq_len + max_new_tokens                   
            input_ids = torch.full((bsz, seq_len), 0, dtype=torch.long).to(device)

        tok_idx = None
        past_key_values = None

        x = F.pad(input_ids, (0, max_length - input_ids.shape[1]), value=mask_token_id) 

        timesteps = torch.linspace(1, eps, steps + 1, device=x.device)
        x = generation_tokens_hook_func(None, x, None)

        if step_ratio is not None:
            steps = int(max_new_tokens * step_ratio)

        if add_boa_token:
            bos_index = int((x.shape[1] - (x == mask_token_id).sum()) + (x == mask_token_id).sum() * 0.2)
            x[:, bos_index] = 151684 # <|begin_of_audio|>

        input_x = x.clone()
        total_steps = steps
        block_num = (x == mask_token_id).sum() // block_size
        if block_num * block_size < (x == mask_token_id).sum(): block_num += 1
        input_length = input_ids.shape[-1]

        task = None
        if "task" in kwargs: task = kwargs['task']
        if cfg > 0:
            import random
            empty_prompt = ""
            if task == "S2I":
                empty_prompt = "<|im_start|>system\nPlease generate an image based on the input audio.<|im_end|>\n"
                empty_prompt += "<|im_start|>user\n<|im_end|>\n<|im_start|>assistant\n"
                un_x = kwargs['tokenizer'].encode(empty_prompt)
            elif task == "T2I":
                empty_prompt = "<|im_start|>user\nGenerate an image based on the provided text description.\n"
                empty_prompt += "<|im_end|>\n<|im_start|>assistant\n"
                first_audio_token = kwargs['tokenizer'].encode("<|begin_of_audio|>")[0]
                un_x_text = random.sample([_ for _ in range(first_audio_token)], 
                                          input_ids.shape[1] - len(kwargs['tokenizer'].encode(empty_prompt)))
                un_x = kwargs['tokenizer'].encode("<|im_start|>user\nGenerate an image based on the provided \
                                                  text description.\n") 
                un_x = un_x + un_x_text + kwargs['tokenizer'].encode("<|im_end|>\n<|im_start|>assistant\n")

        for block_idx in range(block_num):
            block_mask = torch.zeros([x.shape[-1]]).to(torch.bool).to(x.device)
            block_mask[input_length + block_idx * block_size: input_length + (block_idx + 1) * block_size] = True
            steps = int(block_mask.sum() / (x.shape[-1] - input_length) * total_steps)
            timesteps = torch.linspace(1, eps, steps + 1, device=x.device)
            for i in tqdm(range(steps)):
                mask_index = (x == mask_token_id)
                if mask_index.sum() == 0: break
                inputs_embeds_curr = self.model.embed_tokens(x)

                if inputs_embeds is not None:
                    inputs_embeds_curr[:, :inputs_embeds.shape[1]] = inputs_embeds

                if cfg > 0:
                    input_un_x = torch.tensor(un_x).unsqueeze(0).to(x.dtype).to(x.device)
                    input_un_x = torch.cat([input_un_x, x[:, input_ids.shape[1]:]], dim=1)
                    un_inpus_embeds = self.model.embed_tokens(input_un_x)

                    attention_mask_cond = torch.ones([1, inputs_embeds_curr.shape[1], inputs_embeds_curr.shape[1]])
                    attention_mask_cond = attention_mask_cond.to(torch.bool).to(inputs_embeds_curr.device)
                    attention_mask_uncond = torch.zeros([1, inputs_embeds_curr.shape[1], inputs_embeds_curr.shape[1]])
                    attention_mask_uncond[:, :un_inpus_embeds.shape[1], :un_inpus_embeds.shape[1]] = 1
                    attention_mask_uncond = attention_mask_uncond.to(torch.bool).to(inputs_embeds.device)
                    attention_mask = torch.cat([attention_mask_cond, attention_mask_uncond])
                    attention_mask = attention_mask.unsqueeze(1)

                    if inputs_embeds_curr.shape[1] != un_inpus_embeds.shape[1]:
                        un_inpus_embeds = torch.cat([un_inpus_embeds, 
                            torch.zeros_like(inputs_embeds_curr[:, :inputs_embeds_curr.shape[1] - 
                                                                un_inpus_embeds.shape[1], :])], dim=1)
                    input_inputs_embeds_curr = torch.cat([inputs_embeds_curr, un_inpus_embeds])

                    model_logits = self.forward_dream(None, attention_mask, tok_idx, 
                                                      inputs_embeds=input_inputs_embeds_curr).logits 
                    logits = model_logits[:1]; un_logits = model_logits[1:]
                    logits = un_logits + (cfg + 1) * (logits - un_logits)
                    logits = torch.cat([logits[:, :1], logits[:, :-1]], dim=1)

                else:
                    logits = self.forward_dream(None, attention_mask, tok_idx, 
                                                inputs_embeds=inputs_embeds_curr).logits
                    logits = torch.cat([logits[:, :1], logits[:, :-1]], dim=1)

                logits = generation_logits_hook_func(i, x, logits)

                mask_logits = logits[mask_index]
                if i == 0:
                    input_index = torch.where(mask_index[0]==True)[0][0]

                t = timesteps[i]
                s = timesteps[i + 1]

                if alg == 'origin':
                    p_transfer = 1 - s / t if i < steps - 1 else 1  
                    x0 = torch.zeros_like(x[mask_index], device=self.device, dtype=torch.long) + mask_token_id
                    transfer_index_t_s = torch.rand(*x0.shape, device=self.device) < p_transfer
                    _, x0[transfer_index_t_s] = sample_tokens(
                        mask_logits[transfer_index_t_s],
                        temperature=temperature,
                        top_p=top_p,
                        top_k=top_k,
                        max_position_penalty=max_position_penalty,
                    )
                    x[mask_index] = x0.clone()

                else:
                    if alg == 'maskgit_plus':
                        confidence, x0 = sample_tokens(mask_logits, temperature=temperature, top_p=top_p, top_k=top_k, 
                                                       max_position_penalty=max_position_penalty)
                    elif alg == 'topk_margin':
                        confidence, x0 = sample_tokens(
                            mask_logits,
                            temperature=temperature,
                            top_p=top_p,
                            top_k=top_k,
                            margin_confidence=True,
                            max_position_penalty=max_position_penalty,
                        )
                    elif alg == 'entropy':
                        confidence, x0 = sample_tokens(
                            mask_logits,
                            temperature,
                            top_p=top_p,
                            top_k=top_k,
                            neg_entropy=True,
                            max_position_penalty=max_position_penalty,
                        )
                    elif alg == "entropy-penalty": 
                        confidence, x0 = sample_tokens(
                            mask_logits,
                            temperature,
                            top_p=top_p,
                            top_k=top_k,
                            neg_entropy=True,
                            repeat_penalty=repeat_penalty if len(histories) != 0 else 1.0,
                            past_x=histories[-1] if len(histories) != 0 else [],
                            mask_id=mask_token_id,
                            max_position_penalty=max_position_penalty,
                        )
                    else:
                        raise RuntimeError(f"Unknown alg: {alg}")

                    block_mask_1 = block_mask[mask_index[0]]
                    confidence = confidence + torch.where(block_mask_1, 0, -torch.inf).to(confidence.device)

                    num_mask_token = mask_index.sum()
                    num_mask_token = (x[:, block_mask] == mask_token_id).sum()
                    number_transfer_tokens = int(num_mask_token * (1 - s / t)) if i < steps - 1 else num_mask_token
                    number_transfer_tokens = max(number_transfer_tokens, 1)

                    if number_transfer_tokens > 0:
                        if alg_temp is None or alg_temp == 0:
                            _, transfer_index = torch.topk(confidence, number_transfer_tokens)
                        else:
                            confidence = confidence / alg_temp
                            confidence = F.softmax(confidence, dim=-1)
                            transfer_index = torch.multinomial(confidence, num_samples=number_transfer_tokens)

                        x0_ = torch.zeros_like(x0, device=self.device, dtype=torch.long) + mask_token_id
                        x0_[transfer_index] = x0[transfer_index].clone()
                        x[mask_index] = x0_

                        logit,indic = torch.max(torch.softmax(logits.clone(),dim=-1),-1)
                        logit = logit[0][x[0]!=0]
                        indic = indic[0][x[0]!=0]
                        temp_X = x[0][x[0]!=0]

                x = generation_tokens_hook_func(i, x, logits)

                if histories is not None:
                    histories.append(x.clone())
                    all_logit.append(torch.max(logits.clone(),-1)[-1])

        return (x, histories)