如何實現一個分詞器

在開發代碼補全插件的過程中，根據項目需要，我實現了一個分詞器，本文將介紹分詞器的具體實現細節。

一、什麼是分詞器？

分詞器是 NLP（natural language processing，自然語言處理）領域的一個重要部分，它可以把一段文本轉換爲小的單元，稱爲 token 。token 可以是單詞、字符、標點符號等。在基於 Transformer 的 LLM （Large Language Model，大語言模型）中，如 BERT 或 GPT 等，分詞器扮演着更重要的角色。大模型通過不斷學習來發現 token 之間的聯繫，從而能夠預測下一個 token，實現與人類對話的效果。

OpenAI 官方提供了一個測試頁面，來幫助我們理解一段文本會怎樣被拆分爲一系列 token ，訪問該鏈接即可查看相關內容。計算一段文本會被拆分爲多少個 token 非常重要，因爲模型能夠一次能夠讀取的 token 數量是有限的。

二、分詞器在代碼補全領域的應用

代碼補全是 LLM 的一個應用場景，在編輯器中安裝 Copilot 插件後，在編寫代碼時，就能自動獲取實時的補全建議，提高開發速度。Copilot 插件的工作原理如下：

（1）理解上下文： 當開發者編寫代碼的時候，Copilot 會持續地分析當前的上下文，包括當前正在輸入的代碼、註釋、和文件中的其他代碼。同時，還會分析項目整體結構和用到的庫。

（2）向模型發送請求： 基於當前上下文，Copilot 會給 LLM 發送一個請求。該請求會包含相關的代碼上下文，例如：光標前的代碼、光標後的代碼、函數名稱和註釋等。

（3）LLM 生成補全建議： LLM 收到請求信息後，基於它從大量公共代碼庫學到的知識和當前項目的具體上下文信息，生成多個代碼建議。

（4）展示補全建議： Copilot 會把多個代碼建議顯示到編輯器中，開發者可以採納、拒絕或修改補全代碼。

在代碼補全插件中，編碼器是一個核心的功能模塊，它會把代碼分割爲 token，包括：關鍵詞、運算符、單詞、標點符號、空格等。假設我們輸入一行代碼：

let name = 'julian'

編碼器會把代碼拆分爲幾個獨立部分：

再把每個部分進行編碼，具體結果爲：

[1169, 836, 284, 364, 73, 360, 1122, 6]

由上圖可以看出，分詞器能夠計算出每段代碼的 token 數量，在向 LLM 發送請求前，Copilot 會先檢查當前 Prompt 的 token 數量是否超過了 LLM 的閾值，如果超出，就需要對 Prompt 進行截取，避免由於 token 超出指定範圍而導致代碼補全失敗。因此，分詞器的作用不言而喻，如果不能精確的計算 token 數量，會影響開發者的使用體驗。

三、BPE 算法

目前主流的 LLM 都基於 BPE 算法（Byte-Pair Encoding，字節對編碼）來實現分詞器，如 OpenAI 的 GPT 模型。

BPE 算法訓練過程

（1）從語料庫中獲取用於編寫所有單詞的符號來構建詞彙表，每個符號即爲一個 token，假設我們的語料庫是一個字符串man woman，那麼我們就可以得到一個詞彙表：

const text = 'man woman'

const tokens = {
  ' ': 0, 
  'a': 1, 
  'm': 3, 
  'n': 4, 
  'o': 5, 
  'w': 6,
}

（2）把 text 轉換爲 token：

const tokenized_text = [
  'm', 'a', 'n', ' ', 'w', 'o', 'm', 'a', 'n'
]

（3）把每兩個相鄰的 token 進行合併，計算每個字節對出現的次數；

（4）把出現次數最多的字節作爲一個新的 token 加入詞表；

（5）重複上述過程，直到沒有字節對可以合併。

具體的字節對合並過程如下：

（a）第 1 次合併

mergeObj = {
 'ma': 2,
 'an': 2,
 ' w': 1,
 'wo': 1,
 'om': 1,
}

// 將'm'和'a'進行合併，作爲新的token
tokens = {
  ' ': 0, 
  'a': 1, 
  'm': 2, 
  'n': 3, 
  'o': 4, 
  'w': 5,
  'ma': 6
}

// 重新把text轉爲token
tokenized_text = [
  'ma', 'n', ' ', 'w', 'o', 'ma', 'n'
]

（b）第 2 次合併

mergeObj = {
 'man': 2,
 ' w': 1,
 'wo': 1,
 'oma': 1,
}

// 將'ma'和'n'進行合併，作爲新的token
tokens = {
  ' ': 0, 
  'a': 1, 
  'm': 2, 
  'n': 3, 
  'o': 4, 
  'w': 5,
  'ma': 6,
  'man': 7
}

// 重新把text轉爲token
tokenized_text = [
  'man', ' ', 'w', 'o', 'man'
]

（c）第 3 次合併

mergeObj = {
 ' w': 1,
 'wo': 1,
 'oman': 1,
}

// 將'w'和'o'進行合併，作爲新的token
tokens = {
  ' ': 0, 
  'a': 1, 
  'm': 2, 
  'n': 3, 
  'o': 4, 
  'w': 5,
  'ma': 6,
  'man': 7,
  'wo': 8
}

// 重新把text轉爲token
tokenized_text = [
  'man', ' ', 'wo', 'man'
]

（d）第 4 次合併

mergeObj = {
 ' wo': 1,
 'woman': 1,
}

// 將'wo'和'man'進行合併，作爲新的token
tokens = {
  ' ': 0, 
  'a': 1, 
  'm': 2, 
  'n': 3, 
  'o': 4, 
  'w': 5,
  'ma': 6,
  'man': 7,
  'wo': 8,
  'woman': 9
}

// 重新把text轉爲token
tokenized_text = [
  'man', ' ', 'woman'
]

最終，會產生一個詞表文件tokenizer.json和一個BPE文件。tokenizer.json文件會包含完整的詞彙表映射，即每個 token（詞或子詞）到其唯一 ID 的對應關係。BPE文件記錄了子詞的合併規則和順序，模型在需要對新文本進行分詞時會根據這些規則進行處理。

// tokenizer.json 文件
{
  ' ': 0, 
  'a': 1, 
  'm': 2, 
  'n': 3, 
  'o': 4, 
  'w': 5,
  'ma': 6,
  'man': 7,
  'wo': 8,
  'woman': 9
}

// BPE文件
m a
ma n
w o
wo man

四、爲什麼不直接使用社區中的分詞器？

當我們向分詞器中輸入一段文本後，分詞器會讀取tokenizer.json文件和BPE文件，對文本進行編碼。假設我們輸入man woman後，分詞器會把該文本編碼爲[7, 0, 9]。由於每個模型經過訓練產生的詞表文件不同，如果直接使用社區中的分詞器，可能導致分詞結果不準確。因此，很有必要基於當前模型產生的tokenizer.json文件和BPE文件，來實現一個分詞器。

五、實現分詞器

1. 實現思路：

將文本塊轉換爲字節數組，這是編碼的第一步，例如，將字符串 "let" 轉換爲字節數組 ['l', 'e', 't']。然後，獲取字符對，即文本中相鄰字符的組合，如上述字節數組會得到 [ ['l', 'e'], ['e', 't'] ]。

如果沒有字符對（通常是輸入文本長度爲 1），則直接返回編碼後的字節。否則，進入一個循環，不斷合併最頻繁的字符對，直到不能再合併爲止。這是 BPE 算法的核心，通過合併頻繁出現的字符對來減少文本的長度。在每次循環中，找出當前最頻繁的字符對，並將它們合併。合併後，更新字節數組並繼續下一輪合併，直到字節數組長度爲 1 或沒有更多字符對可以合併。

最後，將合併後的字節數組轉換爲 tokens，並將結果緩存，這樣相同的輸入在下次處理時可以直接從緩存中獲取結果，提高效率。通過這些步驟，分詞器能夠高效地處理文本塊，將其轉換爲 tokens，同時利用緩存避免重複計算。

2. 前期準備

vocab.bpe: 記錄字符合並的順序。

tokenizer.json: 包含編碼的映射關係。

3. 工具函數

（1）dictZip函數的作用是將兩個數組x和y組合成一個Map對象。對於每個索引i，x數組中的元素將作爲鍵，y數組中相應的元素將作爲值。這樣，每個x中的元素都會與y中相應位置的元素配對，形成鍵值對。最終，函數返回這個包含了所有鍵值對的Map對象。

const dictZip = (x, y) => {
  let result = new Map();
  x.forEach((_, i) => {
    result.set(x[i], y[i]);
  });
  return result;
};

（2）get_char_pairs 函數的作用是接收一個字符串作爲參數，然後生成並返回一個包含該字符串中所有相鄰字符對的集合。

function get_char_pairs(word) {
  // 初始化一個空的Set用於存儲字符對
  let pairs = new Set(),
    prev_char = word[0]; // 存儲前一個字符，初始爲單詞的第一個字符
  for (let i = 1; i < word.length; i++) { // 從第二個字符開始遍歷單詞
    let char = word[i]; // 當前字符
    pairs.add([prev_char, char]), (prev_char = char); // 將前一個字符和當前字符組成的對添加到集合中，並更新前一個字符
  }
  return pairs; // 返回包含所有字符對的集合
}

（3）mutatingConcat 可以將源數組 (src) 的元素添加到目標數組 (dest) 的末尾，並返回修改後的目標數組。

function mutatingConcat(dest, src) {
    for (let i = 0; i < src.length; i++) dest.push(src[i]);
    return dest;
}

4. 讀取tokenizer.json文件和BPE文件

// 讀取 "tokenizer.json" 文件並解析其內容
let encoder_text = fs.readFileSync(path.resolve(__dirname, "tokenizer.json"));
let encoder_json = JSON.parse(encoder_text.toString());

// 讀取 "vocab.bpe" 文件
let bpe_file = fs.readFileSync(path.resolve(__dirname, "vocab.bpe"), "utf-8");

// 把編碼文件中的內容存入 Map 中
// encoder = {
//   "A": 32,
//   "B": 33,
//   "C": 34,
//   "D": 35,
// }
this.encoder = new Map(Object.entries(encoder_json));

// 創建 decoder Map，通過交換 encoder 的鍵和值來實現
// decoder = {
//   '32': 'A',
//   '33': 'B',
//   '34': 'C',
//   '35': 'D'
// }
for (let [key, value] of this.encoder) {
  this.decoder.set(value, key);
}

// 將 bpe_file 按行拆分，過濾掉空行，得到 bpe_merges。
// 使用 dictZip 函數將 bpe_merges 和其索引創建一個字典 bpe_ranks。
// 如 {
//    'Ġ Ġ' => 0,
//    'Ġ t' => 1,
//    'Ġ a' => 2,
//    'i n' => 3,
//    }
let bpe_merges = bpe_file
  .split(/\r?\n/)
  .slice(1)
  .filter((l) => l.trim().length > 0);
this.bpe_ranks = dictZip(bpe_merges, range(0, bpe_merges.length));

5. 分詞器的核心是使用 BPE 算法不斷合併出現最頻繁的字符對，將輸入的文本塊轉換爲 tokens，具體過程如下：

// 假設輸入的文本是"let"

bpe(chunk) {
    // 檢查緩存中是否已有處理結果，如果有，則直接返回緩存的結果，避免重複計算
    if (this.cache.has(chunk)) {
      return this.cache.get(chunk);
    }

    // 將文本塊轉換爲字節數組，這是編碼的第一步
    // 例如，對於字符串 "let"，輸出將是 [ 'l', 'e', 't' ]
    let bytes = this.byteEncodeStr(chunk);

    // 獲取字節對（即字符對），這是爲了找出文本中相鄰字符的組合
    // 例如，對於上面的字節數組，輸出將是 [ ['l', 'e'], ['e', 't'] ]
    let pairs = get_char_pairs(bytes);

    // 如果沒有字符對，則直接返回編碼後的字節
    // 這種情況通常發生在輸入文本長度爲1時
    if (!pairs) {
      return bytes.map((x) => this.encoder.get(x));
    }

    // 不斷合併最頻繁的字符對，直到不能再合併爲止
    // 這是BPE算法的核心，通過合併頻繁出現的字符對來減少文本的長度
    while (true) {
      let minPairs = new Map();

      // 找出當前最頻繁的字符對
      // 這裏使用一個映射來記錄每個字符對的頻率（或排名）
      pairs.forEach((pair) => {
        let joined_pair = pair.join(" ");
        let rank = this.bpe_ranks.get(joined_pair);
        minPairs.set(rank === undefined || isNaN(rank) ? 1e11 : rank, pair);
      });

      // 獲取在 bpe 文件中最先出現的字符對，後續會合並該字節對
      // 之所以先合併該字節對，是因爲後續的合併依賴於前面的合併結果
      let minPairsKeys = Array.from(minPairs.keys()).map((x) => Number(x));
      let bigram = minPairs.get(Math.min(...minPairsKeys));

      // 如果沒有更多字符對可以合併，跳出循環
      if (!bigram || !this.bpe_ranks.has(bigram.join(" "))) {
        break;
      }

      let first = bigram[0];
      let second = bigram[1];
      let new_bytes = [];
      let i = 0;

      // 合併字符對
      // 這個循環遍歷字節數組，尋找併合並指定的字符對
      while (i < bytes.length) {

        let j = bytes.indexOf(first, i);
        if (j === -1) {
          // 如果找不到字符對中的第一個字符，則將剩餘的所有字符添加到新的字節數組中
          this.mutatingConcat(new_bytes, bytes.slice(i));
          break;
        }

        // 將當前位置到找到的位置之間的字符添加到新的字節數組中
        this.mutatingConcat(new_bytes, bytes.slice(i, j));
        i = j;

        if (
          bytes[i] === first &&
          i < bytes.length - 1 &&
          bytes[i + 1] === second
        ) {
          // 如果找到了字符對，則將它們合併爲一個字符，並添加到新的字節數組中
          new_bytes.push(first + second);
          i += 2;
        } else {
          // 如果沒有找到字符對，則只添加當前字符
          new_bytes.push(bytes[i]);
          i += 1;
        }
      }

      // 更新字節數組爲合併後的結果，以便進行下一輪合併
      bytes = new_bytes;
      // 如果字節數組長度爲1，則停止合併
      if (bytes.length === 1) {
        break;
      }

      // 重新獲取字符對，以便進行下一輪合併
      pairs = get_char_pairs(bytes);
    }
    // 將合併後的字節數組轉換爲tokens
    // 這裏的轉換是基於一個預定義的編碼器，將每個字節（或字節組合）映射到一個特定的token
    let tokens = bytes.map((x) => this.encoder.get(x));

    // 緩存結果並返回
    // 這樣，相同的輸入在下次處理時可以直接從緩存中獲取結果，提高效率
    this.cache.set(chunk, tokens);
    return tokens;
  }

以上就是實現一個分詞器的具體過程。

總結

本文詳細介紹瞭如何實現一個分詞器，並探討了其在自然語言處理和代碼補全中的應用。通過理解 BPE 算法的原理和實現過程，我們不僅能夠創建自定義的分詞器，還能更好地適配和優化大語言模型的使用。本文提供的分詞器實現方案不僅適用於代碼補全工具，還可以擴展到其他需要文本處理的領域。通過掌握這些技術，我們可以提升模型的準確性和效率，爲開發更智能的應用打下堅實的基礎。

參考鏈接

• OpenAI Cookbook，如何使用 Tiktoken 計算 token 數量：https://cookbook.openai.com/examples/how_to_count_tokens_with_tiktoken

• 字節對編碼：https://towardsdatascience.com/byte-pair-encoding-subword-based-tokenization-algorithm-77828a70bee0

• Tokenizer 功能預覽：https://platform.openai.com/tokenizer

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/OCgkOaJAgHymNovlKYQmGg