Upload model

generation_utils.py

@@ -302,7 +302,7 @@ class DreamGenerationMixin:
         **kwargs,
     ) -> Union[DreamModelOutput, torch.LongTensor]:
         # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call
-        tokenizer = kwargs.pop("tokenizer", None)  # Pull this out first, we only use it for stopping criteria
+        generation_hook_func = kwargs.pop("generation_hook_func", lambda x, logits: (x, logits))
         generation_config = self._prepare_generation_config(generation_config, **kwargs)
 
         # 2. Define model inputs
@@ -355,6 +355,7 @@ class DreamGenerationMixin:
             input_ids,
             attention_mask=attention_mask,
             generation_config=generation_config,
+            generation_hook_func=generation_hook_func
         )
         return result
 
@@ -363,6 +364,7 @@ class DreamGenerationMixin:
         input_ids: torch.LongTensor,
         attention_mask: Optional[torch.LongTensor],
         generation_config: DreamGenerationConfig,
+        generation_hook_func
     ) -> Union[DreamModelOutput, torch.LongTensor]:
         # init values
         output_history = generation_config.output_history
@@ -402,7 +404,7 @@ class DreamGenerationMixin:
             mask_index = (x == mask_token_id)
             logits = self(x, attention_mask, tok_idx).logits
             logits = torch.cat([logits[:,:1], logits[:, :-1]], dim=1)
-            logits = logits[mask_index]
+            mask_logits = logits[mask_index]
             t = timesteps[i]
             s = timesteps[i + 1]
         
@@ -410,15 +412,15 @@ class DreamGenerationMixin:
                 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(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':
-                    confidence, x0 = sample_tokens(logits, temperature=temperature, top_p=top_p, top_k=top_k)
+                    confidence, x0 = sample_tokens(mask_logits, temperature=temperature, top_p=top_p, top_k=top_k)
                 elif alg == 'topk_margin':
-                    confidence, x0 = sample_tokens(logits, temperature=temperature, top_p=top_p, top_k=top_k, margin_confidence=True)
+                    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(logits, temperature, top_p=top_p, top_k=top_k, neg_entropy=True)
+                    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()
@@ -433,6 +435,9 @@ class DreamGenerationMixin:
                     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_
+                    
+            # this allows user-defined control of the intermediate steps
+            x, logits = generation_hook_func(x, logits)
 
             if histories is not None:
                 histories.append(x.clone())