Fix bnf.

Ai00-X · Jun 14, 2024 · 2b1150c · 2b1150c
1 parent 370092a
commit 2b1150c
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 3 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/ai00-core/Cargo.toml b/crates/ai00-core/Cargo.toml
@@ -16,7 +16,7 @@ bytemuck = "1"
 cbor4ii = { version = "0.3.2", features = ["serde1"] }
 fastrand = "2"
 half = "2.4"
-kbnf = "0.1.2"
+kbnf = "0.1.3"
 qp-trie = "0.8"
 rustc-hash = "1.1.0"
 uuid = { version = "1.8.0", features = ["serde", "v4"] }

diff --git a/crates/ai00-core/src/sampler/bnf.rs b/crates/ai00-core/src/sampler/bnf.rs
@@ -15,12 +15,14 @@ impl BnfSampler {
  .token_index_to_bytes()
  .iter()
  .enumerate()
+ .filter(|(_, v)| !v.is_empty())
  .map(|(k, v)| (k as u32, Token(v.clone().into_boxed_slice())))
  .collect();
  let strings = tokenizer
  .token_index_to_bytes()
  .iter()
  .enumerate()
+ .filter(|(_, v)| !v.is_empty())
  .map(|(k, v)| (k as u32, String::from_utf8_lossy(v).to_string()))
  .collect();
  let vocab = Vocabulary::new(tokens, strings)?;
@@ -31,6 +33,7 @@ impl BnfSampler {
 
 impl Transformer for BnfSampler {
  fn transform(&self, output: &mut [f32]) {
+ let output = &mut output[..self.0.vocab().vocab_size()];
  self.0.mask_logits(output).expect("bnf transform error")
  }