[Feat] support for merging generation config with generate kwargs

generation_utils.py

@@ -1,3 +1,8 @@
+# Copyright 2025 Radical Numerics Inc.
+#
+# This source code is licensed under the Apache License, Version 2.0, found in the
+# LICENSE file in the root directory of this source tree.
+
 """
 RND1 Generation Utilities.
 
@@ -6,12 +11,12 @@ including the main GenerationMixin class that integrates with HuggingFace.
 """
 
 import torch
-import torch.nn as nn
 from typing import Optional, Union, Dict, Any
 from transformers import GenerationMixin as HFGenerationMixin
 from transformers.generation import GenerationConfig
 
-from .sampling import diffusion_sample, apply_top_k_filtering, apply_top_p_filtering
+from .generation_config import RND1GenerationConfig
+from .sampling import diffusion_sample
 
 
 class RND1GenerationMixin(HFGenerationMixin):
@@ -41,12 +46,12 @@ class RND1GenerationMixin(HFGenerationMixin):
 
         Args:
             inputs: Input token IDs to use as prefix (standard HF parameter)
-            generation_config: Generation configuration object
+            generation_config: Generation configuration object. Default is RND1GenerationConfig.
             prefix_ids: Alternative to inputs for infilling tasks
             suffix_ids: Optional suffix for infilling tasks
             infill_length: Length of infill region (for infilling)
             return_dict_in_generate: Whether to return GenerateDecoderOnlyOutput
-            **kwargs: Additional arguments (accepted for compatibility)
+            **kwargs: Additional arguments (accepted for compatibility). These will be passed to the config constructor.
 
         Returns:
             Generated token IDs or GenerateDecoderOnlyOutput
@@ -56,7 +61,7 @@ class RND1GenerationMixin(HFGenerationMixin):
             model_kwargs = kwargs.copy()
         else:
             # Only prepare config from kwargs if no config was provided
-            gen_config, model_kwargs = self._prepare_generation_config(None, **kwargs)
+            gen_config, model_kwargs = self._prepare_generation_config(RND1GenerationConfig(), **kwargs)
 
         device = next(self.parameters()).device
 
@@ -71,7 +76,7 @@ class RND1GenerationMixin(HFGenerationMixin):
             suffix_ids = suffix_ids.to(device)
 
         eos_token_id = gen_config.eos_token_id or getattr(self.config, "eos_token_id", 151645)
-        pad_token_id = gen_config.pad_token_id or getattr(self.config, "pad_token_id", None)
+        pad_token_id = gen_config.pad_token_id or getattr(self.config, "pad_token_id", 151643)
         bos_token_id = gen_config.bos_token_id or getattr(self.config, "bos_token_id", None)
         mask_token_id = getattr(gen_config, "mask_token_id", getattr(self.config, "mask_token_id", 151669))
 
@@ -101,12 +106,6 @@ class RND1GenerationMixin(HFGenerationMixin):
         greedy = getattr(gen_config, "greedy",
                         not bool(gen_config.do_sample) if hasattr(gen_config, "do_sample") else True)
 
-        generator = model_kwargs.get("generator", None)
-        if generator is None:
-            seed = getattr(gen_config, 'seed', None)
-            if seed is not None:
-                generator = torch.Generator(device=device)
-                generator.manual_seed(seed)
 
         with torch.inference_mode():
             sequences = diffusion_sample(
@@ -125,7 +124,6 @@ class RND1GenerationMixin(HFGenerationMixin):
                 pad_token_id=pad_token_id,
                 bos_token_id=bos_token_id,
                 device=device,
-                generator=generator,
                 visualizer=model_kwargs.get("visualizer", None),  # Optional visualizer from kwargs
             )
 
@@ -142,7 +140,6 @@ class RND1GenerationMixin(HFGenerationMixin):
         generation_config: Optional[GenerationConfig] = None,
         suffix_ids: Optional[torch.LongTensor] = None,
         infill_length: Optional[int] = None,
-        generator: Optional[torch.Generator] = None,
         **kwargs,
     ) -> torch.LongTensor:
         """
@@ -157,7 +154,6 @@ class RND1GenerationMixin(HFGenerationMixin):
             generation_config: Generation configuration object
             suffix_ids: Optional suffix token IDs
             infill_length: Length of infill region
-            generator: Random generator for reproducibility
             **kwargs: Additional arguments for backward compatibility
 
         Returns:
@@ -171,7 +167,6 @@ class RND1GenerationMixin(HFGenerationMixin):
             generation_config=generation_config,
             suffix_ids=suffix_ids,
             infill_length=infill_length,
-            generator=generator,
             visualizer=visualizer,
             return_dict_in_generate=False,
             **kwargs,