Add bindings for Trainer in Python

This commit is contained in:
Anthony MOI
2019-12-03 15:54:15 -05:00
parent 310a2af76b
commit eaafb22511
5 changed files with 74 additions and 1 deletions

View File

@ -7,6 +7,7 @@ use super::decoders::Decoder;
use super::models::Model;
use super::pre_tokenizers::PreTokenizer;
use super::token::Token;
use super::trainers::Trainer;
#[pyclass]
pub struct Tokenizer {
@ -97,5 +98,15 @@ impl Tokenizer {
fn id_to_token(&self, id: u32) -> Option<String> {
self.tokenizer.id_to_token(id)
}
fn train(&mut self, trainer: &Trainer, files: Vec<String>) -> PyResult<()> {
trainer.trainer.execute(|trainer| {
if let Err(e) = self.tokenizer.train(trainer, files) {
Err(exceptions::Exception::py_err(format!("{}", e)))
} else {
Ok(())
}
})
}
}