mirror of
https://github.com/jingyaogong/minimind.git
synced 2026-09-26 21:07:22 +00:00
Linxiushen
cb08b9300b
get_texts() 只解析 data.get('conversations')，但主线预训练数据
pretrain_t2t(_mini).jsonl 的每行是 {"text": ...}（README「数据集下载」一节）。
把 DATA_PATH 指向预训练语料时，生成器一条文本都不产出，BpeTrainer 仅基于
initial_alphabet 训练，得到一个没有任何 merge 的词表，全程不报错。
复现（改 DATA_PATH 为任意 pretrain_t2t*.jsonl，或直接调用）：
from train_tokenizer import get_texts
sum(1 for _ in get_texts('../dataset/pretrain_t2t_mini.jsonl'))
# 修复前: 0
# 修复后: 1270238
改动：
- get_texts 同时兼容 {"text": ...} 与 {"conversations": [...]} 两种格式
- 产出 0 条时 raise ValueError，而不是让调用方拿到一个空词表
- 把硬编码的 `if i >= 10000: break` 提为模块级 MAX_LINES（默认 0 = 全量）。
原先这一行也会让 SFT 语料只用前 10000 行，且无从关闭；
另外它按「行号」计数而非「有效文本条数」，遇到空内容或解析失败的行会提前耗尽配额。
验证：用 pretrain_t2t 的前 3000 篇跑 train_tokenizer(vocab_size=2000)，
得到 vocab=2000 / merges=1708；修复前同样输入得到 0 条 merge。
两种格式各 20000 行的合成数据上，产出条数均为 20000（修复前分别是 0 和 10000）。
190 lines
15 KiB
Python
190 lines
15 KiB
Python
# 注：不建议再重复训练tokenizer（“词典”），MiniMind已自带，此脚本仅供学习和参考。基于不同词典训练的模型将导致输出完全不统一，降低社区的模型复用性
|
||
# Note: It is not recommended to re-train the tokenizer. MiniMind already includes one. This script is for learning and reference only. Training models with different tokenizers will lead to inconsistent outputs and reduce model reusability in the community.
|
||
import os
|
||
import json
|
||
from tokenizers import decoders, models, pre_tokenizers, trainers, Tokenizer
|
||
|
||
DATA_PATH = '../dataset/sft_t2t_mini.jsonl'
|
||
TOKENIZER_DIR = '../model_learn_tokenizer/'
|
||
VOCAB_SIZE = 6400
|
||
SPECIAL_TOKENS_NUM = 36
|
||
MAX_LINES = 0 # 0 表示读取全部；设为正数则只取前 N 条（快速试跑）
|
||
|
||
def get_texts(data_path, max_lines=MAX_LINES):
|
||
"""逐行产出用于训练 BPE 的文本，兼容两种主线数据格式。
|
||
|
||
pretrain_t2t(_mini).jsonl 是 {"text": ...}，
|
||
sft_t2t(_mini).jsonl 是 {"conversations": [{"role":..., "content":...}, ...]}。
|
||
原先只解析 conversations，所以把 DATA_PATH 指向预训练语料时一条文本都取不到，
|
||
BPE 仅基于 initial_alphabet 训练，产出一个没有任何 merge 的词表，且全程不报错。
|
||
"""
|
||
used = 0
|
||
with open(data_path, 'r', encoding='utf-8', errors='ignore') as f:
|
||
for line in f:
|
||
if max_lines and used >= max_lines:
|
||
break
|
||
try:
|
||
data = json.loads(line)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
if 'text' in data:
|
||
text = str(data['text'])
|
||
else:
|
||
contents = [item.get('content') for item in data.get('conversations', []) if item.get('content')]
|
||
text = "\n".join(contents) if contents else ''
|
||
if text.strip():
|
||
used += 1
|
||
yield text
|
||
if used == 0:
|
||
raise ValueError(
|
||
f'{data_path} 中没有可用文本：每行应为 {{"text": ...}} 或 '
|
||
f'{{"conversations": [{{"role":..., "content":...}}, ...]}}。'
|
||
f'继续训练只会得到一个空词表。'
|
||
)
|
||
|
||
def train_tokenizer(data_path, tokenizer_dir, vocab_size, special_tokens_num=SPECIAL_TOKENS_NUM):
|
||
tokenizer = Tokenizer(models.BPE())
|
||
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
|
||
|
||
special_tokens_list = [
|
||
"<|endoftext|>", "<|im_start|>", "<|im_end|>",
|
||
"<|object_ref_start|>", "<|object_ref_end|>", "<|box_start|>", "<|box_end|>", "<|quad_start|>", "<|quad_end|>",
|
||
"<|vision_start|>", "<|vision_end|>", "<|vision_pad|>", "<|image_pad|>", "<|video_pad|>",
|
||
"<|audio_start|>", "<|audio_end|>", "<|audio_pad|>", "<tts_pad>", "<tts_text_bos>", "<tts_text_eod>", "<tts_text_bos_single>"
|
||
]
|
||
|
||
additional_tokens_list = [
|
||
"<tool_call>", "</tool_call>",
|
||
"<tool_response>", "</tool_response>",
|
||
"<think>", "</think>"
|
||
]
|
||
num_buffer = special_tokens_num - len(special_tokens_list + additional_tokens_list)
|
||
buffer_tokens = [f"<|buffer{i}|>" for i in range(1, num_buffer + 1)] # 预留一定数量的token位置
|
||
all_special_tokens = special_tokens_list + additional_tokens_list + buffer_tokens
|
||
trainer = trainers.BpeTrainer(
|
||
vocab_size=vocab_size,
|
||
show_progress=True,
|
||
initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
|
||
special_tokens=all_special_tokens
|
||
)
|
||
texts = get_texts(data_path)
|
||
tokenizer.train_from_iterator(texts, trainer=trainer)
|
||
tokenizer.decoder = decoders.ByteLevel()
|
||
tokenizer.add_special_tokens(special_tokens_list)
|
||
|
||
os.makedirs(tokenizer_dir, exist_ok=True)
|
||
tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json"))
|
||
tokenizer.model.save(tokenizer_dir)
|
||
tokenizer_json_path = os.path.join(tokenizer_dir, "tokenizer.json")
|
||
with open(tokenizer_json_path, 'r', encoding='utf-8') as f:
|
||
tokenizer_data = json.load(f)
|
||
for token_info in tokenizer_data.get('added_tokens', []):
|
||
if token_info['content'] not in special_tokens_list:
|
||
token_info['special'] = False
|
||
with open(tokenizer_json_path, 'w', encoding='utf-8') as f:
|
||
json.dump(tokenizer_data, f, ensure_ascii=False, indent=2)
|
||
|
||
added_tokens_decoder = {}
|
||
for i, token in enumerate(all_special_tokens):
|
||
idx = tokenizer.token_to_id(token)
|
||
added_tokens_decoder[str(idx)] = {
|
||
"content": token,
|
||
"lstrip": False,
|
||
"normalized": False,
|
||
"rstrip": False,
|
||
"single_word": False,
|
||
"special": True if token in special_tokens_list else False
|
||
}
|
||
|
||
config = {
|
||
"add_bos_token": False,
|
||
"add_eos_token": False,
|
||
"add_prefix_space": False,
|
||
"added_tokens_decoder": added_tokens_decoder,
|
||
"additional_special_tokens": [t for t in special_tokens_list if t not in ["<|endoftext|>"]],
|
||
"bos_token": "<|im_start|>",
|
||
"clean_up_tokenization_spaces": False,
|
||
"eos_token": "<|im_end|>",
|
||
"legacy": True,
|
||
"model_max_length": 131072,
|
||
"pad_token": "<|endoftext|>",
|
||
"sp_model_kwargs": {},
|
||
"spaces_between_special_tokens": False,
|
||
"unk_token": "<|endoftext|>",
|
||
"image_token": "<|image_pad|>",
|
||
"audio_token": "<|audio_pad|>",
|
||
"video_token": "<|video_pad|>",
|
||
"vision_bos_token": "<|vision_start|>",
|
||
"vision_eos_token": "<|vision_end|>",
|
||
"audio_bos_token": "<|audio_start|>",
|
||
"audio_eos_token": "<|audio_end|>",
|
||
"chat_template": "{%- if tools %}\n {{- '<|im_start|>system\\n' }}\n {%- if messages[0].role == 'system' %}\n {{- messages[0].content + '\\n\\n' }}\n {%- endif %}\n {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n {%- for tool in tools %}\n {{- \"\\n\" }}\n {{- tool | tojson }}\n {%- endfor %}\n {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n {%- if messages[0].role == 'system' %}\n {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n {%- set index = (messages|length - 1) - loop.index0 %}\n {%- if ns.multi_step_tool and message.role == \"user\" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n {%- set ns.multi_step_tool = false %}\n {%- set ns.last_query_index = index %}\n {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n {%- if message.content is string %}\n {%- set content = message.content %}\n {%- else %}\n {%- set content = '' %}\n {%- endif %}\n {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n {%- elif message.role == \"assistant\" %}\n {%- set reasoning_content = '' %}\n {%- if message.reasoning_content is string %}\n {%- set reasoning_content = message.reasoning_content %}\n {%- else %}\n {%- if '</think>' in content %}\n {%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n {%- set content = content.split('</think>')[-1].lstrip('\\n') %}\n {%- endif %}\n {%- endif %}\n {%- if true %}\n {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n {%- endif %}\n {%- if message.tool_calls %}\n {%- for tool_call in message.tool_calls %}\n {%- if (loop.first and content) or (not loop.first) %}\n {{- '\\n' }}\n {%- endif %}\n {%- if tool_call.function %}\n {%- set tool_call = tool_call.function %}\n {%- endif %}\n {{- '<tool_call>\\n{\"name\": \"' }}\n {{- tool_call.name }}\n {{- '\", \"arguments\": ' }}\n {%- if tool_call.arguments is string %}\n {{- tool_call.arguments }}\n {%- else %}\n {{- tool_call.arguments | tojson }}\n {%- endif %}\n {{- '}\\n</tool_call>' }}\n {%- endfor %}\n {%- endif %}\n {{- '<|im_end|>\\n' }}\n {%- elif message.role == \"tool\" %}\n {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n {{- '<|im_start|>user' }}\n {%- endif %}\n {{- '\\n<tool_response>\\n' }}\n {{- content }}\n {{- '\\n</tool_response>' }}\n {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n {{- '<|im_end|>\\n' }}\n {%- endif %}\n {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n {{- '<|im_start|>assistant\\n' }}\n {%- if open_thinking is defined and open_thinking is true %}\n {{- '<think>\\n' }}\n {%- else %}\n {{- '<think>\\n\\n</think>\\n\\n' }}\n {%- endif %}\n{%- endif %}",
|
||
"tokenizer_class": "PreTrainedTokenizerFast"
|
||
}
|
||
|
||
with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w", encoding="utf-8") as f:
|
||
json.dump(config, f, ensure_ascii=False, indent=4)
|
||
print("Tokenizer training completed.")
|
||
|
||
def eval_tokenizer(tokenizer_dir):
|
||
from transformers import AutoTokenizer
|
||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir)
|
||
messages = [
|
||
{"role": "system", "content": "你是一个优秀的聊天机器人，总是给我正确的回应！"},
|
||
{"role": "user", "content": '你来自哪里？'},
|
||
{"role": "assistant", "content": '我来自月球'},
|
||
{"role": "user", "content": '你到底来自哪里？'},
|
||
{"role": "assistant", "content": '我来自地球'}
|
||
]
|
||
new_prompt = tokenizer.apply_chat_template(
|
||
messages,
|
||
tokenize=False
|
||
)
|
||
print('-'*100)
|
||
print(new_prompt)
|
||
print('-'*100)
|
||
print('tokenizer词表长度：', len(tokenizer))
|
||
model_inputs = tokenizer(new_prompt)
|
||
print('encoder长度：', len(model_inputs['input_ids']))
|
||
response = tokenizer.decode(model_inputs['input_ids'], skip_special_tokens=False)
|
||
print('decoder一致性：', response == new_prompt, "\n")
|
||
print('-'*100)
|
||
print('压缩率测试（Chars/Tokens）：')
|
||
test_texts = [
|
||
# 中文样本 (约200字)
|
||
"人工智能是计算机科学的一个分支，它企图了解智能的实质，并生产出一种新的能以人类智能相似的方式做出反应的智能机器，该领域的研究包括机器人、语言识别、图像识别、自然语言处理和专家系统等。人工智能从诞生以来，理论和技术日益成熟，应用领域也不断扩大，可以设想，未来人工智能带来的科技产品，将会是人类智慧的“容器”。人工智能可以对人的意识、思维的信息过程的模拟。人工智能不是人的智能，但能像人那样思考、也可能超过人的智能。",
|
||
"星际航行是指在星系内甚至星系间的空间中进行的航行。由于宇宙空间极其广阔，传统的化学火箭动力在恒星间航行时显得力不从心。科学家们提出了多种方案，包括离子推进器、核热火箭、甚至是利用反物质作为能源的设想。此外，曲率驱动和虫洞旅行等科幻概念也在理论物理研究中被反复探讨。尽管目前人类的足迹仅限于月球，但随着核聚变技术和材料科学的突破，前往火星乃至更遥远的太阳系边缘将成为可能。",
|
||
# 英文样本 (约200词/字符)
|
||
"Large language models (LLMs) are a type of artificial intelligence (AI) trained on vast amounts of text data to understand and generate human-like language. These models use deep learning techniques, specifically transformers, to process and predict the next word in a sequence. LLMs like GPT-4, Llama, and Claude have demonstrated remarkable capabilities in coding, translation, and creative writing. However, they also face challenges such as hallucinations, where the model generates factually incorrect information, and the need for significant computational resources.",
|
||
"The development of sustainable energy is crucial for the future of our planet. As climate change continues to impact global weather patterns, transitioning from fossil fuels to renewable sources like solar, wind, and hydroelectric power has become an urgent priority. Innovations in battery storage technology and smart grid management are essential to ensure a reliable energy supply. International cooperation and policy frameworks are also necessary to drive the global shift towards a greener economy and reduce carbon emissions.",
|
||
# 混合样本
|
||
"Python 是一种高级编程语言，以其简洁的语法和强大的生态系统而闻名。It is widely used in data science, machine learning, and web development. 开发者可以利用 NumPy, Pandas, and PyTorch 等库快速构建复杂的应用。学习 Python 的过程非常愉快，因为它的代码读起来就像英语一样。Whether you are a beginner or an expert, Python offers something for everyone.",
|
||
]
|
||
|
||
total_compression = 0
|
||
for i, text in enumerate(test_texts):
|
||
encoded = tokenizer.encode(text)
|
||
token_count = len(encoded)
|
||
char_count = len(text)
|
||
compression_ratio = char_count / token_count
|
||
total_compression += compression_ratio
|
||
print(f"样本 {i+1} | 字符数: {char_count:4} | Tokens: {token_count:3} | 压缩率: {compression_ratio:.2f}")
|
||
|
||
print(f"平均压缩率: {total_compression / len(test_texts):.2f}")
|
||
print('-'*100)
|
||
print('流式解码（字节缓冲）测试：')
|
||
input_ids = model_inputs['input_ids']
|
||
token_cache = []
|
||
for tid in input_ids:
|
||
token_cache.append(tid)
|
||
current_decode = tokenizer.decode(token_cache)
|
||
if current_decode and '\ufffd' not in current_decode:
|
||
display_ids = token_cache[0] if len(token_cache) == 1 else token_cache
|
||
raw_tokens = [tokenizer.convert_ids_to_tokens(int(t)) for t in (token_cache if isinstance(token_cache, list) else [token_cache])]
|
||
print(f'Token ID: {str(display_ids):15} -> Raw: {str(raw_tokens):20} -> Decode Str: {current_decode}')
|
||
token_cache = []
|
||
|
||
if __name__ == '__main__':
|
||
train_tokenizer(DATA_PATH, TOKENIZER_DIR, VOCAB_SIZE)
|
||
eval_tokenizer(TOKENIZER_DIR)
|