Upload generation_utils.py with huggingface_hub

generation_utils.py

@@ -138,20 +138,12 @@ class RND1GenerationMixin(HFGenerationMixin):
     def generate_with_visualization(
         self,
         tokenizer,
-        prefix_ids: Optional[torch.LongTensor] = None,
+        inputs: Optional[torch.LongTensor] = None,
+        generation_config: Optional[GenerationConfig] = None,
         suffix_ids: Optional[torch.LongTensor] = None,
         infill_length: Optional[int] = None,
-        seq_len: int = 256,
-        num_steps: int = 256,
-        mask_token_id: int = 151669,
-        temperature: float = 1.0,
-        top_k: Optional[int] = None,
-        top_p: Optional[float] = None,
-        greedy: bool = True,
-        eos_token_id: int = 151645,
-        pad_token_id: Optional[int] = None,
-        bos_token_id: Optional[int] = None,
         generator: Optional[torch.Generator] = None,
+        **kwargs,
     ) -> torch.LongTensor:
         """
         Generate with live visualization (for demos).
@@ -161,20 +153,12 @@ class RND1GenerationMixin(HFGenerationMixin):
 
         Args:
             tokenizer: Tokenizer for decoding tokens to text
-            prefix_ids: Optional prefix token IDs
+            inputs: Input token IDs to use as prefix
+            generation_config: Generation configuration object
             suffix_ids: Optional suffix token IDs
             infill_length: Length of infill region
-            seq_len: Target sequence length
-            num_steps: Number of diffusion steps
-            mask_token_id: Mask token ID
-            temperature: Sampling temperature
-            top_k: Top-k filtering
-            top_p: Top-p filtering
-            greedy: Whether to use greedy sampling
-            eos_token_id: End of sequence token ID
-            pad_token_id: Padding token ID
-            bos_token_id: Beginning of sequence token ID
             generator: Random generator for reproducibility
+            **kwargs: Additional arguments for backward compatibility
 
         Returns:
             Generated token IDs as LongTensor
@@ -182,33 +166,15 @@ class RND1GenerationMixin(HFGenerationMixin):
         from .terminal_visualizer import TerminalVisualizer
         visualizer = TerminalVisualizer(tokenizer, show_visualization=True)
 
-        max_new_tokens = None
-        if seq_len is not None and prefix_ids is not None:
-            max_new_tokens = seq_len - prefix_ids.shape[1]
-
-        from .generation_config import RND1GenerationConfig
-        gen_config = RND1GenerationConfig(
-            max_length=seq_len,
-            max_new_tokens=max_new_tokens,
-            num_diffusion_steps=num_steps,
-            mask_token_id=mask_token_id,
-            temperature=temperature,
-            top_k=top_k,
-            top_p=top_p,
-            greedy=greedy,
-            bos_token_id=bos_token_id,
-            eos_token_id=eos_token_id,
-            pad_token_id=pad_token_id,
-        )
-
         return self.generate(
-            inputs=prefix_ids,
+            inputs=inputs,
+            generation_config=generation_config,
             suffix_ids=suffix_ids,
             infill_length=infill_length,
-            generation_config=gen_config,
             generator=generator,
             visualizer=visualizer,
             return_dict_in_generate=False,
+            **kwargs,
         )
 
     def prepare_inputs_for_generation(