特殊标记设计：从 [CLS] 到 ChatML 的 Prompt 格式演化

特殊标记的核心功能:
1. 序列边界标记: [BOS] [EOS] - 标识序列开始/结束
2. 任务分隔标记: [SEP] - 分隔不同输入段
3. 分类聚合标记: [CLS] - 聚合全局语义
4. 填充对齐标记: [PAD] - 批处理中对齐长度
5. 未知词标记:   [UNK] - 处理未登录词
6. 掩码标记:     [MASK] - 遮蔽语言模型训练
7. 角色标记:     <|im_start|> - 标识对话角色
8. 系统标记:     <|system|> - 系统指令

# 不同模型系列的特殊标记 ID 对比
special_token_ids = {
    "BERT": {
        "[PAD]": 0, "[UNK]": 1, "[CLS]": 2, "[SEP]": 3, "[MASK]": 103,
    },
    "GPT-2": {
        "<|endoftext|>": 50256,  # 单一结束标记
    },
    "LLaMA": {
        "<unk>": 0, "<s>": 1, "</s>": 2,
        "<pad>": 32000,
    },
    "Llama-3": {
        "<|begin_of_text|>": 128000,
        "<|end_of_text|>": 128001,
        "<|start_header_id|>": 128006,
        "<|end_header_id|>": 128007,
        "<|eot_id|>": 128009,
    },
    "ChatML (OpenAI)": {
        "<|im_start|>": 100264,
        "<|im_end|>": 100265,
        "<|endoftext|>": 100257,
    },
}

# BERT 输入格式示例
# 单句分类任务:
input_text = "[CLS] 这部电影太精彩了 [SEP]"
# Token ID: [2, 6821, 3175, 6629, 6885, 4638, 3, ...]

# 句对任务 (NLI):
input_text = "[CLS] 猫在垫子上 [SEP] 有一只动物在休息 [SEP]"
# Token ID: [2, ..., 3, ..., 3]

# [CLS] 的 hidden state 用于分类
# 最终输出: pooled_output = tanh(W * hidden_states[0] + b)

class BertInputFormatter:
    """BERT 输入格式化器"""
    def __init__(self, tokenizer, max_length=512):
        self.tokenizer = tokenizer
        self.max_length = max_length
    
    def format_single_sentence(self, text):
        """单句格式: [CLS] text [SEP]"""
        tokens = self.tokenizer.tokenize(text)
        tokens = ['[CLS]'] + tokens[:self.max_length-2] + ['[SEP]']
        input_ids = self.tokenizer.convert_tokens_to_ids(tokens)
        attention_mask = [1] * len(input_ids)
        token_type_ids = [0] * len(input_ids)
        return {'input_ids': input_ids, 'attention_mask': attention_mask, 'token_type_ids': token_type_ids}
    
    def format_sentence_pair(self, text_a, text_b):
        """句对格式: [CLS] A [SEP] B [SEP]"""
        tokens_a = self.tokenizer.tokenize(text_a)
        tokens_b = self.tokenizer.tokenize(text_b)
        max_len = self.max_length - 3
        while len(tokens_a) + len(tokens_b) > max_len:
            if len(tokens_a) > len(tokens_b):
                tokens_a.pop()
            else:
                tokens_b.pop()
        tokens = ['[CLS]'] + tokens_a + ['[SEP]'] + tokens_b + ['[SEP]']
        token_type_ids = [0] * (len(tokens_a) + 2) + [1] * (len(tokens_b) + 1)
        input_ids = self.tokenizer.convert_tokens_to_ids(tokens)
        attention_mask = [1] * len(input_ids)
        return {'input_ids': input_ids, 'attention_mask': attention_mask, 'token_type_ids': token_type_ids}

class MLMMasker:
    """BERT 的遮蔽语言模型掩码策略"""
    def __init__(self, tokenizer, mlm_probability=0.15):
        self.tokenizer = tokenizer
        self.mlm_probability = mlm_probability
    
    def mask_tokens(self, inputs):
        """掩码策略: 80% [MASK], 10% 随机, 10% 不变"""
        labels = inputs.clone()
        probability_matrix = torch.full(labels.shape, self.mlm_probability)
        special_tokens_mask = [
            self.tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True)
            for val in labels.tolist()
        ]
        probability_matrix.masked_fill_(torch.tensor(special_tokens_mask, dtype=torch.bool), 0)
        masked_indices = torch.bernoulli(probability_matrix).bool()
        labels[~masked_indices] = -100
        # 80% -> [MASK]
        indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices
        inputs[indices_replaced] = self.tokenizer.mask_token_id
        # 10% -> 随机 token
        indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced
        random_words = torch.randint(len(self.tokenizer), labels.shape, dtype=torch.long)
        inputs[indices_random] = random_words[indices_random]
        return inputs, labels

# GPT-2 训练数据的组织方式
training_corpus = [
    "文档1内容...<|endoftext|>",
    "文档2内容...<|endoftext|>",
    "文档3内容...<|endoftext|>",
]
# 模型学到: 在 <|endoftext|> 之后可以开始全新话题

class GPT2InputFormatter:
    def __init__(self, tokenizer, max_length=1024):
        self.tokenizer = tokenizer
        self.max_length = max_length
    def format_prompt(self, prompt):
        """GPT-2 的简单 prompt 格式 (无系统角色)"""
        input_ids = self.tokenizer.encode(prompt)
        return input_ids
    def format_training(self, documents):
        """训练数据打包: 多文档拼接"""
        all_tokens = []
        for doc in documents:
            tokens = self.tokenizer.encode(doc)
            tokens.append(self.tokenizer.eos_token_id)
            all_tokens.extend(tokens)
        blocks = []
        for i in range(0, len(all_tokens) - self.max_length, self.max_length):
            blocks.append(all_tokens[i:i + self.max_length])
        return blocks

# InstructGPT 的 Prompt 模板
INSTRUCTGPT_TEMPLATE = """
{instruction}

### Response:
{response}
"""
# 问题和答案的边界由 "### Response:" 标识
# 这是一种"伪特殊标记" - 用文本模式而非专用 token

# ChatML 格式示例
chatml_example = """<|im_start|>system
You are a helpful assistant.<|im_end|>
<|im_start|>user
What is the capital of France?<|im_end|>
<|im_start|>assistant
The capital of France is Paris.<|im_end|>
<|im_start|>assistant
"""

# 核心设计:
# 1. <|im_start|> 标记消息开始
# 2. 紧跟角色名 (system/user/assistant)
# 3. 换行后是消息内容
# 4. <|im_end|> 标记消息结束
# 5. 新消息以 <|im_start|> 开始

class ChatMLFormatter:
    """ChatML 格式化器"""
    IM_START = "<|im_start|>"
    IM_END = "<|im_end|>"
    
    def format_messages(self, messages, add_generation_prompt=False):
        """将消息列表格式化为 ChatML 文本"""
        formatted = ""
        for message in messages:
            role = message['role']
            content = message['content']
            formatted += f"{self.IM_START}{role}\n{content}{self.IM_END}\n"
        if add_generation_prompt:
            formatted += f"{self.IM_START}assistant\n"
        return formatted
    
    def format_for_training(self, messages, train_on_responses_only=True):
        """训练数据格式化 - 只在 assistant 回复部分计算 loss"""
        formatted = ""
        loss_mask = []
        for message in messages:
            role = message['role']
            content = message['content']
            if role == "assistant":
                segment = f"{self.IM_START}{role}\n{content}{self.IM_END}\n"
                formatted += segment
                loss_mask.extend([1] * len(segment.split()))
            else:
                segment = f"{self.IM_START}{role}\n{content}{self.IM_END}\n"
                formatted += segment
                loss_mask.extend([0] * len(segment.split()))
        return formatted, loss_mask

# Prompt Injection 对比
# InstructGPT 风格 (易受注入):
# 模型可能混淆 "### Response:" 是新的回复还是用户输入
# ChatML (更安全):
# 用户输入中的角色标记不会被当作真正的角色标记
# (因为 <|im_start|> 是专用 token, 用户文本中不会产生)

# === LLaMA-2 Chat ===
LLAMA2_TEMPLATE = """<s>[INST] <<SYS>>
{system_prompt}
<</SYS>>
{user_message} [/INST] {assistant_response} </s>"""

# === LLaMA-3 ===
LLAMA3_TEMPLATE = """<|begin_of_text|><|start_header_id|>system<|end_header_id|>
{system_prompt}<|eot_id|><|start_header_id|>user<|end_header_id|>
{user_message}<|eot_id|><|start_header_id|>assistant<|end_header_id|>
{assistant_response}<|eot_id|>"""

# === Mistral / Mixtral ===
MISTRAL_TEMPLATE = """<s>[INST] {user_message}[/INST] {assistant_response}</s>"""
MISTRAL_MULTI = """<s>[INST] {msg1}[/INST] {resp1}</s>[INST] {msg2}[/INST] {resp2}</s>"""

# === Qwen ===
QWEN_TEMPLATE = """<|im_start|>system
{system_prompt}<|im_end|>
<|im_start|>user
{user_message}<|im_end|>
<|im_start|>assistant
{assistant_response}<|im_end|>"""

# === Claude ===
CLAUDE_TEMPLATE = "\n\nHuman: {user_message}\n\nAssistant: {assistant_response}"

# === Phi-3 ===
PHI3_TEMPLATE = """<|system|>
{system_prompt}<|end|>
<|user|>
{user_message}<|end|>
<|assistant|>
{assistant_response}<|end|>"""

from transformers import AutoTokenizer
import torch

def add_special_tokens_to_model(model, tokenizer, special_tokens_list):
    """向已有模型添加新的特殊标记"""
    num_added = tokenizer.add_special_tokens({
        'additional_special_tokens': special_tokens_list
    })
    model.resize_token_embeddings(len(tokenizer))
    with torch.no_grad():
        existing_mean = model.get_input_embeddings().weight[:-num_added].mean(dim=0)
        model.get_input_embeddings().weight[-num_added:] = existing_mean
        model.get_output_embeddings().weight[-num_added:] = existing_mean
    return num_added

# 示例: 为模型添加工具调用标记
special_tokens = ["<|tool_call|>", "<|/tool_call|>", "<|tool_result|>", "<|/tool_result|>"]
num_added = add_special_tokens_to_model(model, tokenizer, special_tokens)

class SpecialTokenAwareLoss:
    """特殊标记感知的损失计算"""
    def __init__(self, tokenizer, ignore_roles=None):
        self.tokenizer = tokenizer
        self.ignore_roles = ignore_roles or ['system', 'user']
    def compute_loss_mask(self, input_ids, labels):
        """只在 assistant 回复部分计算 loss"""
        im_start_id = self.tokenizer.convert_tokens_to_ids("<|im_start|>")
        im_end_id = self.tokenizer.convert_tokens_to_ids("<|im_end|>")
        loss_mask = torch.zeros_like(input_ids, dtype=torch.bool)
        for batch_idx in range(input_ids.size(0)):
            ids = input_ids[batch_idx].tolist()
            i = 0
            while i < len(ids):
                if ids[i] == im_start_id:
                    j = i + 1
                    while j < len(ids) and ids[j] != im_end_id:
                        if self._is_assistant_token(ids, j):
                            start = j + 2
                            end = self._find_next(ids, im_end_id, start)
                            loss_mask[batch_idx, start:end+1] = True
                            break
                        j += 1
                i += 1
        labels[~loss_mask] = -100
        return labels

class SpecialTokenStreamHandler:
    """流式生成中的特殊标记处理"""
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
        self.eos_tokens = {tokenizer.eos_token_id}
        stop_strings = ["<|im_end|>", "</s>", "<|eot_id|>"]
        self.stop_token_ids = set()
        for s in stop_strings:
            ids = tokenizer.encode(s, add_special_tokens=False)
            if len(ids) == 1:
                self.stop_token_ids.add(ids[0])
        self.buffer = ""
    def process_token(self, token_id):
        """处理单个新 token"""
        if token_id in self.eos_tokens or token_id in self.stop_token_ids:
            return None, True  # 停止生成
        token_text = self.tokenizer.decode([token_id])
        self.buffer += token_text
        for stop in self.stop_strings:
            if stop in self.buffer:
                idx = self.buffer.index(stop)
                output = self.buffer[:idx]
                self.buffer = ""
                return output, True
        for stop in self.stop_strings:
            for i in range(1, len(stop)):
                if self.buffer.endswith(stop[:i]):
                    return "", False  # 可能是不完整的停止标记
        output = self.buffer
        self.buffer = ""
        return output, False

1. 唯一性: 不与自然文本冲突
   GOOD: <|im_start|> - 不会出现在正常文本中
   BAD:  "###" - 可能出现在代码或文档中
2. 不可伪造: 用户无法通过输入生成
   GOOD: <|im_start|> 作为单 token - 用户文本不会产生此 token
   BAD:  [INST] - 用户可以输入这些字符
3. 简洁性: 不浪费太多 token
   GOOD: <|im_end|> - 1 token
   BAD:  "<|start_of_assistant_message|>" - 多个 token
4. 可扩展: 支持添加新角色/功能
   GOOD: ChatML: <|im_start|>{role}\n...<|im_end|>
   BAD:  固定模板: [INST]...[/INST] 只支持 user/assistant
5. 可检测: 容易在流式生成中检测
   GOOD: 单 token 的 <|im_end|> - 一检测到就停止
   BAD:  多 token 的停止串 - 需要 buffer 处理

# 策略1: 词表开头 (SentencePiece 传统)
# 优点: 快速访问, 缺点: 占用低 ID
# 代表: LLaMA (<unk>=0, <s>=1, </s>=2)
# 策略2: 词表末尾 (BPE 扩展)
# 优点: 不影响原始词表, 缺点: 需要扩展 embedding
# 代表: GPT-2 (eos=50256)
# 策略3: 保留区间
# 优点: 系统化管理, 缺点: 浪费部分 ID 空间
# 代表: LLaMA-3 (128000+ 为特殊标记)
RESERVED_SPECIAL_TOKENS = {
    "pad": {"id": 0, "token": "<pad>"},
    "unk": {"id": 1, "token": "<unk>"},
    "bos": {"id": 2, "token": "<s>"},
    "eos": {"id": 3, "token": "</s>"},
}

# 多模态模型需要更多特殊标记
multimodal_special_tokens = [
    "<image>", "</image>",  # 图像
    "<audio>", "</audio>",  # 音频
    "<video>", "</video>",  # 视频
    "<|tool_call|>", "<|/tool_call|>",  # 工具调用
    "<tool_result>", "</tool_result>",  # 工具结果
    "<|thinking|>", "<|/thinking|>",  # 思考过程
]
# 多模态格式示例:
# <|im_start|>user
# <image>What is in this image?<|im_end|>
# <|im_start|>assistant
# I can see a cat sitting on a windowsill.<|im_end|>

演化时间线:
2018 - BERT: [CLS] text [SEP]
2019 - GPT-2: text <|endoftext|>
2022 - InstructGPT: ### Instruction / ### Response
2023 - ChatML: <|im_start|>role\n...<|im_end|>
2024 - LLaMA-3: <|start_header_id|>role<|end_header_id|>\n...<|eot_id|>
2025 - 统一化趋势: 向 ChatML 靠拢
未来方向:
- 跨模型统一格式 (减少适配成本)
- 原生多模态标记标准化
- 可验证的 Prompt 格式 (防注入)
- 动态角色扩展 (插件系统)