diff --git a/tokenizers/tk-serialize/src/from_json/model.rs b/tokenizers/tk-serialize/src/from_json/model.rs index ee400d21b2..ffe29eabf0 100644 --- a/tokenizers/tk-serialize/src/from_json/model.rs +++ b/tokenizers/tk-serialize/src/from_json/model.rs @@ -55,6 +55,10 @@ pub(super) fn read_bpe(cfg: &Json<'_>) -> Result<(Vocab, Merges, BpeConfig)> { .field("ignore_merges") .and_then(Json::as_bool) .unwrap_or_default(), + cache_capacity: cfg + .field("cache_capacity") + .and_then(Json::as_usize) + .unwrap_or_else(|| BpeConfig::default().cache_capacity), byte_level, ..BpeConfig::default() }; diff --git a/tokenizers/tk-serialize/src/from_json/tests.rs b/tokenizers/tk-serialize/src/from_json/tests.rs index 6a332c6980..9d7f266614 100644 --- a/tokenizers/tk-serialize/src/from_json/tests.rs +++ b/tokenizers/tk-serialize/src/from_json/tests.rs @@ -103,6 +103,21 @@ fn every_component_encodes_what_it_should() { } } +#[test] +fn bpe_cache_capacity_is_configurable() { + let explicit = Json::parse( + r#"{"type": "BPE", "byte_level": false, "cache_capacity": 0, + "vocab": {"a": 0, "b": 1, "ab": 2}, "merges": [["a", "b"]]}"#, + ) + .expect("valid BPE JSON"); + let (_, _, options) = read_bpe(&explicit).expect("the BPE reads"); + assert_eq!(options.cache_capacity, 0); + + let absent = Json::parse(TINY_BPE).expect("valid BPE JSON"); + let (_, _, options) = read_bpe(&absent).expect("the BPE reads"); + assert_eq!(options.cache_capacity, BpeConfig::default().cache_capacity); +} + /// `(slot, json, in the message)`: a legacy or malformed shape, and what the refusal has to name so /// the caller knows what to convert. Inferring a model kind from its keys, rewriting `"a b"` merges, /// and folding a `Metaspace` or `ByteLevel` pre-tokenizer are all tk-convert's job now.