bring over some of the code for testing

Author: aduarterengifo

pyproject.toml

@@ -12,8 +12,12 @@ classifiers = [
 ]
 dynamic = ["version"]
 dependencies = [
+    "jaxtyping>=0.2.19",
     "maturin>=1.8.7",
+    "numpy>=1.24.4",
+    "pytest>=8.3.5",
     "regex>=2024.11.6",
+    "torch>=2.5.1",
 ]
 [tool.maturin]
 features = ["pyo3/extension-module"]

python/rusty_tokey/__init__.py

@@ -0,0 +1,8 @@
+from .rusty_tokey import (
+    rusty_full_merge,
+    rusty_merge,
+    rusty_pre_tok,
+    rusty_get_pre_toks,
+    sum_as_string,
+    simpl,
+)

python/rusty_tokey/rlly_serious.py

@@ -1,15 +1,9 @@
-import itertools
 import os
 import regex as re
 from typing import BinaryIO
-import multiprocessing as mp
 import cProfile
 import pstats
-from collections import Counter
-import heapq
-# 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+"""
@@ -58,132 +52,7 @@ def find_chunk_boundaries(file: BinaryIO, desired_num_chunks: int, split_special
     # Make sure all boundaries are unique, but might be fewer than desired_num_chunks
     return sorted(set(chunk_boundaries))
 
-
-def pre_tokenize_chunk(input_path: str, start: int, end: int, pattern: re.Pattern[str]):
-    with open(input_path, "rb") as f:
-        f.seek(start)
-        chunk = f.read(end - start).decode("utf-8", errors="ignore")
-        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 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")
-
-#         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
-
-
-def get_all_simple_pairs(
-    pre_tok_dic: Counter[tuple[bytes], int],
-) -> tuple[dict[tuple[bytes, bytes], int], dict[tuple[bytes, bytes], set[tuple[bytes]]]]:
-    pair_to_count = Counter()
-    pair_to_tokens = {}
-    for pre_token_key, count in pre_tok_dic.items():
-        if len(pre_token_key) < 2:
-            continue
-        for i in range(len(pre_token_key) - 1):
-            pair = (pre_token_key[i], pre_token_key[i + 1])
-            pair_to_count[pair] += count
-            if pair not in pair_to_tokens:
-                pair_to_tokens[pair] = set()
-            pair_to_tokens[pair].add(pre_token_key)
-    return (pair_to_count, pair_to_tokens)
-
-
-# not optimized for now.
-def merge(pre_tok_dic: Counter[tuple[bytes], int], stopping_condition: int):
-    max_pairs: list[tuple[bytes, bytes]] = []
-    (pair_to_count, pair_to_tokens) = get_all_simple_pairs(pre_tok_dic)
-
-    # build heap from pair_to_count, to efficiently find max_pair
-    heap = [(-count, pair) for pair, count in pair_to_count.items()]
-    heapq.heapify(heap)
-
-    while len(max_pairs) < stopping_condition and heap:
-        # neg_count, max_pair = heapq.heappop(heap)
-        # if max_pair not in pair_to_count or -neg_count != pair_to_count[max_pair]:
-        #     continue
-        max_pair = max(pair_to_count, key=lambda pair: (pair_to_count[pair], pair))
-        # max_pair = alt_max_pair
-        max_pairs.append(max_pair)
-
-        # INSERT_YOUR_CODE
-        # Print the pair at the top of the heap (heap max) and the pair with the max count (max max)
-        # print("heap max:", alt_max_pair, -neg_count, "max max:", max_pair, pair_to_count[max_pair])
-
-        # ------------ EFFICIENTLY UPDATE ------------
-        affected_pre_toks = set(pair_to_tokens[max_pair])
-        for pre_tok in affected_pre_toks:
-            count = pre_tok_dic[pre_tok]
-            for i in range(len(pre_tok) - 1):
-                # calculate pair
-                pair = (pre_tok[i], pre_tok[i + 1])
-                # decrement count from pair
-                pair_to_count[pair] -= count
-                # ------ HEAP MAINTENANCE ------
-                heapq.heappush(heap, (-pair_to_count[pair], pair))
-                # ------ HEAP MAINTENANCE ------
-                # remove token from pairs_to_tokens
-                pair_to_tokens[pair].discard(pre_tok)
-                # if pair doesn't occur anymore remove all together.
-                if pair_to_count[pair] == 0:
-                    del pair_to_count[pair]
-                    del pair_to_tokens[pair]
-            # construct new merged_token.
-            new_pre_tok = []
-            i = 0
-            while i < len(pre_tok):
-                if i < len(pre_tok) - 1 and (pre_tok[i], pre_tok[i + 1]) == max_pair:
-                    new_pre_tok.append(pre_tok[i] + pre_tok[i + 1])
-                    i += 2
-                else:
-                    new_pre_tok.append(pre_tok[i])
-                    i += 1
-            new_pre_tok = tuple(new_pre_tok)
-
-            # update dic.
-            pre_tok_dic[new_pre_tok] += count
-            pre_tok_dic[pre_tok] -= count
-
-            # remove all together if appropriate
-            if pre_tok_dic[pre_tok] == 0:
-                del pre_tok_dic[pre_tok]
-
-            for i in range(len(new_pre_tok) - 1):
-                pair = (new_pre_tok[i], new_pre_tok[i + 1])
-                pair_to_count[pair] += count
-                # ------ HEAP MAINTENANCE ------
-                heapq.heappush(heap, (-pair_to_count[pair], pair))
-                # ------ HEAP MAINTENANCE ------
-                if pair not in pair_to_tokens:
-                    pair_to_tokens[pair] = set()
-                pair_to_tokens[pair].add(new_pre_tok)
-
-        # ------------ EFFICIENTLY UPDATE ------------
-    return max_pairs
-
-
-def train_bpe(
+def rusty_train_bpe(
     input_path: str, vocab_size: int, special_tokens: list[str]
 ) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
     special_tokens_len = len(special_tokens)
@@ -212,8 +81,9 @@ def train_bpe(
 if __name__ == "__main__":
     profiler = cProfile.Profile()
     profiler.enable()
-    (vocab, max_pairs) = train_bpe("./python/rusty_tokey/data/tinystories_sample_5M.txt", 400, ["<|endoftext|>"])
+    (vocab, max_pairs) = rusty_train_bpe("./python/rusty_tokey/data/tinystories_sample_5M.txt", 400, ["<|endoftext|>"])
     print("max_pairs", max_pairs)
+    print("vocab", vocab)
     profiler.disable()
     stats = pstats.Stats(profiler).sort_stats("cumtime")
     stats.print_stats(20)  # Show top 20 slowest functions