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

Merge f88153065f into 9a71d13688

Browse Source
This commit is contained in:
adamnemecek 2025-11-13 20:46:53 +02:00 committed by GitHub
commit e23096e827
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 18 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
rustbpe/src/lib.rs
View File

@ -27,6 +27,13 @@ pub struct Tokenizer {
// ------------------------ internal helpers ------------------------
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(i8)]
enum Delta {
Rem = -1,
Ins = 1,
}
#[derive(Clone, Debug)]
struct Word {
ids: Vec<u32>,
@ -48,7 +55,7 @@ impl Word {
/// -1 for removed pairs, +1 for newly created pairs.
///
/// 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 n = self.ids.len();
if n < 2 {
@ -56,7 +63,7 @@ impl Word {
}
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;
while i < n {
@ -66,13 +73,13 @@ impl Word {
// remove old pairs
if let Some(x) = left {
deltas.push(((x, a), -1));
deltas.push(((x, new_id), 1));
deltas.push(((x, a), Delta::Rem));
deltas.push(((x, new_id), Delta::Ins));
}
deltas.push(((a, b), -1));
deltas.push(((a, b), Delta::Rem));
if let Some(y) = right {
deltas.push(((b, y), -1));
deltas.push(((new_id, y), 1));
deltas.push(((b, y), Delta::Rem));
deltas.push(((new_id, y), Delta::Ins));
}
// write merged token
@ -112,12 +119,10 @@ impl PartialOrd for MergeJob {
impl Ord for MergeJob {
fn cmp(&self, other: &Self) -> Ordering {
// Max-heap by count; tie-break to ascending pair order (deterministic)
if self.count != other.count {
self.count.cmp(&other.count)
} else {
self.count.cmp(&other.count).then_with(||
// ascending order on the pair when counts tie
other.pair.cmp(&self.pair)
}
)
}
}
@ -217,10 +222,10 @@ impl Tokenizer {
let changes = words[word_idx].merge_pair(top.pair, new_id);
// Update global pair counts based on this word's count
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 {
*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);
}
}