Skip to content

Commit 83ddf6f

Browse files
aduarterengifoaduarterengifo
committed
0.59s with partially interned strings
1 parent b0aae55 commit 83ddf6f

1 file changed

Lines changed: 142 additions & 129 deletions

File tree

‎src/lib.rs‎

Lines changed: 142 additions & 129 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,29 @@ use rayon::prelude::*;
1111

1212
const PAT: &str = r"'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+";
1313

14-
// (NVM) TODO: maybe start accepting &[u8] instead? keep in mind all hell might break loose.
14+
struct TokenInterner {
15+
tokens: Vec<Vec<u8>>,
16+
token_to_id: HashMap<Vec<u8>, u32>
17+
}
18+
19+
type TokenId = u32;
20+
21+
impl TokenInterner {
22+
fn intern(&mut self, token: Vec<u8>) -> TokenId {
23+
if let Some(&id) = self.token_to_id.get(&token) {
24+
id
25+
} else {
26+
let id = self.tokens.len() as TokenId;
27+
self.token_to_id.insert(token.clone(), id);
28+
self.tokens.push(token);
29+
id
30+
}
31+
}
1532

16-
// TODO: actually use Cow<u8> instead of Vec<u8>
33+
fn get_bytes(&self, id:TokenId) -> &[u8] {
34+
&self.tokens[id as usize]
35+
}
36+
}
1737

1838
static RE: Lazy<Regex> = Lazy::new(|| Regex::new(PAT).unwrap());
1939

@@ -50,14 +70,16 @@ fn decrement_or_remove<T: std::cmp::Eq + Hash>(map: &mut HashMap<T, usize>, key:
5070
// borrows:
5171
// tok_to_count: hashmap
5272
fn get_pairs(
53-
tok_to_count: &HashMap<Vec<Vec<u8>>, usize>,
73+
tok_to_count: &HashMap<Vec<TokenId>, usize>,
74+
interner: &mut TokenInterner
5475
) -> (
55-
HashMap<(Vec<u8>, Vec<u8>), usize>,
56-
HashMap<(Vec<u8>, Vec<u8>), HashSet<Vec<Vec<u8>>>>,
76+
HashMap<(TokenId, TokenId), usize>,
77+
HashMap<(TokenId, TokenId), HashSet<Vec<TokenId>>>,
5778
BinaryHeap<PairHeapEntry>
5879
) {
59-
let mut pair_to_count: HashMap<(Vec<u8>, Vec<u8>), usize> = HashMap::new();
60-
let mut pair_to_toks: HashMap<(Vec<u8>, Vec<u8>), HashSet<Vec<Vec<u8>>>> = HashMap::new();
80+
let estimated_pairs = tok_to_count.len() * 2;
81+
let mut pair_to_count: HashMap<(TokenId, TokenId), usize> = HashMap::with_capacity(estimated_pairs);
82+
let mut pair_to_toks: HashMap<(TokenId, TokenId), HashSet<Vec<TokenId>>> = HashMap::with_capacity(estimated_pairs);
6183
let mut heap = BinaryHeap::new();
6284

6385
for (tok, count) in tok_to_count {
@@ -68,24 +90,19 @@ fn get_pairs(
6890
for i in 0..(tok.len() - 1) {
6991
let pair = (tok[i].clone(), tok[i+1].clone());
7092

71-
let new_count = match pair_to_count.entry(pair.clone()) {
72-
std::collections::hash_map::Entry::Occupied(mut e) => {
73-
*e.get_mut() += count; // Increment by counter if exists
74-
*e.get()
75-
}
76-
std::collections::hash_map::Entry::Vacant(e) => {
77-
e.insert(*count); // Insert counter if doesn't exist
78-
*count
79-
}
93+
let new_count = {
94+
let entry = pair_to_count.entry(pair.clone()).or_insert(0);
95+
*entry += count;
96+
*entry
8097
};
8198

8299
heap.push(PairHeapEntry {
83100
count: new_count,
84-
pair: pair.clone(),
101+
pair: (interner.get_bytes(pair.0).to_vec(), interner.get_bytes(pair.1).to_vec()),
85102
});
86103

87104
pair_to_toks
88-
.entry(pair.clone())
105+
.entry(pair)
89106
.or_insert(HashSet::new())
90107
.insert(tok.clone());
91108
}
@@ -94,59 +111,117 @@ fn get_pairs(
94111
(pair_to_count, pair_to_toks, heap)
95112
}
96113

114+
fn rusty_get_chunk_pre_toks(filepath: &str, start: u64, end: u64, special_tokens:Vec<String>) -> io::Result<HashMap<Vec<u8>, usize>> {
115+
let mut tok_to_count: HashMap<Vec<u8>, usize> = HashMap::new();
116+
let pat_special_toks = special_tokens.iter().map(|x: &String| regex::escape(x)).collect::<Vec<String>>().join("|");
117+
let re_special_toks: Regex = Regex::new(&pat_special_toks).unwrap();
118+
let mut file = File::open(filepath)?;
119+
file.seek(SeekFrom::Start(start))?;
120+
let mut buffer = vec![0u8; (end - start) as usize];
121+
let bytes_read = file.read(&mut buffer)?;
122+
let chunk = String::from_utf8_lossy(&buffer[..bytes_read]);
123+
for regex_match in re_special_toks.split(&chunk).flat_map(|subchunk| RE.find_iter(subchunk.unwrap())) {
124+
let key: Vec<u8> = regex_match.unwrap().as_str().as_bytes().to_vec();
125+
// += 1
126+
*tok_to_count.entry(key).or_default() += 1;
127+
}
128+
Ok(tok_to_count)
129+
}
97130

98-
/// we take a map from vector of bytes and a max vocab size.
99-
/// and we return?
100-
#[pyfunction]
101-
fn rusty_merge(mut tok_to_count: HashMap<Vec<Vec<u8>>, usize>, max: usize) -> PyResult<Vec<(Vec<u8>, Vec<u8>)>> {
102-
println!("we are inside rusty_merge");
103-
// println!("{:?}", tok_to_count);
104-
let mut max_pairs = vec![(vec![0; 0], vec![0; 0]); 0];
105-
let (mut pair_to_count, mut pair_to_toks, mut heap) = get_pairs(&tok_to_count);
131+
fn rusty_get_pre_toks(filepath: &str, boundaries: Vec<u64>, special_tokens:Vec<String>) -> io::Result<(HashMap<Vec<TokenId>, usize>, TokenInterner)> {
132+
let token_to_id = HashMap::new();
133+
let tokens = Vec::new();
134+
let mut interner = TokenInterner {tokens, token_to_id};
135+
// let pat_special_toks = special_tokens.iter().map(|x: &String| regex::escape(x)).collect::<Vec<String>>().join("|");
136+
let r: Vec<(u64, u64)> = boundaries.windows(2).map(|x| (x[0],x[1])).collect();
106137

107-
// let max_count_2 = pair_to_count[&("w".as_bytes().to_vec(),"w".as_bytes().to_vec())];
108-
// println!("max_count {}", max_count_2);
109-
// TODO: maintain heap of max-pairs, instead of finding the max_pair on every loop.
138+
let intermediate = r
139+
.par_iter()
140+
.map(|(start,end)| {
141+
rusty_get_chunk_pre_toks(filepath,*start,*end, special_tokens.clone()).unwrap()
142+
})
143+
.collect::<Vec<_>>();
144+
145+
let mut toks :HashSet<Vec<u8>> = HashSet::new();
146+
147+
for chunkmap in &intermediate {
148+
for raw in chunkmap.keys() {
149+
toks.insert(raw.clone());
150+
for byte in raw {
151+
toks.insert(vec![*byte]);
152+
}
153+
}
154+
}
155+
156+
let mut sorted_toks: Vec<Vec<u8>> = toks.into_iter().collect();
157+
sorted_toks.sort();
158+
159+
for tok in sorted_toks {
160+
interner.intern(tok);
161+
}
162+
163+
let mut result = HashMap::new();
164+
165+
for chunk_map in intermediate {
166+
for (raw, count) in chunk_map {
167+
let token_ids: Vec<TokenId> = raw
168+
.chunks(1)
169+
.map(|byte| interner.intern(byte.to_vec()))
170+
.collect();
171+
172+
*result.entry(token_ids).or_default() += count;
173+
}
174+
}
175+
176+
Ok((result, interner))
177+
}
178+
179+
fn rusty_merge(mut tok_to_count: HashMap<Vec<TokenId>, usize>, max: usize, interner: &mut TokenInterner) -> PyResult<Vec<(Vec<u8>, Vec<u8>)>> {
180+
let mut max_pairs = vec![(vec![0; 0], vec![0; 0]); 0];
181+
let (mut pair_to_count, mut pair_to_toks, mut heap) = get_pairs(&tok_to_count, interner);
110182

111183
while max_pairs.len() < max {
112184
// pop from heap.
113185
if let Some (heap_entry) = heap.pop() {
114-
if pair_to_count.get(&heap_entry.pair) == Some(&heap_entry.count) {
115-
let max_pair = heap_entry.pair;
116-
117-
max_pairs.push(max_pair.clone());
186+
let max_pair = (interner.intern(heap_entry.pair.0), interner.intern(heap_entry.pair.1));
187+
if pair_to_count.get(&max_pair) == Some(&heap_entry.count) {
118188

119-
let max_pair_to_toks = &pair_to_toks[&max_pair];
189+
max_pairs.push((interner.get_bytes(max_pair.0).to_vec(),interner.get_bytes(max_pair.1).to_vec()));
120190

191+
let tokens_to_process: Vec<Vec<TokenId>> = pair_to_toks[&max_pair].iter().cloned().collect();
121192
// for every tok that contains max_pair
122-
for tok in max_pair_to_toks.clone() {
123-
match tok_to_count.entry(tok.clone()) {
193+
for tok in &tokens_to_process {
194+
match tok_to_count.entry(tok.to_vec()) {
124195
Entry::Occupied(_) => (),
125196
Entry::Vacant(_) => continue, // Skip if doesn't exist
126197
};
127198

128-
let tok_count = tok_to_count[&tok.clone()];
199+
let tok_count = if let Some(&count) = tok_to_count.get(tok) {
200+
count
201+
} else {
202+
continue;
203+
};
129204

130205
// for every pair in tok
131206
for i in 0..tok.len() - 1 {
132-
let pair = (tok[i].clone(), tok[i+1].clone());
207+
let pair = (tok[i], tok[i+1]);
133208

134-
match pair_to_toks.entry(pair.clone()) {
209+
match pair_to_toks.entry(pair) {
135210
std::collections::hash_map::Entry::Occupied(mut e) => {
136211
let set = e.get_mut();
137212

138213
// remove from pair's toks.
139-
set.remove(&tok);
214+
set.remove(tok);
140215
}
141216
std::collections::hash_map::Entry::Vacant(_) => {}
142217
}
143218

144-
match pair_to_count.entry(pair.clone()) {
219+
match pair_to_count.entry(pair) {
145220
std::collections::hash_map::Entry::Occupied(mut e) => {
146221
*e.get_mut() = e.get_mut().saturating_sub(tok_count); // remove tok_count from pair count.
147222

148223
if *e.get() > 0 {
149-
heap.push(PairHeapEntry { count: *e.get(), pair });
224+
heap.push(PairHeapEntry { count: *e.get(), pair:(interner.get_bytes(pair.0).to_vec(), interner.get_bytes(pair.1).to_vec()) });
150225
} else {
151226
e.remove();
152227
}
@@ -157,132 +232,70 @@ fn rusty_merge(mut tok_to_count: HashMap<Vec<Vec<u8>>, usize>, max: usize) -> Py
157232
}
158233

159234
// construct new_tok
160-
let mut new_tok: Vec<Vec<u8>>= Vec::new();
235+
let mut new_tok: Vec<TokenId>= Vec::new();
161236

162237
let mut i = 0;
163238

164239
while i < tok.len() {
165-
if i < tok.len() - 1 && (tok[i].clone(), tok[i+1].clone()) == max_pair {
166-
let new_vocab = [&tok[i][..], &tok[i+1][..]].concat();
167-
// println!("new_vocab {:?}", String::from_utf8(new_vocab.clone()));
168-
new_tok.push(new_vocab);
240+
if i < tok.len() - 1 && (tok[i], tok[i+1]) == max_pair {
241+
242+
let bytes1 = interner.get_bytes(tok[i]);
243+
let bytes2 = interner.get_bytes(tok[i+1]);
244+
245+
let merged_bytes = [bytes1, bytes2].concat();
246+
247+
let new_token_id = interner.intern(merged_bytes);
248+
249+
new_tok.push(new_token_id);
169250
i += 2;
170251
} else {
171-
new_tok.push(tok[i].clone());
252+
new_tok.push(tok[i]);
172253
i += 1
173254
}
174255
}
175256

176-
// println!("new_tok {:?}", new_tok.clone().iter().map(|x| String::from_utf8(x.clone()).unwrap()).collect::<Vec<String>>().join("") );
177-
178257
// increment new_tok count by tok_count
179258
*tok_to_count.entry(new_tok.clone()).or_default() += tok_count;
180259

181260
// decrement tok count by tok_count
182261
// if tok_count is zero -> remove tok entry all together.
183-
decrement_or_remove(&mut tok_to_count, tok.clone(), tok_count);
262+
decrement_or_remove(&mut tok_to_count, tok.to_vec(), tok_count);
184263

185264
// for every pair in new_tok
186-
for i in 0..new_tok.clone().len() - 1 {
187-
let pair = (new_tok[i].clone(), new_tok[i+1].clone());
265+
for i in 0..new_tok.len() - 1 {
266+
let pair = (new_tok[i], new_tok[i+1]);
267+
268+
pair_to_toks
269+
.entry(pair)
270+
.or_insert_with(HashSet::new)
271+
.insert(new_tok.clone());
188272

189-
match pair_to_toks.entry(pair.clone()) {
190-
// old_pair
191-
std::collections::hash_map::Entry::Occupied(mut e) => {
192-
// insert new_tok
193-
e.get_mut().insert(new_tok.clone());
194-
}
195-
// new pair
196-
std::collections::hash_map::Entry::Vacant(e) => {
197-
e.insert(HashSet::from_iter([new_tok.clone()]));
198-
}
199-
}
200273

201-
*pair_to_count.entry(pair.clone()).or_default() += tok_count;
274+
*pair_to_count.entry(pair).or_default() += tok_count;
202275

203276
heap.push(PairHeapEntry {
204277
count: pair_to_count[&pair],
205-
pair: pair.clone(),
278+
pair: (interner.get_bytes(pair.0).to_vec(), interner.get_bytes(pair.1).to_vec()),
206279
});
207280

208281
}
209282
}
210283
}
211284
}
212285
};
213-
// println!("{:?}", max_pairs);
214286
Ok(max_pairs)
215287
}
216288

217-
#[pyfunction]
218-
fn rusty_full_merge(filepath: &str, boundaries: Vec<u64>, special_tokens:Vec<String>, max: usize) -> PyResult<Vec<(Vec<u8>, Vec<u8>)>> {
219-
let tok_to_count = rusty_get_pre_toks(filepath, boundaries, special_tokens);
220-
// println!("tok_to_count {:?}", tok_to_count);
221-
return rusty_merge(tok_to_count.unwrap(), max)
222-
}
223-
224-
225-
#[pyfunction]
226-
fn rusty_pre_tok(chunk:&str, special_tokens:Vec<String>) -> HashMap<Vec<Vec<u8>>, usize> {
227-
let mut tok_to_count: HashMap<Vec<Vec<u8>>, usize> = HashMap::new();
228-
let pat_special_toks = special_tokens.iter().map(|x: &String| regex::escape(x)).collect::<Vec<String>>().join("|");
229-
let re_special_toks: Regex = Regex::new(&pat_special_toks).unwrap();
230-
231-
for regex_match in re_special_toks.split(chunk).flat_map(|subchunk| RE.find_iter(subchunk.unwrap())) {
232-
let key: Vec<Vec<u8>> = regex_match.unwrap().as_str().as_bytes().iter().map(|&byte| vec![byte]).collect();
233-
234-
// += 1
235-
*tok_to_count.entry(key).or_default() += 1;
236-
237-
}
238-
239-
tok_to_count
240-
}
241-
242-
fn rusty_get_chunk_pre_toks(filepath: &str, start: u64, end: u64, special_tokens:Vec<String>) -> io::Result<HashMap<Vec<Vec<u8>>, usize>> {
243-
let mut tok_to_count: HashMap<Vec<Vec<u8>>, usize> = HashMap::new();
244-
let pat_special_toks = special_tokens.iter().map(|x: &String| regex::escape(x)).collect::<Vec<String>>().join("|");
245-
let re_special_toks: Regex = Regex::new(&pat_special_toks).unwrap();
246-
let mut file = File::open(filepath)?;
247-
file.seek(SeekFrom::Start(start))?;
248-
let mut buffer = vec![0u8; (end - start) as usize];
249-
let bytes_read = file.read(&mut buffer)?;
250-
let chunk = String::from_utf8_lossy(&buffer[..bytes_read]);
251-
for regex_match in re_special_toks.split(&chunk).flat_map(|subchunk| RE.find_iter(subchunk.unwrap())) {
252-
let key: Vec<Vec<u8>> = regex_match.unwrap().as_str().as_bytes().iter().map(|&byte| vec![byte]).collect();
253-
// += 1
254-
*tok_to_count.entry(key).or_default() += 1;
255-
}
256-
Ok(tok_to_count)
257-
}
258289

259290
#[pyfunction]
260-
fn rusty_get_pre_toks(filepath: &str, boundaries: Vec<u64>, special_tokens:Vec<String>) -> io::Result<HashMap<Vec<Vec<u8>>, usize>> {
261-
262-
// let pat_special_toks = special_tokens.iter().map(|x: &String| regex::escape(x)).collect::<Vec<String>>().join("|");
263-
let r: Vec<(u64, u64)> = boundaries.windows(2).map(|x| (x[0],x[1])).collect();
264-
let map = r
265-
.par_iter()
266-
.map(|(start,end)| {
267-
rusty_get_chunk_pre_toks(filepath,*start,*end, special_tokens.clone())
268-
})
269-
.collect::<Vec<_>>()
270-
.into_iter()
271-
.fold(HashMap::new(), |mut acc, chunk_map| {
272-
for (key, value) in chunk_map.unwrap() {
273-
*acc.entry(key).or_default() += value;
274-
}
275-
acc
276-
});
277-
Ok(map)
291+
fn rusty_full_merge(filepath: &str, boundaries: Vec<u64>, special_tokens:Vec<String>, max: usize) -> PyResult<Vec<(Vec<u8>, Vec<u8>)>> {
292+
let (tok_to_count,mut interner) = rusty_get_pre_toks(filepath, boundaries, special_tokens).unwrap();
293+
return rusty_merge(tok_to_count, max, &mut interner);
278294
}
279295

280296
/// A Python module implemented in Rust.
281297
#[pymodule]
282298
fn rusty_tokey(m: &Bound<'_, PyModule>) -> PyResult<()> {
283-
m.add_function(wrap_pyfunction!(rusty_merge, m)?)?;
284299
m.add_function(wrap_pyfunction!(rusty_full_merge, m)?)?;
285-
m.add_function(wrap_pyfunction!(rusty_pre_tok, m)?)?;
286-
m.add_function(wrap_pyfunction!(rusty_get_pre_toks, m)?)?;
287300
Ok(())
288301
}

0 commit comments

Comments
 (0)