-
Notifications
You must be signed in to change notification settings - Fork 30.1k
Add tokenizer_kwargs
argument to the text generation pipeline
#40364
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Add tokenizer_kwargs
argument to the text generation pipeline
#40364
Conversation
The test failure seems to be an unrelated flake:
Pushing an empty commit to re-run the CI. |
A disjoint set of tests have failed in the re-run. |
@Rocketknight1 Please review this PR when you have a chance. The CI failures seem to be caused by unrelated, flaky tests. |
92b4d49
to
2fe0979
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This LGTM! cc @gante just in case you have opinions about the max_length
generate kwarg clash.
Also @Joshua-Chin you may need to rebase to fix some conflicts before we can merge the PR! That should also clear up the CI issues. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
One question about variable names, otherwise lgtm :)
@@ -285,6 +289,9 @@ def __call__(self, text_inputs, **kwargs): | |||
- `None` : default strategy where nothing in particular happens | |||
- `"hole"`: Truncates left of input, and leaves a gap wide enough to let generation happen (might | |||
truncate a lot of the prompt and not suitable when generation exceed the model capacity) | |||
tokenizer_kwargs (`dict`, *optional*): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
perhaps tokenizer_encode_kwargs
? There are also kwargs used at decode time, and we don't want to mix the two
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@gante I updated the argument to tokenizer_encode_kwargs
. Please take another look when you have a chance.
2fe0979
to
e484dbb
Compare
The CI is currently failing because of the following test, added by a recently merged change (HunYuan opensource #39606):
|
d80a814
to
f1d1dc1
Compare
What does this PR do?
This PR adds a
tokenizer_kwargs
argument to theTextGenerationPipeline
, allowing users to pass arbitrary arguments to the tokenizer during preprocessing. In particular, this lets users set chat template arguments, such as theenable_thinking
flag for Qwen3 or SmolLM3.Before submitting
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.