import os import re import time import heapq from collections import defaultdict from collections import Counter from multiprocessing import Pool, cpu_count from tqdm import tqdm from typing import Dict, Tuple, Union, Pattern PAT_COMPILED = re.compile(r"\S+") def pre_tokenize_and_count( args: Tuple[bytes, Dict[str, int], Union[Pattern, None]] ) -> Counter: """ Pre-tokenize a chunk of bytes into token ids, handling special tokens. Returns a Counter of tokens. """ chunk_bytes, special_token_to_id, delimiter_pattern_compiled = args chunk = chunk_bytes.decode("utf-8", errors="ignore") special_tokens_set = set(special_token_to_id.keys()) words_list = [] if delimiter_pattern_compiled: sub_chunks = delimiter_pattern_compiled.split(chunk) else: sub_chunks = [chunk] for sub_chunk in tqdm(sub_chunks, desc="Pre-tokenizing subchunks"): if not sub_chunk: continue if sub_chunk in special_tokens_set: token_id = special_token_to_id[sub_chunk] words_list.append((token_id,)) else: for word_str in PAT_COMPILED.findall(sub_chunk): if word_str: byte_sequence = word_str.encode("utf-8") id_sequence = tuple(byte_sequence) words_list.append(id_sequence) return Counter(words_list) # Placeholder, should be replaced with your chunk splitter logic def find_chunk_boundaries(f, num_chunks, marker_bytes): f.seek(0, 2) file_size = f.tell() chunk_size = file_size // num_chunks boundaries = [0] for i in range(1, num_chunks): boundaries.append(min(i * chunk_size, file_size)) boundaries.append(file_size) return boundaries class Node: """表示词内一个 token 节点，便于链表原地更新。""" def __init__(self, value, word_freq): self.value = value self.word_freq = word_freq # 共享引用，节省内存 self.prev = None self.next = None class PQItem: """定义优先队列元素，实现自定义比较：频率优先，其次按字典序逆序。""" def __init__(self, freq, id_pair, byte_pair): self.freq = freq self.id_pair = id_pair self.byte_pair = byte_pair def __lt__(self, other): if self.freq != other.freq: return self.freq > other.freq # 频率高的先出 return self.byte_pair > other.byte_pair # 字典序大的先出 def run_train_bpe( input_path: str | os.PathLike, vocab_size: int, special_tokens: list[str], num_chunks: int = 4, num_processes: int = None, **kwargs, ) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]: before_pretokenization_time = time.time() # 1. Set up initial byte vocab and special tokens vocab = {i: bytes([i]) for i in range(256)} for token in special_tokens: token_bytes = token.encode("utf-8") if token_bytes not in vocab.values(): vocab[len(vocab)] = token_bytes byte_to_token_id = {v: k for k, v in vocab.items()} special_token_to_id = { token: byte_to_token_id[token.encode("utf-8")] for token in special_tokens } # 2. Prepare special tokens regex delimiter delimiter_pattern_compiled = None if special_tokens: # Sort by length descending for proper greedy match special_tokens_sorted = sorted( [t.encode("utf-8") for t in special_tokens], key=len, reverse=True ) escaped_tokens = [re.escape(t.decode("utf-8")) for t in special_tokens_sorted] delimiter_re = "|".join(escaped_tokens) if delimiter_re: delimiter_pattern_compiled = re.compile(f"({delimiter_re})") # 3. Read file, split into chunks for multiprocessing with open(input_path, "rb") as f: boundaries = find_chunk_boundaries( f, num_chunks, "<|endoftext|>".encode("utf-8") ) chunk_args = [] for start, end in zip(boundaries[:-1], boundaries[1:]): f.seek(start) chunk_bytes = f.read(end - start) chunk_args.append( ( chunk_bytes, special_token_to_id, delimiter_pattern_compiled, ) ) # 4. Determine number of processes processes_to_use = num_processes if processes_to_use is None: processes_to_use = min(cpu_count(), 8) processes_to_use = min(processes_to_use, len(chunk_args)) elapsed = time.time() - before_pretokenization_time print(f"Time taken before pretokenization: {elapsed:.2f} seconds") # 5. Multiprocess: Counter aggregation all_word_freqs = Counter() start_time = time.time() with Pool(processes=processes_to_use) as pool: print( f"Starting pre-tokenization with {processes_to_use} processes on {len(chunk_args)} chunks..." ) results_iterator = pool.imap_unordered(pre_tokenize_and_count, chunk_args) for chunk_counter in tqdm( results_iterator, total=len(chunk_args), desc="Processing chunks", leave=True ): all_word_freqs.update(chunk_counter) print(f"Pre-tokenization and initial counting time: {time.time() - start_time:.2f} seconds") ### Pre-tokenization 结束 pair_to_nodes = defaultdict(set) for word_tuple, count in tqdm(all_word_freqs.items(), desc="Building", leave=True): if len(word_tuple) < 2: continue # 所有链表节点共享 word_freq 引用，节省内存 word_freq = {'count': count} head = Node(word_tuple[0], word_freq) prev_node = head for i in range(1, len(word_tuple)): curr_node = Node(word_tuple[i], word_freq) prev_node.next = curr_node curr_node.prev = prev_node pair = (prev_node.value, curr_node.value) pair_to_nodes[pair].add(prev_node) prev_node = curr_node pair_freqs = Counter() for pair, nodes in tqdm(pair_to_nodes.items(), desc="Counting pairs", leave=True): # 某个 pair 的出现次数 = 其所有节点所对应 word 的词频累加 pair_freqs[pair] = sum(node.word_freq['count'] for node in nodes) pq = [ PQItem(freq, p, (vocab[p[0]], vocab[p[1]])) for p, freq in pair_freqs.items() ] heapq.heapify(pq) ### BPE 开始 merges = [] num_merges = vocab_size - len(vocab) pbar = tqdm(total=num_merges, desc="Performing BPE merges") start_time = time.time() for _ in range(num_merges): if not pq: break # 取出频率最高的 pair，处理优先队列惰性删除的过期元素 best_pair = None while pq: item = heapq.heappop(pq) if item.id_pair not in pair_freqs: continue # 已经被合并删除 if pair_freqs[item.id_pair] == item.freq: best_pair = item.id_pair break if best_pair is None: break p1, p2 = best_pair # 合成新 token，添加到 merges/vocab new_token_id = len(vocab) merged_token_bytes = vocab[p1] + vocab[p2] merges.append((vocab[p1], vocab[p2])) vocab[new_token_id] = merged_token_bytes # 逐个更新包含改 pair 的词 nodes_to_process = list(pair_to_nodes[best_pair]) for node1 in nodes_to_process: node2 = node1.next if node2 is None: continue word_freq = node1.word_freq['count'] # 更新左侧相邻 pair 的频率及映射关系 if node1.prev: left = node1.prev old_left_pair = (left.value, node1.value) pair_freqs[old_left_pair] -= word_freq heapq.heappush(pq, PQItem(pair_freqs[old_left_pair], old_left_pair, (vocab[old_left_pair[0]], vocab[old_left_pair[1]]))) pair_to_nodes[old_left_pair].discard(left) new_left_pair = (left.value, new_token_id) pair_to_nodes[new_left_pair].add(left) pair_freqs[new_left_pair] += word_freq heapq.heappush(pq, PQItem(pair_freqs[new_left_pair], new_left_pair, (vocab[new_left_pair[0]], vocab[new_left_pair[1]]))) # 更新右侧相邻 pair 的频率及映射关系 if node2.next: right = node2.next old_right_pair = (node2.value, right.value) pair_freqs[old_right_pair] -= word_freq heapq.heappush(pq, PQItem(pair_freqs[old_right_pair], old_right_pair, (vocab[old_right_pair[0]], vocab[old_right_pair[1]]))) new_right_pair = (new_token_id, right.value) pair_to_nodes[old_right_pair].discard(node2) pair_to_nodes[new_right_pair].add(node1) pair_freqs[new_right_pair] += word_freq heapq.heappush(pq, PQItem(pair_freqs[new_right_pair], new_right_pair, (vocab[new_right_pair[0]], vocab[new_right_pair[1]]))) # 链表合并：node1、node2合成 new_token_id node1.value = new_token_id node1.next = node2.next if node2.next: node2.next.prev = node1 # 删除被合并 pair 的所有统计 del pair_freqs[best_pair] del pair_to_nodes[best_pair] pbar.update(1) end_time = time.time() print(f"Merge time: {end_time - start_time:.2f} seconds") pbar.close() return vocab, merges