Python - Add decoders

This commit is contained in:
Anthony MOI
2019-11-22 21:08:57 -05:00
parent e44f52024c
commit 8fbe3c2662
3 changed files with 44 additions and 2 deletions

View File

@@ -1,3 +1,4 @@
mod decoders;
mod models;
mod pre_tokenizers;
mod token;
@@ -9,7 +10,7 @@ use pyo3::wrap_pymodule;
/// Models Module
#[pymodule]
pub fn models(_py: Python, m: &PyModule) -> PyResult<()> {
fn models(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_class::<models::Model>()?;
m.add_class::<models::BPE>()?;
Ok(())
@@ -17,16 +18,24 @@ pub fn models(_py: Python, m: &PyModule) -> PyResult<()> {
/// PreTokenizers Module
#[pymodule]
pub fn pre_tokenizers(_py: Python, m: &PyModule) -> PyResult<()> {
fn pre_tokenizers(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_class::<pre_tokenizers::ByteLevel>()?;
Ok(())
}
/// Decoders Module
#[pymodule]
fn decoders(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_class::<decoders::ByteLevel>()?;
Ok(())
}
/// Tokenizers Module
#[pymodule]
fn tokenizers(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_class::<tokenizer::Tokenizer>()?;
m.add_wrapped(wrap_pymodule!(models))?;
m.add_wrapped(wrap_pymodule!(pre_tokenizers))?;
m.add_wrapped(wrap_pymodule!(decoders))?;
Ok(())
}