revise generate

modeling.py

@@ -6,6 +6,7 @@ from torch import nn
 import torch.nn.functional as F
 from functools import partial
 
+from transformers.generation.utils import GenerateDecoderOnlyOutput  
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache
 from transformers.generation import GenerationMixin
@@ -643,7 +644,7 @@ class Fast_dLLM_QwenForCausalLM(Fast_dLLM_QwenPreTrainedModel, GenerationMixin):
             loss=loss,
             logits=logits,
             past_key_values=outputs.past_key_values,
-            hidden_states=outputs.hidden_states,
+            hidden_states=hidden_states,
             attentions=outputs.attentions,
             block_past_key_values=outputs.block_past_key_values,
         )
@@ -652,7 +653,9 @@ class Fast_dLLM_QwenForCausalLM(Fast_dLLM_QwenPreTrainedModel, GenerationMixin):
     def generate(
         self,
         input_ids,
-        max_new_tokens, 
+        max_new_tokens=None,
+        max_length=None,
+        tokenizer=None,
         mask_id=151665,
         threshold=1,
         small_block_size=8,
@@ -662,14 +665,36 @@ class Fast_dLLM_QwenForCausalLM(Fast_dLLM_QwenPreTrainedModel, GenerationMixin):
         top_p=0.95,
         temperature=0,
         use_block_cache=False,
+        return_dict_in_generate=False,
+        output_scores=False,
+        output_hidden_states=False,
         **kwargs
     ):
+        if max_new_tokens is None and max_length is None:
+            raise ValueError("Either max_new_tokens or max_length must be specified")
+        if max_new_tokens is None:
+            max_new_tokens = max_length - input_ids.shape[1]
+        
+        scores_list = [] if output_scores else None
+        decoder_hidden_states = [] if output_hidden_states else None
+        
         num_blocks = max_new_tokens // block_size
         original_input_length = input_ids.shape[1]
 
         if input_ids.shape[1] > block_size:
-            output = self.forward(input_ids=input_ids[:, :(input_ids.shape[1] // block_size * block_size)], use_cache=True, update_past_key_values=True, block_size=block_size)
+            output = self.forward(
+                input_ids=input_ids[:, :(input_ids.shape[1] // block_size * block_size)], 
+                use_cache=True, 
+                update_past_key_values=True, 
+                block_size=block_size
+            )
             logits, past_key_values = output.logits, output.past_key_values
+            
+            if output_scores:
+                scores_list.append(logits)
+            if output_hidden_states and hasattr(output, 'hidden_states'):
+                decoder_hidden_states.append(output.hidden_states)
+            
             if input_ids.shape[1] % block_size == 0:
                 next_token = logits[:, -1:, :].argmax(dim=-1)
                 input_ids = torch.cat([input_ids, next_token], dim=1)
@@ -683,30 +708,51 @@ class Fast_dLLM_QwenForCausalLM(Fast_dLLM_QwenPreTrainedModel, GenerationMixin):
                 break
             prompt_length = input_ids.shape[1]
             # Initialize x_init with mask_id
-            x_init = mask_id * torch.ones((input_ids.shape[0], block_size-prompt_length%block_size), device=self.device, dtype=torch.long)
+            x_init = mask_id * torch.ones(
+                (input_ids.shape[0], block_size-prompt_length%block_size), 
+                device=self.device, 
+                dtype=torch.long
+            )
             x_init = torch.cat([input_ids, x_init], dim=1)
 
             x_t = x_init.clone()
             block_past_key_values = None
+            
             while True:
                 if stop_token in x_t[:, prompt_length:]:
                     stop_token_idx = (x_t[:, prompt_length:] == stop_token).nonzero()[0][1]
                     if (x_t[:, prompt_length:prompt_length+stop_token_idx] == mask_id).sum() == 0:
                         break
                 mask_idx = (x_t[:, -block_size:] == mask_id)
+                
                 # Decode a complete block, update cache, and generate the next token
                 if mask_idx.sum() == 0:
-                    output = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=True, block_size=block_size)
+                    output = self.forward(
+                        input_ids=x_t[:, -block_size:], 
+                        use_cache=True, 
+                        past_key_values=past_key_values, 
+                        update_past_key_values=True, 
+                        block_size=block_size
+                    )
                     logits, past_key_values = output.logits, output.past_key_values
+                    
+                    # 收集输出信息
+                    if output_scores:
+                        scores_list.append(logits)
+                    if output_hidden_states and hasattr(output, 'hidden_states'):
+                        decoder_hidden_states.append(output.hidden_states)
+                    
                     next_token = logits[:, -1:, :].argmax(dim=-1)
                     x_t = torch.cat([x_t, next_token], dim=1)
                     break
+                    
                 for small_block_idx in range(num_small_blocks):
                     small_block_start_idx = small_block_idx * small_block_size
                     small_block_end_idx = small_block_start_idx + small_block_size
 
                     start = -block_size + small_block_start_idx
                     end = None if block_size == small_block_end_idx else -block_size + small_block_end_idx
+                    
                     while True:
                         mask_idx = (x_t[:, -block_size:] == mask_id)
                         if mask_idx[:, start:end].sum() == 0:
@@ -718,18 +764,43 @@ class Fast_dLLM_QwenForCausalLM(Fast_dLLM_QwenPreTrainedModel, GenerationMixin):
 
                         if use_block_cache:
                             if block_past_key_values is None or (x_t[:, -block_size+small_block_start_idx] == mask_id).any():
-                                output = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, use_block_cache=True)
+                                output = self.forward(
+                                    input_ids=x_t[:, -block_size:], 
+                                    use_cache=True, 
+                                    past_key_values=past_key_values, 
+                                    update_past_key_values=False, 
+                                    use_block_cache=True,
+                                )
                                 logits, block_past_key_values = output.logits, output.block_past_key_values
                                 logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                                 logits = logits[:, start:end]
                             else:
-                                logits = self.forward(input_ids=x_t[:,start:end], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, use_block_cache=True, block_past_key_values=block_past_key_values, replace_position=small_block_start_idx).logits
+                                output = self.forward(
+                                    input_ids=x_t[:,start:end], 
+                                    use_cache=True, 
+                                    past_key_values=past_key_values, 
+                                    update_past_key_values=False, 
+                                    use_block_cache=True, 
+                                    block_past_key_values=block_past_key_values, 
+                                    replace_position=small_block_start_idx
+                                )
+                                logits = output.logits
                                 logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                         else:
-                            logits = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False).logits
+                            output = self.forward(
+                                input_ids=x_t[:, -block_size:], 
+                                use_cache=True, 
+                                past_key_values=past_key_values, 
+                                update_past_key_values=False
+                            )
+                            logits = output.logits
                             logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                             logits = logits[:, start:end]
 
+                        if output_scores:
+                            scores_list.append(logits)
+                        if output_hidden_states and hasattr(output, 'hidden_states'):
+                            decoder_hidden_states.append(output.hidden_states)
 
                         x_1, p_1t = self.sample_with_top_p(logits, top_p=top_p, temperature=temperature)
                         # Select tokens with probability greater than threshold from p_1t
@@ -744,11 +815,21 @@ class Fast_dLLM_QwenForCausalLM(Fast_dLLM_QwenPreTrainedModel, GenerationMixin):
                         x_t[:, start:end][unmask_idx] = x_1[unmask_idx]
 
             input_ids = x_t
+            
         # Truncate stop_token
         if stop_token in input_ids[:, original_input_length:]:
             stop_token_idx = (input_ids[:, original_input_length:] == stop_token).nonzero()[0][1]
             input_ids = input_ids[:, :stop_token_idx+original_input_length+1]
-        return input_ids
+        
+        if return_dict_in_generate:
+            return GenerateDecoderOnlyOutput(
+                sequences=input_ids,
+                scores=tuple(scores_list) if output_scores and scores_list else None,
+                hidden_states=tuple(decoder_hidden_states) if output_hidden_states and decoder_hidden_states else None,
+            )
+        else:
+            return input_ids
+
 
     def sample_with_top_p(self, logits, top_p=0.95, temperature=1.0):
         # Calculate probabilities
@@ -782,4 +863,4 @@ class Fast_dLLM_QwenForCausalLM(Fast_dLLM_QwenPreTrainedModel, GenerationMixin):
         p_1t = normalized_probs
         x_1 = torch.multinomial(p_1t[0], num_samples=1).unsqueeze(0).squeeze(-1)
 
-        return x_1, p_1t
+        return x_1, p_1t