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

Option to skip special tokens while decoding

Browse Source
This commit is contained in:
Anthony MOI
2019-12-19 20:03:02 -05:00
parent a8d68d516d
commit b7040e0412
3 changed files with 22 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
bindings/python/src/tokenizer.rs
View File

@ -192,12 +192,16 @@ impl Tokenizer {
.into()
}
fn decode(&self, ids: Vec<u32>) -> PyResult<String> {
ToPyResult(self.tokenizer.decode(ids)).into()
fn decode(&self, ids: Vec<u32>, skip_special_tokens: bool) -> PyResult<String> {
ToPyResult(self.tokenizer.decode(ids, skip_special_tokens)).into()
}
fn decode_batch(&self, sentences: Vec<Vec<u32>>) -> PyResult<Vec<String>> {
ToPyResult(self.tokenizer.decode_batch(sentences)).into()
fn decode_batch(
&self,
sentences: Vec<Vec<u32>>,
skip_special_tokens: bool,
) -> PyResult<Vec<String>> {
ToPyResult(self.tokenizer.decode_batch(sentences, skip_special_tokens)).into()
}
fn token_to_id(&self, token: &str) -> Option<u32> {

2
tokenizers/src/cli.rs
View File

@ -53,7 +53,7 @@ fn shell(matches: &ArgMatches) -> Result<()> {
println!("Offsets:\t{:?}", encoded.get_offsets());
println!(
"Decoded:\t{}",
tokenizer.decode(encoded.get_ids().to_vec()).unwrap()
tokenizer.decode(encoded.get_ids().to_vec(), true).unwrap()
);
println!("Tokenized in {:?}", elapsed);
}

18
tokenizers/src/tokenizer/mod.rs
View File

@ -342,15 +342,19 @@ impl Tokenizer {
}
/// Decode the given ids, back to a String
pub fn decode(&self, ids: Vec<u32>) -> Result<String> {
pub fn decode(&self, ids: Vec<u32>, skip_special_tokens: bool) -> Result<String> {
let tokens = ids
.into_iter()
.map(|id| {
if let Some(token) = self.added_tokens_r.get(&id) {
let token = if let Some(token) = self.added_tokens_r.get(&id) {
Some(token.content.to_owned())
} else {
self.model.id_to_token(id)
}
};
token.filter(|token| {
!skip_special_tokens || !self.special_tokens.contains_key(token)
})
})
.filter(|token| token.is_some())
.map(|id| id.unwrap())
@ -364,10 +368,14 @@ impl Tokenizer {
}
/// Decode all sentences in parallel
pub fn decode_batch(&self, sentences: Vec<Vec<u32>>) -> Result<Vec<String>> {
pub fn decode_batch(
&self,
sentences: Vec<Vec<u32>>,
skip_special_tokens: bool,
) -> Result<Vec<String>> {
sentences
.into_par_iter()
.map(|sentence| self.decode(sentence))
.map(|sentence| self.decode(sentence, skip_special_tokens))
.collect()
}