Skip to content

Commit 0547793

Browse files
chore(trainer): add and improve trainer signature
1 parent 7ccf93c commit 0547793

File tree

1 file changed

+8
-3
lines changed

1 file changed

+8
-3
lines changed

bindings/python/src/trainers.rs

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -312,7 +312,10 @@ impl PyBpeTrainer {
312312
}
313313

314314
#[new]
315-
#[pyo3(signature = (**kwargs), text_signature = None)]
315+
#[pyo3(
316+
signature = (**kwargs),
317+
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={})"
318+
)]
316319
pub fn new(kwargs: Option<&Bound<'_, PyDict>>) -> PyResult<(Self, PyTrainer)> {
317320
let mut builder = tk::models::bpe::BpeTrainer::builder();
318321
if let Some(kwargs) = kwargs {
@@ -518,7 +521,7 @@ impl PyWordPieceTrainer {
518521
#[new]
519522
#[pyo3(
520523
signature = (** kwargs),
521-
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)"
524+
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)"
522525
)]
523526
pub fn new(kwargs: Option<&Bound<'_, PyDict>>) -> PyResult<(Self, PyTrainer)> {
524527
let mut builder = tk::models::wordpiece::WordPieceTrainer::builder();
@@ -659,7 +662,9 @@ impl PyWordLevelTrainer {
659662
}
660663

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

0 commit comments

Comments
 (0)