Adding support to pass arguments to custom beam search scorer

Author: GM07

src/transformers/generation/configuration_utils.py

@@ -136,6 +136,9 @@ class GenerationConfig(PushToHubMixin):
         beam_search_scorer_class: (`class`, *optional*, defaults to `None`):
             Which class to use as a beam search scorer. If `None`, it will use the default `BeamSearchScorer` class.
             The type must inherit from `BeamSearchScorer`.
+        beam_search_scorer_args: (`dict`, *optional*, defaults to `None`)
+            Arguments that will be passed when creating the beam search scorer. When this argument is specified,
+            `beam_search_scorer_class` must not be `None`.
 
         > Parameters for manipulation of the model output logits
 
@@ -357,6 +360,7 @@ def __init__(self, **kwargs):
         self.penalty_alpha = kwargs.pop("penalty_alpha", None)
         self.use_cache = kwargs.pop("use_cache", True)
         self.beam_search_scorer_class = kwargs.pop("beam_search_scorer_class", None)
+        self.beam_search_scorer_args = kwargs.pop("beam_search_scorer_args", None)
 
         # Parameters for manipulation of the model output logits
         self.temperature = kwargs.pop("temperature", 1.0)
@@ -652,6 +656,11 @@ def validate(self, is_init=False):
                     UserWarning,
                 )
 
+            if self.beam_search_scorer_class is None and self.beam_search_scorer_args is not None:
+                raise ValueError(
+                    "The initialization arguments for the beam search scorer class were provided, but the class was not",
+                )
+
         # 3. detect incorrect paramaterization specific to advanced beam modes
         else:
             # constrained beam search

src/transformers/generation/utils.py

@@ -2004,7 +2004,6 @@ def generate(
                 if generation_config.do_sample
                 else None
             )
-
             # 12. prepare beam search scorer
             if generation_config.beam_search_scorer_class is None:
                 beam_scorer = BeamSearchScorer(
@@ -2017,6 +2016,11 @@ def generate(
                     max_length=generation_config.max_length,
                 )
             else:
+                args = (
+                    generation_config.beam_search_scorer_args
+                    if generation_config.beam_search_scorer_args is not None
+                    else {}
+                )
                 beam_scorer = generation_config.beam_search_scorer_class(
                     batch_size=batch_size,
                     num_beams=generation_config.num_beams,
@@ -2026,6 +2030,7 @@ def generate(
                     num_beam_hyps_to_keep=generation_config.num_return_sequences,
                     num_beam_groups=generation_config.num_beam_groups,
                     max_length=generation_config.max_length,
+                    **args,
                 )
 
             # 13. interleave input_ids with `num_beams` additional sequences per batch
@@ -2061,6 +2066,11 @@ def generate(
                     max_length=generation_config.max_length,
                 )
             else:
+                args = (
+                    generation_config.beam_search_scorer_args
+                    if generation_config.beam_search_scorer_args is not None
+                    else {}
+                )
                 beam_scorer = generation_config.beam_search_scorer_class(
                     batch_size=batch_size,
                     num_beams=generation_config.num_beams,
@@ -2070,6 +2080,7 @@ def generate(
                     num_beam_hyps_to_keep=generation_config.num_return_sequences,
                     num_beam_groups=generation_config.num_beam_groups,
                     max_length=generation_config.max_length,
+                    **args,
                 )
             # 12. interleave input_ids with `num_beams` additional sequences per batch
             input_ids, model_kwargs = self._expand_inputs_for_generation(

tests/generation/test_utils.py

@@ -586,8 +586,13 @@ class CustomBeamSearchScorer(BeamSearchScorer):
             finalize_called = False
             process_called = False
 
-            def __init__(self, *args, **kwargs):
+            def __init__(inside_self, test_args, *args, **kwargs):
                 super().__init__(*args, **kwargs)
+                inside_self.test_args = test_args
+                self.assertTrue(
+                    inside_self.test_args,
+                    "The argument `test_args` should " "have been passed to the beam search scorer init function",
+                )
 
             def process(self, *args, **kwargs):
                 results = super().process(*args, **kwargs)
@@ -620,6 +625,7 @@ def finalize(self, *args, **kwargs):
                 do_sample=False,
                 max_new_tokens=self.max_new_tokens,
                 beam_search_scorer_class=CustomBeamSearchScorer,
+                beam_search_scorer_args={"test_args": True},
                 output_scores=False,
                 output_logits=False,
                 output_attentions=False,
@@ -2343,8 +2349,13 @@ class CustomBeamSearchScorer(BeamSearchScorer):
             finalize_called = False
             process_called = False
 
-            def __init__(self, *args, **kwargs):
+            def __init__(inside_self, test_args, *args, **kwargs):
                 super().__init__(*args, **kwargs)
+                inside_self.test_args = test_args
+                self.assertTrue(
+                    inside_self.test_args,
+                    ("The argument `test_args` should " "have been passed to the beam search scorer init function"),
+                )
 
             def process(self, *args, **kwargs):
                 results = super().process(*args, **kwargs)
@@ -2371,6 +2382,7 @@ def finalize(self, *args, **kwargs):
             min_length=5,
             eos_token_id=model.config.eos_token_id,
             beam_search_scorer_class=CustomBeamSearchScorer,
+            beam_search_scorer_args={"test_args": True},
             **model_kwargs,
         )
         outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)