mii/tokenizers
1
0
Fork 0
mirror of https://github.com/mii443/tokenizers.git synced 2025-08-23 00:35:35 +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()
.collect();
let mut merges_file = File::create(&merges_path)?;
let mut merges: Vec<(Pair, u32)> = self
let mut merges: Vec<(&Pair, &u32)> = self
.merges
.iter()
.map(|(pair, (rank, _))| (*pair, *rank))
.map(|(pair, (rank, _))| (pair, rank))
.collect();
merges.sort_unstable_by_key(|k| k.1);
merges.sort_unstable_by_key(|k| *k.1);
merges_file.write_all(
&merges
.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()
.collect::<Vec<_>>()[..],
)?;