LLM Tokenizer 方法总结

# 伪代码：BPE 编码
def bpe_encode(word, merge_rules):
    tokens = list(word) + ['</w>']       # 字符级初始化
    for pair in merge_rules:              # 按训练顺序应用合并规则
        tokens = apply_merge(tokens, pair)
    return tokens

from collections import Counter, defaultdict
from typing import Dict, List, Tuple


# ═══════════════════════════════════════════════════
#  BPE 训练：从语料 → 词表 + 合并规则
# ═══════════════════════════════════════════════════
def get_vocab(corpus: List[str]) -> Dict[str, int]:
    """把语料转换为字符级词频词典（每个词加 </w> 结尾标记）"""
    vocab = Counter()
    for text in corpus:
        for word in text.split():
            # 字符间加空格，末尾加 </w>
            vocab[' '.join(list(word)) + ' </w>'] += 1
    return dict(vocab)


def get_stats(vocab: Dict[str, int]) -> Dict[Tuple[str, str], int]:
    """统计词典中所有相邻 bigram 的频率"""
    pairs = Counter()
    for word, freq in vocab.items():
        symbols = word.split()
        for i in range(len(symbols) - 1):
            pairs[(symbols[i], symbols[i + 1])] += freq
    return dict(pairs)


def merge_vocab(pair: Tuple[str, str], vocab: Dict[str, int]) -> Dict[str, int]:
    """将词典中所有词里的 pair 合并为一个新 token"""
    bigram = ' '.join(pair)
    replacement = ''.join(pair)
    return {
        word.replace(bigram, replacement): freq
        for word, freq in vocab.items()
    }


def train_bpe(
    corpus: List[str],
    target_vocab_size: int
) -> Tuple[Dict[str, int], List[Tuple[str, str]]]:
    """
    BPE 训练主函数
    Returns:
        token2id: 词表（token → id）
        merge_rules: 有序合并规则列表
    """
    vocab = get_vocab(corpus)
    
    # 初始词表：所有出现过的字符 + 特殊符号
    all_chars = set()
    for word in vocab:
        all_chars.update(word.split())
    token2id = {ch: i for i, ch in enumerate(sorted(all_chars))}
    
    merge_rules = []
    
    while len(token2id) < target_vocab_size:
        pairs = get_stats(vocab)
        if not pairs:
            break
        
        # 选频率最高的 pair
        best_pair = max(pairs, key=pairs.get)
        
        # 合并并更新词典
        vocab = merge_vocab(best_pair, vocab)
        
        # 记录合并规则 + 新 token
        new_token = ''.join(best_pair)
        merge_rules.append(best_pair)
        token2id[new_token] = len(token2id)
    
    return token2id, merge_rules


# ═══════════════════════════════════════════════════
#  BPE 编码：文本 → token id 序列
# ═══════════════════════════════════════════════════
def bpe_tokenize_word(word: str, merge_rules: List[Tuple[str, str]]) -> List[str]:
    """对单个词应用 BPE 合并规则，返回 token 字符串列表"""
    tokens = list(word) + ['</w>']
    
    for pair in merge_rules:
        i = 0
        new_tokens = []
        while i < len(tokens):
            # 如果当前位置匹配目标 pair，则合并
            if i < len(tokens) - 1 and (tokens[i], tokens[i+1]) == pair:
                new_tokens.append(tokens[i] + tokens[i+1])
                i += 2
            else:
                new_tokens.append(tokens[i])
                i += 1
        tokens = new_tokens
    
    return tokens


def bpe_encode(
    text: str,
    token2id: Dict[str, int],
    merge_rules: List[Tuple[str, str]]
) -> List[int]:
    """文本 → token id 列表（未知 token 用 UNK id=0）"""
    unk_id = token2id.get('<unk>', 0)
    ids = []
    for word in text.split():
        tokens = bpe_tokenize_word(word, merge_rules)
        ids.extend(token2id.get(t, unk_id) for t in tokens)
    return ids


# ═══════════════════════════════════════════════════
#  BPE 解码：token id 序列 → 文本
# ═══════════════════════════════════════════════════
def bpe_decode(
    ids: List[int],
    id2token: Dict[int, str]
) -> str:
    """token id 列表 → 原始文本（去掉 </w> 标记，恢复空格）"""
    tokens = [id2token.get(i, '<unk>') for i in ids]
    # </w> 代表词尾，后面应该跟空格（除了最后一个 token）
    text = ''.join(tokens).replace('</w>', ' ').strip()
    return text


# ═══════════════════════════════════════════════════
#  使用示例
# ═══════════════════════════════════════════════════
if __name__ == '__main__':
    corpus = [
        "low low low low low",
        "lower lower",
        "newest newest newest newest newest newest",
        "widest widest widest",
    ]
    
    token2id, merge_rules = train_bpe(corpus, target_vocab_size=30)
    id2token = {v: k for k, v in token2id.items()}
    
    # 编码
    text = "low newest"
    ids = bpe_encode(text, token2id, merge_rules)
    print(f"encode('{text}') = {ids}")
    # 例如: encode('low newest') = [14, 22, 23]
    
    # 解码
    decoded = bpe_decode(ids, id2token)
    print(f"decode({ids}) = '{decoded}'")
    # 例如: decode([14, 22, 23]) = 'low newest'
    
    print(f"学到的合并规则（前10条）: {merge_rules[:10]}")
    print(f"词表大小: {len(token2id)}")

from typing import Dict, List, Optional


# ═══════════════════════════════════════════════════
#  WordPiece 训练：从语料构建词表（PMI 打分版）
# ═══════════════════════════════════════════════════
def train_wordpiece(
    corpus: List[str],
    target_vocab_size: int,
    special_tokens: List[str] = None
) -> List[str]:
    """
    简化版 WordPiece 训练
    Returns: 词表列表（按 token 字符串排列）
    """
    if special_tokens is None:
        special_tokens = ['[PAD]', '[UNK]', '[CLS]', '[SEP]', '[MASK]']
    
    # Step 1：构建字符级初始词表
    char_counts = Counter()
    word_counts = Counter()
    for text in corpus:
        for word in text.lower().split():
            word_counts[word] += 1
            for i, ch in enumerate(word):
                # 非首字符加 ## 前缀
                symbol = ch if i == 0 else '##' + ch
                char_counts[symbol] += 1
    
    # 词表从特殊 token + 字符开始
    vocab = set(special_tokens) | set(char_counts.keys())
    
    # Step 2：迭代合并（PMI 打分）
    # 将每个词表示为 symbol 序列
    word_symbols = {}
    for word in word_counts:
        symbols = [word[0]] + ['##' + ch for ch in word[1:]]
        word_symbols[word] = symbols
    
    while len(vocab) < target_vocab_size:
        # 计算每个 symbol 的出现频率
        sym_freq = Counter()
        pair_freq = Counter()
        for word, count in word_counts.items():
            syms = word_symbols[word]
            for s in syms:
                sym_freq[s] += count
            for i in range(len(syms) - 1):
                pair_freq[(syms[i], syms[i+1])] += count
        
        if not pair_freq:
            break
        
        total = sum(sym_freq.values())
        # PMI 打分：P(AB) / (P(A) * P(B))
        best_pair, best_score = None, -1
        for (a, b), ab_count in pair_freq.items():
            score = (ab_count / total) / ((sym_freq[a] / total) * (sym_freq[b] / total))
            if score > best_score:
                best_score, best_pair = score, (a, b)
        
        if best_pair is None:
            break
        
        # 合并 best_pair → 新 token
        a, b = best_pair
        # 去掉 b 开头的 ## 前缀（如果 b 是 ##x 则合并后是 ax）
        new_token = a + (b[2:] if b.startswith('##') else b)
        vocab.add(new_token)
        
        # 更新所有词的 symbol 序列
        for word in word_symbols:
            syms = word_symbols[word]
            merged = []
            i = 0
            while i < len(syms):
                if i < len(syms) - 1 and syms[i] == a and syms[i+1] == b:
                    merged.append(new_token)
                    i += 2
                else:
                    merged.append(syms[i])
                    i += 1
            word_symbols[word] = merged
    
    return sorted(vocab)


# ═══════════════════════════════════════════════════
#  WordPiece 推理：最长匹配（BERT 官方算法）
# ═══════════════════════════════════════════════════
class WordPieceTokenizer:
    """BERT 风格的 WordPiece tokenizer"""
    
    def __init__(self, vocab: List[str], unk_token: str = '[UNK]'):
        self.vocab = set(vocab)
        self.token2id = {t: i for i, t in enumerate(vocab)}
        self.id2token = {i: t for i, t in enumerate(vocab)}
        self.unk_token = unk_token
        self.unk_id = self.token2id.get(unk_token, 0)
    
    def _tokenize_word(self, word: str) -> List[str]:
        """
        对单个词用最长匹配分词
        - 首字符从头开始匹配
        - 非首子串加 ## 前缀匹配
        """
        tokens = []
        start = 0
        word_len = len(word)
        is_bad = False
        
        while start < word_len:
            end = word_len
            found = None
            
            while start < end:
                substr = word[start:end]
                # 非开头子串加 ## 前缀
                candidate = substr if start == 0 else '##' + substr
                
                if candidate in self.vocab:
                    found = candidate
                    break
                end -= 1
            
            if found is None:
                # 没有任何匹配 → 整个词用 [UNK]
                is_bad = True
                break
            
            tokens.append(found)
            start = end
        
        return [self.unk_token] if is_bad else tokens
    
    def encode(self, text: str, add_special_tokens: bool = True) -> List[int]:
        """文本 → token id 列表（可选加 [CLS]/[SEP]）"""
        tokens = []
        if add_special_tokens:
            tokens.append('[CLS]')
        
        for word in text.lower().split():
            tokens.extend(self._tokenize_word(word))
        
        if add_special_tokens:
            tokens.append('[SEP]')
        
        return [self.token2id.get(t, self.unk_id) for t in tokens]
    
    def decode(self, ids: List[int], skip_special_tokens: bool = True) -> str:
        """token id 列表 → 文本（合并 ## 前缀）"""
        special = {'[CLS]', '[SEP]', '[PAD]', '[MASK]'}
        tokens = [self.id2token.get(i, '[UNK]') for i in ids]
        
        words = []
        current = ''
        for token in tokens:
            if skip_special_tokens and token in special:
                if current:
                    words.append(current)
                    current = ''
                continue
            if token.startswith('##'):
                current += token[2:]      # ## 前缀直接拼接，不加空格
            else:
                if current:
                    words.append(current)
                current = token
        if current:
            words.append(current)
        
        return ' '.join(words)
    
    def tokenize(self, text: str) -> List[str]:
        """文本 → token 字符串列表（调试用）"""
        result = []
        for word in text.lower().split():
            result.extend(self._tokenize_word(word))
        return result


# 使用示例
if __name__ == '__main__':
    # 用自定义词表创建 tokenizer
    vocab = [
        '[PAD]', '[UNK]', '[CLS]', '[SEP]', '[MASK]',
        'un', '##aff', '##ord', '##able',
        'hello', 'world', '##ing', '##ed',
        'a', 'b', 'c', 'd', 'e', 'f', 'g'
    ]
    tokenizer = WordPieceTokenizer(vocab)
    
    # tokenize
    tokens = tokenizer.tokenize("unaffordable hello")
    print(f"tokenize: {tokens}")
    # → ['un', '##aff', '##ord', '##able', 'hello']
    
    # encode（含 [CLS]/[SEP]）
    ids = tokenizer.encode("hello world")
    print(f"encode: {ids}")
    
    # decode
    text = tokenizer.decode(ids)
    print(f"decode: '{text}'")
    # → 'hello world'

import math
import random
from typing import Dict, List, Tuple, Optional


# ═══════════════════════════════════════════════════
#  Unigram 词表格式：token → log_prob（对数概率）
# ═══════════════════════════════════════════════════
class UnigramTokenizer:
    """
    Unigram LM Tokenizer
    vocab: dict[token_str → log_prob]（log_prob 均为负数或 0）
    """
    
    def __init__(self, vocab: Dict[str, float], unk_token: str = '<unk>'):
        self.vocab = vocab
        self.unk_token = unk_token
        self.token2id = {t: i for i, t in enumerate(sorted(vocab.keys()))}
        self.id2token = {i: t for t, i in self.token2id.items()}
        if unk_token not in self.token2id:
            uid = len(self.token2id)
            self.token2id[unk_token] = uid
            self.id2token[uid] = unk_token
    
    def _viterbi_segment(self, text: str) -> List[str]:
        """
        Viterbi 算法：找使 Unigram 概率最大的分词方案
        时间复杂度 O(n²)，n 为文本长度
        """
        n = len(text)
        NEG_INF = float('-inf')
        
        # best_score[i]：前 i 个字符的最优对数概率
        best_score = [NEG_INF] * (n + 1)
        best_score[0] = 0.0
        # back[i]：最优路径的回溯指针
        back = [0] * (n + 1)
        
        for i in range(1, n + 1):
            for j in range(i):
                chunk = text[j:i]
                if chunk in self.vocab:
                    score = best_score[j] + self.vocab[chunk]
                    if score > best_score[i]:
                        best_score[i] = score
                        back[i] = j
        
        # 回溯路径
        tokens = []
        i = n
        while i > 0:
            j = back[i]
            tokens.append(text[j:i])
            i = j
        
        return tokens[::-1]
    
    def _sample_segment(self, text: str, alpha: float = 0.1) -> List[str]:
        """
        子词正则化：按概率分布随机采样一种分词方案（用于训练数据增强）
        alpha：温度系数，越大越随机；alpha=0 等价于贪心最优
        """
        n = len(text)
        # 用前向算法计算每个位置的所有可能分词及其分数
        candidates = []
        
        def _segment_recursive(pos, path):
            if pos == n:
                candidates.append((path[:], sum(self.vocab.get(t, -20) for t in path)))
                return
            for end in range(pos + 1, n + 1):
                chunk = text[pos:end]
                if chunk in self.vocab:
                    path.append(chunk)
                    _segment_recursive(end, path)
                    path.pop()
        
        # 只采样少量候选以控制复杂度（实际生产用 lattice sampling）
        _segment_recursive(0, [])
        
        if not candidates:
            return self._viterbi_segment(text)  # fallback
        
        # softmax 采样
        scores = [s * (1 / (alpha + 1e-9)) for _, s in candidates]
        max_s = max(scores)
        weights = [math.exp(s - max_s) for s in scores]
        total = sum(weights)
        probs = [w / total for w in weights]
        
        r = random.random()
        cumsum = 0.0
        for (tokens, _), p in zip(candidates, probs):
            cumsum += p
            if r < cumsum:
                return tokens
        return candidates[-1][0]
    
    def tokenize(self, text: str, sample: bool = False) -> List[str]:
        """文本 → token 字符串列表；sample=True 启用子词正则化"""
        result = []
        # 按空格预切分（实际 SentencePiece 版本不需要这步）
        for word in text.split():
            if sample:
                result.extend(self._sample_segment(word))
            else:
                result.extend(self._viterbi_segment(word))
        return result
    
    def encode(self, text: str, sample: bool = False) -> List[int]:
        """文本 → token id 列表"""
        tokens = self.tokenize(text, sample=sample)
        unk_id = self.token2id[self.unk_token]
        return [self.token2id.get(t, unk_id) for t in tokens]
    
    def decode(self, ids: List[int]) -> str:
        """token id 列表 → 文本"""
        return ' '.join(self.id2token.get(i, self.unk_token) for i in ids)


# 使用示例
if __name__ == '__main__':
    # 模拟词表（token → log_prob）
    vocab = {
        'un': -2.0, 'afford': -3.5, 'able': -2.5,
        'unafford': -6.0, 'unaffordable': -8.0,
        'u': -5.0, 'n': -5.0, 'a': -4.0,
        'hello': -2.0, 'world': -2.1,
        '<unk>': -20.0,
    }
    tokenizer = UnigramTokenizer(vocab)
    
    text = "unaffordable"
    
    # 最优分词（Viterbi）
    best = tokenizer.tokenize(text, sample=False)
    print(f"Viterbi 最优: {best}")
    # → ['un', 'afford', 'able']
    
    # 子词正则化（随机采样，训练时用）
    samples = [tokenizer.tokenize(text, sample=True) for _ in range(5)]
    print(f"随机采样 5 次: {samples}")
    # 可能输出: [['un','afford','able'], ['u','n','afford','able'], ...]
    
    # encode / decode
    ids = tokenizer.encode("hello world")
    print(f"encode: {ids}")
    print(f"decode: {tokenizer.decode(ids)}")

import json
from typing import Dict, List, Union
from pathlib import Path


# ═══════════════════════════════════════════════════
#  通用词表匹配器
# ═══════════════════════════════════════════════════
class VocabMatcher:
    """
    通用 Token ↔ ID 双向映射器
    支持从 dict / json 文件 / txt 文件加载词表
    """
    
    def __init__(
        self,
        vocab: Union[Dict[str, int], None] = None,
        unk_token: str = '[UNK]',
        pad_token: str = '[PAD]',
        bos_token: str = '[BOS]',
        eos_token: str = '[EOS]',
    ):
        self.unk_token = unk_token
        self.pad_token = pad_token
        self.bos_token = bos_token
        self.eos_token = eos_token
        
        if vocab is not None:
            self.token2id = vocab
        else:
            self.token2id = {}
        
        self.id2token = {v: k for k, v in self.token2id.items()}
        
        # 确保特殊 token 存在
        for tok in [pad_token, unk_token, bos_token, eos_token]:
            if tok not in self.token2id:
                new_id = len(self.token2id)
                self.token2id[tok] = new_id
                self.id2token[new_id] = tok
    
    # ─── 加载词表 ───────────────────────────────────────
    
    @classmethod
    def from_json(cls, path: str, **kwargs) -> 'VocabMatcher':
        """从 vocab.json 加载（格式：{"token": id, ...}）"""
        with open(path, 'r', encoding='utf-8') as f:
            vocab = json.load(f)
        return cls(vocab=vocab, **kwargs)
    
    @classmethod
    def from_txt(cls, path: str, **kwargs) -> 'VocabMatcher':
        """
        从 vocab.txt 加载（格式：每行一个 token，行号即 id）
        BERT 词表就是这种格式
        """
        with open(path, 'r', encoding='utf-8') as f:
            tokens = [line.rstrip('\n') for line in f]
        vocab = {tok: i for i, tok in enumerate(tokens)}
        return cls(vocab=vocab, **kwargs)
    
    @classmethod
    def from_list(cls, tokens: List[str], **kwargs) -> 'VocabMatcher':
        """从 token 列表构建（顺序即 id）"""
        vocab = {tok: i for i, tok in enumerate(tokens)}
        return cls(vocab=vocab, **kwargs)
    
    # ─── 属性 ───────────────────────────────────────────
    
    @property
    def vocab_size(self) -> int:
        return len(self.token2id)
    
    @property
    def unk_id(self) -> int:
        return self.token2id[self.unk_token]
    
    @property
    def pad_id(self) -> int:
        return self.token2id[self.pad_token]
    
    # ─── 单条编解码 ─────────────────────────────────────
    
    def token_to_id(self, token: str) -> int:
        """单个 token → id（OOV 返回 unk_id）"""
        return self.token2id.get(token, self.unk_id)
    
    def id_to_token(self, id: int) -> str:
        """单个 id → token（越界返回 unk_token）"""
        return self.id2token.get(id, self.unk_token)
    
    def encode(
        self,
        tokens: List[str],
        add_bos: bool = False,
        add_eos: bool = False,
        max_length: Optional[int] = None,
        padding: bool = False,
    ) -> List[int]:
        """
        token 列表 → id 列表
        - tokens: 已经由分词器切好的 token 字符串列表
        - add_bos/add_eos: 是否在首尾加特殊 token
        - max_length: 截断长度（None 不截断）
        - padding: 是否补 PAD 到 max_length
        """
        ids = [self.token_to_id(t) for t in tokens]
        
        if add_bos:
            ids = [self.token2id[self.bos_token]] + ids
        if add_eos:
            ids = ids + [self.token2id[self.eos_token]]
        
        if max_length is not None:
            ids = ids[:max_length]
            if padding and len(ids) < max_length:
                ids += [self.pad_id] * (max_length - len(ids))
        
        return ids
    
    def decode(
        self,
        ids: List[int],
        skip_special_tokens: bool = True,
        join_str: str = ' ',
    ) -> str:
        """
        id 列表 → token 字符串（空格连接）
        - skip_special_tokens: 跳过 PAD/BOS/EOS/UNK 等
        - join_str: token 间分隔符（BPE 解码后通常不需要空格，根据算法决定）
        """
        special = {self.pad_token, self.unk_token, self.bos_token, self.eos_token}
        tokens = [
            self.id_to_token(i) for i in ids
            if not (skip_special_tokens and self.id_to_token(i) in special)
        ]
        return join_str.join(tokens)
    
    # ─── 批量编码 ───────────────────────────────────────
    
    def batch_encode(
        self,
        batch_tokens: List[List[str]],
        max_length: int,
        padding: bool = True,
        add_eos: bool = True,
    ) -> Dict[str, List[List[int]]]:
        """
        批量编码（返回 input_ids + attention_mask）
        batch_tokens: 多条样本的 token 列表
        """
        input_ids = []
        attention_masks = []
        
        for tokens in batch_tokens:
            ids = self.encode(
                tokens,
                add_eos=add_eos,
                max_length=max_length,
                padding=padding,
            )
            # attention_mask：非 PAD 位置为 1，PAD 位置为 0
            mask = [0 if i == self.pad_id else 1 for i in ids]
            input_ids.append(ids)
            attention_masks.append(mask)
        
        return {'input_ids': input_ids, 'attention_mask': attention_masks}
    
    # ─── 工具方法 ───────────────────────────────────────
    
    def save_vocab(self, path: str):
        """保存词表到 JSON 文件"""
        with open(path, 'w', encoding='utf-8') as f:
            json.dump(self.token2id, f, ensure_ascii=False, indent=2)
    
    def contains(self, token: str) -> bool:
        """检查 token 是否在词表中"""
        return token in self.token2id
    
    def oov_rate(self, tokens: List[str]) -> float:
        """计算 OOV 比例"""
        oov = sum(1 for t in tokens if t not in self.token2id)
        return oov / len(tokens) if tokens else 0.0
    
    def __repr__(self):
        return f"VocabMatcher(vocab_size={self.vocab_size}, unk='{self.unk_token}')"


# ═══════════════════════════════════════════════════
#  使用示例
# ═══════════════════════════════════════════════════
if __name__ == '__main__':

    # ── 方式 1：从 dict 构建 ──────────────────────────
    vocab_dict = {
        '[PAD]': 0, '[UNK]': 1, '[CLS]': 2, '[SEP]': 3,
        'hello': 4, 'world': 5, 'un': 6, '##aff': 7,
        '##ord': 8, '##able': 9, 'token': 10,
    }
    matcher = VocabMatcher(vocab=vocab_dict)
    print(matcher)                    # VocabMatcher(vocab_size=13, unk='[UNK]')
    
    # ── 方式 2：从 vocab.txt 加载（BERT 格式）──────────
    # matcher = VocabMatcher.from_txt('/path/to/vocab.txt')
    
    # ── 方式 3：从 vocab.json 加载（HuggingFace 格式）──
    # matcher = VocabMatcher.from_json('/path/to/vocab.json')
    
    # ── 单条 encode ───────────────────────────────────
    tokens = ['un', '##aff', '##ord', '##able']
    ids = matcher.encode(tokens)
    print(f"encode: {ids}")            # [6, 7, 8, 9]
    
    # 带截断和 padding
    ids_padded = matcher.encode(tokens, add_eos=True, max_length=8, padding=True)
    print(f"encode(padded): {ids_padded}")   # [6, 7, 8, 9, <eos_id>, 0, 0, 0]
    
    # ── decode ────────────────────────────────────────
    text = matcher.decode(ids)
    print(f"decode: '{text}'")           # 'un ##aff ##ord ##able'
    
    # ── 批量 encode ──────────────────────────────────
    batch = [['hello', 'world'], ['un', '##aff', '##ord']]
    result = matcher.batch_encode(batch, max_length=6, padding=True)
    print("batch input_ids:", result['input_ids'])
    print("batch attention_mask:", result['attention_mask'])
    # input_ids:   [[4, 5, 3, 0, 0, 0], [6, 7, 8, 3, 0, 0]]
    # attention:   [[1, 1, 1, 0, 0, 0], [1, 1, 1, 1, 0, 0]]
    
    # ── OOV 率检查 ───────────────────────────────────
    oov_tokens = ['hello', 'xyz123', 'world', 'foobar']
    print(f"OOV 率: {matcher.oov_rate(oov_tokens):.1%}")   # 50.0%

import tiktoken


# ═══════════════════════════════════════════════════
#  tiktoken 完整使用指南
# ═══════════════════════════════════════════════════

# 1. 加载编码器
enc = tiktoken.get_encoding("cl100k_base")      # GPT-4 / GPT-3.5
# enc = tiktoken.get_encoding("p50k_base")          # GPT-3
# enc = tiktoken.get_encoding("r50k_base")          # GPT-2

# 2. 基础编解码
text = "Hello, 你好世界！🌍"
ids = enc.encode(text)
print(f"token ids : {ids}")
# [9906, 11, 220, 57668, 53901, 11375, 6447, 127, 246, 9468, 242, 101]

decoded = enc.decode(ids)
print(f"decoded   : {decoded}")                # 完全还原

# 3. 查看每个 token 对应的文本
tokens_bytes = [enc.decode_single_token_bytes(t) for t in ids]
print(f"token bytes: {tokens_bytes}")
# [b'Hello', b',', b' ', b'\xe4\xbd\xa0', ...]

# 4. 不切分特殊 token（用于 Chat 格式）
ids_no_special = enc.encode(text, disallowed_special=())
# 阻止 <|endoftext|> 等被误识别为特殊 token

# 5. 计算 token 数（不实际编码，更快）
num_tokens = len(enc.encode(text))
print(f"token 数: {num_tokens}")

# 6. 模型专属 encoding（自动对应正确词表）
enc_gpt4 = tiktoken.encoding_for_model("gpt-4")
enc_gpt35 = tiktoken.encoding_for_model("gpt-3.5-turbo")


# ═══════════════════════════════════════════════════
#  实用工具：计算 Chat Messages 的 token 数
#  参考 OpenAI 官方文档 "How to count tokens with tiktoken"
# ═══════════════════════════════════════════════════
def num_tokens_from_messages(messages: list, model: str = "gpt-4") -> int:
    """
    计算 messages 列表的 token 数（兼容 gpt-3.5-turbo / gpt-4）
    messages 格式：[{"role": "user", "content": "..."}]
    """
    try:
        encoding = tiktoken.encoding_for_model(model)
    except KeyError:
        encoding = tiktoken.get_encoding("cl100k_base")
    
    # gpt-3.5-turbo / gpt-4 格式
    tokens_per_message = 3  # <|im_start|> role \n content <|im_end|>
    tokens_per_name = 1
    
    num_tokens = 0
    for message in messages:
        num_tokens += tokens_per_message
        for key, value in message.items():
            num_tokens += len(encoding.encode(str(value)))
            if key == "name":
                num_tokens += tokens_per_name
    num_tokens += 3  # 每次回复前有 <|im_start|>assistant
    return num_tokens


# 使用示例
messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user",   "content": "Hello, how are you?"},
]
n = num_tokens_from_messages(messages)
print(f"这段对话占 {n} tokens")    # ~22 tokens


# ═══════════════════════════════════════════════════
#  实用工具：截断文本到指定 token 数
# ═══════════════════════════════════════════════════
def truncate_text(text: str, max_tokens: int, model: str = "gpt-4") -> str:
    """截断文本到 max_tokens（完整 token 边界截断，不会截断 UTF-8 字节）"""
    enc = tiktoken.encoding_for_model(model)
    ids = enc.encode(text)
    if len(ids) <= max_tokens:
        return text
    return enc.decode(ids[:max_tokens])

# 超长文本截断到 100 tokens
long_text = "这是一段很长的文本 " * 100
truncated = truncate_text(long_text, max_tokens=100)
print(f"截断后 token 数: {len(tiktoken.encoding_for_model('gpt-4').encode(truncated))}")