LLM/BPE.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a0cf1513",
   "metadata": {},
   "source": [
    "# 原理和代码\n",
    "- BPE (Byte-Pair Encoding) 算法训练流程:\n",
    "1. 准备语料库，将文本拆分为基本单位(通常是字符或字节)\n",
    "2. 统计所有相邻单元对的频率\n",
    "3. 选择最高频率的单元对合并为新的单元\n",
    "4. 更新语料库，替换所有该单元对为新单元\n",
    "5. 重复步骤2-4直到达到预设的词表大小或合并次数\n",
    "\n",
    "- BPE 编码流程:\n",
    "1. 把文本先拆成最小单位(比如单个字母或字节)\n",
    "2. 识别文本中的特殊token(如<s>, </s>)，这些特殊token会直接映射到对应ID\n",
    "3. 对于普通文本部分:\n",
    "    a. 初始化为基本单位序列\n",
    "    b. 扫描当前序列中的所有相邻单位对\n",
    "    c. 查找这些单位对是否在训练时学到的\"合并规则表\"中\n",
    "    d. 优先选择合并后ID值最小的单位对进行合并(即优先合并更基础、更短的单位)\n",
    "    e. 合并后继续重复b-d步骤，直到没有可合并的单位对\n",
    " 4. 将特殊token和处理好的普通文本部分组合，得到完整的token序列\n",
    " 5. 最后得到的单位序列就是文本的token编码结果\n",
    "\n",
    "- 参考：\n",
    "1. https://github.com/jingyaogong/minimind/blob/master/scripts/train_tokenizer.py\n",
    "2. https://github.com/SmartFlowAI/EmoLLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3eb0a804",
   "metadata": {},
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "unterminated string literal (detected at line 154) (261106877.py, line 154)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;36m  Cell \u001b[0;32mIn[1], line 154\u001b[0;36m\u001b[0m\n\u001b[0;31m    {\"role\": \"assistant\", \"content\": '你能告诉我你怎么了吗？}\u001b[0m\n\u001b[0m                                     ^\u001b[0m\n\u001b[0;31mSyntaxError\u001b[0m\u001b[0;31m:\u001b[0m unterminated string literal (detected at line 154)\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import os\n",
    "import random\n",
    "\n",
    "from tokenizers import (Tokenizer, decoders, models, normalizers,\n",
    "                        pre_tokenizers, processors, trainers)\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "# 设置随机种子以确保结果可复现\n",
    "random.seed(42)\n",
    "\n",
    "# BPE (Byte-Pair Encoding) 算法训练流程:\n",
    "# 1. 准备语料库，将文本拆分为基本单位(通常是字符或字节)\n",
    "# 2. 统计所有相邻单元对的频率\n",
    "# 3. 选择最高频率的单元对合并为新的单元\n",
    "# 4. 更新语料库，替换所有该单元对为新单元\n",
    "# 5. 重复步骤2-4直到达到预设的词表大小或合并次数\n",
    "\n",
    "# BPE 编码流程:\n",
    "# 1. 把文本先拆成最小单位(比如单个字母或字节)\n",
    "# 2. 识别文本中的特殊token(如<s>, </s>)，这些特殊token会直接映射到对应ID\n",
    "# 3. 对于普通文本部分:\n",
    "#    a. 初始化为基本单位序列\n",
    "#    b. 扫描当前序列中的所有相邻单位对\n",
    "#    c. 查找这些单位对是否在训练时学到的\"合并规则表\"中\n",
    "#    d. 优先选择合并后ID值最小的单位对进行合并(即优先合并更基础、更短的单位)\n",
    "#    e. 合并后继续重复b-d步骤，直到没有可合并的单位对\n",
    "# 4. 将特殊token和处理好的普通文本部分组合，得到完整的token序列\n",
    "# 5. 最后得到的单位序列就是文本的token编码结果\n",
    "\n",
    "\n",
    "def train_tokenizer():\n",
    "    # 读取JSONL文件并提取文本数据\n",
    "    def read_texts_from_jsonl(file_path):\n",
    "        with open(file_path, 'r', encoding='utf-8') as f:\n",
    "            for line in f:\n",
    "                data = json.loads(line)\n",
    "                yield data['text']\n",
    "\n",
    "    data_path = './bpe.jsonl'\n",
    "\n",
    "    # 初始化tokenizer\n",
    "    tokenizer = Tokenizer(models.BPE())\n",
    "    # 预分词器负责将原始文本进行初步分割，为后续的 BPE 合并操作做准备add_prefix_space=False 表示不在每个序列前添加空格\n",
    "    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)\n",
    "\n",
    "    # 定义特殊token\n",
    "    special_tokens = [\"<unk>\", \"<s>\", \"</s>\"]\n",
    "\n",
    "    # 设置训练器并添加特殊token\n",
    "    trainer = trainers.BpeTrainer(\n",
    "        vocab_size=6400,\n",
    "        special_tokens=special_tokens,  # 确保这三个token被包含\n",
    "        show_progress=True,\n",
    "        # 使用字节级别的基础字符集作为初始词汇，它包含了256个可能的字节值（0-255）对应的Unicode字符表示\n",
    "        initial_alphabet=pre_tokenizers.ByteLevel.alphabet()\n",
    "    )\n",
    "\n",
    "    # 读取文本数据\n",
    "    texts = read_texts_from_jsonl(data_path)\n",
    "\n",
    "    # 训练tokenizer\n",
    "    tokenizer.train_from_iterator(texts, trainer=trainer)\n",
    "\n",
    "    # 设置解码器\n",
    "    tokenizer.decoder = decoders.ByteLevel()\n",
    "\n",
    "    # 检查特殊token的索引\n",
    "    # 确保三个特殊token被正确分配了预期的ID：\n",
    "    # <unk>应该是ID 0 - 用于表示未知词汇\n",
    "    assert tokenizer.token_to_id(\"<unk>\") == 0\n",
    "    # <s>应该是ID 1 - 用于表示文本开始\n",
    "    assert tokenizer.token_to_id(\"<s>\") == 1\n",
    "    # </s>应该是ID 2 - 用于表示文本结束\n",
    "    assert tokenizer.token_to_id(\"</s>\") == 2\n",
    "\n",
    "    # 保存tokenizer\n",
    "    tokenizer_dir = \"./EmoLLM_tokenizer\"\n",
    "    os.makedirs(tokenizer_dir, exist_ok=True)\n",
    "    tokenizer.save(os.path.join(tokenizer_dir, \"tokenizer.json\"))\n",
    "    tokenizer.model.save(\"./EmoLLM_tokenizer\")\n",
    "\n",
    "    # 手动创建配置文件\n",
    "    config = {\n",
    "        \"add_bos_token\": False,\n",
    "        \"add_eos_token\": False,\n",
    "        \"add_prefix_space\": False,\n",
    "        \"added_tokens_decoder\": {\n",
    "            # 配置 <unk> 标记 (未知词标记)\n",
    "            # 索引为0，用于表示词表外的词或字符\n",
    "            \"0\": {\n",
    "                \"content\": \"<unk>\",  # 标记的实际内容\n",
    "                \"lstrip\": False,     # 是否从左侧移除空白\n",
    "                \"normalized\": False,  # 是否规范化\n",
    "                \"rstrip\": False,     # 是否从右侧移除空白\n",
    "                \"single_word\": False,  # 是否作为单个词处理\n",
    "                \"special\": True      # 标记为特殊标记，在解码时会特殊处理\n",
    "            },\n",
    "            # 配置 <s> 标记 (序列开始标记)\n",
    "            # 索引为1，用于标识文本序列的开始\n",
    "            \"1\": {\n",
    "                \"content\": \"<s>\",\n",
    "                \"lstrip\": False,\n",
    "                \"normalized\": False,\n",
    "                \"rstrip\": False,\n",
    "                \"single_word\": False,\n",
    "                \"special\": True\n",
    "            },\n",
    "            # 配置 </s> 标记 (序列结束标记)\n",
    "            # 索引为2，用于标识文本序列的结束\n",
    "            \"2\": {\n",
    "                \"content\": \"</s>\",\n",
    "                \"lstrip\": False,\n",
    "                \"normalized\": False,\n",
    "                \"rstrip\": False,\n",
    "                \"single_word\": False,\n",
    "                \"special\": True\n",
    "            }\n",
    "        },\n",
    "        \"additional_special_tokens\": [],\n",
    "        \"bos_token\": \"<s>\",\n",
    "        \"clean_up_tokenization_spaces\": False,\n",
    "        \"eos_token\": \"</s>\",\n",
    "        \"legacy\": True,\n",
    "        \"model_max_length\": 32768,\n",
    "        \"pad_token\": \"<unk>\",\n",
    "        \"sp_model_kwargs\": {},\n",
    "        \"spaces_between_special_tokens\": False,\n",
    "        \"tokenizer_class\": \"PreTrainedTokenizerFast\",\n",
    "        \"unk_token\": \"<unk>\",\n",
    "        # 配置聊天模板 - 使用Jinja2模板格式定义模型输入的格式化方式\n",
    "        \"chat_template\": \"{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{{ '<s>system\\\\n' + system_message + '</s>\\\\n' }}{% else %}{{ '<s>system\\\\n你是 EmoLLM，是一个完全开源的心理健康大模型。</s>\\\\n' }}{% endif %}{% for message in messages %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<s>user\\\\n' + content + '</s>\\\\n<s>assistant\\\\n' }}{% elif message['role'] == 'assistant' %}{{ content + '</s>' + '\\\\n' }}{% endif %}{% endfor %}\"\n",
    "        # 聊天模板说明:\n",
    "        # 1. 如果第一条消息是系统消息，将其作为系统指令；否则使用默认系统消息\n",
    "        # 2. 用户消息格式: <s>user\\n{用户内容}</s>\\n<s>assistant\\n\n",
    "        # 3. 助手消息格式: {助手内容}</s>\\n\n",
    "        # 4. 特殊标记<s>和</s>用于标记消息的开始和结束\n",
    "    }\n",
    "\n",
    "    # 保存配置文件\n",
    "    with open(os.path.join(tokenizer_dir, \"tokenizer_config.json\"), \"w\", encoding=\"utf-8\") as config_file:\n",
    "        json.dump(config, config_file, ensure_ascii=False, indent=4)\n",
    "\n",
    "    print(\"Tokenizer training completed and saved.\")\n",
    "\n",
    "\n",
    "def eval_tokenizer():\n",
    "    # 加载预训练的tokenizer\n",
    "    tokenizer = AutoTokenizer.from_pretrained(\"./EmoLLM_tokenizer\")\n",
    "\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": \"你是EmoLLM，优秀的心理健康大模型。\"},\n",
    "        {\"role\": \"user\", \"content\": '我最近很苦恼'},\n",
    "        {\"role\": \"assistant\", \"content\": '你能告诉我你怎么了吗？'}\n",
    "    ]\n",
    "    new_prompt = tokenizer.apply_chat_template(\n",
    "        messages,\n",
    "        tokenize=False\n",
    "    )\n",
    "    print(new_prompt)\n",
    "\n",
    "    # 获取实际词汇表长度（包括特殊符号）\n",
    "    actual_vocab_size = len(tokenizer)\n",
    "    print('tokenizer实际词表长度：', actual_vocab_size)\n",
    "\n",
    "    model_inputs = tokenizer(new_prompt)\n",
    "    print('encoder长度：', len(model_inputs['input_ids']))\n",
    "\n",
    "    input_ids = model_inputs['input_ids']\n",
    "    response = tokenizer.decode(input_ids, skip_special_tokens=False)\n",
    "    print('decoder和原始文本是否一致：', response == new_prompt)\n",
    "\n",
    "\n",
    "def main():\n",
    "    train_tokenizer()\n",
    "    eval_tokenizer()\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "wxz",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}