tacit/nanochat
1
0
Fork 0
mirror of https://github.com/karpathy/nanochat.git synced 2025-12-06 04:12:13 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

use a Delta enum

Browse Source
This commit is contained in:
adamnemecek 2025-10-14 16:36:34 -07:00
parent a13c9ca6ae
commit f88153065f
1 changed files with 16 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
rustbpe/src/lib.rs
View File

@ -27,6 +27,13 @@ pub struct Tokenizer {
// ------------------------ internal helpers ------------------------ // ------------------------ internal helpers ------------------------
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(i8)]
enum Delta {
Rem = -1,
Ins = 1,
}
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
struct Word { struct Word {
ids: Vec<u32>, ids: Vec<u32>,
@ -48,7 +55,7 @@ impl Word {
/// -1 for removed pairs, +1 for newly created pairs. /// -1 for removed pairs, +1 for newly created pairs.
/// ///
/// NOTE: this version deliberately avoids a HashMap in the hot loop. /// NOTE: this version deliberately avoids a HashMap in the hot loop.
fn merge_pair(&mut self, pair: Pair, new_id: u32) -> Vec<(Pair, i32)> { fn merge_pair(&mut self, pair: Pair, new_id: u32) -> Vec<(Pair, Delta)> {
let (a, b) = pair; let (a, b) = pair;
let n = self.ids.len(); let n = self.ids.len();
if n < 2 { if n < 2 {
@ -56,7 +63,7 @@ impl Word {
} }
let mut out: Vec<u32> = Vec::with_capacity(n); let mut out: Vec<u32> = Vec::with_capacity(n);
let mut deltas: Vec<(Pair, i32)> = Vec::with_capacity(6); let mut deltas: Vec<(Pair, Delta)> = Vec::with_capacity(6);
let mut i = 0; let mut i = 0;
while i < n { while i < n {
@ -66,13 +73,13 @@ impl Word {
// remove old pairs // remove old pairs
if let Some(x) = left { if let Some(x) = left {
deltas.push(((x, a), -1)); deltas.push(((x, a), Delta::Rem));
deltas.push(((x, new_id), 1)); deltas.push(((x, new_id), Delta::Ins));
} }
deltas.push(((a, b), -1)); deltas.push(((a, b), Delta::Rem));
if let Some(y) = right { if let Some(y) = right {
deltas.push(((b, y), -1)); deltas.push(((b, y), Delta::Rem));
deltas.push(((new_id, y), 1)); deltas.push(((new_id, y), Delta::Ins));
} }
// write merged token // write merged token
@ -215,10 +222,10 @@ impl Tokenizer {
let changes = words[word_idx].merge_pair(top.pair, new_id); let changes = words[word_idx].merge_pair(top.pair, new_id);
// Update global pair counts based on this word's count // Update global pair counts based on this word's count
for (pair, delta) in changes { for (pair, delta) in changes {
let delta_total = delta * counts[word_idx]; let delta_total = (delta as i32) * counts[word_idx];
if delta_total != 0 { if delta_total != 0 {
*pair_counts.entry(pair).or_default() += delta_total; *pair_counts.entry(pair).or_default() += delta_total;
if delta > 0 { if delta == Delta::Ins {
local_pos_updates.entry(pair).or_default().insert(word_idx); local_pos_updates.entry(pair).or_default().insert(word_idx);
} }
} }