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

Implement __new__ on Trainers

Browse Source 
__new__ allows Trainers to be initialized in the normal python
fashion.
This commit is contained in:
Bjarte Johansen
2020-02-07 10:55:37 +01:00
parent 0e5d81b400
commit 4971e9608d
6 changed files with 79 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -9,7 +9,7 @@ pub struct Trainer {
pub trainer: Container<dyn tk::tokenizer::Trainer>,
}
#[pyclass]
#[pyclass(extends=Trainer)]
pub struct BpeTrainer {}
#[pymethods]
impl BpeTrainer {
@ -17,9 +17,9 @@ impl BpeTrainer {
/// --
///
/// Create a new BpeTrainer with the given configuration
#[staticmethod]
#[new]
#[args(kwargs = "**")]
pub fn new(kwargs: Option<&PyDict>) -> PyResult<Trainer> {
pub fn new(obj: &PyRawObject, kwargs: Option<&PyDict>) -> PyResult<()> {
let mut builder = tk::models::bpe::BpeTrainer::builder();
if let Some(kwargs) = kwargs {
for (key, val) in kwargs {
@ -49,10 +49,9 @@ impl BpeTrainer {
};
}
}
Ok(Trainer {
Ok(obj.init(Trainer {
trainer: Container::Owned(Box::new(builder.build())),
})
}))
}
}
@ -64,9 +63,9 @@ impl WordPieceTrainer {
/// --
///
/// Create a new BpeTrainer with the given configuration
#[staticmethod]
#[new]
#[args(kwargs = "**")]
pub fn new(kwargs: Option<&PyDict>) -> PyResult<Trainer> {
pub fn new(obj: &PyRawObject, kwargs: Option<&PyDict>) -> PyResult<()> {
let mut builder = tk::models::wordpiece::WordPieceTrainer::builder();
if let Some(kwargs) = kwargs {
for (key, val) in kwargs {
@ -97,8 +96,8 @@ impl WordPieceTrainer {
}
}
Ok(Trainer {
Ok(obj.init(Trainer {
trainer: Container::Owned(Box::new(builder.build())),
})
}))
}
}