stopping criteria for TextGenerationPipeline

[geronimi73]

Feature request

pass stopping criteria or string to TextGenerationPipeline

Motivation

it does not exist, have not found any way to do it at least, but would be very useful

Your contribution

none

[ArthurZucker]

Hey! You can pass generation_kwargs to the pipeline, which are usually used for stopping criteria such as max_length, max_new_tokens, max_time . You can also pass some stopping_criteria argument to the generate function using `generation_kwargs = {stopping_criteria = StoppingCriteriaList: [MaxTimeCriteria(32)] }.

from transformers import pipeline, StoppingCriteriaList, MaxTimeCriteria

# Initialize the text generation pipeline
generator = pipeline("text-generation")

# Define the stopping criteria using MaxTimeCriteria
stopping_criteria = StoppingCriteriaList([MaxTimeCriteria(32)])

# Define the generation_kwargs with stopping criteria
generation_kwargs = {
    "max_length": 100,  # Maximum length of the generated text
    "max_new_tokens": 10,  # Maximum number of new tokens to generate
    "generation_kwargs": {"stopping_criteria": stopping_criteria}  # Add stopping criteria to generation_kwargs
}

# Pass the generation_kwargs to the pipeline
generated_text = generator(
    "Hey!  How are you able.",
    **generation_kwargs
)

# Print the generated text
print(generated_text[0]["generated_text"])
>>> Hey!  How are you able.  Do you have a job or do you have

[ArthurZucker]

This should probably be added to the documentation!

[LysandreJik]

cc @MKhalusova who is currently working on the generate docs!

[MKhalusova]

The max_length and max_new_tokens are mentioned in nearly every doc on text generation:

• https://huggingface.co/docs/transformers/llm_tutorial#generated-output-is-too-shortlong
• https://huggingface.co/docs/transformers/generation_strategies#customize-text-generation
• https://huggingface.co/docs/transformers/main_classes/text_generation
• upcoming LLM prompting doc as well

However, I see a couple of issues:

• the MaxTimeCriteria is much less discoverable, as it is only mentioned in the Utilities for Generation.
• A more critical issue from my point of view is the lack of connection between the pipeline documentation and the text generation docs. The TextGenerationPipeline API doc lists parameters that control how pipeline is instantiated, but doesn't link to text generation docs.

I can address these.

[Tanman2001]

Is this a duplicate of #17562?

If so, it seems the only reason that older issue is still open is because of the missing documentation. If the documentation is fixed then both can be closed.

[peng-yiwen]

According to my experiments, setting

generation_kwargs = {
    "max_new_tokens": 1000,  # Maximum number of new tokens to generate
    "generation_kwargs": {"stopping_criteria": stopping_criteria}  # Add stopping criteria to generation_kwargs
}

is not feasible for passing a stopping criteria into a pipeline. Instead, if we let "stopping_criteria" out of "generation_kwargs", it then works well. As below:

generation_kwargs = {
    "max_new_tokens": 1000,  # Maximum number of new tokens to generate
    "stopping_criteria": stopping_criteria,
}