Skip to content
Merged
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 27 additions & 17 deletions src/transformers/generation/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -1177,21 +1177,35 @@ def _merge_criteria_processor_list(
default_list: Union[LogitsProcessorList, StoppingCriteriaList],
custom_list: Union[LogitsProcessorList, StoppingCriteriaList],
) -> Union[LogitsProcessorList, StoppingCriteriaList]:
"""
Copy link
Contributor Author

@gante gante Mar 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The changes in this function are secondary to the main change:

  • whisper breaks because it both sets custom logits processors AND has the default flags in the generation config to instantiate them
  • after the original change (inherit defaults from the model's generation config), we were throwing an exception here
  • after this secondary change, we only throw a warning and discard the logits processor instance created inside .generate() (i.e. assumes the user knows what they are doing when the pass logits_processors to .generate() instead of crashing)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A tiny comment wouldn't hurt for future us, isn't very easy to get why we do this without reading PR description.

Also I am not sure if this is required, aren't we restricting custom logits processors to only those that cannot be configured by generation config? Something that is only defined by users for their use-case

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also I am not sure if this is required, aren't we restricting custom logits processors to only those that cannot be configured by generation config? Something that is only defined by users for their use-case

The user always had the option of unsetting a flag and passing the corresponding processor. This change makes it less restricting: if they pass both a flag and the corresponding processor, we keep the processor and ignore the flag (previously we would throw an exception)

I'm not super fan of the new behavior, I think the exception less ambiguous (and thus preferable). However, it's needed to coexist with the main change in this PR, which is (IMO) more important.

I'll add a comment to clarify the evolution of this function, in case we need to trace back a related change :)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes, same feeling here. I would even restrict it to only custom logits processors in v5 to free us from the burden of maintaining correct priority/ordering etc. Looks like pandora's box what is being fixed 😿

Copy link
Contributor Author

@gante gante Mar 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would even restrict it to only custom logits processors in v5 to free us from the burden of maintaining correct priority/ordering etc.

I definitely don't want this, generate would become very uncomfortable to use 😅 A user always has the option of disabling generation_config flags and passing the processors in the order they want. But most users don't want that level of control, and yet may want an external processor

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The reason transformers is so popular is because we enable complex use-cases from a few lines of code

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I agree that this gives users more freedom to order processors the way they want by disabling generation config. Though I feel like it is not very clear to me as a user what happens under the hood, when I pass my own Temperature processor or use config. IMO we need a better docs page for advanced usage, if we allow that much freedom and expect users to know what they are doing

Users almost never 100% know what they are doing, thus open issues on GH 😆

Merge user-defined processors/criteria with the ones instantiated inside `generate`. In case the same
processor/criteria is present on both lists, use the user-defined one.
"""
if len(custom_list) == 0:
return default_list

final_list = type(default_list)()
for default in default_list:
using_custom = False
for custom in custom_list:
if type(custom) is type(default):
object_type = "stopping criteria" if isinstance(custom, StoppingCriteria) else "logits processor"
raise ValueError(
f"A custom {object_type} of type {type(custom)} with values {custom} has been passed to"
f" `.generate()`, but it has already been created with the values {default}. {default} has been"
" created by passing the corresponding arguments to generate or by the model's config default"
f" values. If you just want to change the default values of {object_type} consider passing"
f" them as arguments to `.generate()` instead of using a custom {object_type}."
logger.warning_once(
f"A custom {object_type} of type {type(custom)} has been passed to `.generate()`, but it "
f"was also created in `.generate()`, given its parameterization. The custom {type(custom)} "
f"will take precedence. Please check the docstring of {type(custom)} to see related "
"`.generate()` flags."
)
default_list.extend(custom_list)
return default_list
final_list.append(custom)
using_custom = True
break
if not using_custom:
final_list.append(default)

for custom in custom_list:
if custom not in final_list:
final_list.append(custom)
return final_list

def compute_transition_scores(
self,
Expand Down Expand Up @@ -1574,16 +1588,12 @@ def _prepare_generation_config(
if not is_torchdynamo_compiling():
generation_config = copy.deepcopy(generation_config)
model_kwargs = generation_config.update(**kwargs)
# If `generation_config` is provided, let's fallback ALL special tokens to the default values for the model
# If `generation_config` is provided, let's fallback ALL default values to the model's generation config
if not using_model_generation_config:
if generation_config.bos_token_id is None:
generation_config.bos_token_id = self.generation_config.bos_token_id
if generation_config.eos_token_id is None:
generation_config.eos_token_id = self.generation_config.eos_token_id
if generation_config.pad_token_id is None:
generation_config.pad_token_id = self.generation_config.pad_token_id
if generation_config.decoder_start_token_id is None:
generation_config.decoder_start_token_id = self.generation_config.decoder_start_token_id
default_generation_config = GenerationConfig()
for key, default_value in default_generation_config.__dict__.items():
if getattr(generation_config, key) == default_value:
setattr(generation_config, key, getattr(self.generation_config, key))
else:
model_kwargs = kwargs

Expand Down
Loading