Tokenizer 全面对比：BPE vs WordPiece vs Unigram vs SentencePiece

英文 "Hello World" 的不同分词:
  字符级:  H e l l o   W o r l d          → 11 tokens
  词级:    Hello World                     → 2 tokens (但词表爆炸)
  BPE:    Hello World                      → 2 tokens (高频词保留)
  字节级:  Hello World                      → 2 tokens (256基础, 递归合并)

中文 "你好世界" 的不同分词:
  字符级:  你 好 世 界                      → 4 tokens
  BPE:    你 好 世 界                      → 4 tokens (中文常见)
  词级:    你好世界                         → 1 token (理想但不可达)

from collections import Counter, defaultdict

class BPETokenizer:
    def __init__(self):
        self.vocab = []  # token 列表
        self.merges = []  # 合并规则列表
    
    def _get_word_freqs(self, corpus):
        """统计词频，每个词表示为字符元组"""
        word_freqs = Counter()
        for text in corpus:
            words = text.split()
            for word in words:
                # 每个词用字符元组表示，末尾加 </w> 标记词边界
                word_tuple = tuple(word) + ('</w>',)
                word_freqs[word_tuple] += 1
        return word_freqs
    
    def _get_pair_freqs(self, word_freqs):
        """统计相邻 token 对的频率"""
        pair_freqs = Counter()
        for word, freq in word_freqs.items():
            for i in range(len(word) - 1):
                pair = (word[i], word[i + 1])
                pair_freqs[pair] += freq
        return pair_freqs
    
    def _merge(self, word_freqs, pair):
        """执行一次合并"""
        new_word_freqs = {}
        for word, freq in word_freqs.items():
            new_word = []
            i = 0
            while i < len(word):
                if i < len(word) - 1 and (word[i], word[i + 1]) == pair:
                    new_word.append(word[i] + word[i + 1])
                    i += 2
                else:
                    new_word.append(word[i])
                    i += 1
            new_word_freqs[tuple(new_word)] = freq
        return new_word_freqs
    
    def train(self, corpus, vocab_size=10000):
        # 初始化词表为所有字符
        word_freqs = self._get_word_freqs(corpus)
        
        # 初始词表 = 所有出现的字符
        char_set = set()
        for word in word_freqs:
            for char in word:
                char_set.add(char)
        self.vocab = sorted(char_set)
        
        # 迭代合并
        num_merges = vocab_size - len(self.vocab)
        for i in range(num_merges):
            pair_freqs = self._get_pair_freqs(word_freqs)
            if not pair_freqs:
                break
            
            # 选择频率最高的 pair
            best_pair = pair_freqs.most_common(1)[0][0]
            
            # 合并
            word_freqs = self._merge(word_freqs, best_pair)
            new_token = best_pair[0] + best_pair[1]
            self.vocab.append(new_token)
            self.merges.append(best_pair)
        
        return self.vocab, self.merges
    
    def tokenize(self, text):
        """使用训练好的 merges 对文本分词"""
        words = text.split()
        tokens = []
        for word in words:
            word_tuple = tuple(word) + ('</w>',)
            
            # 贪心地应用 merge 规则
            while len(word_tuple) > 1:
                # 找到优先级最高的 merge
                pairs = [(word_tuple[i], word_tuple[i+1]) 
                         for i in range(len(word_tuple) - 1)]
                
                applicable = [(pair, self.merges.index(pair)) 
                             for pair in pairs if pair in self.merges]
                
                if not applicable:
                    break
                
                # 选择最早学到的 merge（优先级最高）
                best_pair = min(applicable, key=lambda x: x[1])[0]
                
                # 执行合并
                new_word = []
                i = 0
                while i < len(word_tuple):
                    if i < len(word_tuple) - 1 and \
                       (word_tuple[i], word_tuple[i+1]) == best_pair:
                        new_word.append(word_tuple[i] + word_tuple[i+1])
                        i += 2
                    else:
                        new_word.append(word_tuple[i])
                        i += 1
                word_tuple = tuple(new_word)
            
            tokens.extend(word_tuple)
        
        return tokens

# Byte-level BPE 的关键区别
def text_to_bytes(text):
    """将文本编码为字节序列，避免 OOV"""
    # UTF-8 编码后映射到可打印字符范围
    raw_bytes = text.encode('utf-8')
    # GPT-2 的字节到 Unicode 映射
    bs = list(range(ord("!"), ord("~") + 1)) + \
         list(range(ord("¡"), ord("¬") + 1)) + \
         list(range(ord("®"), ord("ÿ") + 1))
    cs = bs[:]
    n = 0
    for b in range(256):
        if b not in bs:
            bs.append(b)
            cs.append(256 + n)
            n += 1
    cs = [chr(c) for c in cs]
    byte_to_unicode = dict(zip(bs, cs))
    
    return ''.join([byte_to_unicode[b] for b in raw_bytes])

# 优势: 永远不会出现 OOV (任何文本都是字节的组合)

class WordPieceTokenizer:
    def __init__(self):
        self.vocab = set()
        self.unk_token = '[UNK]'
    
    def train(self, corpus, vocab_size=30000):
        # 预处理: 统计所有词的频率
        word_freqs = Counter()
        for text in corpus:
            for word in text.split():
                word_freqs[word] += 1
        
        # 初始化: 将每个词拆分为字符，首字符 + ##后续字符
        # ## 前缀表示该 token 是词的延续部分
        splits = {}
        for word, freq in word_freqs.items():
            split = [word[0]] + ['##' + c for c in word[1:]]
            splits[word] = split
            for s in split:
                self.vocab.add(s)
        
        # 迭代合并
        while len(self.vocab) < vocab_size:
            # 计算所有相邻 pair 的分数
            pair_scores = {}
            pair_freqs = Counter()
            single_freqs = Counter()
            
            for word, split in splits.items():
                freq = word_freqs[word]
                if len(split) == 1:
                    single_freqs[split[0]] += freq
                    continue
                for i in range(len(split) - 1):
                    pair = (split[i], split[i + 1])
                    pair_freqs[pair] += freq
                    single_freqs[split[i]] += freq
                single_freqs[split[-1]] += freq
            
            for pair, freq in pair_freqs.items():
                score = freq / (single_freqs[pair[0]] * single_freqs[pair[1]])
                pair_scores[pair] = score
            
            if not pair_scores:
                break
            
            # 选择分数最高的 pair
            best_pair = max(pair_scores, key=pair_scores.get)
            
            # 合并
            merged = best_pair[0] + best_pair[1].replace('##', '')
            self.vocab.add(merged)
            
            # 更新所有词的分割
            for word in splits:
                split = splits[word]
                if len(split) == 1:
                    continue
                i = 0
                while i < len(split) - 1:
                    if (split[i], split[i + 1]) == best_pair:
                        split = split[:i] + [merged] + split[i + 2:]
                    else:
                        i += 1
                splits[word] = split
        
        return self.vocab
    
    def tokenize(self, text):
        """贪心最长匹配分词"""
        tokens = []
        for word in text.split():
            # 贪心: 从最长开始匹配
            start = 0
            sub_tokens = []
            while start < len(word):
                end = len(word)
                cur_substr = None
                while start < end:
                    substr = word[start:end]
                    if start > 0:
                        substr = '##' + substr
                    if substr in self.vocab:
                        cur_substr = substr
                        break
                    end -= 1
                
                if cur_substr is None:
                    sub_tokens = [self.unk_token]
                    break
                sub_tokens.append(cur_substr)
                start = end
            
            tokens.extend(sub_tokens)
        return tokens

import math

class UnigramTokenizer:
    def __init__(self):
        self.vocab = {}  # token -> log_prob
    
    def train(self, corpus, vocab_size=8000, shrinking_factor=0.75):
        # Step 1: 枚举所有可能的子串作为候选词表
        word_freqs = Counter()
        for text in corpus:
            for word in text.split():
                word_freqs[word] += 1
        
        # 生成所有子串候选
        substring_freqs = Counter()
        for word, freq in word_freqs.items():
            for i in range(len(word)):
                for j in range(i + 1, len(word) + 1):
                    substring = word[i:j]
                    substring_freqs[substring] += freq
        
        # 初始化词表（取高频子串）
        candidates = substring_freqs.most_common(vocab_size * 10)
        self.vocab = {token: -math.log(freq / sum(substring_freqs.values())) 
                      for token, freq in candidates}
        
        # Step 2: EM 迭代优化
        for iteration in range(100):
            # E-step: 计算每个词的最优分词
            word_segments = {}
            total_loss = 0
            for word, freq in word_freqs.items():
                best_seg, best_loss = self._best_segmentation(word)
                word_segments[word] = best_seg
                total_loss += best_loss * freq
            
            # M-step: 更新 token 概率
            token_freqs = Counter()
            for word, freq in word_freqs.items():
                for token in word_segments[word]:
                    token_freqs[token] += freq
            
            total = sum(token_freqs.values())
            for token in self.vocab:
                if token in token_freqs:
                    self.vocab[token] = -math.log(token_freqs[token] / total)
            
            # Step 3: 删除低贡献 token
            if len(self.vocab) > vocab_size:
                # 计算删除每个 token 后的似然变化
                sorted_tokens = sorted(self.vocab.items(), 
                                       key=lambda x: x[1])  # log_prob 越大越罕见
                keep_count = int(len(self.vocab) * shrinking_factor)
                self.vocab = dict(sorted_tokens[:keep_count])
        
        return self.vocab
    
    def _best_segmentation(self, word):
        """使用动态规划找到最优分词（最大化似然）"""
        n = len(word)
        # dp[i] = (最大log_prob, 最优分词)
        dp = [(0, [])] * (n + 1)
        
        for i in range(1, n + 1):
            best_prob = float('-inf')
            best_seg = []
            for j in range(i):
                substr = word[j:i]
                if substr in self.vocab:
                    prob = dp[j][0] + self.vocab[substr]
                    if prob > best_prob:
                        best_prob = prob
                        best_seg = dp[j][1] + [substr]
            dp[i] = (best_prob, best_seg)
        
        return dp[n][1], dp[n][0]
    
    def tokenize(self, text):
        """使用 Viterbi 算法找最优分词"""
        tokens = []
        for word in text.split():
            # 通过 forward-backward 采样获得多种分词可能
            seg, _ = self._best_segmentation(word)
            tokens.extend(seg)
        return tokens

# SentencePiece 的核心创新:
# 1. 将空格编码为特殊字符 ▁ (U+2581)
# 2. 在原始字节序列上直接分词，不依赖空格分词
# 3. 支持 BPE 和 Unigram 两种算法
# 4. 可逆: detokenize(tokenize(text)) == text

# 安装: pip install sentencepiece
import sentencepiece as spm

# 训练
spm.SentencePieceTrainer.train(
    input='corpus.txt',
    model_prefix='spm_model',
    vocab_size=32000,
    # 算法选择
    model_type='unigram',  # 或 'bpe', 'char', 'word'
    # 特殊标记
    pad_id=0, pad_piece='<pad>',
    unk_id=1, unk_piece='<unk>',
    bos_id=2, bos_piece='<s>',
    eos_id=3, eos_piece='</s>',
    # 字节回退 (处理OOV)
    byte_fallback=True,
    # 归一化
    normalization_rule_name='nmt_nfkc',
)

# 使用
sp = spm.SentencePieceProcessor()
sp.load('spm_model.model')

# 编码
tokens = sp.encode('Hello World!', out_type=str)
ids = sp.encode('Hello World!', out_type=int)

# 解码
text = sp.decode(ids)  # 逆操作，恢复原始文本

# 空格处理的差异示例
text = "Hello  World"  # 注意双空格

# WordPiece (BERT): 先按空格分词，空格信息丢失
# → ['Hello', 'World']  (双空格信息丢失)

# SentencePiece: 将空格替换为 ▁
# → ['▁Hello', '▁', '▁World']  (保留了双空格信息)

# Byte-level BPE (GPT-2): 将空格编码为 Ġ
# → ['Hello', 'Ġ', 'ĠWorld']  (通过字节映射保留)

文本: "The transformer architecture revolutionized NLP. Transformer架构彻底改变了自然语言处理领域。"

| Tokenizer | Token 数 | 压缩率 (字符/token) | 模型 |
|-----------|----------|---------------------|------|
| GPT-2 BPE | 24 | 2.92 | tiktoken |
| GPT-4 BPE | 18 | 3.89 | tiktoken |
| LLaMA SPM | 21 | 3.33 | sentencepiece |
| BERT WP   | 28 | 2.50 | wordpiece |
| Qwen BPE  | 16 | 4.38 | tiktoken |

code = "def fibonacci(n):\n    if n <= 1:\n        return n\n    return fib(n-1) + fib(n-2)"

# 不同 tokenizer 对代码的处理:
# GPT-4:  "def", " fib", "onacci", "(", "n", "):", "\n", "    ", "if", ...
#         → 代码token较多但合理

# LLaMA:  "def", "▁fib", "onacci", "(", "n", "):", "\n", "▁▁▁▁", "if", ...
#         → ▁ 编码空格，缩进信息保留

# BERT:   "def", "fib", "##ona", "##cci", "(", "n", ")", ":", ...
#         → WordPiece 拆分过细，## 前缀干扰

是否需要多语言支持?
├─ 是 → SentencePiece (Unigram)
│       └─ 是否需要字节回退? → byte_fallback=True
└─ 否 → 是否以英文为主?
        ├─ 是 → Byte-level BPE (GPT-4 style)
        │       └─ 词表: 32K(小模型) ~ 128K(大模型)
        └─ 否 → 考虑领域特性
                ├─ 中文为主 → Qwen/GLM 的 BPE 变体
                └─ 代码为主 → GPT-4 的 BPE (代码覆盖好)

# 1. 训练数据应与模型训练数据同分布
# 2. 注意数据采样比例（多语言场景）
# 3. 保留特殊标记的 ID 空间

# SentencePiece 多语言训练示例
spm.SentencePieceTrainer.train(
    input=['en_corpus.txt', 'zh_corpus.txt', 'code_corpus.txt'],
    model_prefix='multilingual_spm',
    vocab_size=64000,
    model_type='unigram',
    # 多语言采样比例
    input_sentence_size=5000000,
    shuffle_input_sentence=True,
    # 特殊标记
    user_defined_symbols=['<sep>', '<mask>', '<code>', '</code>'],
    # 字节回退
    byte_fallback=True,
    # 限制 token 长度
    max_sentence_length=16384,
    # 控制字符处理
    split_by_unicode_script=True,
    split_by_whitespace=True,
    split_by_number=True,
    max_sentencepiece_length=16,
)