diff --git a/bindings/python/tests/bindings/test_models.py b/bindings/python/tests/bindings/test_models.py index c6a50ce86..063698384 100644 --- a/bindings/python/tests/bindings/test_models.py +++ b/bindings/python/tests/bindings/test_models.py @@ -69,6 +69,10 @@ def test_can_modify(self): model.byte_fallback = True assert model.byte_fallback == True + def test_dropout_zero(self): + model = BPE(dropout=0.0) + assert model.dropout == 0.0 + class TestWordPiece: def test_instantiate(self, bert_files): diff --git a/tokenizers/Cargo.toml b/tokenizers/Cargo.toml index 07cc85d1b..038486880 100644 --- a/tokenizers/Cargo.toml +++ b/tokenizers/Cargo.toml @@ -75,6 +75,8 @@ unstable_wasm = ["fancy-regex", "getrandom/js"] criterion = "0.5" tempfile = "3.10" assert_approx_eq = "1.1" +tracing = "0.1" +tracing-subscriber = "0.3.18" [profile.release] lto = "fat" diff --git a/tokenizers/src/models/bpe/mod.rs b/tokenizers/src/models/bpe/mod.rs index 51e7b6ca0..f0d40b2df 100644 --- a/tokenizers/src/models/bpe/mod.rs +++ b/tokenizers/src/models/bpe/mod.rs @@ -31,7 +31,7 @@ pub enum Error { #[error("Unk token `{0}` not found in the vocabulary")] UnkTokenOutOfVocabulary(String), /// Dropout not between 0 and 1. - #[error("Dropout should be between 0 and 1")] + #[error("Dropout should be between 0 and 1, inclusive")] InvalidDropout, } diff --git a/tokenizers/src/models/bpe/model.rs b/tokenizers/src/models/bpe/model.rs index 618f42b47..8d22ab52d 100644 --- a/tokenizers/src/models/bpe/model.rs +++ b/tokenizers/src/models/bpe/model.rs @@ -136,7 +136,7 @@ impl BpeBuilder { pub fn build(mut self) -> Result { // Validate dropout. if let Some(p) = self.config.dropout { - if p <= 0.0 || p > 1.0 { + if !(0.0..=1.0).contains(&p) { return Err(Error::InvalidDropout.into()); } } @@ -214,7 +214,7 @@ pub struct BPE { pub(crate) merges: MergeMap, /// Contains the cache for optimizing the encoding step. cache: Option>, - /// Dropout probability for merges. 0 = no dropout is the default. At 1.0, tokenization will + /// Dropout probability for merges. 0.0 = no dropout is the default. At 1.0, tokenization will /// perform no merges, so the result will just be characters. pub dropout: Option, /// The unknown token to be used when we encounter an unknown char @@ -493,7 +493,7 @@ impl Model for BPE { return Ok(vec![]); } - if self.dropout.is_none() { + if self.dropout.is_none() || self.dropout == Some(0.0) { self.tokenize_with_cache(sequence) } else { let word = self.merge_word(sequence)?; @@ -685,6 +685,11 @@ mod tests { let tokens = bpe.tokenize("unrelated").unwrap(); assert_eq!(tokens, vec![Token::new(15u32, "unrelated".into(), (0, 9))]); + // With dropout = 0.0 (equivalent to dropout == none) + bpe.dropout = Some(0.0); + let tokens = bpe.tokenize("unrelated").unwrap(); + assert_eq!(tokens, vec![Token::new(15u32, "unrelated".into(), (0, 9))]); + // Now set dropout to 1.0. Result should be no merges performed. bpe.dropout = Some(1.0); let tokens = bpe.tokenize("unrelated").unwrap(); @@ -739,6 +744,13 @@ mod tests { assert_eq!(bpe.vocab.get("ab").unwrap(), &3u32); } + #[test] + // Ensure BPEBuilder with dropout = 0.0 doesn't error + fn test_bpe_with_dropout_0() { + let bpe = BPE::builder().dropout(0.0).build().unwrap(); + assert_eq!(bpe.dropout, Some(0.0)); + } + #[test] // Ensure `BPE::from_file` works as expected. fn test_bpe_with_continuing_subword_prefix() { diff --git a/tokenizers/src/tokenizer/serialization.rs b/tokenizers/src/tokenizer/serialization.rs index f08783dd3..c3ad5b410 100644 --- a/tokenizers/src/tokenizer/serialization.rs +++ b/tokenizers/src/tokenizer/serialization.rs @@ -155,17 +155,15 @@ where for token in &tokens { // Warn the user if the id is different than expected let received_id = tokenizer.token_to_id(&token.token.content); - if received_id != Some(token.id) { - warn!( - "Warning: Token '{}' was expected to have ID '{}' but was given ID '{}'", - token.token.content, - token.id, - if let Some(rid) = received_id { + if let Some(rid) = received_id { + if rid != token.id { + warn!( + "Warning: Token '{}' was expected to have ID '{}' but was given ID '{}'", + token.token.content, + token.id, rid.to_string() - } else { - "None".to_string() - } - ); + ); + } } } let added_tokens: Vec<_> = tokens.into_iter().map(|token| token.token).collect(); @@ -179,6 +177,7 @@ where mod tests { use crate::tokenizer::Tokenizer; use std::str::FromStr; + use tracing_subscriber::fmt; #[test] fn test_deserialization_serialization_invariant() { @@ -233,4 +232,15 @@ mod tests { // It should be exactly the same as above assert_eq!(tok_str, tok_json); } + + #[cfg(feature = "http")] + #[test] + fn test_from_pretrained() { + fmt() + .with_max_level(tracing::Level::DEBUG) + .with_target(false) + .init(); + let _ = Tokenizer::from_pretrained("Qwen/Qwen2-7B-Instruct", None); + warn!("This should be the first warning"); + } } diff --git a/tokenizers/tests/serialization.rs b/tokenizers/tests/serialization.rs index 54fa9053d..4d51d4281 100644 --- a/tokenizers/tests/serialization.rs +++ b/tokenizers/tests/serialization.rs @@ -229,6 +229,21 @@ fn tokenizer() { assert_eq!(serde_json::to_string(&de).unwrap(), ser); } +#[test] +fn bpe_with_dropout_serde() { + let mut bpe = BPE::default(); + bpe.dropout = Some(0.1); + let ser = serde_json::to_string(&bpe).unwrap(); + let de = serde_json::from_str(&ser).unwrap(); + assert_eq!(bpe, de); + + // set dropout to 0.0 (which is analogous to None) and reserialize + bpe.dropout = Some(0.0); + let ser = serde_json::to_string(&bpe).unwrap(); + let de = serde_json::from_str(&ser).unwrap(); + assert_eq!(bpe, de); +} + #[test] fn test_deserialize_long_file() { let _tokenizer = Tokenizer::from_file("data/albert-base-v1-tokenizer.json").unwrap();