Update custom_generate/generate.py

custom_generate/generate.py

@@ -17,12 +17,12 @@ from transformers.generation.utils import (
 )
 
 
-def generate(
+def _deepconf_generate(
     model: Any,
     input_ids: torch.LongTensor,
-    logits_processor: Optional[LogitsProcessorList] = None,
-    stopping_criteria: Optional[StoppingCriteriaList] = None,
-    generation_config: Optional[GenerationConfig] = None,
+    logits_processor: Optional[LogitsProcessorList],
+    stopping_criteria: Optional[StoppingCriteriaList],
+    generation_config: Optional[GenerationConfig],
     synced_gpus: bool = False,
     streamer: Optional[Any] = None,
     **model_kwargs,
@@ -44,12 +44,6 @@ def generate(
         depending on `return_dict_in_generate` and model type.
     """
 
-    # Ensure processors/criteria are defined
-    if logits_processor is None:
-        logits_processor = LogitsProcessorList()
-    if stopping_criteria is None:
-        stopping_criteria = StoppingCriteriaList()
-
     # Get DeepCONF parameters from generation_config or set defaults
     enable_conf = getattr(generation_config, "enable_conf", False)
     enable_early_stopping = getattr(generation_config, "enable_early_stopping", True)  # NEW: Allow disabling early stopping
@@ -75,14 +69,7 @@ def generate(
 
     # Initialize values
     # Handle pad token properly (following HF best practices)
-    pad_token_id = generation_config.pad_token_id
-    if pad_token_id is None and hasattr(generation_config, "_pad_token_tensor"):
-        pad_token_id = generation_config._pad_token_tensor
-    if pad_token_id is None and hasattr(model.config, "pad_token_id"):
-        pad_token_id = model.config.pad_token_id
-    if pad_token_id is None and generation_config.eos_token_id is not None:
-        # Use eos token as pad token if not set
-        pad_token_id = generation_config.eos_token_id
+    pad_token_id = generation_config._pad_token_tensor
 
     output_attentions = generation_config.output_attentions
     output_hidden_states = generation_config.output_hidden_states
@@ -383,3 +370,21 @@ def generate(
             return output
     else:
         return input_ids
+
+
+def generate(model, *args, **kwargs):
+    """Custom generate function for group beam search decoding.
+    Args:
+        model (`PreTrainedModel`):
+            The model to generate from.
+        num_beams (`int`): The number of beams to use for beam search.
+        num_beam_groups (`int`): The number of beam groups to use for beam search.
+        length_penalty (`float`): The length penalty to use for beam search.
+        early_stopping (`bool`): Whether to stop beam search when sufficient beams have finished.
+        num_return_sequences (`int`): The number of sequences to return.
+        max_length (`int`): The maximum length of the generated sequence.
+    """
+    generation_outputs = GenerationMixin.generate(
+        model, *args, custom_generate=_deepconf_generate, **kwargs
+    )
+    return generation_outputs