GPU 別再空等 CPU：tokenizers v1 怎麼把切詞加速數十倍，還保證結果一模一樣

# tokenizers 四階段示範：正規化 -> 預切詞 -> 模型(BPE) -> 後處理
# 現場用幾句話訓練一個小 BPE，全程離線，不需要下載任何模型
import json
import sys
from pathlib import Path

import tokenizers
from tokenizers import Tokenizer, decoders, models, normalizers, pre_tokenizers, processors, trainers

# Windows 主控台預設編碼不一定是 UTF-8；先強制改，避免印出 byte-level 符號時噴錯
sys.stdout.reconfigure(encoding='utf-8')

HERE = Path(__file__).parent

CORPUS = [
    'A tokenizer converts text into the list of integers a model reads.',
    'Normalization applies operations such as lowercasing or Unicode normalization.',
    'Pre-tokenization splits the text into smaller pieces called pre-tokens.',
    'The model turns each pre-token into tokens and maps them to IDs in its vocabulary.',
    'Post-processing adds any special tokens the model expects.',
    'BPE repeatedly joins the highest ranked adjacent pair until no ranked pair remains.',
    'Your GPUs should never sit idle waiting for the CPU to complete its tokenization.',
    'Tokenization should be light and should scale with your workflow.',
]
EXTRA = ['Café GPUs never sit idle!', 'Tokenization should scale with your workflow.']


def build_tokenizer():
    tok = Tokenizer(models.BPE())
    # 階段一 正規化：統一 Unicode 寫法並轉小寫（有損，解碼回來會是小寫）
    tok.normalizer = normalizers.Sequence([normalizers.NFKC(), normalizers.Lowercase()])
    # 階段二 預切詞：byte-level 讓任何字元都能拆成位元組，不會出現認不得的字
    tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tok.decoder = decoders.ByteLevel()
    trainer = trainers.BpeTrainer(
        vocab_size=400,
        special_tokens=['[CLS]', '[SEP]'],
        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
    )
    # 階段三 模型：訓練時學出合併排名，編碼時照排名合併
    tok.train_from_iterator(CORPUS, trainer)
    # 階段四 後處理：句首句尾補特殊標記（ID 要等訓練完才知道，所以最後才設）
    tok.post_processor = processors.TemplateProcessing(
        single='[CLS] $A [SEP]',
        special_tokens=[('[CLS]', tok.token_to_id('[CLS]')), ('[SEP]', tok.token_to_id('[SEP]'))],
    )
    return tok


def show_stages(tok, text):
    print('原文        :', text)
    norm = tok.normalizer.normalize_str(text)
    print('1 正規化    :', norm)
    pieces = tok.pre_tokenizer.pre_tokenize_str(norm)
    print('2 預切詞    :', [p for p, _ in pieces])
    core = tok.encode(text, add_special_tokens=False)
    print('3 模型(BPE) :', core.tokens)
    print('  對應 ID   :', core.ids)
    full = tok.encode(text)
    print('4 後處理    :', full.tokens)
    back = tok.decode(full.ids)
    print('解碼還原    :', back)
    # 正規化是有損的，所以拿正規化後的文字比對才公平
    print('還原檢查    :', '一致' if back == norm else '不一致！')
    print()


def main():
    tok = build_tokenizer()
    print('tokenizers 版本:', tokenizers.__version__)
    print()
    for text in EXTRA:
        show_stages(tok, text)
    tok.save(str(HERE / 'tokenizer.json'))
    # 快照：訓練當下每句話的 token ID，升級後用來驗證『輸出沒變』
    snapshot = {t: tok.encode(t).ids for t in CORPUS + EXTRA}
    (HERE / 'ids_snapshot.json').write_text(
        json.dumps(snapshot, ensure_ascii=False, indent=1), encoding='utf-8'
    )
    print('已存檔：tokenizer.json、ids_snapshot.json（升級前後都用同一份，才是公平比較）')


if __name__ == '__main__':
    main()

# 量測：逐筆 encode vs 批次 encode_batch vs decode_batch，並確認升級後 token ID 沒變
# 同一份 tokenizer.json 在不同版本載入，才能在『同樣輸入、同樣輸出』下比較速度
import json
import os
import sys
import time
from pathlib import Path

import tokenizers
from tokenizers import Tokenizer

from stages import CORPUS, EXTRA

sys.stdout.reconfigure(encoding='utf-8')

HERE = Path(__file__).parent
REPEAT = 5000  # 把句子重複多次，讓總量夠大，計時才不會被雜訊淹沒
ROUNDS = 3


def best_of(fn):
    # 取最快的一次：排除首次載入、快取與作業系統排程造成的偶發變慢
    best = None
    for _ in range(ROUNDS):
        t0 = time.perf_counter()
        fn()
        dt = time.perf_counter() - t0
        best = dt if best is None else min(best, dt)
    return best


def check_ids(tok):
    snap_path = HERE / 'ids_snapshot.json'
    if not snap_path.exists():
        print('找不到 ids_snapshot.json，請先執行 python stages.py')
        return False
    snap = json.loads(snap_path.read_text(encoding='utf-8'))
    bad = [t for t, ids in snap.items() if tok.encode(t).ids != ids]
    if bad:
        print(f'ID 檢查：有 {len(bad)} 句和快照不同！例如：{bad[0]}')
    else:
        print(f'ID 檢查：{len(snap)} 句的 token ID 全部和快照相同（升級沒有改變輸出）')
    return not bad


def show_history():
    files = sorted(HERE.glob('results_*.json'))
    if len(files) < 2:
        print('提示：升級版本後再跑一次，這裡就會出現新舊對照表')
        return
    print('--- 版本對照 ---')
    for f in files:
        r = json.loads(f.read_text(encoding='utf-8'))
        print('{version:>12} | 逐筆 {single_mb_s:>7.1f} MB/s | 批次 {batch_mb_s:>7.1f} MB/s | 解碼 {decode_m_tok_s:>6.2f} M tok/s | ID 相同：{ids_ok}'.format(**r))


def main():
    tok = Tokenizer.from_file(str(HERE / 'tokenizer.json'))
    version = tokenizers.__version__
    print('tokenizers 版本:', version, '| CPU 邏輯核心數:', os.cpu_count())
    ids_ok = check_ids(tok)

    texts = (CORPUS + EXTRA) * REPEAT
    total_mb = sum(len(t.encode('utf-8')) for t in texts) / 1e6
    ids_list = [e.ids for e in tok.encode_batch(texts)]
    n_tokens = sum(len(ids) for ids in ids_list)

    # 逐筆：在 Python 迴圈裡一句一句呼叫，含 Python 呼叫本身的開銷
    t_single = best_of(lambda: [tok.encode(t) for t in texts])
    # 批次：一次交出整批文字，函式庫有機會把工作分給多個 CPU 核心
    t_batch = best_of(lambda: tok.encode_batch(texts))
    t_decode = best_of(lambda: tok.decode_batch(ids_list))

    single_mb_s = total_mb / t_single
    batch_mb_s = total_mb / t_batch
    decode_m_tok_s = n_tokens / t_decode / 1e6
    print(f'資料量：{len(texts)} 句 / {total_mb:.2f} MB / {n_tokens} tokens')
    print(f'逐筆 encode  ：{t_single:.3f} 秒（{single_mb_s:.1f} MB/s）')
    print(f'encode_batch ：{t_batch:.3f} 秒（{batch_mb_s:.1f} MB/s）')
    print(f'decode_batch ：{t_decode:.3f} 秒（{decode_m_tok_s:.2f} M tokens/s）')
    print(f'批次相對逐筆的加速比：{t_single / t_batch:.1f} 倍')

    result = {
        'version': version,
        'single_mb_s': round(single_mb_s, 2),
        'batch_mb_s': round(batch_mb_s, 2),
        'decode_m_tok_s': round(decode_m_tok_s, 3),
        'ids_ok': ids_ok,
    }
    out = HERE / f'results_{version}.json'
    out.write_text(json.dumps(result, ensure_ascii=False, indent=1), encoding='utf-8')
    print()
    show_history()


if __name__ == '__main__':
    main()