save

Author: aduarterengifo

Cargo.lock

@@ -9,8 +9,10 @@ name = "rusty_tokey"
 crate-type = ["cdylib"]
 
 [dependencies]
+fancy-regex = "0.14.0"
 once_cell = "1.21.3"
 rand = "0.9.0"
+rayon = "1.10.0"
 regex = "1.11.1"
 
 [dependencies.pyo3]

Cargo.toml

@@ -0,0 +1 @@
+stub_only

python/rusty_tokey/py.typed

@@ -7,8 +7,9 @@
 import pstats
 from collections import Counter
 import heapq
-from rusty_tokey import rusty_merge
-from rusty_tokey import rusty_pre_tok
+# from rusty_tokey import rusty_merge
+from rusty_tokey import rusty_full_merge
+# from rusty_tokey import rusty_pre_tok
 
 
 PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
@@ -73,23 +74,23 @@ def pre_tokenize_chunk(input_path: str, start: int, end: int, pattern: re.Patter
             store[res_bytes] += 1
         return store
 
-def rusty_pre_tokenize_chunk(input_path: str, start: int, end: int, special_tokens: list[str]) -> dict[list[]]:
-    with open(input_path, "rb") as f:
-        f.seek(start)
-        chunk = f.read(end - start).decode("utf-8", errors="ignore")
+# def rusty_pre_tokenize_chunk(input_path: str, start: int, end: int, pattern: re.Pattern[str], special_tokens: list[str]) -> dict[list[list[int]],int]:
+#     with open(input_path, "rb") as f:
+#         f.seek(start)
+#         chunk = f.read(end - start).decode("utf-8", errors="ignore")
 
-        return rusty_pre_tok(chunk,special_tokens)
+#         result = rusty_pre_tok(chunk,special_tokens)
 
-        # sub_chunks = pattern.split(chunk)
-        # # run tokenization for each of the sub_chunks.
-        # iterators = [PAT_RE.finditer(sub_chunk) for sub_chunk in sub_chunks]
-        # flattened_iterators = itertools.chain(*iterators)
-        # store = Counter()
-        # for match in flattened_iterators:
-        #     res = match.group()
-        #     res_bytes = tuple(bytes([c]) for c in res.encode("utf-8"))
-        #     store[res_bytes] += 1
-        # return store
+#         sub_chunks = pattern.split(chunk)
+#         # run tokenization for each of the sub_chunks.
+#         iterators = [PAT_RE.finditer(sub_chunk) for sub_chunk in sub_chunks]
+#         flattened_iterators = itertools.chain(*iterators)
+#         store = Counter()
+#         for match in flattened_iterators:
+#             res = match.group()
+#             res_bytes = tuple(bytes([c]) for c in res.encode("utf-8"))
+#             store[res_bytes] += 1
+#         return store
 
 
 def get_all_simple_pairs(
@@ -192,19 +193,22 @@ def train_bpe(
     with open(input_path, "rb") as f:
         # find the chunk boundaries
         boundaries = find_chunk_boundaries(f, 16, "<|endoftext|>".encode("utf-8"))
-        num_workers = min(mp.cpu_count(), len(boundaries) - 1)
-        with mp.Pool(processes=num_workers) as pool:
-            results = pool.starmap(
-                rusty_pre_tokenize_chunk,
-                [(input_path, start, end, special_tokens) for start, end in zip(boundaries[:-1], boundaries[1:])],
-            )
+        # num_workers = min(mp.cpu_count(), len(boundaries) - 1)
+        # with mp.Pool(processes=num_workers) as pool:
+        #     results = pool.starmap(
+        #         rusty_pre_tokenize_chunk,
+        #         [(input_path, start, end, pattern, special_tokens) for start, end in zip(boundaries[:-1], boundaries[1:])],
+        #     )
         # results = [pre_tokenize_chunk(chunk) for chunk in chunks]
-        combined = Counter()
+        # combined = Counter()
+
+        # for d in results:
+        #     combined.update(d)
 
-        for d in results:
-            combined.update(d)
+        # max_pairs = rusty_merge(combined, stopping_condition)
 
-        max_pairs = rusty_merge(combined, stopping_condition)
+        max_pairs = rusty_full_merge(input_path, boundaries, special_tokens, vocab_size)
+        # print('max_pairs', max_pairs)
         vocab: dict[int, bytes] = {}
         # single-bytes
         for i in range(0, 256):

python/rusty_tokey/rlly_serious.py

@@ -1,13 +1,12 @@
 use pyo3::prelude::*;
 use std::collections::{HashMap, HashSet};
-use std::fmt::Error;
 use std::hash::Hash;
 use std::collections::hash_map::Entry;
-use std::num::ParseIntError;
-use regex::Regex;
+use fancy_regex::Regex;
 use once_cell::sync::Lazy;
 use std::fs::File;
-use std::io::{self, Read, BufReader};
+use std::io::{self, Read, Seek, SeekFrom};
+use rayon::prelude::*;
 
 const PAT: &str = r"'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+";
 
@@ -83,7 +82,7 @@ fn get_pairs(
 #[pyfunction]
 fn rusty_merge(mut tok_to_count: HashMap<Vec<Vec<u8>>, usize>, max: usize) -> PyResult<Vec<(Vec<u8>, Vec<u8>)>> {
     println!("we are inside rusty_merge");
-    println!("{:?}", tok_to_count);
+    // println!("{:?}", tok_to_count);
     let mut max_pairs = vec![(vec![0; 0], vec![0; 0]); max];
     let (mut pair_to_count, mut pair_to_toks) = get_pairs(&tok_to_count);
     // TODO: maintain heap of max-pairs, instead of finding the max_pair on every loop.
@@ -194,18 +193,26 @@ fn rusty_merge(mut tok_to_count: HashMap<Vec<Vec<u8>>, usize>, max: usize) -> Py
 
 
     };
-    println!("{:?}", max_pairs);
+    // println!("{:?}", max_pairs);
     Ok(max_pairs)
 }
 
+#[pyfunction]
+fn rusty_full_merge(filepath: &str, boundaries: Vec<u64>, special_tokens:Vec<String>, max: usize) -> PyResult<Vec<(Vec<u8>, Vec<u8>)>> {
+    let tok_to_count = rusty_get_pre_toks(filepath, boundaries, special_tokens);
+    // println!("tok_to_count {:?}", tok_to_count);
+    return rusty_merge(tok_to_count.unwrap(), max)
+}
+
+
 #[pyfunction]
 fn rusty_pre_tok(chunk:&str, special_tokens:Vec<String>) -> HashMap<Vec<Vec<u8>>, usize> {
     let mut tok_to_count: HashMap<Vec<Vec<u8>>, usize> = HashMap::new();
     let pat_special_toks = special_tokens.iter().map(|x: &String| regex::escape(x)).collect::<Vec<String>>().join("|");
     let re_special_toks: Regex = Regex::new(&pat_special_toks).unwrap();
 
-    for regex_match in re_special_toks.split(chunk).flat_map(|subchunk| RE.find_iter(subchunk)) {
-        let key: Vec<Vec<u8>>  = regex_match.as_str().chars().map(|c| c.to_string().as_bytes().to_vec()).collect();
+    for regex_match in re_special_toks.split(chunk).flat_map(|subchunk| RE.find_iter(subchunk.unwrap())) {
+        let key: Vec<Vec<u8>>  = regex_match.unwrap().as_str().chars().map(|c| c.to_string().as_bytes().to_vec()).collect();
 
         // += 1
         *tok_to_count.entry(key).or_insert(1) += 1;
@@ -215,18 +222,56 @@ fn rusty_pre_tok(chunk:&str, special_tokens:Vec<String>) -> HashMap<Vec<Vec<u8>>
     tok_to_count
 }
 
+fn rusty_get_chunk_pre_toks(filepath: &str, start: u64, end: u64, special_tokens:Vec<String>) ->  io::Result<HashMap<Vec<Vec<u8>>, usize>> {
+    let mut tok_to_count: HashMap<Vec<Vec<u8>>, usize> = HashMap::new();
+    let pat_special_toks = special_tokens.iter().map(|x: &String| regex::escape(x)).collect::<Vec<String>>().join("|");
+    let re_special_toks: Regex = Regex::new(&pat_special_toks).unwrap();
+    let mut file = File::open(filepath)?;
+    file.seek(SeekFrom::Start(start))?;
+    let mut buffer = vec![0u8; (end - start) as usize];
+    let bytes_read = file.read(&mut buffer)?;
+    let chunk = String::from_utf8_lossy(&buffer[..bytes_read]);
+    println!("------ chunk ------");
+    println!("{}", chunk);
+    println!("------ chunk ------");
+    for regex_match in re_special_toks.split(&chunk).flat_map(|subchunk| RE.find_iter(subchunk.unwrap())) {
+        println!("match {}", regex_match.clone().unwrap().as_str());
+        let key: Vec<Vec<u8>>  = regex_match.unwrap().as_str().chars().map(|c| c.to_string().as_bytes().to_vec()).collect();
+
+        // += 1
+        *tok_to_count.entry(key).or_insert(1) += 1;
+        
+    }
+    Ok(tok_to_count)   
+}
+
 #[pyfunction]
-fn rusty_get_pre_toks(filepath: String) -> Result<(), std::io::Error>{
-    let file = File::open(filepath)?;
-    
-    Ok(())
+fn rusty_get_pre_toks(filepath: &str, boundaries: Vec<u64>, special_tokens:Vec<String>) -> io::Result<HashMap<Vec<Vec<u8>>, usize>> {
+
+    // let pat_special_toks = special_tokens.iter().map(|x: &String| regex::escape(x)).collect::<Vec<String>>().join("|");
+    let r: Vec<(u64, u64)> = boundaries.windows(2).map(|x| (x[0],x[1])).collect();
+    let map = r
+        .par_iter()
+        .map(|(start,end)| {
+            rusty_get_chunk_pre_toks(filepath,*start,*end, special_tokens.clone())
+        })
+        .collect::<Vec<_>>()
+        .into_iter()
+        .fold(HashMap::new(), |mut acc, chunk_map| {
+            for (key, value) in chunk_map.unwrap() {
+                *acc.entry(key).or_insert(0) += value;
+            }
+            acc
+        });
+    Ok(map)
 }
 
 /// A Python module implemented in Rust.
 #[pymodule]
 fn rusty_tokey(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
     m.add_function(wrap_pyfunction!(rusty_merge, m)?)?;
+    m.add_function(wrap_pyfunction!(rusty_full_merge, m)?)?;
     m.add_function(wrap_pyfunction!(rusty_pre_tok, m)?)?;
     m.add_function(wrap_pyfunction!(rusty_get_pre_toks, m)?)?;
     m.add_function(wrap_pyfunction!(simpl, m)?)?;