Skip to content

chore(trainer): add and improve trainer signature #1838

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.

Sign up for GitHub

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

Merged
merged 7 commits into from
Sep 16, 2025
+27 −4
Merged

chore(trainer): add and improve trainer signature #1838

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
98ac665
chore(trainers): add __init__ to fix python type check errors
shenxiangzhuang Jul 31, 2025
7ccf93c
restore
shenxiangzhuang Aug 28, 2025
0547793
chore(trainer): add and improve trainer signature
shenxiangzhuang Aug 29, 2025
3d77532
clean fix
shenxiangzhuang Aug 29, 2025
69a2814
Merge branch 'huggingface:main' into chore/fix_trainer_types
shenxiangzhuang Aug 29, 2025
e127013
Merge branch 'main' into chore/fix_trainer_types
ArthurZucker Aug 29, 2025
8da5806
chore(fmt): fix cargo fmt error
shenxiangzhuang Sep 12, 2025
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
17 changes: 17 additions & 0 deletions bindings/python/py_src/tokenizers/trainers/__init__.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,20 @@ class BpeTrainer(Trainer):
highly repetitive tokens like `======` for wikipedia

"""
def __init__(
self,
vocab_size=30000,
min_frequency=0,
show_progress=True,
special_tokens=[],
limit_alphabet=None,
initial_alphabet=[],
continuing_subword_prefix=None,
end_of_word_suffix=None,
max_token_length=None,
words={},
):
pass

class UnigramTrainer(Trainer):
"""
Expand Down Expand Up @@ -85,6 +99,7 @@ class UnigramTrainer(Trainer):
vocab_size=8000,
show_progress=True,
special_tokens=[],
initial_alphabet=[],
shrinking_factor=0.75,
unk_token=None,
max_piece_length=16,
Expand All @@ -109,6 +124,8 @@ class WordLevelTrainer(Trainer):
special_tokens (:obj:`List[Union[str, AddedToken]]`):
A list of special tokens the model should know of.
"""
def __init__(self, vocab_size=30000, min_frequency=0, show_progress=True, special_tokens=[]):
pass

class WordPieceTrainer(Trainer):
"""
Expand Down
14 changes: 10 additions & 4 deletions bindings/python/src/trainers.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -312,7 +312,10 @@ impl PyBpeTrainer {
}

#[new]
#[pyo3(signature = (**kwargs), text_signature = None)]
#[pyo3(
signature = (**kwargs),
text_signature = "(self, vocab_size=30000, min_frequency=0, show_progress=True, special_tokens=[], limit_alphabet=None, initial_alphabet=[], continuing_subword_prefix=None, end_of_word_suffix=None, max_token_length=None, words={})"
Copy link
Collaborator

Choose a reason for hiding this comment

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

Suggested change
text_signature = "(self, vocab_size=30000, min_frequency=0, show_progress=True, special_tokens=[], limit_alphabet=None, initial_alphabet=[], continuing_subword_prefix=None, end_of_word_suffix=None, max_token_length=None, words={})"
text_signature = "(self, vocab_size=30000, min_frequency=0, show_progress=True, special_tokens=None, limit_alphabet=None, initial_alphabet=None, continuing_subword_prefix=None, end_of_word_suffix=None, max_token_length=None, words=None)"

default to mutable is a mess in python

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Well, I also notice this. There is a question, I find that text_signature in WordPieceTrainer already use []:

#[new]
    #[pyo3(
        signature = (** kwargs),
        text_signature = "(self, vocab_size=30000, min_frequency=0, show_progress=True, special_tokens=[], limit_alphabet=None, initial_alphabet= [],continuing_subword_prefix=\"##\", end_of_word_suffix=None)"
    )]
    pub fn new(kwargs: Option<&Bound<'_, PyDict>>) -> PyResult<(Self, PyTrainer)> {

And, I tested it in python, and find that the default values are just same as the text_signature:

In [5]: import tokenizers; tokenizers.trainers.WordPieceTrainer()
Out[5]: WordPieceTrainer(WordPieceTrainer(bpe_trainer=BpeTrainer(min_frequency=0, vocab_size=30000, show_progress=True, special_tokens=[], limit_alphabet=None, initial_alphabet=[], continuing_subword_prefix="##", end_of_word_suffix=None, max_token_length=None, words={})))

And the doc for WordPieceTrainer is consistent with the real default values:

Init signature:
tokenizers.trainers.WordPieceTrainer(
    self,
    vocab_size=30000,
    min_frequency=0,
    show_progress=True,
    special_tokens=[],
    limit_alphabet=None,
    initial_alphabet=[],
    continuing_subword_prefix='##',
    end_of_word_suffix=None,
)
Docstring:     
Trainer capable of training a WordPiece model

...(ignored)

So maybe we should keep the text_signature same the real default value to avoid introduce break changes and confuse the users(due to the inconsistence) too.

Copy link
Collaborator

Choose a reason for hiding this comment

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

Ok no worries then

)]
pub fn new(kwargs: Option<&Bound<'_, PyDict>>) -> PyResult<(Self, PyTrainer)> {
let mut builder = tk::models::bpe::BpeTrainer::builder();
if let Some(kwargs) = kwargs {
Expand Down Expand Up @@ -518,7 +521,7 @@ impl PyWordPieceTrainer {
#[new]
#[pyo3(
signature = (** kwargs),
text_signature = "(self, vocab_size=30000, min_frequency=0, show_progress=True, special_tokens=[], limit_alphabet=None, initial_alphabet= [],continuing_subword_prefix=\"##\", end_of_word_suffix=None)"
text_signature = "(self, vocab_size=30000, min_frequency=0, show_progress=True, special_tokens=[], limit_alphabet=None, initial_alphabet=[], continuing_subword_prefix=\"##\", end_of_word_suffix=None)"
)]
pub fn new(kwargs: Option<&Bound<'_, PyDict>>) -> PyResult<(Self, PyTrainer)> {
let mut builder = tk::models::wordpiece::WordPieceTrainer::builder();
Expand Down Expand Up @@ -659,7 +662,10 @@ impl PyWordLevelTrainer {
}

#[new]
#[pyo3(signature = (**kwargs), text_signature = None)]
#[pyo3(
signature = (**kwargs),
text_signature = "(self, vocab_size=30000, min_frequency=0, show_progress=True, special_tokens=[])"
)]
pub fn new(kwargs: Option<&Bound<'_, PyDict>>) -> PyResult<(Self, PyTrainer)> {
let mut builder = tk::models::wordlevel::WordLevelTrainer::builder();

Expand Down Expand Up @@ -826,7 +832,7 @@ impl PyUnigramTrainer {
#[new]
#[pyo3(
signature = (**kwargs),
text_signature = "(self, vocab_size=8000, show_progress=True, special_tokens=[], shrinking_factor=0.75, unk_token=None, max_piece_length=16, n_sub_iterations=2)"
text_signature = "(self, vocab_size=8000, show_progress=True, special_tokens=[], initial_alphabet=[], shrinking_factor=0.75, unk_token=None, max_piece_length=16, n_sub_iterations=2)"
)]
pub fn new(kwargs: Option<Bound<'_, PyDict>>) -> PyResult<(Self, PyTrainer)> {
let mut builder = tk::models::unigram::UnigramTrainer::builder();
Expand Down
Loading