程式碼Tokenizer訓練與模型建構

– 9 分鐘閱讀

預訓練的Tokenizer通常針對自然語言文字進行最佳化，對於程式碼中的特殊符號、縮排和關鍵字處理能力不足。為提升程式碼相關任務的效能，本文示範如何使用Python和Transformers函式庫訓練客製化Tokenizer，並以GPT-2模型為例，建構程式碼自動完成模型。首先，從程式碼資料集中選取部分資料訓練Tokenizer，並調整詞彙大小以平衡效能和資源消耗。接著，透過分析詞彙表，驗證Tokenizer對程式碼元素的理解能力。後續使用簡單的Python程式碼測試Tokenizer的效能，並比較不同詞彙量和資料集大小對結果的影響。模型建構部分，則示範如何初始化GPT-2模型，並調整組態以適應新的Tokenizer。為了提升訓練效率，設計了客製化的資料載入器，將多個程式碼片段連線成固定長度的序列，並以特殊符號區隔。最後，文章也說明如何將訓練好的Tokenizer上傳至Hugging Face Hub，方便社群分享和重複使用。

建立適合程式碼的Tokenizer

在訓練Transformer模型處理程式碼任務時，選擇適當的Tokenizer至關重要。預設的Tokenizer可能無法有效地處理程式碼的結構和語法，因此需要建立一個新的Tokenizer。

為什麼需要自定義Tokenizer

預設的Tokenizer，如GPT-2使用的Tokenizer，是針對自然語言文字進行最佳化的，可能無法有效地處理程式碼中的特殊字元、縮排和關鍵字。因此，需要建立一個新的Tokenizer，使其能夠更好地理解和處理程式碼。

訓練新的Tokenizer

首先，需要從資料集中選取一部分資料來訓練新的Tokenizer。在這個例子中，選取了約1-2 GB的資料，大約10萬個檔案。

from tqdm.auto import tqdm

length = 100000
dataset_name = 'transformersbook/codeparrot-train'
dataset = load_dataset(dataset_name, split="train", streaming=True)
iter_dataset = iter(dataset)

def batch_iterator(batch_size=10):
    for _ in tqdm(range(0, length, batch_size)):
        yield [next(iter_dataset)['content'] for _ in range(batch_size)]

new_tokenizer = tokenizer.train_new_from_iterator(batch_iterator(),
                                                  vocab_size=12500,
                                                  initial_alphabet=base_vocab)

檢視新的Tokenizer的詞彙

訓練完成後，需要檢視新的Tokenizer的詞彙，以確保它能夠有效地處理程式碼。

tokens = sorted(new_tokenizer.vocab.items(), key=lambda x: x[1], reverse=False)
print([f'{new_tokenizer.convert_tokens_to_string(t)}' for t, _ in tokens[257:280]])

輸出結果顯示，新的Tokenizer能夠有效地處理程式碼中的縮排、空格和常見的Python關鍵字。

檢視不常見的Token

print([f'{new_tokenizer.convert_tokens_to_string(t)}' for t, _ in tokens[-12:]])

輸出結果顯示，一些不常見的Token可能是來自註解或特定的程式碼。

測試新的Tokenizer

使用一個簡單的Python程式碼範例來測試新的Tokenizer。

python_code = "def say_hello():\n    print('Hello, World!')\n    # Print it\nsay_hello()"
print(new_tokenizer(python_code).tokens())

輸出結果顯示，新的Tokenizer能夠有效地處理程式碼中的關鍵字、縮排和字串。

比較不同大小的詞彙

嘗試使用更大的詞彙量（32,768）和更多的資料（20萬個檔案）來訓練新的Tokenizer。

length = 200000
new_tokenizer_larger = tokenizer.train_new_from_iterator(batch_iterator(),
                                                         vocab_size=32768,
                                                         initial_alphabet=base_vocab)

檢視更大的詞彙

tokens = sorted(new_tokenizer_larger.vocab.items(), key=lambda x: x[1], reverse=False)
print([f'{new_tokenizer_larger.convert_tokens_to_string(t)}' for t, _ in tokens[-12:]])

輸出結果顯示，更大的詞彙量能夠包含更多的程式碼相關的Token。

測試更大的Tokenizer

print(new_tokenizer_larger(python_code).tokens())

輸出結果顯示，更大的Tokenizer能夠更有效地處理程式碼中的關鍵字和字串。

內容解密：
  1. 本段主要介紹瞭如何為Transformer模型建立一個適合處理程式碼任務的Tokenizer。
  2. 說明瞭預設Tokenizer的侷限性，以及為什麼需要自定義Tokenizer。
  3. 詳細描述瞭如何訓練新的Tokenizer，包括資料的選取、訓練過程和詞彙的檢視。
  4. 使用簡單的Python程式碼範例測試了新的Tokenizer，驗證了其有效性。
  5. 比較了不同大小的詞彙對Tokenizer效能的影響，並選擇了最適合的Tokenizer。

未來的工作可以進一步最佳化Tokenizer，例如嘗試不同的詞彙量、不同的訓練資料集，或者結合多種程式語言的資料來訓練更通用的Tokenizer。同時，也可以評估不同的Tokenizer對模型效能的影響，以選擇最優的解決方案。

內容解密：

  1. 未來的工作將著重於進一步最佳化Tokenizer，以提高其處理不同程式語言和任務的能力。
  2. 將嘗試不同的詞彙量和訓練資料集，以評估其對模型效能的影響。
  3. 結合多種程式語言的資料來訓練更通用的Tokenizer，以提高模型的通用性和準確性。

在Hugging Face Hub上儲存自定義的Tokenizer

既然我們的tokenizer已經訓練完畢，我們應該儲存它。最簡單的儲存方式是將其推播到Hugging Face Hub，以便稍後可以從任何地方存取它。當我們使用單獨的訓練伺服器時，這將特別有用。

要建立一個私有模型儲存函式庫並將我們的tokenizer儲存為其中的第一個檔案，我們可以直接使用tokenizer的push_to_hub()方法。由於我們已經使用huggingface-cli login驗證了我們的帳戶，我們可以簡單地推播tokenizer，如下所示：

model_ckpt = "codeparrot"
org = "transformersbook"
new_tokenizer_larger.push_to_hub(model_ckpt, organization=org)

如果我們不想推播到組織，可以簡單地省略organization引數。這將在我們的名稱空間中建立一個名為codeparrot的儲存函式庫，任何人都可以使用以下命令載入它：

reloaded_tokenizer = AutoTokenizer.from_pretrained(org + "/" + model_ckpt)
print(reloaded_tokenizer(python_code).tokens())

輸出結果將與我們剛才看到的完全相同。我們還可以在Hub上調查其檔案和儲存的詞彙表。為了重現性，讓我們也儲存我們的較小的tokenizer：

new_tokenizer.push_to_hub(model_ckpt + "-small-vocabulary", organization=org)

這是為特定使用案例構建tokenizer的探討。接下來，我們將建立一個新的模型並從頭開始訓練它。

從頭開始訓練模型

這可能是您一直在等待的部分：模型訓練。在本文中，我們將決定哪種架構最適合該任務，初始化一個沒有預訓練權重的全新模型，設定自定義資料載入類別，並建立一個可擴充套件的訓練迴圈。在最後，我們將訓練具有1.11億和15億引數的小型和大型GPT-2模型！但讓我們不要過於超前。首先，我們需要決定哪種架構最適合程式碼自動補全任務。

在本文中，我們將實作一個比平常更長的指令碼來在分散式基礎設施上訓練模型。因此，您不應該獨立執行每個程式碼片段，而應該下載Transformers儲存函式庫中提供的指令碼。按照隨附的說明使用Accelerate在您的硬體上執行指令碼。

預訓練目標的故事

現在我們已經可以存取大規模預訓練語料函式庫和高效的tokenizer，我們可以開始思考如何預訓練轉換器模型。有了像圖10-1所示的程式碼片段這樣的大型程式碼函式庫，我們可以處理多個任務。我們選擇哪個任務將影響我們的預訓練目標的選擇。讓我們來看看三個常見的任務。

因果語言建模

文字資料的一個自然任務是為模型提供程式碼樣本的開頭，並要求它生成可能的補全。這是一個自監督的訓練目標，我們可以在沒有註解的情況下使用資料集。這應該聽起來很熟悉：這是我們在第5章中遇到的因果語言建模任務。一個直接相關的下游任務是程式碼自動補全，因此我們肯定會將此模型列入候選名單。像GPT系列模型這樣的解碼器架構通常最適合此任務，如圖10-2所示。

遮罩語言建模

一個相關但略有不同的任務是為模型提供一個嘈雜的程式碼樣本，例如用隨機或遮罩的單詞替換程式碼指令，並要求它重建原始的乾淨樣本，如圖10-3所示。這也是一個自監督的訓練目標，通常稱為遮罩語言建模或去噪目標。很難想到與去噪直接相關的下游任務，但去噪通常是一個很好的預訓練任務，可以學習一般的表示形式，以便稍後用於下游任務。我們在前幾章中使用的許多模型（如BERT和XLM-RoBERTa）都是以這種方式預訓練的。因此，在大型語料函式庫上訓練遮罩語言模型可以與在下游任務上對模型進行微調結合起來，使用有限數量的標註示例。

序列到序列訓練

一個替代任務是使用正規表示式等啟發式方法將註解或檔案字串與程式碼分開，並構建一個大規模的（程式碼，註解）配對資料集，可以用作註解資料集。然後，訓練任務是一個監督訓練目標，其中一個類別（程式碼或註解）用作模型的輸入，另一個類別（註解或程式碼）用作標籤。這是一種帶有（輸入，標籤）對的有監督學習，如圖10-4所示。使用大型、乾淨、多樣化的資料集以及具有足夠容量的模型，我們可以嘗試訓練一個模型，學習將註解轉錄為程式碼，反之亦然。與此監督訓練任務直接相關的下游任務是從程式碼生成檔案或從檔案生成程式碼，具體取決於我們如何設定輸入/輸出。在這種情況下，序列被翻譯成另一個序列，這是像T5、BART和PEGASUS這樣的編碼器-解碼器架構所擅長的。

從零開始訓練程式碼自動完成模型

在訓練一個能夠自動完成程式碼的模型時，我們將採用編碼器-解碼器（encoder-decoder）架構來處理序列到序列（sequence-to-sequence）的任務。由於我們的目標是建立一個程式碼自動完成的模型，因此我們將選擇第一個目標並採用GPT架構。

初始化模型

這是本文第一次不使用from_pretrained()方法來載入模型，而是初始化一個新的模型。我們將載入gpt2-xl的組態，以便使用相同的超引數，並僅調整新標記器的詞彙大小。然後，我們使用from_config()方法初始化一個新的模型：

from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
config = AutoConfig.from_pretrained("gpt2-xl", vocab_size=len(tokenizer))
model = AutoModelForCausalLM.from_config(config)

內容解密：

  1. 載入預訓練的標記器（tokenizer），以便將輸入文字轉換為模型可以理解的格式。
  2. 使用gpt2-xl的組態，並調整詞彙大小以匹配新的標記器。
  3. 初始化一個新的因果語言模型（Causal Language Model），該模型適用於文字生成任務。

讓我們檢查模型的實際大小：

print(f'GPT-2 (xl) size: {model_size(model)/1000**2:.1f}M parameters')

輸出結果顯示，GPT-2（xl）模型具有15.3億個引數。由於模型較大，我們還將建立一個較小的版本，以確保在擴充套件之前一切正常運作。

實作資料載入器

為了高效地訓練模型，我們需要確保輸入資料的序列長度與模型的上下文長度相匹配。我們將採用一個小技巧，即將多個示例的標記化序列連線起來，並以特殊的結束符號（EOS token）分隔，然後將這個長序列分割成相同大小的區塊，如圖10-5所示。

圖示：準備不同長度的序列進行因果語言建模

此圖示展示瞭如何透過連線多個標記化的示例，並以EOS token分隔，然後將其分割成固定大小的區塊。

首先，我們需要估計資料集中每個標記的平均字元長度：

examples, total_characters, total_tokens = 500, 0, 0
dataset = load_dataset('transformersbook/codeparrot-train', split='train', streaming=True)

for _, example in tqdm(zip(range(examples), iter(dataset)), total=examples):
    total_characters += len(example['content'])
    total_tokens += len(tokenizer(example['content']).tokens())

characters_per_token = total_characters / total_tokens
print(characters_per_token)

輸出結果顯示，每個標記的平均字元長度約為3.62。

建立可迭代資料集

我們將建立一個自定義的可迭代資料集（IterableDataset），以便準備固定長度的輸入資料供模型使用。

import torch
from torch.utils.data import IterableDataset

class ConstantLengthDataset(IterableDataset):
    def __init__(self, tokenizer, dataset, seq_length=1024, num_of_sequences=1024, chars_per_token=3.6):
        self.tokenizer = tokenizer
        self.concat_token_id = tokenizer.eos_token_id
        self.dataset = dataset
        self.seq_length = seq_length
        self.input_characters = seq_length * chars_per_token * num_of_sequences

    def __iter__(self):
        iterator = iter(self.dataset)
        more_examples = True
        while more_examples:
            buffer, buffer_len = [], 0
            while True:
                if buffer_len >= self.input_characters:
                    print(f"Buffer full: {buffer_len} >= {self.input_characters:.0f}")
                    break
                try:
                    print(f"Fill buffer: {buffer_len} < {self.input_characters:.0f}")
                    buffer.append(next(iterator)["content"])
                    buffer_len += len(buffer[-1])
                except StopIteration:
                    iterator = iter(self.dataset)
            # 將緩衝區中的內容轉換為標記ID並進行處理

內容解密：

  1. ConstantLengthDataset類別繼承自IterableDataset，並在初始化時接受標記器、資料集、序列長度等多個引數。
  2. __iter__方法定義瞭如何迭代資料集中的元素，將多個示例連線起來直到達到所需的輸入字元數。
  3. 當緩衝區滿時，將其內容轉換為標記ID並進行進一步處理，以準備固定長度的輸入資料。
玄貓 BlackCat

玄貓 BlackCat

技術愛好者，專注於分享程式開發、雲端技術與 AI 應用的心得體會。