1- from rusty_tokey import sum_as_string
2- from rusty_tokey import simpl
1+ import itertools
2+ import os
3+ import regex as re
4+ from typing import BinaryIO
5+ import multiprocessing as mp
6+ import cProfile
7+ import pstats
8+ from collections import Counter
9+ import heapq
310
4- print (sum_as_string (2 ,2 ))
5- print (simpl ())
11+
12+ PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
13+ PAT_RE = re .compile (PAT )
14+
15+
16+ def find_chunk_boundaries (file : BinaryIO , desired_num_chunks : int , split_special_token : bytes ) -> list [int ]:
17+ """
18+ Chunk the file into parts that can be counted independently.
19+ May return fewer chunks if the boundaries end up overlapping.
20+ """
21+ assert isinstance (split_special_token , bytes ), "Must represent special token as a bytestring"
22+
23+ # Get total file size in bytes
24+ file .seek (0 , os .SEEK_END )
25+ file_size = file .tell ()
26+ file .seek (0 )
27+
28+ chunk_size = file_size // desired_num_chunks
29+
30+ # Initial guesses for chunk boundary locations, uniformly spaced
31+ # Chunks start on previous index, don't include last index
32+ chunk_boundaries = [i * chunk_size for i in range (desired_num_chunks + 1 )]
33+ chunk_boundaries [- 1 ] = file_size
34+
35+ mini_chunk_size = 4096 # Read ahead by 4k bytes at a time
36+
37+ for bi in range (1 , len (chunk_boundaries ) - 1 ):
38+ initial_position = chunk_boundaries [bi ]
39+ file .seek (initial_position ) # Start at boundary guess
40+ while True :
41+ mini_chunk = file .read (mini_chunk_size ) # Read a mini chunk
42+
43+ # If EOF, this boundary should be at the end of the file
44+ if mini_chunk == b"" :
45+ chunk_boundaries [bi ] = file_size
46+ break
47+
48+ # Find the special token in the mini chunk
49+ found_at = mini_chunk .find (split_special_token )
50+ if found_at != - 1 :
51+ chunk_boundaries [bi ] = initial_position + found_at
52+ break
53+ initial_position += mini_chunk_size
54+
55+ # Make sure all boundaries are unique, but might be fewer than desired_num_chunks
56+ return sorted (set (chunk_boundaries ))
57+
58+
59+ def pre_tokenize_chunk (input_path : str , start : int , end : int , pattern : re .Pattern [str ]):
60+ with open (input_path , "rb" ) as f :
61+ f .seek (start )
62+ chunk = f .read (end - start ).decode ("utf-8" , errors = "ignore" )
63+ sub_chunks = pattern .split (chunk )
64+ # run tokenization for each of the sub_chunks.
65+ iterators = [PAT_RE .finditer (sub_chunk ) for sub_chunk in sub_chunks ]
66+ flattened_iterators = itertools .chain (* iterators )
67+ store = Counter ()
68+ for match in flattened_iterators :
69+ res = match .group ()
70+ res_bytes = tuple (bytes ([c ]) for c in res .encode ("utf-8" ))
71+ store [res_bytes ] += 1
72+ return store
73+
74+
75+ def get_all_simple_pairs (
76+ pre_tok_dic : Counter [tuple [bytes ], int ],
77+ ) -> tuple [dict [tuple [bytes , bytes ], int ], dict [tuple [bytes , bytes ], set [tuple [bytes ]]]]:
78+ pair_to_count = Counter ()
79+ pair_to_tokens = {}
80+ for pre_token_key , count in pre_tok_dic .items ():
81+ if len (pre_token_key ) < 2 :
82+ continue
83+ for i in range (len (pre_token_key ) - 1 ):
84+ pair = (pre_token_key [i ], pre_token_key [i + 1 ])
85+ pair_to_count [pair ] += count
86+ if pair not in pair_to_tokens :
87+ pair_to_tokens [pair ] = set ()
88+ pair_to_tokens [pair ].add (pre_token_key )
89+ return (pair_to_count , pair_to_tokens )
90+
91+
92+ # not optimized for now.
93+ def merge (pre_tok_dic : Counter [tuple [bytes ], int ], stopping_condition : int ):
94+ max_pairs : list [tuple [bytes , bytes ]] = []
95+ (pair_to_count , pair_to_tokens ) = get_all_simple_pairs (pre_tok_dic )
96+
97+ # build heap from pair_to_count, to efficiently find max_pair
98+ heap = [(- count , pair ) for pair , count in pair_to_count .items ()]
99+ heapq .heapify (heap )
100+
101+ while len (max_pairs ) < stopping_condition and heap :
102+ # neg_count, max_pair = heapq.heappop(heap)
103+ # if max_pair not in pair_to_count or -neg_count != pair_to_count[max_pair]:
104+ # continue
105+ max_pair = max (pair_to_count , key = lambda pair : (pair_to_count [pair ], pair ))
106+ if len (max_pairs ) == 0 :
107+ print ("max_count_py" , pair_to_count [max_pair ], max_pair )
108+ # max_pair = alt_max_pair
109+ max_pairs .append (max_pair )
110+
111+ # INSERT_YOUR_CODE
112+ # Print the pair at the top of the heap (heap max) and the pair with the max count (max max)
113+ # print("heap max:", alt_max_pair, -neg_count, "max max:", max_pair, pair_to_count[max_pair])
114+
115+ # ------------ EFFICIENTLY UPDATE ------------
116+ affected_pre_toks = set (pair_to_tokens [max_pair ])
117+ for pre_tok in affected_pre_toks :
118+ count = pre_tok_dic [pre_tok ]
119+ for i in range (len (pre_tok ) - 1 ):
120+ # calculate pair
121+ pair = (pre_tok [i ], pre_tok [i + 1 ])
122+ # decrement count from pair
123+ pair_to_count [pair ] -= count
124+ # ------ HEAP MAINTENANCE ------
125+ heapq .heappush (heap , (- pair_to_count [pair ], pair ))
126+ # ------ HEAP MAINTENANCE ------
127+ # remove token from pairs_to_tokens
128+ pair_to_tokens [pair ].discard (pre_tok )
129+ # if pair doesn't occur anymore remove all together.
130+ if pair_to_count [pair ] == 0 :
131+ del pair_to_count [pair ]
132+ del pair_to_tokens [pair ]
133+ # construct new merged_token.
134+ new_pre_tok = []
135+ i = 0
136+ while i < len (pre_tok ):
137+ if i < len (pre_tok ) - 1 and (pre_tok [i ], pre_tok [i + 1 ]) == max_pair :
138+ new_pre_tok .append (pre_tok [i ] + pre_tok [i + 1 ])
139+ i += 2
140+ else :
141+ new_pre_tok .append (pre_tok [i ])
142+ i += 1
143+ new_pre_tok = tuple (new_pre_tok )
144+
145+ # update dic.
146+ pre_tok_dic [new_pre_tok ] += count
147+ pre_tok_dic [pre_tok ] -= count
148+
149+ # remove all together if appropriate
150+ if pre_tok_dic [pre_tok ] == 0 :
151+ del pre_tok_dic [pre_tok ]
152+
153+ for i in range (len (new_pre_tok ) - 1 ):
154+ pair = (new_pre_tok [i ], new_pre_tok [i + 1 ])
155+ pair_to_count [pair ] += count
156+ # ------ HEAP MAINTENANCE ------
157+ heapq .heappush (heap , (- pair_to_count [pair ], pair ))
158+ # ------ HEAP MAINTENANCE ------
159+ if pair not in pair_to_tokens :
160+ pair_to_tokens [pair ] = set ()
161+ pair_to_tokens [pair ].add (new_pre_tok )
162+
163+ # ------------ EFFICIENTLY UPDATE ------------
164+ return max_pairs
165+
166+
167+ def train_bpe (
168+ input_path : str , vocab_size : int , special_tokens : list [str ]
169+ ) -> tuple [dict [int , bytes ], list [tuple [bytes , bytes ]]]:
170+ special_tokens_len = len (special_tokens )
171+ stopping_condition = vocab_size - 256 - special_tokens_len
172+ pattern = re .compile ("|" .join ([re .escape (tok ) for tok in special_tokens ]))
173+ # open the input path
174+ with open (input_path , "rb" ) as f :
175+ # find the chunk boundaries
176+ boundaries = find_chunk_boundaries (f , 16 , "<|endoftext|>" .encode ("utf-8" ))
177+ num_workers = min (mp .cpu_count (), len (boundaries ) - 1 )
178+ with mp .Pool (processes = num_workers ) as pool :
179+ results = pool .starmap (
180+ pre_tokenize_chunk ,
181+ [(input_path , start , end , pattern ) for start , end in zip (boundaries [:- 1 ], boundaries [1 :])],
182+ )
183+ # results = [pre_tokenize_chunk(chunk) for chunk in chunks]
184+ combined = Counter ()
185+
186+ for d in results :
187+ combined .update (d )
188+
189+ max_pairs = merge (combined , stopping_condition )
190+ vocab : dict [int , bytes ] = {}
191+ # single-bytes
192+ for i in range (0 , 256 ):
193+ vocab [i ] = bytes ([i ])
194+ # special_tokens
195+ for i , token in enumerate (special_tokens ):
196+ vocab [i + 256 ] = token .encode ("utf-8" )
197+
198+ # vocab
199+ for i , (a , b ) in enumerate (max_pairs ):
200+ vocab [i + 256 + special_tokens_len ] = a + b
201+ return (vocab , max_pairs )
202+
203+
204+ if __name__ == "__main__" :
205+ profiler = cProfile .Profile ()
206+ profiler .enable ()
207+ (vocab , max_pairs ) = train_bpe ("./tests/fixtures/tinystories_sample_5M.txt" , 400 , ["<|endoftext|>" ])
208+ print ("max_pairs" , max_pairs )
209+ profiler .disable ()
210+ stats = pstats .Stats (profiler ).sort_stats ("cumtime" )
211+ stats .print_stats (20 ) # Show top 20 slowest functions
0 commit comments