Skip to content

Commit 3a14632

Browse files
hendrikvanantwerpenhendrikvanantwerpen
authored
Merge pull request #16 from github/add-o200k-tokens
Add support for o200k tokenization
2 parents 2112a0e + 2a4deef commit 3a14632

7 files changed

Lines changed: 166 additions & 86 deletions

File tree

‎.gitignore‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,4 +2,4 @@ Cargo.lock
22
/target/
33
/crates/*/target/
44
/crates/*/Cargo.lock
5-
.vscode/
5+
.vscode/

‎crates/bpe/benches/counting.rs‎

Lines changed: 88 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -6,85 +6,100 @@ use criterion::{criterion_group, criterion_main, Criterion};
66
use rand::{thread_rng, Rng};
77

88
fn counting_benchmark(c: &mut Criterion) {
9-
let bpe = BytePairEncoding::cl100k();
10-
let text = create_test_bytes(&bpe, 20000);
9+
for (name, bpe) in [
10+
("cl100k", BytePairEncoding::cl100k()),
11+
("o200k", BytePairEncoding::o200k()),
12+
] {
13+
let text = create_test_bytes(&bpe, 20000);
14+
let fast = IntervalEncoding::new(&bpe, &text);
1115

12-
let fast = IntervalEncoding::new(&bpe, &text);
13-
14-
for bytes in [10, 100, 1000, 10000] {
15-
let mut group = c.benchmark_group(format!("bytes-{bytes}"));
16-
group.bench_function("hybrid counting", |b| {
17-
b.iter_batched(
18-
|| thread_rng().gen_range(0..text.len() - bytes),
19-
|start| fast.count(start..start + bytes),
20-
criterion::BatchSize::SmallInput,
21-
)
22-
});
23-
group.bench_function("backtrack counting", |b| {
24-
b.iter_batched(
25-
|| thread_rng().gen_range(0..text.len() - bytes),
26-
|start| bpe.count(&text[start..start + bytes]),
27-
criterion::BatchSize::SmallInput,
28-
)
29-
});
16+
for bytes in [10, 100, 1000, 10000] {
17+
let mut group = c.benchmark_group(format!("bpe-{name}-bytes-{bytes}"));
18+
group.bench_function("hybrid counting", |b| {
19+
b.iter_batched(
20+
|| thread_rng().gen_range(0..text.len() - bytes),
21+
|start| fast.count(start..start + bytes),
22+
criterion::BatchSize::SmallInput,
23+
)
24+
});
25+
group.bench_function("backtrack counting", |b| {
26+
b.iter_batched(
27+
|| thread_rng().gen_range(0..text.len() - bytes),
28+
|start| bpe.count(&text[start..start + bytes]),
29+
criterion::BatchSize::SmallInput,
30+
)
31+
});
32+
}
3033
}
3134
}
3235

3336
fn encoding_benchmark(c: &mut Criterion) {
34-
let bpe = BytePairEncoding::cl100k();
35-
let tiktoken = tiktoken_rs::cl100k_base().unwrap();
36-
let text = create_test_string(&bpe, 20000);
37-
let input = text.as_bytes();
37+
for (name, bpe, tiktoken) in [
38+
(
39+
"cl100k",
40+
BytePairEncoding::cl100k(),
41+
tiktoken_rs::cl100k_base().unwrap(),
42+
),
43+
(
44+
"o200k",
45+
BytePairEncoding::o200k(),
46+
tiktoken_rs::o200k_base().unwrap(),
47+
),
48+
] {
49+
let text = create_test_string(&bpe, 20000);
50+
let input = text.as_bytes();
3851

39-
for bytes in [10, 100, 1000, 10000] {
40-
let mut group = c.benchmark_group(format!("bytes-{bytes}"));
41-
group.bench_function("backtracking", |b| {
42-
b.iter_batched(
43-
|| thread_rng().gen_range(0..input.len() - bytes),
44-
|start| bpe.encode_via_backtracking(&input[start..start + bytes]),
45-
criterion::BatchSize::SmallInput,
46-
)
47-
});
48-
group.bench_function("heap", |b| {
49-
b.iter_batched(
50-
|| thread_rng().gen_range(0..input.len() - bytes),
51-
|start| bpe.encode_via_bitfield(&input[start..start + bytes]),
52-
criterion::BatchSize::SmallInput,
53-
)
54-
});
55-
group.bench_function("dynamic programming", |b| {
56-
b.iter_batched(
57-
|| thread_rng().gen_range(0..input.len() - bytes),
58-
|start| bpe.encode_via_table(&input[start..start + bytes]),
59-
criterion::BatchSize::SmallInput,
60-
)
61-
});
62-
group.bench_function("greedy", |b| {
63-
b.iter_batched(
64-
|| thread_rng().gen_range(0..input.len() - bytes),
65-
|start| bpe.encode_greedy(&input[start..start + bytes]),
66-
criterion::BatchSize::SmallInput,
67-
)
68-
});
69-
group.bench_function("minimal", |b| {
70-
b.iter_batched(
71-
|| thread_rng().gen_range(0..input.len() - bytes),
72-
|start| bpe.encode_minimal(&input[start..start + bytes]),
73-
criterion::BatchSize::SmallInput,
74-
)
75-
});
76-
group.bench_function("tiktoken", |b| {
77-
b.iter_batched(
78-
|| loop {
79-
let start = thread_rng().gen_range(0..input.len() - bytes - 1);
80-
if is_char_boundary(input[start]) && is_char_boundary(input[start + bytes]) {
81-
return start;
82-
}
83-
},
84-
|start| tiktoken.encode_ordinary(&text[start..start + bytes]),
85-
criterion::BatchSize::SmallInput,
86-
)
87-
});
52+
for bytes in [10, 100, 1000, 10000] {
53+
let mut group = c.benchmark_group(format!("bpe-{name}-bytes-{bytes}"));
54+
group.bench_function("backtracking", |b| {
55+
b.iter_batched(
56+
|| thread_rng().gen_range(0..input.len() - bytes),
57+
|start| bpe.encode_via_backtracking(&input[start..start + bytes]),
58+
criterion::BatchSize::SmallInput,
59+
)
60+
});
61+
group.bench_function("heap", |b| {
62+
b.iter_batched(
63+
|| thread_rng().gen_range(0..input.len() - bytes),
64+
|start| bpe.encode_via_bitfield(&input[start..start + bytes]),
65+
criterion::BatchSize::SmallInput,
66+
)
67+
});
68+
group.bench_function("dynamic programming", |b| {
69+
b.iter_batched(
70+
|| thread_rng().gen_range(0..input.len() - bytes),
71+
|start| bpe.encode_via_table(&input[start..start + bytes]),
72+
criterion::BatchSize::SmallInput,
73+
)
74+
});
75+
group.bench_function("greedy", |b| {
76+
b.iter_batched(
77+
|| thread_rng().gen_range(0..input.len() - bytes),
78+
|start| bpe.encode_greedy(&input[start..start + bytes]),
79+
criterion::BatchSize::SmallInput,
80+
)
81+
});
82+
group.bench_function("minimal", |b| {
83+
b.iter_batched(
84+
|| thread_rng().gen_range(0..input.len() - bytes),
85+
|start| bpe.encode_minimal(&input[start..start + bytes]),
86+
criterion::BatchSize::SmallInput,
87+
)
88+
});
89+
group.bench_function("tiktoken", |b| {
90+
b.iter_batched(
91+
|| loop {
92+
let start = thread_rng().gen_range(0..input.len() - bytes - 1);
93+
if is_char_boundary(input[start]) && is_char_boundary(input[start + bytes])
94+
{
95+
return start;
96+
}
97+
},
98+
|start| tiktoken.encode_ordinary(&text[start..start + bytes]),
99+
criterion::BatchSize::SmallInput,
100+
)
101+
});
102+
}
88103
}
89104
}
90105

‎crates/bpe/src/byte_pair_encoding.rs‎

Lines changed: 74 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,11 @@ static BPE_CL100K: Lazy<BytePairEncoding> = Lazy::new(|| {
1919
rmp_serde::from_slice(bytes).expect("")
2020
});
2121

22+
static BPE_O200K: Lazy<BytePairEncoding> = Lazy::new(|| {
23+
let bytes = include_bytes!("data/bpe_o200k.dict");
24+
rmp_serde::from_slice(bytes).expect("")
25+
});
26+
2227
/// Representation of the byte pair dictionary.
2328
/// This struct provides various conversions.
2429
/// We put all of them into a single struct so that they can be reused by different implementations.
@@ -153,11 +158,15 @@ fn token_bytes<'a>(all_tokens: &'a [u8], token_starts: &[u32], token_id: u32) ->
153158
}
154159

155160
fn hash_bytes(bytes: &[u8]) -> u32 {
161+
hash_bytes_with_factor(bytes, 17846336922010275747)
162+
}
163+
164+
fn hash_bytes_with_factor(bytes: &[u8], factor: u64) -> u32 {
156165
let mut hasher = FnvHasher::default();
157166
bytes.hash(&mut hasher);
158167
// Note: we save 1/3 of space for the hashmap by only using the most significant bits of the hash.
159168
// To make them unique for the given tokens, we have to add unfortunately another multiplication.
160-
((hasher.finish().wrapping_mul(37493864257)) >> 32) as u32
169+
((hasher.finish().wrapping_mul(factor)) >> 32) as u32
161170
}
162171

163172
fn find_token_by_bytes(
@@ -180,6 +189,10 @@ impl BytePairEncoding {
180189
&BPE_CL100K
181190
}
182191

192+
pub fn o200k() -> &'static Self {
193+
&BPE_O200K
194+
}
195+
183196
/// Construct a BytePairEncoding instance frmo a tiktoken dictionary.
184197
pub fn from_tiktoken(tiktoken_bpe: &CoreBPE, num_tokens: usize) -> Self {
185198
Self::from_dictionary((0..num_tokens).map(|i| tiktoken_bpe._decode_native(&[i])))
@@ -492,13 +505,11 @@ pub fn create_test_bytes(bpe: &BytePairEncoding, tokens: usize) -> Vec<u8> {
492505

493506
#[cfg(test)]
494507
mod tests {
495-
use std::fs::File;
496-
use std::path::PathBuf;
508+
497509
use std::time::Instant;
498510

499511
use itertools::Itertools;
500-
use serde::Serialize;
501-
use tiktoken_rs::{cl100k_base, cl100k_base_singleton};
512+
use tiktoken_rs::cl100k_base_singleton;
502513

503514
use crate::byte_pair_encoding::{create_test_bytes, BytePairEncoding};
504515

@@ -541,19 +552,72 @@ mod tests {
541552
}
542553
}
543554
}
555+
}
544556

545-
// TODO: Move the generation of the dictionary into some build procedure?
557+
#[cfg(test)]
558+
mod data {
559+
use std::collections::HashSet;
560+
use std::fs::File;
561+
use std::path::PathBuf;
562+
563+
use rand::Rng;
564+
use serde::Serialize;
565+
use tiktoken_rs::{cl100k_base, o200k_base};
566+
567+
use super::*;
568+
569+
const BPE_CL100K_LEN: usize = 100256;
570+
const BPE_O200K_LEN: usize = 199998;
571+
572+
/// Use this to find a hashing factor for [`hash_bytes`] that prevents collisions.
573+
/// 1. Ensure all supported tokenizers are in the list.
574+
/// 2. Update the hash factor in [`hash_bytes`].
575+
/// 3. Run [`update_token_dicts`] tests below to update data files.
546576
#[test]
547-
fn test_serialize() {
577+
#[ignore = "run manually to find a suitable hash factor"]
578+
fn find_hash_factor() {
579+
let bpes: &mut [(CoreBPE, usize)] = &mut [
580+
(cl100k_base().unwrap(), BPE_CL100K_LEN),
581+
(o200k_base().unwrap(), BPE_O200K_LEN),
582+
];
583+
let mut rnd = rand::thread_rng();
584+
loop {
585+
let factor: u64 = rnd.gen();
586+
if bpes.iter().all(|(bpe, len)| {
587+
let mut seen = HashSet::with_capacity(*len);
588+
(0..*len)
589+
.all(|i| seen.insert(hash_bytes_with_factor(&bpe._decode_native(&[i]), factor)))
590+
}) {
591+
println!("hash factor: {factor}");
592+
return;
593+
}
594+
}
595+
}
596+
597+
#[test]
598+
fn update_token_dicts() {
599+
serialize_tokens(
600+
&cl100k_base().expect("tiktoken initialization must not fail!"),
601+
BPE_CL100K_LEN,
602+
"cl100k",
603+
);
604+
serialize_tokens(
605+
&o200k_base().expect("tiktoken initialization must not fail!"),
606+
BPE_O200K_LEN,
607+
"o200k",
608+
);
609+
}
610+
611+
#[track_caller]
612+
fn serialize_tokens(dict: &CoreBPE, num_tokens: usize, name: &str) {
548613
let path = PathBuf::from(file!());
549614
let dir = path.parent().unwrap();
550-
let data_file = dir.join("data/bpe_cl100k.dict");
615+
let data_file = dir.join(format!("data/bpe_{name}.dict"));
551616
let current_dir = std::env::current_dir().unwrap();
552617
let abs_path = current_dir.parent().unwrap().parent().unwrap();
553618
let file = File::create(abs_path.join(data_file)).unwrap();
554619
let mut serializer = rmp_serde::Serializer::new(file);
555-
let cl100_dict = cl100k_base().expect("tiktoken initialization must not fail!");
556-
BytePairEncoding::from_tiktoken(&cl100_dict, 100256)
620+
BytePairEncoding::from_tiktoken(dict, num_tokens)
557621
.serialize(&mut serializer)
558622
.unwrap();
559623
}

‎crates/bpe/src/data/bpe_cl100k.dict‎

0 Bytes
Binary file not shown.

‎crates/bpe/src/data/bpe_o200k.dict‎

30.7 MB
Binary file not shown.

‎crates/geo_filters/src/config/bitchunks.rs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -217,9 +217,8 @@ pub(crate) fn count_ones_from_bitchunks<T: IsBucketType>(
217217
let mut total = take_ref(&mut ones, max_msb_len - 1).count();
218218
let smallest_msb = ones
219219
.next()
220-
.map(|bucket| {
220+
.inspect(|_| {
221221
total += 1;
222-
bucket
223222
})
224223
.unwrap_or_default();
225224

‎criterion.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
# save report in this directory, even if a custom target directory is set
2+
criterion_home = "./target/criterion"

0 commit comments

Comments
 (0)