Generate: model_kwargs can also be an input to prepare_inputs_for_generation

src/transformers/generation/flax_utils.py

@@ -194,9 +194,9 @@ def _validate_model_kwargs(self, model_kwargs: Dict[str, Any]):
         """Validates model kwargs for generation. Generate argument typos will also be caught here."""
         unused_model_args = []
         model_args = set(inspect.signature(self.prepare_inputs_for_generation).parameters)
-        # `kwargs` if often used to handle optional forward pass inputs like `attention_mask`. If
+        # `kwargs`/`model_kwargs` is often used to handle optional forward pass inputs like `attention_mask`. If
         # `prepare_inputs_for_generation` doesn't accept `kwargs`, then a stricter check can be made ;)
-        if "kwargs" in model_args:
+        if "kwargs" in model_args or "model_kwargs" in model_args:
             model_args |= set(inspect.signature(self.__call__).parameters)
         for key, value in model_kwargs.items():
             if value is not None and key not in model_args:

src/transformers/generation/tf_utils.py

@@ -1445,9 +1445,9 @@ def _validate_model_kwargs(self, model_kwargs: Dict[str, Any]):
 
         unused_model_args = []
         model_args = set(inspect.signature(self.prepare_inputs_for_generation).parameters)
-        # `kwargs` if often used to handle optional forward pass inputs like `attention_mask`. If
+        # `kwargs`/`model_kwargs` is often used to handle optional forward pass inputs like `attention_mask`. If
         # `prepare_inputs_for_generation` doesn't accept `kwargs`, then a stricter check can be made ;)
-        if "kwargs" in model_args:
+        if "kwargs" in model_args or "model_kwargs" in model_args:
             model_args |= set(inspect.signature(self.call).parameters)
         for key, value in model_kwargs.items():
             if value is not None and key not in model_args:

src/transformers/generation/utils.py

@@ -981,9 +981,9 @@ def _validate_model_kwargs(self, model_kwargs: Dict[str, Any]):
 
         unused_model_args = []
         model_args = set(inspect.signature(self.prepare_inputs_for_generation).parameters)
-        # `kwargs` if often used to handle optional forward pass inputs like `attention_mask`. If
+        # `kwargs`/`model_kwargs` is often used to handle optional forward pass inputs like `attention_mask`. If
         # `prepare_inputs_for_generation` doesn't accept `kwargs`, then a stricter check can be made ;)
-        if "kwargs" in model_args:
+        if "kwargs" in model_args or "model_kwargs" in model_args:
             model_args |= set(inspect.signature(self.forward).parameters)
         for key, value in model_kwargs.items():
             if value is not None and key not in model_args: