Python - Fix cases where str expected instead of AddedToken

This commit is contained in:
Anthony MOI
2020-03-25 19:22:37 -04:00
parent d25eb075c8
commit f8d54edcdd
4 changed files with 33 additions and 19 deletions

View File

@ -3,6 +3,7 @@ extern crate tokenizers as tk;
use pyo3::exceptions;
use pyo3::prelude::*;
use pyo3::types::*;
use pyo3::PyObjectProtocol;
use super::decoders::Decoder;
use super::encoding::Encoding;
@ -45,6 +46,19 @@ impl AddedToken {
Ok(())
}
}
#[pyproto]
impl PyObjectProtocol for AddedToken {
fn __str__(&'p self) -> PyResult<&'p str> {
Ok(&self.token.content)
}
fn __repr__(&self) -> PyResult<String> {
Ok(format!(
"AddedToken(\"{}\", rstrip={}, lstrip={}, single_word={})",
self.token.content, self.token.rstrip, self.token.lstrip, self.token.single_word
))
}
}
#[pyclass(dict)]
pub struct Tokenizer {