BPE 分词器算法原理与实现

import regex

PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""

regex.findall(PAT, "你好世界, Hello World!")  # [' 你好世界 ', ',', ' Hello', ' World', '!']

# utils.py

class myHeap:

    def __init__(self, mode="max", init_list=None):

        self.mode = mode

        self._init_data(init_list or [])

    def _init_data(self, init_list) -> list:

        self.data = list(init_list)

        self.size = len(self.data)

        if self.size == 0:

            return

        for i in range((len(self.data) - 1) // 2, -1, -1):

            self._shift_down(i)

    def _compare(self, a, b):

        if self.mode == "max":

            return a > b

        else:

            return a < b

    def _shift_down(self, index):

        """Shift down the element at index to maintain heap property"""

        while True:

            left = 2 * index + 1

            right = 2 * index + 2

            maxIndex = index

            if left < self.size and self._compare(self.data[left], self.data[maxIndex]):

                maxIndex = left

            if right < self.size and self._compare(self.data[right], self.data[maxIndex]):

                maxIndex = right

            if maxIndex != index:

                self.data[index], self.data[maxIndex] = self.data[maxIndex], self.data[index]

                index = maxIndex

            else:

                break

    def push(self, item):

        """Push an item onto the heap"""

        self.data.append(item)

        self.size += 1

        index = self.size - 1

        while index > 0:

            parent = (index - 1) // 2

            if self._compare(self.data[index], self.data[parent]):

                self.data[index], self.data[parent] = self.data[parent], self.data[index]

                index = parent

            else:

                break

    def pop(self) -> any:

        """Pop the top item off the heap and return it"""

        if self.size == 0:

            return None

        top_item = self.data[0]

        self.data[0] = self.data[-1]

        self.data.pop()

        self.size -= 1

        if self.size > 0:

            self._shift_down(0)

        return top_item

    def top(self) -> any:

        """Return the top item of the heap without removing it"""

        return self.data[0] if self.size > 0 else None

# utils.py

class BytePair(NamedTuple):

    left: bytes

    right: bytes

    @property

    def merged_bytes(self):

        return self.left + self.right

    def __str__(self) -> str:

        return f"({self.left.decode('utf-8', errors='ignore')} {self.right.decode('utf-8', errors='ignore')})"

    def __repr__(self) -> str:

        return self.__str__()

# utils.py

@dataclass

class WordRef:

    tokens: list[bytes]

    count: int

    @property

    def byte_pairs_dict(self) -> defaultdict[BytePair, int]:

        bp_to_count = defaultdict(int)

        for i in range(len(self.tokens) - 1):

            bp = BytePair(self.tokens[i], self.tokens[i + 1])

            bp_to_count[bp] += self.count

        return bp_to_count

    def merge(self, bp_merged: BytePair):

        """merge the given byte pair in this word"""

        new_tokens: list[bytes] = []

        i = 0

        while i < len(self.tokens):

            if i < len(self.tokens) - 1 and self.tokens[i] == bp_merged.left and self.tokens[i + 1] == bp_merged.right:

                new_tokens.append(bp_merged.merged_bytes)

                i += 2

            else:

                new_tokens.append(self.tokens[i])

                i += 1

        self.tokens = new_tokens

    def __str__(self) -> str:

        return f"{self.tokens} : {self.count}"

    def __repr__(self) -> str:

        return self.__str__()

# pretokenize.py

import os

import re

import regex

import multiprocessing

import io

from typing import BinaryIO, Dict, Tuple, List, Iterable

from collections import Counter

from cs336_basics.bpe.utils import split_bytes, PAT

def find_chunk_boundaries(

    file: BinaryIO,

    desired_num_chunks: int,

    split_special_token: bytes,

) -> list[int]:

    """

    Chunk the file into parts that can be counted independently.

    May return fewer chunks if the boundaries end up overlapping.

    """

    assert isinstance(split_special_token, bytes), "Must represent special token as a bytestring"

    # Get total file size in bytes

    file.seek(0, os.SEEK_END)

    file_size = file.tell()

    file.seek(0)

    chunk_size = file_size // desired_num_chunks

    # Initial guesses for chunk boundary locations, uniformly spaced

    # Chunks start on previous index, don't include last index

    chunk_boundaries = [i * chunk_size for i in range(desired_num_chunks + 1)]

    chunk_boundaries[-1] = file_size

    mini_chunk_size = 4096  # Read ahead by 4k bytes at a time

    for bi in range(1, len(chunk_boundaries) - 1):

        initial_position = chunk_boundaries[bi]

        file.seek(initial_position)  # Start at boundary guess

        while True:

            mini_chunk = file.read(mini_chunk_size)  # Read a mini chunk

            # If EOF, this boundary should be at the end of the file

            if mini_chunk == b"":

                chunk_boundaries[bi] = file_size

                break

            # Find the special token in the mini chunk

            found_at = mini_chunk.find(split_special_token)

            if found_at != -1:

                chunk_boundaries[bi] = initial_position + found_at

                break

            initial_position += mini_chunk_size

    # Make sure all boundaries are unique, but might be fewer than desired_num_chunks

    return sorted(set(chunk_boundaries))

def pretokenize_file(file: BinaryIO, special_tokens: List[str] = None) -> Counter[Tuple[bytes, ...], int]:

    """

    Pre-tokenize a file into words and count their occurrences, using multiprocessing.

    The primary function to execute the parallel pre-tokenization process.

    """

    # 1. prepare special tokens

    split_special_token = b"<|endoftext|>"

    # 2. find chunk boundaries

    num_processes = 16

    chunk_max_size = 1024 * 1024 * 500  # 500 MB

    desired_chunks = max(num_processes, file.seek(0, os.SEEK_END) // chunk_max_size + 1)

    boundaries = find_chunk_boundaries(file, desired_num_chunks=desired_chunks, split_special_token=split_special_token)

    # 3. pretokenize chunks in parallel

    def chunk_generator():

        for start, end in zip(boundaries[:-1], boundaries[1:]):

            file.seek(start)

            chunk_data = file.read(end - start)

            yield (chunk_data, special_tokens)

    final_words_counter: Counter[Tuple[bytes, ...], int] = Counter()

    try:

        with multiprocessing.Pool(processes=num_processes) as pool:

            for result in pool.imap_unordered(pretokenize_chunk, chunk_generator()):

                final_words_counter.update(result)

    except Exception as e:

        print(f"Multiprocessing failed, falling back to single process: {e}")

        for chunk_data, special_tokens in chunk_generator():

            counter = pretokenize_chunk((chunk_data, special_tokens))

            final_words_counter.update(counter)

    return final_words_counter

def pretokenize_chunk(args) -> Counter[Tuple[bytes, ...], int]:

    """

    Pre-tokenize a chunk, replace special tokens, and count word occurrences.

    This function runs in parallel worker processes.

    """

    chunk, special_tokens = args

    chunk_str = chunk.decode("utf-8", errors="ignore")

    words_counter = Counter()

    last_end = 0

    # pretokenize between special tokens

    if special_tokens and len(special_tokens) > 0:

        for match in regex.finditer("|".join(re.escape(st) for st in special_tokens), chunk_str):

            pre_token_text = chunk_str[last_end : match.start()]

            for pre_token in regex.finditer(PAT, pre_token_text):

                word_bytes = pre_token.group(0).encode("utf-8")

                words_counter[split_bytes(word_bytes)] += 1

                del word_bytes  # free memory

            last_end = match.end()

    # pretokenize after last special token

    pre_token_text = chunk_str[last_end:]

    for pre_token in regex.finditer(PAT, pre_token_text):

        word_bytes = pre_token.group(0).encode("utf-8")

        words_counter[split_bytes(word_bytes)] += 1

        del word_bytes

    return words_counter

# train.py

import pdb

import pickle

from tqdm import tqdm

from collections import Counter, defaultdict

from cs336_basics.bpe.pretokenize import pretokenize_file

from cs336_basics.bpe.utils import myHeap, BytePair, WordRef

def train_bpe(

    input_path: str,

    vocab_size: int,

    special_tokens: list[str] = None,

):

    """Train a BPE tokenizer on the given input file."""

    # 1. initialize byte-level vocab and add special tokens

    vocab = {i: bytes([i]) for i in range(256)}  # int to bytes

    for special_token in special_tokens or []:

        vocab[len(vocab)] = special_token.encode("utf-8", errors="ignore")

    # 2. pretokenize input file into words and count occurrences

    with open(input_path, "rb") as f:

        word_counts = pretokenize_file(f, special_tokens)  # Counter[Tuple[bytes, ...], int]

    # 3. create WordRef objects and initialize byte pair stats

    word_refs: list[WordRef] = []

    pair_to_count: defaultdict[BytePair, int] = defaultdict(int)

    pair_to_word: defaultdict[BytePair, set[int]] = defaultdict(set)

    for word_tuple, count in word_counts.items():

        word_ref = WordRef(tokens=list(word_tuple), count=count)

        word_refs.append(WordRef(tokens=list(word_tuple), count=count))

        for bp, cnt in word_ref.byte_pairs_dict.items():

            pair_to_count[bp] += cnt

            pair_to_word[bp].add(len(word_refs) - 1)

    pair_heap = myHeap()

    for bp, cnt in pair_to_count.items():

        pair_heap.push((cnt, bp))

    pair_heap_to_delete = dict()  # heap lazy deletion map

    # 4. merge BPEs until reaching the desired vocab size

    merges = []

    with tqdm(total=vocab_size - len(vocab), desc="BPE Merging") as pbar:

        while len(vocab) < vocab_size:

            # Lazy deletion: pop until we find a valid top pair

            while pair_heap.size > 0:

                top_cnt, top_bp = pair_heap.pop()

                if top_bp in pair_heap_to_delete:

                    top_cnt -= pair_heap_to_delete[top_bp]

                    if top_cnt > 0:

                        pair_heap.push((top_cnt, top_bp))

                    del pair_heap_to_delete[top_bp]

                else:

                    break

            # merge best pair

            assert top_bp is not None, "No more pairs to merge"

            merges.append((top_bp.left, top_bp.right))

            new_token = top_bp.merged_bytes

            vocab[len(vocab)] = new_token

            # update word counts and pair stats

            pair_to_count_change: defaultdict[BytePair, int] = defaultdict(int)

            pair_heap_to_add = defaultdict(int)  # cumulative count changes for lazy deletion

            for word_ref_id in pair_to_word[top_bp]:

                word_ref = word_refs[word_ref_id]

                old_bp_dict = word_ref.byte_pairs_dict

                word_ref.merge(top_bp)

                new_bp_dict = word_ref.byte_pairs_dict

                # update old pairs

                for bp, cnt in old_bp_dict.items():

                    if bp not in new_bp_dict and bp != top_bp:

                        pair_to_word[bp].remove(word_ref_id)

                    change = new_bp_dict.get(bp, 0) - cnt

                    pair_to_count_change[bp] += change

                    if change < 0:

                        pair_heap_to_delete[bp] = pair_heap_to_delete.get(bp, 0) - change

                # update new pairs

                for bp, cnt in new_bp_dict.items():

                    if bp not in old_bp_dict:

                        pair_to_word[bp].add(word_ref_id)

                        pair_to_count_change[bp] += cnt

                        pair_heap_to_add[bp] += cnt

            # apply count changes

            for bp, change in pair_to_count_change.items():

                pair_to_count[bp] += change

                if pair_to_count[bp] <= 0:

                    del pair_to_count[bp]

                    del pair_to_word[bp]

            for bp, add_cnt in pair_heap_to_add.items():

                pair_heap.push((add_cnt, bp))  # push updated count to heap

            del pair_heap_to_delete[top_bp]  # remove merged pair from lazy deletion map

            assert top_bp not in pair_to_count and top_bp not in pair_to_word

            # update progress bar

            pbar.update(1)

    return vocab, merges

# pretokenize.py

def pretokenize_text(text: str, special_tokens: List[str] = None) -> Iterable[bytes]:

    """

    Pre-tokenize a string into words, replacing special tokens.

    Yields pre-tokenized byte sequences.

    """

    last_end = 0  # Track the end of the last match

    if special_tokens and len(special_tokens) > 0:

        special_tokens.sort(key=len, reverse=True)  # Longer tokens first

        for match in re.finditer("|".join(re.escape(st) for st in special_tokens or []), text):

            special_token = match.group()

            # yield pre-tokens before the special token

            for pre_token in regex.finditer(PAT, text[last_end : match.start()]):

                yield pre_token.group().encode("utf-8")

            yield special_token.encode("utf-8")

            last_end = match.end()

    for pre_token in regex.finditer(PAT, text[last_end:]):

        yield pre_token.group().encode("utf-8")

# tokenizer.py

import pickle

import time

from typing import Dict, Tuple, List, Union, Iterable

from cs336_basics.bpe.pretokenize import pretokenize_text

from cs336_basics.bpe.utils import split_bytes, WordRef, BytePair

class Tokenizer:

    vocab: Dict[int, bytes]

    inv_vocab: Dict[bytes, int]

    merges: List[Tuple[bytes, bytes]]

    special_tokens: List[str]

    encode_cache: Dict[bytes, List[int]] = {}

    def __init__(self, vocab: Dict[int, bytes], merges: List[Tuple[bytes, bytes]], special_tokens: List[str] = None):

        self.vocab = vocab

        self.inv_vocab = {v: k for k, v in vocab.items()}

        self.merges = merges

        self.special_tokens = special_tokens or []

    @classmethod

    def from_file(cls, vocab_filepath: str, merges_filepath: str, special_tokens: List[str] | None = None):

        with open(vocab_filepath, "rb") as vf:

            vocab = pickle.load(vf)

        with open(merges_filepath, "rb") as mf:

            merges = pickle.load(mf)

        return cls(vocab=vocab, merges=merges, special_tokens=special_tokens)

    def _pretoken_to_ids(self, text: bytes) -> Iterable[int]:

        """Encode a byte string into a list of token IDs using the BPE merges"""

        if text in self.encode_cache:

            yield from self.encode_cache[text]

            return

        word_ref = WordRef(tokens=list(split_bytes(text)), count=1)

        byte_pairs_dict = word_ref.byte_pairs_dict

        for merge_left, merge_right in self.merges:

            bp = BytePair(merge_left, merge_right)

            if bp not in byte_pairs_dict:

                continue

            word_ref.merge(bp)

            byte_pairs_dict = word_ref.byte_pairs_dict

            if len(word_ref.tokens) == 1:

                break

        token_ids = []

        for token in word_ref.tokens:

            yield self.inv_vocab[token]

            token_ids.append(self.inv_vocab[token])

        self.encode_cache[text] = token_ids

    def encode(self, text: str) -> List[int]:

        """Encode a string into a list of token IDs using the BPE merges"""

        encoded_ids = []

        # reset cache with special tokens

        self.encode_cache = {st.encode("utf-8"): [self.inv_vocab[st.encode("utf-8")]] for st in self.special_tokens}

        for pre_toeken in pretokenize_text(text, self.special_tokens):

            for token_id in self._pretoken_to_ids(pre_toeken):

                encoded_ids.append(token_id)

        return encoded_ids

    def encode_iterable(self, iterable: Iterable[str]) -> Iterable[int]:

        """Encode an iterable of strings into a flat iterable of token IDs using the BPE merges"""

        self.encode_cache = {st.encode("utf-8"): [self.inv_vocab[st.encode("utf-8")]] for st in self.special_tokens}

        for text in iterable:

            for pre_toeken in pretokenize_text(text, self.special_tokens):

                for token_id in self._pretoken_to_ids(pre_toeken):

                    yield token_id

    def decode(self, ids: List[int]) -> str:

        bytes_seq = []

        for i in ids:

            if i >= len(self.vocab):

                Warning(f"Token ID {i} not in vocab, replacing with b''")

            bytes_seq.append(self.vocab.get(i, b""))

        seq = b"".join(self.vocab.get(i) for i in ids).decode("utf-8", errors="replace")

        return seq