Update generation_utils.py

generation_utils.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2024 The Dream team, HKUNLP Group and the HuggingFace Inc. team. All rights reserved.
+# Copyright 2024 The Dream team, HKUNLP Group and the HuggingFace Inc. team.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -22,14 +22,8 @@ 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,
-)
+from transformers.generation.configuration_utils import GenerationConfig
+from transformers.utils import ModelOutput, is_torchdynamo_compiling, logging
 
 logger = logging.get_logger(__name__)
 
@@ -47,6 +41,7 @@ def top_p_logits(logits, top_p=None):
     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
@@ -56,7 +51,6 @@ def top_k_logits(logits, top_k=None):
 
 
 def sample_tokens(logits, temperature=0.0, top_p=None, top_k=None, margin_confidence=False, neg_entropy=False):
-
     if temperature > 0:
         logits = logits / temperature
     if top_p is not None and top_p < 1:
@@ -69,24 +63,22 @@ def sample_tokens(logits, temperature=0.0, top_p=None, top_k=None, margin_confid
         try:
             x0 = dists.Categorical(probs=probs).sample()
             confidence = torch.gather(probs, -1, x0.unsqueeze(-1)).squeeze(-1)
-        except:
+        except Exception:
             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)
-        # Extract top1 and top2 probabilities
-        top1_probs = sorted_probs[:, 0] 
-        top2_probs = sorted_probs[:, 1] 
-        # Calculate confidence as top1 - top2
-        confidence = top1_probs - top2_probs 
-    
+        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
 
 
@@ -109,6 +101,11 @@ class DreamGenerationConfig(GenerationConfig):
         self.alg: str = kwargs.pop("alg", 'origin')
         self.alg_temp: Optional[float] = kwargs.pop("alg_temp", None)
 
+        # === RCR 相关参数（新增；默认不影响原逻辑） ===
+        self.rcr: bool = kwargs.pop("rcr", False)
+        # 仅在 rcr=True 时用于选择置信度算法；rcr=False 不读取它
+        self.conf_alg: str = kwargs.pop("conf_alg", 'maskgit_plus')
+
         # 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)
@@ -123,16 +120,12 @@ class DreamGenerationConfig(GenerationConfig):
         # 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.
+        # 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)
@@ -140,12 +133,12 @@ class DreamGenerationConfig(GenerationConfig):
                     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(
@@ -153,9 +146,6 @@ class DreamGenerationMixin:
         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:
@@ -164,16 +154,63 @@ class DreamGenerationMixin:
             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"""
+    # === 新增：RCR 逻辑，仅在 rcr=True 时被调用；不改动非 RCR 分支 ===
+    def _apply_rcr_logic(self, x, x0, confidence, mask_index, overtime_confidence,
+                         mask_token_id, step, total_steps, s, t):
+        """
+        在 Dream 的“maskgit”采样骨架上，执行 Running Confidence Remasking：
+        - 本步采用 Dream 原调度：global_k = num_mask_token * (1 - s/t)
+        - 先以当前置信度将 top-k token 从 [MASK] 转为预测 token，并累计它们的置信度
+        - 再施加“目标累计”约束：截至本步应累计生成 target_cum = num_mask_token * (1 - s/t)
+          若当前累计 > 目标，则把最低置信度的那些 token 反遮盖回 [MASK]
+        说明：只影响 rcr=True 的路径；rcr=False 时完全不调用本函数。
+        """
+        device = x.device
+        B = x.shape[0]
+
+        # 与 Dream 一致的 num_mask_token（按 batch 平均）
+        num_mask_token = mask_index.sum() / mask_index.shape[0]
+        # 本步的转移数量（按 Dream 调度）
+        number_transfer_tokens = int(num_mask_token * (1 - s / t)) if step < total_steps - 1 else int(num_mask_token)
+
+        # 构造全长置信度和候选值（非 mask 位置分别设为 -inf / mask_token_id）
+        full_conf = torch.full_like(x, -torch.inf, device=device, dtype=confidence.dtype)
+        x_temp = torch.zeros_like(x, device=device, dtype=torch.long) + mask_token_id
+        full_conf[mask_index] = confidence
+        x_temp[mask_index] = x0.clone()
+
+        for j in range(B):
+            # 逐样本 clamp，避免 batch 均值带来越界
+            masked_j = int(mask_index[j].sum().item())
+            k_j = min(number_transfer_tokens, masked_j)
+
+            # 先按置信度选出本步 top-k_j
+            if k_j > 0:
+                _, select_idx = torch.topk(full_conf[j], k=k_j, largest=True)
+                x[j, select_idx] = x_temp[j, select_idx]
+                overtime_confidence[j, select_idx] = full_conf[j, select_idx].clone().float()
+
+            # 目标累计约束：截至本步应累计的生成数
+            if step < total_steps - 1:
+                target_cum = int(num_mask_token * (1 - s / t))
+                gen_mask = overtime_confidence[j] > 0
+                current_gen = int(gen_mask.sum().item())
+                # 若超额，则按最低置信度回遮
+                to_remask = max(0, current_gen - target_cum)
+                if to_remask > 0:
+                    gen_indices = torch.where(gen_mask)[0]
+                    if gen_indices.numel() > 0:
+                        gen_conf = overtime_confidence[j, gen_indices]
+                        to_remask = min(to_remask, int(gen_indices.numel()))
+                        _, local_low = torch.topk(gen_conf, k=to_remask, largest=False)
+                        low_global = gen_indices[local_low]
+                        x[j, low_global] = mask_token_id
+                        overtime_confidence[j, low_global] = 0.0
 
-        # Can't throw warnings/exceptions during compilation
+    def _validate_generated_length(self, generation_config, input_ids_length, has_default_max_length):
         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 "
@@ -188,14 +225,7 @@ class DreamGenerationMixin:
                 " 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"""
-
+    def _prepare_generated_length(self, generation_config, has_default_max_length, input_ids_length):
         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(
@@ -212,29 +242,17 @@ class DreamGenerationMixin:
                 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)
+    def _prepare_generation_config(self, generation_config: Optional[DreamGenerationConfig], **kwargs: Dict) -> DreamGenerationConfig:
         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
@@ -247,25 +265,10 @@ class DreamGenerationMixin:
 
         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 _prepare_special_tokens(self, generation_config: DreamGenerationConfig, device: Optional[Union[torch.device, str]] = None):
         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)
@@ -276,19 +279,13 @@ class DreamGenerationMixin:
         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
@@ -301,19 +298,16 @@ class DreamGenerationMixin:
         generation_config: Optional[DreamGenerationConfig] = None,
         **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)
 
-        # 2. Define model inputs
         assert inputs is not None
         input_ids = inputs
         device = input_ids.device
         attention_mask = kwargs.pop("attention_mask", None)
         self._prepare_special_tokens(generation_config, device=device)
 
-        # 3. Prepare `max_length`.
         input_ids_length = input_ids.shape[-1]
         has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
         generation_config = self._prepare_generated_length(
@@ -323,8 +317,7 @@ class DreamGenerationMixin:
         )
 
         self._validate_generated_length(generation_config, input_ids_length, has_default_max_length)
-        
-        # 4. Check input_ids
+
         if not is_torchdynamo_compiling() and self.device.type != input_ids.device.type:
             warnings.warn(
                 "You are calling .generate() with the `input_ids` being on a device type different"
@@ -336,9 +329,9 @@ class DreamGenerationMixin:
                 UserWarning,
             )
         if (
-            hasattr(generation_config, "pad_token_id") and
-            torch.any(input_ids == generation_config.pad_token_id) and 
-            attention_mask is None
+            hasattr(generation_config, "pad_token_id")
+            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 "
@@ -349,7 +342,7 @@ class DreamGenerationMixin:
         input_ids, attention_mask = self._expand_inputs_for_generation(
             expand_size=generation_config.num_return_sequences,
             input_ids=input_ids,
-            attention_mask=attention_mask 
+            attention_mask=attention_mask,
         )
 
         result = self._sample(
@@ -357,7 +350,7 @@ class DreamGenerationMixin:
             attention_mask=attention_mask,
             generation_config=generation_config,
             generation_tokens_hook_func=generation_tokens_hook_func,
-            generation_logits_hook_func=generation_logits_hook_func
+            generation_logits_hook_func=generation_logits_hook_func,
         )
         return result
 
@@ -369,7 +362,7 @@ class DreamGenerationMixin:
         generation_tokens_hook_func,
         generation_logits_hook_func
     ) -> Union[DreamModelOutput, torch.LongTensor]:
-        # init values
+        # === 原变量 ===
         output_history = generation_config.output_history
         return_dict_in_generate = generation_config.return_dict_in_generate
         max_length = generation_config.max_length
@@ -382,18 +375,19 @@ class DreamGenerationMixin:
         top_p = generation_config.top_p
         top_k = generation_config.top_k
 
+        # === 新增：RCR 控制变量（不会影响 rcr=False 的路径） ===
+        rcr = generation_config.rcr
+        conf_alg = generation_config.conf_alg
+
         histories = [] if (return_dict_in_generate and output_history) else None
 
         # pad input_ids to max_length
         x = F.pad(input_ids, (0, max_length - input_ids.shape[1]), value=mask_token_id)
 
         if attention_mask is not None and torch.any(attention_mask == 0.0):
-            # we do not mask the [MASK] tokens so value = 1.0
             attention_mask = F.pad(attention_mask, (0, max_length - attention_mask.shape[1]), value=1.0)
             tok_idx = attention_mask.long().cumsum(-1) - 1
             tok_idx.masked_fill_(attention_mask == 0, 1)
-            # attention_mask is of shape [B, N]
-            # broadcast to [B, 1, N, N]
             attention_mask = torch.logical_and(
                 attention_mask.unsqueeze(1).unsqueeze(-2),
                 attention_mask.unsqueeze(1).unsqueeze(-1),
@@ -404,12 +398,15 @@ class DreamGenerationMixin:
 
         timesteps = torch.linspace(1, eps, steps + 1, device=x.device)
 
+        # === 仅在 rcr=True 时分配 Overtime Confidence（不影响 baseline） ===
+        overtime_confidence = torch.zeros_like(x, dtype=torch.float32) if rcr else None
+
         # this allows user-defined token control of the intermediate steps
         x = generation_tokens_hook_func(None, x, None)
         for i in range(steps):
             mask_index = (x == mask_token_id)
             logits = self(x, attention_mask, tok_idx).logits
-            logits = torch.cat([logits[:,:1], logits[:, :-1]], dim=1)
+            logits = torch.cat([logits[:, :1], logits[:, :-1]], dim=1)
 
             # this allows user-defined logits control of the intermediate steps
             logits = generation_logits_hook_func(i, x, logits)
@@ -417,48 +414,82 @@ class DreamGenerationMixin:
             mask_logits = logits[mask_index]
             t = timesteps[i]
             s = timesteps[i + 1]
-        
+
             if alg == 'origin':
+                # === 原版 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)
+                _, x0[transfer_index_t_s] = sample_tokens(
+                    mask_logits[transfer_index_t_s], temperature=temperature, top_p=top_p, top_k=top_k
+                )
                 x[mask_index] = x0.clone()
             else:
-                if alg == 'maskgit_plus':
+                # === 非 origin 分支 ===
+                # rcr=False：保持原有使用 alg 的置信度算法
+                # rcr=True ：使用 conf_alg 指定的置信度算法（不改变 rcr=False 的行为）
+                if (not rcr and alg == 'maskgit_plus') or (rcr and conf_alg == 'maskgit_plus'):
                     confidence, x0 = sample_tokens(mask_logits, temperature=temperature, top_p=top_p, top_k=top_k)
-                elif alg == 'topk_margin':
-                    confidence, x0 = sample_tokens(mask_logits, temperature=temperature, top_p=top_p, top_k=top_k, margin_confidence=True)
-                elif alg == 'entropy':
-                    confidence, x0 = sample_tokens(mask_logits, temperature, top_p=top_p, top_k=top_k, neg_entropy=True)
+                elif (not rcr and alg == 'topk_margin') or (rcr and conf_alg == 'topk_margin'):
+                    confidence, x0 = sample_tokens(
+                        mask_logits, temperature=temperature, top_p=top_p, top_k=top_k, margin_confidence=True
+                    )
+                elif (not rcr and alg == 'entropy') or (rcr and conf_alg == 'entropy'):
+                    confidence, x0 = sample_tokens(
+                        mask_logits, temperature, top_p=top_p, top_k=top_k, neg_entropy=True
+                    )
                 else:
-                    raise RuntimeError(f"Unknown alg: {alg}")
-                num_mask_token = mask_index.sum() / mask_index.shape[0]
-                number_transfer_tokens = int(num_mask_token * (1 - s / t)) if i < steps - 1 else int(num_mask_token)
-                full_confidence = torch.full_like(x, -torch.inf, device=self.device, dtype=logits.dtype)
-                full_confidence[mask_index] = confidence
-                if number_transfer_tokens > 0:
-                    if alg_temp is None or alg_temp == 0:
-                        _, transfer_index = torch.topk(full_confidence, number_transfer_tokens)
+                    # 兼容：如果 rcr=True 但 conf_alg 非上述三者，回退到 alg 指定
+                    if rcr:
+                        if alg == 'maskgit_plus':
+                            confidence, x0 = sample_tokens(mask_logits, temperature=temperature, top_p=top_p, top_k=top_k)
+                        elif alg == 'topk_margin':
+                            confidence, x0 = sample_tokens(
+                                mask_logits, temperature=temperature, top_p=top_p, top_k=top_k, margin_confidence=True
+                            )
+                        elif alg == 'entropy':
+                            confidence, x0 = sample_tokens(
+                                mask_logits, temperature, top_p=top_p, top_k=top_k, neg_entropy=True
+                            )
+                        else:
+                            raise RuntimeError(f"Unknown alg: {alg}")
                     else:
-                        full_confidence = full_confidence / alg_temp
-                        full_confidence = F.softmax(full_confidence, dim=-1)
-                        transfer_index = torch.multinomial(full_confidence, num_samples=number_transfer_tokens)
-                    x_ = torch.zeros_like(x, device=self.device, dtype=torch.long) + mask_token_id
-                    x_[mask_index] = x0.clone()
-                    row_indices = torch.arange(x.size(0), device=self.device).unsqueeze(1).expand_as(transfer_index)
-                    x[row_indices,transfer_index] = x_[row_indices,transfer_index]
+                        raise RuntimeError(f"Unknown alg: {alg}")
+
+                if rcr:
+                    # === 仅在 rcr=True 时：应用 RCR；不会触碰 baseline 分支实现 ===
+                    self._apply_rcr_logic(
+                        x, x0, confidence, mask_index, overtime_confidence,
+                        mask_token_id, i, steps, s, t
+                    )
+                else:
+                    # === 原版 Dream 逻辑：保持不变（包括 device=self.device 等细节） ===
+                    num_mask_token = mask_index.sum() / mask_index.shape[0]
+                    number_transfer_tokens = int(num_mask_token * (1 - s / t)) if i < steps - 1 else int(num_mask_token)
+                    full_confidence = torch.full_like(x, -torch.inf, device=self.device, dtype=logits.dtype)
+                    full_confidence[mask_index] = confidence
+                    if number_transfer_tokens > 0:
+                        if alg_temp is None or alg_temp == 0:
+                            _, transfer_index = torch.topk(full_confidence, number_transfer_tokens)
+                        else:
+                            full_confidence = full_confidence / alg_temp
+                            full_confidence = F.softmax(full_confidence, dim=-1)
+                            transfer_index = torch.multinomial(full_confidence, num_samples=number_transfer_tokens)
+                        x_ = torch.zeros_like(x, device=self.device, dtype=torch.long) + mask_token_id
+                        x_[mask_index] = x0.clone()
+                        row_indices = torch.arange(x.size(0), device=self.device).unsqueeze(1).expand_as(transfer_index)
+                        x[row_indices, transfer_index] = x_[row_indices, transfer_index]
 
             # this allows user-defined token control of the intermediate steps
             x = generation_tokens_hook_func(i, x, logits)
 
             if histories is not None:
                 histories.append(x.clone())
-        
+
         if return_dict_in_generate:
             return DreamModelOutput(
                 sequences=x,
                 history=histories,
             )
         else:
-            return x
+            return x