mii/tokenizers
1
0
Fork 0
mirror of https://github.com/mii443/tokenizers.git synced 2025-08-23 16:49:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

simplify initialization of BpeTrainer

Browse Source
This commit is contained in:
epwalsh
2019-12-23 14:38:54 -08:00
committed by MOI Anthony
parent fab1d4cabc
commit c0ed873c4d
2 changed files with 22 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
bindings/python/src/trainers.rs
View File

@ -20,18 +20,18 @@ impl BpeTrainer {
#[staticmethod]
#[args(kwargs = "**")]
pub fn new(kwargs: Option<&PyDict>) -> PyResult<Trainer> {
let mut config: tk::models::bpe::BpeTrainerConfig = Default::default();
let mut trainer = tk::models::bpe::BpeTrainer::default();
if let Some(kwargs) = kwargs {
for (key, val) in kwargs {
let key: &str = key.extract()?;
match key {
"vocab_size" => {
let size: usize = val.extract()?;
config.set_vocab_size(size);
trainer.vocab_size = size;
}
"min_frequency" => {
let freq: u32 = val.extract()?;
config.set_min_frequency(freq);
trainer.min_frequency = freq;
}
_ => println!("Ignored unknown kwargs option {}", key),
};
@ -39,7 +39,7 @@ impl BpeTrainer {
}
Ok(Trainer {
trainer: Container::Owned(Box::new(tk::models::bpe::BpeTrainer::new(config))),
trainer: Container::Owned(Box::new(trainer)),
})
}
}