代碼補全快餐教程(4) - 訓練語言模型

一個強大的語言模型可以是其他任務的良好基礎。預訓練的模型就為我們提供了一個強大的語言模型基礎，在些基礎上，我們進行微調，就可以實作滿足特殊需求的模型。

我們先做實操，然後再講解相關理論。

代碼資料準備

嚴格來講，進行代碼資料準備需要做代碼的排重，後面講到相關論文時我們會講到。

現在我們就用個最簡單的辦法，将代碼先拼接在一起。

我們寫個小腳本，将transformer庫中的python檔案都讀出來連接配接在一起：

import os


def walkPrograms(dir, datafile, wildcard):
    exts = wildcard.split(" ")
    for root, subdirs, files in os.walk(dir):
        for name in files:
            for ext in exts:
                if name.endswith(ext):
                    print(root)
                    # print(subdirs)
                    print(name)
                    filename = os.path.join(root, name)
                    print(filename)
                    try:
                        f1 = open(filename, 'r', encoding='utf-8')
                        datafile.writelines(f1.readlines())
                    except UnicodeDecodeError:
                        continue
                    break


outfile = open('transformer.data', 'w', encoding='utf-8')
wildcard = '.py'
walkPrograms('/home/xulun/github/transformers/', outfile, wildcard)           

最後會生成一個transformer.data檔案，其中是python檔案的組合。

語言模型fine-tuning

進行訓練之前，我們先安裝下transformer庫，先cd到transformers的下載下傳目錄，然後執行

pip3 install -e . --user           

安裝成功之後，我們就可以使用transformers下的examples中的run_lm_finetuning.py腳本來進行fine-tuning:

python3 run_lm_finetuning.py \
    --output_dir=/home/xulun/out_trans \
    --model_type=gpt2 \
    --model_name_or_path=gpt2 \
    --per_gpu_train_batch_size=1 \
    --do_train \
    --train_data_file=/home/xulun/github/lusinga/localcomplete/server/transformer.data \
    --block_size=512 --save_steps=500 --overwrite_output_dir           

我們來介紹下這些參數的含義：

• output_dir: 最終我們要儲存的是權值，這裡給出儲存權值的目錄
• model_type: 模型的大類，比如gpt2或者其他
• model_name_or_path: 模型的小類，比如gpt2-medium, gpt2-large, gpt2-xl等
• per_gpu_train_batch_size: 多CPU訓練時每個CPU批次的大小
• do_train: 隻有指定了這個才會進行訓練
• train_data_file: 要訓練的檔案名
• block_size: 分塊的大小，如果GPU記憶體大就多選點，我用的是NVidia 2060 GPU，記憶體較小，是以我選了個相對較小的值
• save_steps: 訓練多少步儲存一次，預設值是50，我覺得有點小，這裡改成500
• overwrite_output_dir: 輸出目錄不為空時覆寫之，節省存儲空間

驗證效果

我們做個補全效果測試吧，還是我們之前的代碼，我們先用gpt2試試效果：

import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel

# MODEL = '/home/xulun/out_trans/'
MODEL = 'gpt2'

# 加載詞彙表
tokenizer = GPT2Tokenizer.from_pretrained(MODEL)

# 輸入待補全的文本
text = '    indexed_tokens = tokenizer.'
predicted_text = text

# 加載模型中預訓練好的權值
model = GPT2LMHeadModel.from_pretrained(MODEL)

# 設定為eval模式，這樣就不會執行訓練模式下的Dropout過程
model.eval()
#model.to('cuda')

# 每一個隻能補一個token出來，補一句話需要多次，30次是我拍腦袋的
for i in range(0,30):

    # 以上次預測結果作為本次的輸入，所謂的自回歸
    indexed_tokens = tokenizer.encode(predicted_text)

    # 将讀出的索引标記轉化成PyTorch向量
    tokens_tensor = torch.tensor([indexed_tokens])

    # 使用GPU進行加速，誠實地講速度不太快
    #tokens_tensor = tokens_tensor.to('cuda')

    # 進行推理
    with torch.no_grad():
        outputs = model(tokens_tensor)
        predictions = outputs[0]

    # 擷取預測的下一個子詞
    predicted_index = torch.argmax(predictions[0, -1, :]).item()
    # 解碼成我們都讀懂的文本
    predicted_text = tokenizer.decode(indexed_tokens + [predicted_index])
    # 列印輸入結果
    print(predicted_text)           

輸出如下：

indexed_tokens = tokenizer.get_tokenizer_id(tokenizer.get_tokenizer_id(), tokenizer.get_tokenizer_id(), tokenizer.           

下面我們換成我們剛才訓練的模型，就是讓MODEL從gpt2換成剛才我們訓練好的目錄：

MODEL = '/home/xulun/out_trans/'           

好吧，有同學要完整的：

import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel

MODEL = '/home/xulun/out_trans/'
# MODEL = 'gpt2'

# 加載詞彙表
tokenizer = GPT2Tokenizer.from_pretrained(MODEL)

# 輸入待補全的文本
#text = 'function walk(dir, fn) { if (fs.existsSync(dir)) { let stat ='
#text = 'if (stat.isDirectory()) {fs.readdirSync(dir).'
#text = 'mediaFileText.color ='
#text = 'mediaFileText.top ='
text = '    indexed_tokens = tokenizer.'
predicted_text = text

# 加載模型中預訓練好的權值
model = GPT2LMHeadModel.from_pretrained(MODEL)

# 設定為eval模式，這樣就不會執行訓練模式下的Dropout過程
model.eval()
#model.to('cuda')

# 每一個隻能補一個token出來，補一句話需要多次，30次是我拍腦袋的
for i in range(0,30):

    # 以上次預測結果作為本次的輸入，所謂的自回歸
    indexed_tokens = tokenizer.encode(predicted_text)

    # 将讀出的索引标記轉化成PyTorch向量
    tokens_tensor = torch.tensor([indexed_tokens])

    # 使用GPU進行加速，誠實地講速度不太快
    #tokens_tensor = tokens_tensor.to('cuda')

    # 進行推理
    with torch.no_grad():
        outputs = model(tokens_tensor)
        predictions = outputs[0]

    # 擷取預測的下一個子詞
    predicted_index = torch.argmax(predictions[0, -1, :]).item()
    # 解碼成我們都讀懂的文本
    predicted_text = tokenizer.decode(indexed_tokens + [predicted_index])
    # 列印輸入結果
    print(predicted_text)           

輸出結果如下：

indexed_tokens = tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)           

看起來是比原始模型更懂transformers。我們可以用更多的代碼進行訓練，這樣就能對于寫python代碼的效果更好。

如果要支援其他語言，我們将訓練集換成其他語言就可以了。