CS336_Assignment 1 (basics)

# Unicode example
>>> ord('牛')
29275
>>> chr(29275)
'牛'

# A: What Unicode character does chr(0) return?
# Q: '\x00'
>>> chr(0)
'\x00'
# A: How does this character’s string representation (__repr__()) differ from its printed representa-tion?
# Q: It's a invisible string, so no output.
>>> print(chr(0))

# A: What happens when this character occurs in text? 
# It may be helpful to play around with the following in your Python interpreter and see if it matches your expectations:
# >>> chr(0)
# >>> print(chr(0))
# >>> "this is a test" + chr(0) + "string"
# >>> print("this is a test" + chr(0) + "string")
# Q:
>>> chr(0)
'\x00'
>>> print(chr(0))

>>> "this is a test" + chr(0) + "string"
'this is a test\x00string'
>>> print("this is a test" + chr(0) + "string")
this is a teststring

>>> test_string = "hello! こんにちは!"
>>> utf8_encoded = test_string.encode("utf-8")
>>> print(utf8_encoded)
b'hello! \xe3\x81\x93\xe3\x82\x93\xe3\x81\xab\xe3\x81\xa1\xe3\x81\xaf!'
>>> print(type(utf8_encoded))
<class 'bytes'>
>>> list(utf8_encoded)
[104, 101, 108, 108, 111, 33, 32, 227, 129, 147, 227, 130, 147, 227, 129, 171, 227, 129, 161, 227, 129, 175, 33]
>>> print(len(test_string))
13
>>> print(len(utf8_encoded))
23
>>> print(utf8_encoded.decode("utf-8"))
hello! こんにちは!

# Q: What are some reasons to prefer training our tokenizer on UTF-8 encoded bytes, rather than UTF-16 or UTF-32? 
# It may be helpful to compare the output of these encodings for various input strings.

# A: UTF-8 is a dynamic length code, so it just need a byte for ASCII, making it efficient for English.

# Q: Consider the following (incorrect) function, which is intended to decode a UTF-8 byte string into a Unicode string. 
# Why is this function incorrect? Provide an example of an input byte string that yields incorrect results.

def decode_utf8_bytes_to_str_wrong(bytestring: bytes):
    return "".join([bytes([b]).decode("utf-8") for b in bytestring])
>>> decode_utf8_bytes_to_str_wrong("hello".encode("utf-8"))
'hello'

# A: The functino just support to encode utf-8 byte by byte, but it will go wrong when we decode muti-byte string such as "你好". 

# Q: Give a two byte sequence that does not decode to any Unicode character(s).
# A: b'\xc3\x41' is a illegal two byte sequence.

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

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

# use findall
>>> re.findall(PAT, "some text that i'll pre-tokenize")
['some', ' text', ' that', ' i', "'ll", ' pre', '-', 'tokenize']

# use finditer is recommended, for flexble usage of memory
>>> [match.group(0) for match in re.finditer(PAT, "some text that i'll pre-tokenize")]         
['some', ' text', ' that', ' i', "'ll", ' pre', '-', 'tokenize']

# bpe example
def pre_tokenize(input_text: str, special_tokens: set):
    split_words = input_text.split(' ')
    pre_tokenization = {}
    
    for word in split_words:
        if word in special_tokens:
            pre_tokenization[(word,)] = pre_tokenization.get((word,), 0) + 1
        else:
            word_as_tuple = tuple(word) 
            pre_tokenization[word_as_tuple] = pre_tokenization.get(word_as_tuple, 0) + 1
            
    return pre_tokenization


def get_pair_frequencies(input : dict):
    output = {}
    for word, count in input.items():
        for i in range(len(word) - 1):
            left_char = word[i]
            right_char = word[i + 1]
            pair = (left_char, right_char)
            output[pair] = output.get(pair, 0) + count
    return output

def find_key(input : dict):
    max_value = max(input.values())
    max_key = [k for k, v in input.items() if v == max_value]
    return max(max_key)

def merge(vocabulary : dict, pair):
    new_vocab = {}
    for word_tuple, count in vocabulary.items():
        new_word_tuple = merge_single_word(word_tuple, pair)
        new_vocab[new_word_tuple] = count
    return new_vocab

def merge_single_word(word_tuple, pair_to_merge):
    new_word = list()
    i = 0
    while i < len(word_tuple):
        if i < len(word_tuple) - 1 and (word_tuple[i], word_tuple[i + 1]) == pair_to_merge:
            new_word.append(word_tuple[i] + word_tuple[i + 1])
            i += 2
        else:
            new_word.append(word_tuple[i])
            i += 1
    return tuple(new_word)

def main():
    special_tokens = {'<|endoftext|>'}
    input_text = "low low low low low lower lower widest widest widest newest newest newest newest newest newest"
    num_merges = 6 # 执行6次合并
    
    # 1. 初始化
    vocab = pre_tokenize(input_text, special_tokens) # 你的 pre_tokenization 字典
    
    # 2. 循环执行合并
    for i in range(num_merges):
        print(f"--- 合并第 {i+1} 次 ---")
        
        # 2a. 统计当前词汇表中的字节对频率
        pair_freqs = get_pair_frequencies(vocab)
        
        # 2b. 找到最佳对
        if not pair_freqs: # 如果没有对可以合并了
            break
        best_pair = find_key(pair_freqs)
        
        # 2c. 执行合并，更新词汇表
        vocab = merge(vocab, best_pair)
        
        print(f"合并了: {best_pair}")
        print("合并后的词汇表:")
        print(vocab)

    print("\n--- 最终结果 ---")
    print(vocab)
if __name__ == '__main__':
    main()

from .pretokenization_example import find_chunk_boundaries
import regex as re
import os
from typing import Dict, List, Tuple

def pre_tokenize(input_text: str, special_tokens: List[str]) -> Dict[Tuple[bytes, ...], int]:
    word_counts = {}
    
    # 构建一个能分割文本的、包含所有特殊token的正则表达式
    # re.escape确保了特殊token中的元字符被正确处理
    # f"({ ... })" 创建了一个捕获组，这样 re.split 可以在结果中保留分隔符
    special_pattern = "|".join(re.escape(s) for s in special_tokens)
    special_splitter = re.compile(f'({special_pattern})')

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

    # 1. 按照特殊token分割文本
    text_parts = special_splitter.split(input_text)

    # 2. 遍历分割后的部分
    for i, part in enumerate(text_parts):
        if not part:  # 跳过 re.split 可能产生的空字符串
            continue
            
        if i % 2 == 1:  # 奇数索引是分隔符本身，即特殊token
            token_bytes = part.encode('utf-8')
            word_tuple = (token_bytes,)
            word_counts[word_tuple] = word_counts.get(word_tuple, 0) + 1
        else:  # 偶数索引是普通文本
            # 对普通文本部分应用基础的正则表达式
            for match in base_pattern.finditer(part):
                token_str = match.group(0)
                token_bytes = token_str.encode('utf-8')
                word_tuple = tuple(bytes([b]) for b in token_bytes)
                word_counts[word_tuple] = word_counts.get(word_tuple, 0) + 1
                
    return word_counts

def get_pair_frequencies(input : dict):
    output = {}
    for word, count in input.items():
        for i in range(len(word) - 1):
            left_char = word[i]
            right_char = word[i + 1]
            pair = (left_char, right_char)
            output[pair] = output.get(pair, 0) + count
    return output

def find_key(input : dict):
    max_value = max(input.values())
    max_key = [k for k, v in input.items() if v == max_value]
    return max(max_key)

def merge_single_word(word_tuple, pair_to_merge):
    new_word = list()
    i = 0
    while i < len(word_tuple):
        if i < len(word_tuple) - 1 and (word_tuple[i], word_tuple[i + 1]) == pair_to_merge:
            new_word.append(word_tuple[i] + word_tuple[i + 1])
            i += 2
        else:
            new_word.append(word_tuple[i])
            i += 1
    return tuple(new_word)

def bpe(
    input_path: str,
    vocab_size: int,
    special_tokens: list[str],
) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
    """
    训练一个字节级别的BPE分词器。
    """
    # --- 1. 初始化 ---
    final_vocab: Dict[int, bytes] = {}
    merges: List[Tuple[bytes, bytes]] = []
    
    # 首先添加特殊tokens到词汇表开头
    current_id = 0
    for token_str in special_tokens:
        final_vocab[current_id] = token_str.encode('utf-8')
        current_id += 1
    
    # 然后填充基础词汇表 (0-255字节)
    for i in range(256):
        final_vocab[current_id] = bytes([i])
        current_id += 1
        
    # --- 2. 全局预分词和频率统计 ---
    word_counts: Dict[Tuple[bytes, ...], int] = {}
    
    with open(input_path, "rb") as f:
        boundaries = find_chunk_boundaries(
            f, 32, "<|endoftext|>".encode("utf-8"))
            
        for start, end in zip(boundaries[:-1], boundaries[1:]):
            f.seek(start)
            chunk_text = f.read(end - start).decode("utf-8", errors="ignore")
            
            # 对每个块进行预分词
            chunk_word_counts = pre_tokenize(chunk_text, special_tokens)
            
            # 聚合到全局词频统计中
            for word, count in chunk_word_counts.items():
                word_counts[word] = word_counts.get(word, 0) + count

    # --- 3. BPE 迭代合并 ---
    num_merges_needed = vocab_size - len(final_vocab)

    for i in range(num_merges_needed):
        pair_freqs = get_pair_frequencies(word_counts)
        if not pair_freqs:
            break
        
        best_pair = find_key(pair_freqs)
        merges.append(best_pair)
        
        words_to_update = {}
        # 1. 找出所有包含 best_pair 的词
        for word, count in word_counts.items():
            if best_pair[0] in word and best_pair[1] in word:
                # 这是一个潜在的候选词，但我们还需要检查它们是否相邻
                for j in range(len(word) - 1):
                    if (word[j], word[j+1]) == best_pair:
                        words_to_update[word] = count
                        break # 找到一个匹配就够了，处理下一个词
        
        # 2. 更新这些词
        for word, count in words_to_update.items():
            # 从旧字典中移除
            del word_counts[word]
            # 创建新词并添加到字典中
            new_word = merge_single_word(word, best_pair)
            word_counts[new_word] = word_counts.get(new_word, 0) + count

        new_token_bytes = best_pair[0] + best_pair[1]
        final_vocab[current_id] = new_token_bytes
        current_id += 1
    return final_vocab, merges