use std::cell::RefCell;
use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap};
use std::fmt;
use std::sync::Mutex;
use std::sync::atomic::{AtomicUsize, Ordering};
use daachorse::{DoubleArrayAhoCorasick, DoubleArrayAhoCorasickBuilder};
use serde::Deserialize;
use serde_json::Value;
type TokenId = u32;
type ParsedMergeMap = HashMap<(u32, u32), (u32, u32)>;
type Vocab = HashMap<String, u32>;
const EMPTY_KEY: u64 = u64::MAX;
const INVALID_TOKEN: u32 = u32::MAX;
use crate::byte_level::BYTE_TO_CHAR;
#[inline(always)]
const fn pack_pair(t1: u32, t2: u32) -> u64 {
(t1 as u64) << 32 | t2 as u64
}
#[inline(always)]
const fn fx_hash(key: u64) -> u64 {
key.wrapping_mul(0x517cc1b727220a95)
}
#[inline(always)]
fn fx_hash_bytes(bytes: &[u8]) -> u64 {
let mut h: u64 = bytes.len() as u64;
let mut i = 0;
while i + 8 <= bytes.len() {
let word = u64::from_ne_bytes(bytes[i..i + 8].try_into().unwrap());
h = h.wrapping_add(word).wrapping_mul(0x517cc1b727220a95);
i += 8;
}
while i < bytes.len() {
h = h
.wrapping_add(bytes[i] as u64)
.wrapping_mul(0x517cc1b727220a95);
i += 1;
}
h
}
struct FxBuildHasher;
impl std::hash::BuildHasher for FxBuildHasher {
type Hasher = FxStrHasher;
#[inline(always)]
fn build_hasher(&self) -> FxStrHasher {
FxStrHasher(0)
}
}
#[repr(transparent)]
struct FxStrHasher(u64);
impl std::hash::Hasher for FxStrHasher {
#[inline]
fn finish(&self) -> u64 {
self.0
}
#[inline]
fn write(&mut self, bytes: &[u8]) {
self.0 = fx_hash_bytes(bytes);
}
}
type FxHashMap<K, V> = HashMap<K, V, FxBuildHasher>;
const FLAT_CACHE_BITS: usize = 16;
const FLAT_CACHE_SIZE: usize = 1 << FLAT_CACHE_BITS;
const FLAT_CACHE_MASK: usize = FLAT_CACHE_SIZE - 1;
const EMPTY_SLOT: u64 = 0;
#[derive(Clone, Copy)]
#[repr(C)]
struct CountCacheSlot {
hash: u64,
count: u32,
key_len: u16,
key_offset: u32,
}
const FLAT_CACHE_MAX_LOAD: usize = FLAT_CACHE_SIZE * 3 / 4;
struct CountFlatCache {
bpe_id: usize,
slots: Vec<CountCacheSlot>,
key_pool: Vec<u8>,
entry_count: usize,
}
impl CountFlatCache {
fn new() -> Self {
Self {
bpe_id: 0,
slots: vec![
CountCacheSlot {
hash: EMPTY_SLOT,
count: 0,
key_len: 0,
key_offset: 0,
};
FLAT_CACHE_SIZE
],
key_pool: Vec::with_capacity(512 * 1024),
entry_count: 0,
}
}
fn clear(&mut self) {
for slot in &mut self.slots {
slot.hash = EMPTY_SLOT;
}
self.key_pool.clear();
self.entry_count = 0;
}
#[inline(always)]
fn hash_str(s: &str) -> u64 {
let h = fx_hash_bytes(s.as_bytes());
if h == EMPTY_SLOT { 1 } else { h }
}
#[inline(always)]
fn get(&self, key: &str) -> Option<usize> {
let hash = Self::hash_str(key);
let key_bytes = key.as_bytes();
let mut idx = hash as usize & FLAT_CACHE_MASK;
loop {
let slot = unsafe { self.slots.get_unchecked(idx) };
if slot.hash == hash {
let ks = slot.key_offset as usize;
let ke = ks + slot.key_len as usize;
if unsafe { self.key_pool.get_unchecked(ks..ke) } == key_bytes {
return Some(slot.count as usize);
}
}
if slot.hash == EMPTY_SLOT {
return None;
}
idx = (idx + 1) & FLAT_CACHE_MASK;
}
}
#[inline(always)]
fn insert(&mut self, key: &str, token_count: usize) {
if self.entry_count >= FLAT_CACHE_MAX_LOAD {
self.clear();
}
let hash = Self::hash_str(key);
let key_bytes = key.as_bytes();
let mut idx = hash as usize & FLAT_CACHE_MASK;
loop {
let slot = unsafe { self.slots.get_unchecked(idx) };
let h = slot.hash;
if h == EMPTY_SLOT {
let Ok(key_len) = u16::try_from(key_bytes.len()) else {
return;
};
let Ok(count) = u32::try_from(token_count) else {
return;
};
self.entry_count += 1;
let key_offset = self.key_pool.len() as u32;
self.key_pool.extend_from_slice(key_bytes);
let slot = unsafe { self.slots.get_unchecked_mut(idx) };
slot.hash = hash;
slot.count = count;
slot.key_offset = key_offset;
slot.key_len = key_len;
return;
}
if h == hash {
let ks = slot.key_offset as usize;
let ke = ks + slot.key_len as usize;
if unsafe { self.key_pool.get_unchecked(ks..ke) } == key_bytes {
let Ok(count) = u32::try_from(token_count) else {
return;
};
let slot = unsafe { self.slots.get_unchecked_mut(idx) };
slot.count = count;
return;
}
}
idx = (idx + 1) & FLAT_CACHE_MASK;
}
}
}
thread_local! {
static TL_BPE_COUNT_CACHE: RefCell<CountFlatCache> = RefCell::new(CountFlatCache::new());
static TL_FUSED_COUNT_CACHE: RefCell<CountFlatCache> = RefCell::new(CountFlatCache::new());
}
const CACHE_SHARDS: usize = 64;
const MAX_SHARD_ENTRIES: usize = 4096;
#[repr(transparent)]
struct CountSharedCache {
shards: Vec<Mutex<FxHashMap<String, usize>>>,
}
impl CountSharedCache {
fn new() -> Self {
Self {
shards: (0..CACHE_SHARDS)
.map(|_| Mutex::new(HashMap::with_hasher(FxBuildHasher)))
.collect(),
}
}
#[inline]
fn shard_index(key: &str) -> usize {
fx_hash_bytes(key.as_bytes()) as usize & (CACHE_SHARDS - 1)
}
#[inline]
fn get(&self, key: &str) -> Option<usize> {
let shard = self.shards[Self::shard_index(key)].lock().unwrap();
shard.get(key).copied()
}
fn insert(&self, key: String, value: usize) {
let idx = Self::shard_index(&key);
let mut shard = self.shards[idx].lock().unwrap();
if shard.len() >= MAX_SHARD_ENTRIES {
shard.clear();
}
shard.insert(key, value);
}
}
#[derive(Deserialize)]
struct RawBpe {
#[serde(default)]
vocab: Vocab,
#[serde(default)]
merges: Vec<Value>,
#[serde(default)]
ignore_merges: bool,
}
static BPE_ID_COUNTER: AtomicUsize = AtomicUsize::new(1);
#[derive(Clone, Copy, Eq, PartialEq)]
#[repr(C)]
struct MergeEntry {
key: u64,
val: u64,
}
impl MergeEntry {
#[inline(always)]
fn new(rank: u32, pos: u32, left_c: u32, right_c: u32) -> Self {
Self {
key: (rank as u64) << 32 | pos as u64,
val: (left_c as u64) << 32 | right_c as u64,
}
}
#[inline(always)]
fn pos(&self) -> u32 {
self.key as u32
}
#[inline(always)]
fn left_c(&self) -> u32 {
(self.val >> 32) as u32
}
#[inline(always)]
fn right_c(&self) -> u32 {
self.val as u32
}
}
impl Ord for MergeEntry {
#[inline(always)]
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
self.key.cmp(&other.key)
}
}
impl PartialOrd for MergeEntry {
#[inline(always)]
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
#[derive(Clone, Copy)]
struct MergeSymbol {
c: u32,
prev: i32,
next: i32,
}
struct MergeScratch {
symbols: Vec<MergeSymbol>,
heap: BinaryHeap<Reverse<MergeEntry>>,
}
impl MergeScratch {
const fn new() -> Self {
Self {
symbols: Vec::new(),
heap: BinaryHeap::new(),
}
}
}
thread_local! {
static TL_MERGE_SCRATCH: RefCell<MergeScratch> = const { RefCell::new(MergeScratch::new()) };
}
#[derive(Clone, Copy)]
#[repr(C)]
struct RankedMergeSlot {
key: u64,
rank: u32,
id: u32,
}
#[derive(Clone)]
struct RankedMergeMap {
mask: usize,
slots: Vec<RankedMergeSlot>,
}
impl RankedMergeMap {
fn from_parsed(parsed: &ParsedMergeMap) -> Self {
if parsed.is_empty() {
return Self {
mask: 0,
slots: Vec::new(),
};
}
let capacity = (parsed.len() << 1).next_power_of_two();
let mask = capacity - 1;
let mut slots = vec![
RankedMergeSlot {
key: EMPTY_KEY,
rank: 0,
id: 0
};
capacity
];
for (&(t1, t2), &(rank, merged_id)) in parsed {
let key = pack_pair(t1, t2);
let mut idx = fx_hash(key) as usize & mask;
loop {
if slots[idx].key == EMPTY_KEY {
slots[idx] = RankedMergeSlot {
key,
rank,
id: merged_id,
};
break;
}
idx = (idx + 1) & mask;
}
}
Self { mask, slots }
}
#[inline(always)]
fn get(&self, t1: u32, t2: u32) -> Option<(u32, u32)> {
if self.slots.is_empty() {
return None;
}
let key = pack_pair(t1, t2);
let mut idx = fx_hash(key) as usize & self.mask;
loop {
let slot = unsafe { self.slots.get_unchecked(idx) };
if slot.key == key {
return Some((slot.rank, slot.id));
}
if slot.key == EMPTY_KEY {
return None;
}
idx = (idx + 1) & self.mask;
}
}
}
#[derive(Clone)]
struct MergeAdjacency {
offsets: Vec<u32>,
data: Vec<(u32, u32, u32)>, }
impl MergeAdjacency {
fn from_parsed(parsed: &ParsedMergeMap, vocab_size: usize) -> Self {
let mut counts = vec![0u32; vocab_size];
for &(left, _right) in parsed.keys() {
counts[left as usize] += 1;
}
let mut offsets = Vec::with_capacity(vocab_size + 1);
offsets.push(0u32);
let mut running = 0u32;
for &c in &counts {
running += c;
offsets.push(running);
}
let mut data = vec![(0u32, 0u32, 0u32); running as usize];
let mut write_pos = offsets[..vocab_size].to_vec();
for (&(left, right), &(rank, merged_id)) in parsed {
let idx = write_pos[left as usize] as usize;
data[idx] = (right, rank, merged_id);
write_pos[left as usize] += 1;
}
for i in 0..vocab_size {
let start = offsets[i] as usize;
let end = offsets[i + 1] as usize;
data[start..end].sort_unstable_by_key(|&(neighbor, _, _)| neighbor);
}
Self { offsets, data }
}
#[inline(always)]
fn get(&self, left: u32, right: u32) -> Option<(u32, u32)> {
let start = unsafe { *self.offsets.get_unchecked(left as usize) } as usize;
let end = unsafe { *self.offsets.get_unchecked(left as usize + 1) } as usize;
let slice = unsafe { self.data.get_unchecked(start..end) };
match slice.binary_search_by_key(&right, |&(n, _, _)| n) {
Ok(idx) => {
let entry = unsafe { slice.get_unchecked(idx) };
Some((entry.1, entry.2))
}
Err(_) => None,
}
}
}
#[derive(Deserialize)]
#[serde(try_from = "RawBpe")]
pub struct Bpe {
#[serde(skip)]
id: usize,
daac: DoubleArrayAhoCorasick<TokenId>,
token_lens: Vec<u16>,
count_cache: CountSharedCache,
token_to_id: HashMap<String, u32>,
byte_to_initial_token: [u32; 256],
byte_pair_initial: Vec<(u32, u32)>,
merge_adj: MergeAdjacency,
ignore_merges: bool,
}
impl TryFrom<RawBpe> for Bpe {
type Error = String;
fn try_from(raw: RawBpe) -> Result<Self, String> {
let merge_map = parse_merges(&raw.vocab, &raw.merges)?;
let mut bpe = Self::new(&raw.vocab, merge_map)?;
bpe.ignore_merges = raw.ignore_merges;
Ok(bpe)
}
}
fn parse_merges(vocab: &Vocab, merges: &[Value]) -> Result<ParsedMergeMap, String> {
let mut merge_map = ParsedMergeMap::new();
for (rank, entry) in merges.iter().enumerate() {
let (left, right) = parse_merge_entry(entry)?;
let &left_id = vocab
.get(left)
.ok_or_else(|| format!("merge token not in vocab: {left:?}"))?;
let &right_id = vocab
.get(right)
.ok_or_else(|| format!("merge token not in vocab: {right:?}"))?;
let merged = format!("{left}{right}");
let &merged_id = vocab
.get(&merged)
.ok_or_else(|| format!("merged token not in vocab: {merged:?}"))?;
merge_map.insert((left_id, right_id), (rank as u32, merged_id));
}
Ok(merge_map)
}
fn parse_merge_entry(entry: &Value) -> Result<(&str, &str), String> {
match entry {
Value::String(s) => {
let (left, right) = s
.split_once(' ')
.ok_or_else(|| format!("invalid merge entry (no space): {s:?}"))?;
Ok((left, right))
}
Value::Array(arr) if arr.len() == 2 => {
let left = arr[0]
.as_str()
.ok_or_else(|| format!("merge element not a string: {:?}", arr[0]))?;
let right = arr[1]
.as_str()
.ok_or_else(|| format!("merge element not a string: {:?}", arr[1]))?;
Ok((left, right))
}
_ => Err(format!("unrecognized merge entry format: {entry:?}")),
}
}
impl Bpe {
pub fn new(vocab: &Vocab, merge_map: ParsedMergeMap) -> Result<Self, String> {
if vocab.is_empty() {
return Err("cannot build Bpe with empty vocabulary".into());
}
let vocab_r: std::collections::BTreeMap<u32, &str> =
vocab.iter().map(|(s, &id)| (id, s.as_str())).collect();
let max_token = *vocab_r.keys().max().unwrap();
let daac = DoubleArrayAhoCorasickBuilder::new()
.match_kind(daachorse::MatchKind::LeftmostLongest)
.build_with_values(vocab_r.iter().map(|(&token, pattern)| (pattern, token)))
.map_err(|e| format!("error building DAAC: {e}"))?;
let token_lens: Vec<u16> = (0..=max_token)
.map(|t| {
let s = vocab_r[&t];
u16::try_from(s.len())
.map_err(|_| format!("token {t} length {} exceeds u16::MAX", s.len()))
})
.collect::<Result<Vec<_>, String>>()?;
let ranked_merge_map = RankedMergeMap::from_parsed(&merge_map);
let mut byte_to_initial_token = [INVALID_TOKEN; 256];
for byte_val in 0u16..256 {
let ch = BYTE_TO_CHAR[byte_val as usize];
let mut buf = [0u8; 4];
let s = ch.encode_utf8(&mut buf);
if let Some(&id) = vocab.get(s) {
byte_to_initial_token[byte_val as usize] = id;
}
}
let mut byte_pair_initial = vec![(u32::MAX, 0u32); 65536];
for b1 in 0u16..256 {
let t1 = byte_to_initial_token[b1 as usize];
if t1 == INVALID_TOKEN {
continue;
}
for b2 in 0u16..256 {
let t2 = byte_to_initial_token[b2 as usize];
if t2 == INVALID_TOKEN {
continue;
}
if let Some((rank, new_id)) = ranked_merge_map.get(t1, t2) {
byte_pair_initial[b1 as usize * 256 + b2 as usize] = (rank, new_id);
}
}
}
let vocab_size = (max_token + 1) as usize;
let merge_adj = MergeAdjacency::from_parsed(&merge_map, vocab_size);
Ok(Self {
id: BPE_ID_COUNTER.fetch_add(1, Ordering::Relaxed),
daac,
token_lens,
count_cache: CountSharedCache::new(),
token_to_id: vocab.clone(),
byte_to_initial_token,
byte_pair_initial,
merge_adj,
ignore_merges: false,
})
}
fn next_match(&self, input: &str) -> Option<TokenId> {
let m = self.daac.leftmost_find_iter(input).next()?;
(m.start() == 0).then(|| m.value())
}
#[inline(always)]
pub fn count_tokens(&self, input: &str) -> Result<usize, String> {
if input.is_empty() {
return Ok(0);
}
if let Some(token) = self.next_match(input)
&& self.token_lens[token as usize] as usize == input.len()
{
return Ok(1);
}
let bpe_id = self.id;
let hit = TL_BPE_COUNT_CACHE.with(|c| {
let c = c.borrow();
if c.bpe_id != bpe_id {
return None;
}
c.get(input)
});
if let Some(count) = hit {
return Ok(count);
}
if let Some(count) = self.count_cache.get(input) {
TL_BPE_COUNT_CACHE.with(|c| {
let mut c = c.borrow_mut();
if c.bpe_id != bpe_id {
c.bpe_id = bpe_id;
c.clear();
}
c.insert(input, count);
});
return Ok(count);
}
if self.ignore_merges {
let count = input.chars().count();
TL_BPE_COUNT_CACHE.with(|c| {
let mut c = c.borrow_mut();
if c.bpe_id != bpe_id {
c.bpe_id = bpe_id;
c.clear();
}
c.insert(input, count);
});
self.count_cache.insert(input.to_string(), count);
return Ok(count);
}
let count = self.count_merge_all_encoded(input)?;
TL_BPE_COUNT_CACHE.with(|c| {
let mut c = c.borrow_mut();
if c.bpe_id != bpe_id {
c.bpe_id = bpe_id;
c.clear();
}
c.insert(input, count);
});
self.count_cache.insert(input.to_string(), count);
Ok(count)
}
fn count_merge_all_encoded(&self, input: &str) -> Result<usize, String> {
if input.is_empty() {
return Ok(0);
}
TL_MERGE_SCRATCH.with(|s| {
let mut scratch = s.borrow_mut();
scratch.symbols.clear();
scratch.heap.clear();
let mut n = 0usize;
for ch in input.chars() {
let mut buf = [0u8; 4];
let s = ch.encode_utf8(&mut buf);
let id = self
.token_to_id
.get(s)
.copied()
.ok_or_else(|| format!("character {ch:?} not in vocabulary"))?;
scratch.symbols.push(MergeSymbol {
c: id,
prev: if n == 0 { -1 } else { (n - 1) as i32 },
next: -1,
});
if n > 0 {
scratch.symbols[n - 1].next = n as i32;
}
n += 1;
}
if n == 1 {
return Ok(1);
}
self.init_merge_heap(&mut scratch, n);
self.count_merge_loop(&mut scratch)
})
}
#[inline(always)]
fn init_merge_heap(&self, scratch: &mut MergeScratch, n: usize) {
let symbols = &scratch.symbols;
scratch.heap.extend((0..n - 1).filter_map(|i| {
let left = symbols[i].c;
let right = symbols[i + 1].c;
self.merge_adj
.get(left, right)
.map(|(rank, _new_id)| Reverse(MergeEntry::new(rank, i as u32, left, right)))
}));
}
#[inline(always)]
fn count_merge_loop(&self, scratch: &mut MergeScratch) -> Result<usize, String> {
let symbols = &mut scratch.symbols;
let heap = &mut scratch.heap;
while let Some(Reverse(entry)) = heap.pop() {
let pos = entry.pos() as usize;
let sym = symbols[pos];
let left_c = entry.left_c();
let right_c = entry.right_c();
if sym.c != left_c {
continue;
}
let next_idx = sym.next;
if next_idx < 0 {
continue;
}
let next_idx = next_idx as usize;
let next_sym = symbols[next_idx];
if next_sym.c != right_c {
continue;
}
let new_id = match self.merge_adj.get(left_c, right_c) {
Some((_, nid)) => nid,
None => continue,
};
symbols[pos].c = new_id;
symbols[pos].next = next_sym.next;
if next_sym.next >= 0 {
symbols[next_sym.next as usize].prev = pos as i32;
}
symbols[next_idx].c = INVALID_TOKEN;
if sym.prev >= 0 {
let prev_c = symbols[sym.prev as usize].c;
if let Some((rank, _)) = self.merge_adj.get(prev_c, new_id) {
heap.push(Reverse(MergeEntry::new(
rank,
sym.prev as u32,
prev_c,
new_id,
)));
}
}
let new_next = symbols[pos].next;
if new_next >= 0 {
let next_c = symbols[new_next as usize].c;
if let Some((rank, _)) = self.merge_adj.get(new_id, next_c) {
heap.push(Reverse(MergeEntry::new(rank, pos as u32, new_id, next_c)));
}
}
}
let mut count: usize = 0;
let mut i: i32 = 0;
while i >= 0 {
count += 1;
i = symbols[i as usize].next;
}
Ok(count)
}
}
impl Clone for Bpe {
fn clone(&self) -> Self {
Self {
id: BPE_ID_COUNTER.fetch_add(1, Ordering::Relaxed),
daac: self.daac.clone(),
token_lens: self.token_lens.clone(),
count_cache: CountSharedCache::new(),
token_to_id: self.token_to_id.clone(),
byte_to_initial_token: self.byte_to_initial_token,
byte_pair_initial: self.byte_pair_initial.clone(),
merge_adj: self.merge_adj.clone(),
ignore_merges: self.ignore_merges,
}
}
}
impl fmt::Debug for Bpe {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Bpe")
.field("vocab_size", &self.token_to_id.len())
.field("merges", &self.merge_adj.data.len())
.finish()
}
}
impl PartialEq for Bpe {
fn eq(&self, other: &Self) -> bool {
self.daac == other.daac
&& self.token_lens == other.token_lens
&& self.ignore_merges == other.ignore_merges
}
}