Skip to content

Commit b0aae55

Browse files
aduarterengifoaduarterengifo
committed
heap change, for this size dataset it is *much* slower.
1 parent d8c6e20 commit b0aae55

1 file changed

Lines changed: 112 additions & 98 deletions

File tree

‎src/lib.rs‎

Lines changed: 112 additions & 98 deletions
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,19 @@ const PAT: &str = r"'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s
1717

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

20-
/// Formats the sum of two numbers as string.
21-
#[pyfunction]
22-
fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
23-
Ok((a + b).to_string())
20+
#[derive(Eq, PartialEq)]
21+
#[derive(PartialOrd)]
22+
struct PairHeapEntry {
23+
count: usize,
24+
pair: (Vec<u8>, Vec<u8>),
2425
}
2526

26-
#[pyfunction]
27-
fn simpl() -> PyResult<String> {
28-
Ok("COW".to_string())
27+
impl Ord for PairHeapEntry {
28+
fn cmp(&self, other: &Self) -> Ordering {
29+
// Max heap: higher count = higher priority
30+
self.count.cmp(&other.count)
31+
.then(self.pair.cmp(&other.pair)) // tie-breaking
32+
}
2933
}
3034

3135
fn decrement_or_remove<T: std::cmp::Eq + Hash>(map: &mut HashMap<T, usize>, key: T, amount: usize) -> () {
@@ -50,9 +54,11 @@ fn get_pairs(
5054
) -> (
5155
HashMap<(Vec<u8>, Vec<u8>), usize>,
5256
HashMap<(Vec<u8>, Vec<u8>), HashSet<Vec<Vec<u8>>>>,
57+
BinaryHeap<PairHeapEntry>
5358
) {
5459
let mut pair_to_count: HashMap<(Vec<u8>, Vec<u8>), usize> = HashMap::new();
5560
let mut pair_to_toks: HashMap<(Vec<u8>, Vec<u8>), HashSet<Vec<Vec<u8>>>> = HashMap::new();
61+
let mut heap = BinaryHeap::new();
5662

5763
for (tok, count) in tok_to_count {
5864
// not big enough for a pair.
@@ -62,14 +68,21 @@ fn get_pairs(
6268
for i in 0..(tok.len() - 1) {
6369
let pair = (tok[i].clone(), tok[i+1].clone());
6470

65-
match pair_to_count.entry(pair.clone()) {
71+
let new_count = match pair_to_count.entry(pair.clone()) {
6672
std::collections::hash_map::Entry::Occupied(mut e) => {
6773
*e.get_mut() += count; // Increment by counter if exists
74+
*e.get()
6875
}
6976
std::collections::hash_map::Entry::Vacant(e) => {
7077
e.insert(*count); // Insert counter if doesn't exist
78+
*count
7179
}
72-
}
80+
};
81+
82+
heap.push(PairHeapEntry {
83+
count: new_count,
84+
pair: pair.clone(),
85+
});
7386

7487
pair_to_toks
7588
.entry(pair.clone())
@@ -78,120 +91,123 @@ fn get_pairs(
7891
}
7992
}
8093

81-
(pair_to_count, pair_to_toks)
94+
(pair_to_count, pair_to_toks, heap)
8295
}
8396

97+
8498
/// we take a map from vector of bytes and a max vocab size.
8599
/// and we return?
86100
#[pyfunction]
87101
fn rusty_merge(mut tok_to_count: HashMap<Vec<Vec<u8>>, usize>, max: usize) -> PyResult<Vec<(Vec<u8>, Vec<u8>)>> {
88102
println!("we are inside rusty_merge");
89103
// println!("{:?}", tok_to_count);
90104
let mut max_pairs = vec![(vec![0; 0], vec![0; 0]); 0];
91-
let (mut pair_to_count, mut pair_to_toks) = get_pairs(&tok_to_count);
105+
let (mut pair_to_count, mut pair_to_toks, mut heap) = get_pairs(&tok_to_count);
92106

93107
// let max_count_2 = pair_to_count[&("w".as_bytes().to_vec(),"w".as_bytes().to_vec())];
94108
// println!("max_count {}", max_count_2);
95109
// TODO: maintain heap of max-pairs, instead of finding the max_pair on every loop.
96110

97111
while max_pairs.len() < max {
98-
let max_pair_opt = pair_to_count.iter()
99-
.max_by_key(|(pair, count)| (*count, *pair))
100-
.map(|pair| pair.0.clone());
101-
102-
// if there is a max_pair
103-
if let Some(max_pair) = max_pair_opt {
104-
max_pairs.push(max_pair.clone());
105-
// ---------------- EFFICIENTLY UPDATE ----------------
106-
let max_pair_to_toks = &pair_to_toks[&max_pair];
107-
// println!("max_pair_to_toks: {:?}", max_pair_to_toks);
108-
// for every tok that contains max_pair
109-
for tok in max_pair_to_toks.clone() {
110-
match tok_to_count.entry(tok.clone()) {
111-
Entry::Occupied(_) => (),
112-
Entry::Vacant(_) => continue, // Skip if doesn't exist
113-
};
114-
115-
let tok_count = tok_to_count[&tok.clone()];
116-
117-
// for every pair in tok
118-
for i in 0..tok.len() - 1 {
119-
let pair = (tok[i].clone(), tok[i+1].clone());
120-
121-
match pair_to_toks.entry(pair.clone()) {
122-
std::collections::hash_map::Entry::Occupied(mut e) => {
123-
let set = e.get_mut();
112+
// pop from heap.
113+
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;
124116

125-
// remove from pair's toks.
126-
set.remove(&tok);
127-
}
128-
std::collections::hash_map::Entry::Vacant(_) => {}
129-
}
130-
131-
match pair_to_count.entry(pair.clone()) {
132-
std::collections::hash_map::Entry::Occupied(mut e) => {
133-
*e.get_mut() = e.get_mut().saturating_sub(tok_count); // remove tok_count from pair count.
134-
135-
// if pair's count is zero.
136-
if *e.get() == 0 {
137-
// remove pair's pair_to_count entry
138-
e.remove();
117+
max_pairs.push(max_pair.clone());
139118

140-
// TODO: also remove the pair's pair_to_toks entry
119+
let max_pair_to_toks = &pair_to_toks[&max_pair];
120+
121+
// for every tok that contains max_pair
122+
for tok in max_pair_to_toks.clone() {
123+
match tok_to_count.entry(tok.clone()) {
124+
Entry::Occupied(_) => (),
125+
Entry::Vacant(_) => continue, // Skip if doesn't exist
126+
};
127+
128+
let tok_count = tok_to_count[&tok.clone()];
129+
130+
// for every pair in tok
131+
for i in 0..tok.len() - 1 {
132+
let pair = (tok[i].clone(), tok[i+1].clone());
133+
134+
match pair_to_toks.entry(pair.clone()) {
135+
std::collections::hash_map::Entry::Occupied(mut e) => {
136+
let set = e.get_mut();
137+
138+
// remove from pair's toks.
139+
set.remove(&tok);
141140
}
141+
std::collections::hash_map::Entry::Vacant(_) => {}
142142
}
143-
std::collections::hash_map::Entry::Vacant(_) => {}
143+
144+
match pair_to_count.entry(pair.clone()) {
145+
std::collections::hash_map::Entry::Occupied(mut e) => {
146+
*e.get_mut() = e.get_mut().saturating_sub(tok_count); // remove tok_count from pair count.
147+
148+
if *e.get() > 0 {
149+
heap.push(PairHeapEntry { count: *e.get(), pair });
150+
} else {
151+
e.remove();
152+
}
153+
}
154+
std::collections::hash_map::Entry::Vacant(_) => {}
155+
};
156+
144157
}
145-
}
146-
147-
// construct new_tok
148-
let mut new_tok: Vec<Vec<u8>>= Vec::new();
149-
150-
let mut i = 0;
151-
152-
while i < tok.len() {
153-
if i < tok.len() - 1 && (tok[i].clone(), tok[i+1].clone()) == max_pair {
154-
let new_vocab = [&tok[i][..], &tok[i+1][..]].concat();
155-
// println!("new_vocab {:?}", String::from_utf8(new_vocab.clone()));
156-
new_tok.push(new_vocab);
157-
i += 2;
158-
} else {
159-
new_tok.push(tok[i].clone());
160-
i += 1
158+
159+
// construct new_tok
160+
let mut new_tok: Vec<Vec<u8>>= Vec::new();
161+
162+
let mut i = 0;
163+
164+
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);
169+
i += 2;
170+
} else {
171+
new_tok.push(tok[i].clone());
172+
i += 1
173+
}
161174
}
162-
}
163175

164-
// println!("new_tok {:?}", new_tok.clone().iter().map(|x| String::from_utf8(x.clone()).unwrap()).collect::<Vec<String>>().join("") );
165-
166-
// increment new_tok count by tok_count
167-
*tok_to_count.entry(new_tok.clone()).or_default() += tok_count;
168-
169-
// decrement tok count by tok_count
170-
// if tok_count is zero -> remove tok entry all together.
171-
decrement_or_remove(&mut tok_to_count, tok.clone(), tok_count);
176+
// println!("new_tok {:?}", new_tok.clone().iter().map(|x| String::from_utf8(x.clone()).unwrap()).collect::<Vec<String>>().join("") );
172177

173-
// for every pair in new_tok
174-
for i in 0..new_tok.clone().len() - 1 {
175-
let pair = (new_tok[i].clone(), new_tok[i+1].clone());
178+
// increment new_tok count by tok_count
179+
*tok_to_count.entry(new_tok.clone()).or_default() += tok_count;
176180

177-
match pair_to_toks.entry(pair.clone()) {
178-
// old_pair
179-
std::collections::hash_map::Entry::Occupied(mut e) => {
180-
// insert new_tok
181-
e.get_mut().insert(new_tok.clone());
182-
}
183-
// new pair
184-
std::collections::hash_map::Entry::Vacant(e) => {
185-
e.insert(HashSet::from_iter([new_tok.clone()]));
181+
// decrement tok count by tok_count
182+
// if tok_count is zero -> remove tok entry all together.
183+
decrement_or_remove(&mut tok_to_count, tok.clone(), tok_count);
184+
185+
// 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());
188+
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+
}
186199
}
187-
}
188200

189-
*pair_to_count.entry(pair.clone()).or_default() += tok_count;
201+
*pair_to_count.entry(pair.clone()).or_default() += tok_count;
202+
203+
heap.push(PairHeapEntry {
204+
count: pair_to_count[&pair],
205+
pair: pair.clone(),
206+
});
207+
208+
}
190209
}
191-
}
192-
// ---------------- EFFICIENTLY UPDATE ----------------
193-
} else {
194-
break;
210+
}
195211
}
196212
};
197213
// println!("{:?}", max_pairs);
@@ -264,11 +280,9 @@ fn rusty_get_pre_toks(filepath: &str, boundaries: Vec<u64>, special_tokens:Vec<S
264280
/// A Python module implemented in Rust.
265281
#[pymodule]
266282
fn rusty_tokey(m: &Bound<'_, PyModule>) -> PyResult<()> {
267-
m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
268283
m.add_function(wrap_pyfunction!(rusty_merge, m)?)?;
269284
m.add_function(wrap_pyfunction!(rusty_full_merge, m)?)?;
270285
m.add_function(wrap_pyfunction!(rusty_pre_tok, m)?)?;
271286
m.add_function(wrap_pyfunction!(rusty_get_pre_toks, m)?)?;
272-
m.add_function(wrap_pyfunction!(simpl, m)?)?;
273287
Ok(())
274288
}

0 commit comments

Comments
 (0)