如何在HuggingFace环境中用Tiktoken的ck_100k_base替换BERT分词器？

因为HuggingFace的AutoTokenizer没有直接内置ck_100k_base的支持，你可以通过两种方式实现兼容：

方法一：使用Transformers官方的TiktokenTokenizer #

Transformers提供了TiktokenTokenizer类，可直接对接tiktoken的编码方案，步骤如下：

1. 确保依赖版本达标（transformers>=4.28.0），安装必要包：

pip install tiktoken transformers datasets

2. 初始化兼容HuggingFace接口的分词器：

from transformers import TiktokenTokenizer

# 初始化ck_100k_base分词器，并设置必要的特殊token
tokenizer = TiktokenTokenizer(encoding_name="ck_100k_base")
# 为适配HuggingFace模型，手动指定pad_token（可根据模型需求调整）
tokenizer.pad_token = tokenizer.eos_token

3. 修改预处理函数（逻辑与原BERT分词器一致，直接替换tokenizer即可）：

def preprocess_dataset(examples):
    inputs = [prefix + example[source_lang] for example in examples["translation"]]
    targets = [example[target_lang] for example in examples["translation"]]
    
    # 使用tiktoken分词器处理输入和目标文本
    model_inputs = tokenizer(
        inputs,
        text_target=targets,
        max_length=128,
        truncation=True,
        padding="max_length"  # 可选，根据训练需求添加padding
    )
    return model_inputs

source_lang = "en"
target_lang = "fr"
prefix = "Translate English to French: "
tokenized = my_dataset.map(preprocess_dataset, batched=True)

方法二：自定义封装tiktoken（更灵活） #

如果需要更精细的控制，可以自行封装tiktoken的编码逻辑，适配HuggingFace的数据集处理格式：

import tiktoken
from typing import List, Dict

# 初始化tiktoken编码
enc = tiktoken.get_encoding("ck_100k_base")
# 定义特殊token（根据模型需求调整）
PAD_TOKEN = "<pad>"
EOS_TOKEN = "<|endoftext|>"
# 将特殊token加入编码表
enc.add_tokens([PAD_TOKEN, EOS_TOKEN])
PAD_ID = enc.encode_ordinary(PAD_TOKEN)[0]
EOS_ID = enc.encode_ordinary(EOS_TOKEN)[0]

def custom_tokenize(texts: List[str], max_length: int, is_target: bool = False) -> Dict:
    input_ids = []
    attention_mask = []
    for text in texts:
        # 编码文本
        ids = enc.encode_ordinary(text)
        # 目标文本末尾添加eos token
        if is_target:
            ids.append(EOS_ID)
        # 截断或填充到指定长度
        if len(ids) > max_length:
            ids = ids[:max_length]
        else:
            ids += [PAD_ID] * (max_length - len(ids))
        # 生成attention mask
        mask = [1] * len(ids)
        if len(mask) < max_length:
            mask += [0] * (max_length - len(mask))
        
        input_ids.append(ids)
        attention_mask.append(mask)
    
    return {"input_ids": input_ids, "attention_mask": attention_mask}

def preprocess_dataset(examples):
    inputs = [prefix + example[source_lang] for example in examples["translation"]]
    targets = [example[target_lang] for example in examples["translation"]]
    
    # 处理输入文本
    model_inputs = custom_tokenize(inputs, max_length=128)
    # 处理目标文本并作为labels
    target_outputs = custom_tokenize(targets, max_length=128, is_target=True)
    model_inputs["labels"] = target_outputs["input_ids"]
    
    return model_inputs

source_lang = "en"
target_lang = "fr"
prefix = "Translate English to French: "
tokenized = my_dataset.map(preprocess_dataset, batched=True)

注意事项 #

• 特殊token的设置需要与你使用的模型匹配，若模型未预训练这些token，需在模型初始化时指定pad_token_id、eos_token_id等参数。
• 若无需padding，可移除相关逻辑，仅保留截断处理。

内容的提问来源于stack exchange，提问作者Raptor