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

Fix BPE saving (u32 => String)

Browse Source
This commit is contained in:
Anthony MOI
2019-12-31 23:15:10 -05:00
parent 90dfdc715d
commit f3aef0e4e6
1 changed files with 6 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tokenizers/src/models/bpe/model.rs
View File

@ -276,16 +276,18 @@ impl Model for BPE {
.iter() .iter()
.collect(); .collect();
let mut merges_file = File::create(&merges_path)?; let mut merges_file = File::create(&merges_path)?;
let mut merges: Vec<(Pair, u32)> = self let mut merges: Vec<(&Pair, &u32)> = self
.merges .merges
.iter() .iter()
.map(|(pair, (rank, _))| (*pair, *rank)) .map(|(pair, (rank, _))| (pair, rank))
.collect(); .collect();
merges.sort_unstable_by_key(|k| k.1); merges.sort_unstable_by_key(|k| *k.1);
merges_file.write_all( merges_file.write_all(
&merges &merges
.into_iter() .into_iter()
.map(|(pair, _)| format!("{} {}\n", pair.0, pair.1).into_bytes()) .map(|(pair, _)| {
format!("{} {}\n", self.vocab_r[&pair.0], self.vocab_r[&pair.1]).into_bytes()
})
.flatten() .flatten()
.collect::<Vec<_>>()[..], .collect::<Vec<_>>()[..],
)?; )?;