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.
+# 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.
@@ -85,8 +85,7 @@ def sample_tokens(logits, temperature=0.0, top_p=None, top_k=None, margin_confid
     if neg_entropy:
         epsilon = 1e-10
         log_probs = torch.log(probs + epsilon)
-        # 注意：这里返回的是“负熵”的相反数（越大越自信）
-        confidence = -(probs * log_probs).sum(dim=-1)
+        confidence = torch.sum(probs * log_probs, dim=-1)
     
     return confidence, x0
 
@@ -110,10 +109,6 @@ class DreamGenerationConfig(GenerationConfig):
         self.alg: str = kwargs.pop("alg", 'origin')
         self.alg_temp: Optional[float] = kwargs.pop("alg_temp", None)
 
-        # RCR specific parameters
-        self.rcr: bool = kwargs.pop("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)
@@ -169,58 +164,6 @@ class DreamGenerationMixin:
             attention_mask = attention_mask.repeat_interleave(expand_size, dim=0)
         return input_ids, attention_mask
 
-    def _apply_rcr_logic(self, x, x0, confidence, mask_index, overtime_confidence,
-                        mask_token_id, step, total_steps, s, t):
-        """
-        RCR：在 Dream 原逻辑上做“最小侵入”改动，使其真正生效。
-        - 仍采用 Dream 的调度：本步 global k = num_mask_token * (1 - s/t)
-        - 逐样本 clamp，避免批均值 k 在样本上越界
-        - 目标累计约束：到本步为止累计应已生成 target_cum = num_mask_token * (1 - s/t)。
-          若当前累计 > 目标，按最低置信度反遮盖回 [MASK]。
-        """
-        device = x.device
-        B, L = x.shape
-
-        # 与 Dream 保持一致：使用“批均值”的 num_mask_token 与 (1 - s/t) 调度定义
-        num_mask_token = (mask_index.sum() / mask_index.shape[0]).item()
-        k_global = int(num_mask_token * (1 - (s / t).item())) if step < total_steps - 1 else int(num_mask_token)
-
-        # 构造全长置信度和临时候选（非 mask 位置分别置为 -inf / mask_token）
-        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
-            masked_count_j = int(mask_index[j].sum().item())
-            k_j = min(k_global, masked_count_j)
-            if k_j > 0:
-                # 只在 mask 内选 topk（非 mask 位置 full_conf 为 -inf，不会被选中）
-                _, select_idx = torch.topk(full_conf[j], 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:
-                # Dream 的“到本步为止累计应已生成”的目标数量
-                target_cum = int(num_mask_token * (1 - (s / t).item()))
-                # 当前已生成的数量（overtime_confidence>0 的位置视为已确定）
-                gen_mask = overtime_confidence[j] > 0
-                current_gen = int(gen_mask.sum().item())
-
-                # 若超过目标，反遮盖（remask）最低置信度的那部分，使当前累计 ≈ 目标累计
-                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
-
     def _validate_generated_length(self, generation_config, input_ids_length, has_default_max_length):
         """Performs validation related to the resulting generated length"""
 
@@ -439,10 +382,6 @@ class DreamGenerationMixin:
         top_p = generation_config.top_p
         top_k = generation_config.top_k
 
-        # RCR specific values
-        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
@@ -465,9 +404,6 @@ class DreamGenerationMixin:
 
         timesteps = torch.linspace(1, eps, steps + 1, device=x.device)
 
-        # RCR tracking - initialize overtime confidence tracking
-        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):
@@ -489,38 +425,29 @@ class DreamGenerationMixin:
                 _, 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' or (rcr and conf_alg == 'maskgit_plus'):
+                if alg == 'maskgit_plus':
                     confidence, x0 = sample_tokens(mask_logits, temperature=temperature, top_p=top_p, top_k=top_k)
-                elif alg == 'topk_margin' or (rcr and conf_alg == 'topk_margin'):
+                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' or (rcr and conf_alg == 'entropy'):
+                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}")
-
-                # Apply RCR logic if enabled
-                if rcr:
-                    print(f"[RCR EXEC] Step {i}: RCR logic executed")
-                    self._apply_rcr_logic(x, x0, confidence, mask_index, overtime_confidence,
-                                        mask_token_id, i, steps, s, t)
-                else:
-                    # Original Dream sampling logic
-                    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)
-                    # --------- 仅此处小修：device 用 x.device，避免跨设备 ----------
-                    full_confidence = torch.full_like(x, -torch.inf, device=x.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]
+                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)
@@ -534,4 +461,4 @@ class DreamGenerationMixin:
                 history=histories,
             )
         else:
-            return x
+            return x