sft/add_tokens.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "cc99f1c1-834e-4657-be50-8a7ba81acba7",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7e09040157354b36a629f7adba0e31ab",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "vocabulary size: 128256\n",
      "['[3]', '[2]', '[1]', '[4]', '[5]']\n",
      "After we add 5 tokens\n",
      "vocabulary size: 128261\n",
      "torch.Size([128261, 4096])\n",
      "tensor([[ 0.0027,  0.0267,  0.0343,  ...,  0.0065, -0.0214,  0.0078],\n",
      "        [-0.0338, -0.0048, -0.0302,  ...,  0.0280,  0.0134,  0.0313]],\n",
      "       grad_fn=<SliceBackward0>)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "('new_token/tokenizer_config.json',\n",
       " 'new_token/special_tokens_map.json',\n",
       " 'new_token/tokenizer.json')"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import AutoTokenizer, AutoModel\n",
    "\n",
    "checkpoint = \"llm-research/meta-llama-3-8b\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(checkpoint)\n",
    "model = AutoModel.from_pretrained(checkpoint)\n",
    "print('vocabulary size:', len(tokenizer))\n",
    "\n",
    "new_tokens = ['[5]', '[4]', '[3]', '[2]', '[1]']\n",
    "# check if the tokens are already in the vocabulary\n",
    "new_tokens = set(new_tokens) - set(tokenizer.vocab.keys())\n",
    "new_tokens = list(new_tokens)\n",
    "print(new_tokens)\n",
    "\n",
    "num_added_toks = tokenizer.add_tokens(new_tokens, special_tokens=True)\n",
    "\n",
    "print(\"After we add\", num_added_toks, \"tokens\")\n",
    "print('vocabulary size:', len(tokenizer))\n",
    "\n",
    "# add new, random embeddings for the new tokens\n",
    "model.resize_token_embeddings(len(tokenizer))\n",
    "\n",
    "print(model.embed_tokens.weight.size())\n",
    "\n",
    "# Randomly generated matrix\n",
    "print(model.embed_tokens.weight[-2:, :])\n",
    "\n",
    "# save new tokenizer\n",
    "tokenizer.save_pretrained(\"new_token\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "ee9222d2-46d9-4e61-975d-102c55718dcb",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "vocabulary size: 128261\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"new_token\")\n",
    "\n",
    "print('vocabulary size:', len(tokenizer))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48f48ca9-574b-454e-a40b-4d6684835dfd",
   "metadata": {},
   "source": [
    "在使用预训练模型时，我们有时需要使用一些自定义 token 来增强输入，例如使用 `[ENT_START]` 和 `[ENT_END]` 在文本中标记出实体。由于自定义 token 并不在预训练模型原来的词表中，因此直接运用分词器 (Tokenizer) 处理输入就会出现问题。\n",
    "\n",
    "例如直接使用 BERT 分词器处理下面的句子："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b3a7c733-e529-4f88-96b4-26da7581bc4c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import subprocess\n",
    "import os\n",
    "\n",
    "result = subprocess.run('bash -c \"source /etc/network_turbo && env | grep proxy\"', shell=True, capture_output=True, text=True)\n",
    "output = result.stdout\n",
    "for line in output.splitlines():\n",
    "    if '=' in line:\n",
    "        var, value = line.split('=', 1)\n",
    "        os.environ[var] = value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a84b953e-6b0e-458c-bdbc-30d67bef9aa9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['<|begin_of_text|>', 'Two', 'Ġ[', 'ENT', '_START', ']', 'Ġcars', 'Ġ[', 'ENT', '_END', ']', 'Ġcollided', 'Ġin', 'Ġa', 'Ġ[', 'ENT', '_START', ']', 'Ġtunnel', 'Ġ[', 'ENT', '_END', ']', 'Ġthis', 'Ġmorning', '.']\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"llm-research/meta-llama-3-8b\")\n",
    "\n",
    "sentence = 'Two [ENT_START] cars [ENT_END] collided in a [ENT_START] tunnel [ENT_END] this morning.'\n",
    "print(tokenizer(sentence).tokens())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d576b80c-4271-4cda-8eb2-9b130528b4e8",
   "metadata": {},
   "source": [
    "由于分词器无法识别 `[ENT_START]` 和 `[ENT_END]` ，将它们都当作未知字符处理，例如 `[ENT_END]` 被切分成了 `'['、'E'、'##NT'、'_'、'E'、'##ND'、']'` 七个 token，很明显不符合我们的预期。\n",
    "\n",
    "此外，有时我们还会遇到一些领域相关词汇，例如医学领域的文本通常会包含大量的医学术语，它们可能并不在模型的词表中（例如一些术语是使用多个词语的缩写拼接而成），这时也会出现上面的问题。\n",
    "\n",
    "此时我们就需要将这些新 token 添加到模型的词表中，让分词器与模型可以识别并处理这些 token。\n",
    "\n",
    "## 添加新 token\n",
    "### 添加方法\n",
    "\n",
    "Huggingface 的 Transformers 库提供了两种方式来添加新 token，分别是：\n",
    "\n",
    "- `add_tokens()` 添加普通 token：添加新 token 列表，如果 token 不在词表中，就会被添加到词表的最后。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4569e3c6-255d-4de7-8fd9-2abdc1ce87e5",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "We have added 2 tokens\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"llm-research/meta-llama-3-8b\")\n",
    "  \n",
    "num_added_toks = tokenizer.add_tokens([\"new_tok1\", \"my_new-tok2\"])\n",
    "print(\"We have added\", num_added_toks, \"tokens\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0515f33-85d3-48a1-b59f-48158dfee609",
   "metadata": {},
   "source": [
    "为了防止 token 已经包含在词表中，我们还可以预先对新 token 列表进行过滤："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b022b907-1200-4393-9371-d01a91a6f082",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_tokens = [\"new_tok1\", \"my_new-tok2\"]\n",
    "new_tokens = set(new_tokens) - set(tokenizer.vocab.keys())\n",
    "tokenizer.add_tokens(list(new_tokens))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b7ad757-2797-4b24-a68d-c8546887a620",
   "metadata": {},
   "source": [
    "- `add_special_tokens()` 添加特殊 token：添加包含特殊 token 的字典，键值从 bos_token, eos_token, unk_token, sep_token, pad_token, cls_token, mask_token, additional_special_tokens 中选择。与 add_tokens() 类似，如果 token 不在词表中，就会被添加到词表的最后。添加后，还可以通过特殊属性来访问这些 token，例如 tokenizer.cls_token 就指向 cls token。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "044b7bc5-530a-42d8-a643-1b0718e0bd7a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "We have added 1 tokens\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"llm-research/meta-llama-3-8b\")\n",
    "\n",
    "special_tokens_dict = {\"cls_token\": \"[MY_CLS]\"}\n",
    "  \n",
    "num_added_toks = tokenizer.add_special_tokens(special_tokens_dict)\n",
    "print(\"We have added\", num_added_toks, \"tokens\")\n",
    "  \n",
    "assert tokenizer.cls_token == \"[MY_CLS]\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "501e7b49-f001-4ea9-bb32-352cebbb5314",
   "metadata": {},
   "source": [
    "我们也可以使用 `add_tokens()` 添加特殊 token，只需要额外设置参数 `special_tokens=True`："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "cfc71cbe-d753-476b-b1a2-057f44265dc1",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "We have added 2 tokens\n",
      "['<|begin_of_text|>', '[NEW_tok1]', 'ĠHello', 'Ġ', '[NEW_tok2]', 'Ġ', '[NEW_tok3]', 'ĠWorld', 'Ġ', '[NEW_tok4]', '!']\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"llm-research/meta-llama-3-8b\")\n",
    "  \n",
    "num_added_toks = tokenizer.add_tokens([\"[NEW_tok1]\", \"[NEW_tok2]\"])\n",
    "num_added_toks = tokenizer.add_tokens([\"[NEW_tok3]\", \"[NEW_tok4]\"], special_tokens=True)\n",
    "  \n",
    "print(\"We have added\", num_added_toks, \"tokens\")\n",
    "print(tokenizer('[NEW_tok1] Hello [NEW_tok2] [NEW_tok3] World [NEW_tok4]!').tokens())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b58784a2-edcd-4881-88b1-ad5f1db29225",
   "metadata": {},
   "source": [
    "特殊 token 的标准化 (normalization) 过程与普通 token 有一些不同，比如不会被小写。\n",
    "\n",
    "对于之前的例子，很明显实体标记符 `[ENT_START]` 和 `[ENT_END]` 属于特殊 token，因此按添加特殊 token 的方式进行。如果使用 `add_tokens()` 则需要额外设置 `special_tokens=True`，或者也可以直接使用 `add_special_tokens()`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ef98fbd7-0849-47ed-807c-c6e70badc371",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "We have added 2 tokens\n",
      "['<|begin_of_text|>', 'Two', 'Ġ', '[ENT_START]', 'Ġcars', 'Ġ', '[ENT_END]', 'Ġcollided', 'Ġin', 'Ġa', 'Ġ', '[ENT_START]', 'Ġtunnel', 'Ġ', '[ENT_END]', 'Ġthis', 'Ġmorning', '.']\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"llm-research/meta-llama-3-8b\")\n",
    "\n",
    "num_added_toks = tokenizer.add_tokens(['[ENT_START]', '[ENT_END]'], special_tokens=True)\n",
    "# num_added_toks = tokenizer.add_special_tokens({'additional_special_tokens': ['[ENT_START]', '[ENT_END]']})\n",
    "print(\"We have added\", num_added_toks, \"tokens\")\n",
    "\n",
    "sentence = 'Two [ENT_START] cars [ENT_END] collided in a [ENT_START] tunnel [ENT_END] this morning.'\n",
    "\n",
    "print(tokenizer(sentence).tokens())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fed6b432-4de4-4c5d-b2d2-471775e608ef",
   "metadata": {},
   "source": [
    "可以看到，分词器成功地将 `[ENT_START]` 和 `[ENT_END]` 识别为 token，并且依旧保持大写。\n",
    "\n",
    "### 调整 embedding 矩阵\n",
    "\n",
    "**注意！无论使用哪种方式向词表中添加新 token 后，都需要重置模型 token embedding 矩阵的大小，也就是向矩阵中添加新 token 对应的 embedding，这样模型才可以正常工作（将 token 映射到对应的 embedding）。**\n",
    "\n",
    "该操作通过调用预训练模型的 `resize_token_embeddings()` 函数来实现，例如对于上面的例子："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "a5422526-af9c-4cca-9242-49bfe376ee25",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "68db1d253f7148b0960cfe4074800c1a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "128256\n",
      "We have added 2 tokens\n",
      "128258\n",
      "torch.Size([128258, 4096])\n",
      "tensor([[-0.0065, -0.0182,  0.0124,  ..., -0.0061, -0.0275, -0.0052],\n",
      "        [ 0.0128,  0.0129, -0.0050,  ...,  0.0162,  0.0117, -0.0353]],\n",
      "       grad_fn=<SliceBackward0>)\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoTokenizer, AutoModel\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"llm-research/meta-llama-3-8b\")\n",
    "model = AutoModel.from_pretrained(\"llm-research/meta-llama-3-8b\")\n",
    "\n",
    "print(len(tokenizer))\n",
    "num_added_toks = tokenizer.add_tokens(['[ENT_START]', '[ENT_END]'], special_tokens=True)\n",
    "print(\"We have added\", num_added_toks, \"tokens\")\n",
    "print(len(tokenizer))\n",
    "\n",
    "model.resize_token_embeddings(len(tokenizer))\n",
    "print(model.embed_tokens.weight.size())\n",
    "\n",
    "# Randomly generated matrix\n",
    "print(model.embed_tokens.weight[-2:, :])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cb1ec0c-60d9-4991-bd2e-e98c7cf55f11",
   "metadata": {
    "tags": []
   },
   "source": [
    "可以看到，在添加了特殊 token `[ENT_START]` 和 `[ENT_END]` 之后，分词器的词表大小从 128256 增加到了 128258，并且模型的 token embedding 矩阵大小也成功调整为了 \n",
    "128258×4096。\n",
    "\n",
    "我们还尝试打印出新添加 token 对应的 embedding。因为新 token 会添加在词表的末尾，因此只需打印出矩阵最后两行。如果你重复运行一下上面的代码，就会发现每次打印出的 `[ENT_START]` 和 `[ENT_END]` 的 embedding 是不同的。这是因为在默认情况下，这些新 token 的 embedding 是随机初始化的。\n",
    "\n",
    "如果有充分的训练语料对模型进行微调或者继续预训练，那么将新添加 token 初始化为随机向量没什么问题。但是如果训练语料较少，甚至是只有很少语料的 few-shot learning 场景下，这种做法就可能存在问题。研究表明，在训练数据不够多的情况下，这些新添加 token 的 embedding 只会在初始值附近小幅波动。换句话说，即使经过训练，它们的值事实上还是随机的。\n",
    "\n",
    "因此，在很多情况下，我们需要手工初始化这些新 token 的 embedding。对于 Transformers 库来说，可以通过直接对 embedding 矩阵赋值来实现。例如对于上面的例子，我们将这两个新 token 的 embedding 都初始化为全零向量："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "c9ff4918-5715-4701-b3ad-119c25cb7fd5",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.]], grad_fn=<SliceBackward0>)\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "with torch.no_grad():\n",
    "    model.embed_tokens.weight[-2:, :] = torch.zeros([2, model.config.hidden_size], requires_grad=True)\n",
    "print(model.embed_tokens.weight[-2:, :])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02d07321-f7e0-421d-bcc3-07e11b0383e8",
   "metadata": {},
   "source": [
    "注意，初始化 embedding 的过程并不可导，因此这里通过 `torch.no_grad()` 暂停梯度的计算。\n",
    "\n",
    "还有另外一些比较常见的操作是根据新添加 token 的语义，将其值初始化为训练好 token 的 embedding。例如对于上面的例子，我们可以将 `[ENT_START]` 和 `[ENT_END]` 的值都初始化为“entity”对应的 embedding。因为 token id 就是 token 在矩阵中的索引，因此我们可以直接通过 `weight[token_id]` 取出“entity”对应的 embedding。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "5a4bf56a-35ce-406c-afb8-534043c1d8fb",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3069\n",
      "tensor([[-0.0030,  0.0055,  0.0151,  ..., -0.0129,  0.0041,  0.0093],\n",
      "        [-0.0030,  0.0055,  0.0151,  ..., -0.0129,  0.0041,  0.0093]],\n",
      "       grad_fn=<SliceBackward0>)\n"
     ]
    }
   ],
   "source": [
    "token_id = tokenizer.convert_tokens_to_ids('entity')\n",
    "token_embedding = model.embed_tokens.weight[token_id]\n",
    "print(token_id)\n",
    "\n",
    "with torch.no_grad():\n",
    "    for i in range(1, num_added_toks+1):\n",
    "        model.embed_tokens.weight[-i:, :] = token_embedding.clone().detach().requires_grad_(True)\n",
    "print(model.embed_tokens.weight[-2:, :])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0aaff2d-ce85-4cc7-8bdb-10da7e048da8",
   "metadata": {},
   "source": [
    "可以看到最终结果符合我们的预期，`[ENT_START]` 和 `[ENT_END]` 被初始化为相同的 embedding。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c392663-9744-4442-baac-7eb7c3a36d60",
   "metadata": {
    "tags": []
   },
   "source": [
    "另一种常见的做法是根据新 token 的语义，使用对应的描述文本来完成初始化。例如将值初始化为描述文本中所有 token 的平均值，假设新 token $t_i$ 的描述文本为 $w_{i,1}, w_{i,2}, \\ldots, w_{i,n}$，那么 $t_i$ 的初始化 embedding 为：\n",
    "\n",
    "$E(t_i) = \\frac{1}{n} \\sum_{j=1}^{n} E(w_{i,j})$\n",
    "\n",
    "这里 $E$ 表示预训练模型的 token embedding 矩阵。对于上面的例子，我们可以分别为 `[ENT_START]` 和 `[ENT_END]` 编写对应的描述，然后再到它们的值进行初始化："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "a5d10139-848e-4a7a-9575-07113b04a604",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['end', 'Ġof', 'Ġentity']\n",
      "['start', 'Ġof', 'Ġentity']\n",
      "tensor([[-0.0044, -0.0003,  0.0006,  ..., -0.0017, -0.0066, -0.0005],\n",
      "        [-0.0054, -0.0013,  0.0025,  ...,  0.0010, -0.0039,  0.0031]],\n",
      "       grad_fn=<SliceBackward0>)\n"
     ]
    }
   ],
   "source": [
    "descriptions = ['start of entity', 'end of entity']\n",
    "\n",
    "with torch.no_grad():\n",
    "    for i, token in enumerate(reversed(descriptions), start=1):\n",
    "        tokenized = tokenizer.tokenize(token)\n",
    "        print(tokenized)\n",
    "        tokenized_ids = tokenizer.convert_tokens_to_ids(tokenized)\n",
    "        new_embedding = model.embed_tokens.weight[tokenized_ids].mean(axis=0)\n",
    "        model.embed_tokens.weight[-i, :] = new_embedding.clone().detach().requires_grad_(True)\n",
    "\n",
    "print(model.embed_tokens.weight[-2:, :])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77339a42-554f-4a6b-93a6-d747dd8df30e",
   "metadata": {},
   "source": [
    "可以看到，这里成功地将 `[ENT_START]` 初始化为“start”、“of”、“entity”三个 token embedding 的平均值，将 `[ENT_END]` 初始化为“end”、“of”、“entity” embedding 的平均值。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "405b4fe8-2ce7-4ff3-958b-603b63b2621f",
   "metadata": {},
   "source": [
    "参考：\n",
    "1. https://xiaosheng.blog/2023/01/07/add-new-token\n",
    "2. https://stackoverflow.com/questions/76198051/how-to-add-new-tokens-to-an-existing-huggingface-tokenizer"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}