代碼補(bǔ)全快餐教程(4) - 訓(xùn)練語言模型
一個強(qiáng)大的語言模型可以是其他任務(wù)的良好基礎(chǔ)。預(yù)訓(xùn)練的模型就為我們提供了一個強(qiáng)大的語言模型基礎(chǔ)，在些基礎(chǔ)上，我們進(jìn)行微調(diào)，就可以實現(xiàn)滿足特殊需求的模型。
我們先做實操，然后再講解相關(guān)理論。
代碼數(shù)據(jù)準(zhǔn)備
嚴(yán)格來講，進(jìn)行代碼數(shù)據(jù)準(zhǔn)備需要做代碼的排重，后面講到相關(guān)論文時我們會講到。
現(xiàn)在我們就用個最簡單的辦法，將代碼先拼接在一起。
我們寫個小腳本，將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
進(jìn)行訓(xùn)練之前，我們先安裝下transformer庫，先cd到transformers的下載目錄，然后執(zhí)行
pip3 install -e . --user
安裝成功之后，我們就可以使用transformers下的examples中的run_lm_finetuning.py腳本來進(jìn)行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
在新版中，這個腳本名字改成了：run_language_modeling.py。
我們來介紹下這些參數(shù)的含義：
- output_dir: 最終我們要保存的是權(quán)值，這里給出保存權(quán)值的目錄
- model_type: 模型的大類，比如gpt2或者其他
- model_name_or_path: 模型的小類，比如gpt2-medium, gpt2-large, gpt2-xl等
- per_gpu_train_batch_size: 多CPU訓(xùn)練時每個CPU批次的大小
- do_train: 只有指定了這個才會進(jìn)行訓(xùn)練
- train_data_file: 要訓(xùn)練的文件名
- block_size: 分塊的大小，如果GPU內(nèi)存大就多選點，我用的是NVidia 2060 GPU，內(nèi)存較小，所以我選了個相對較小的值
- save_steps: 訓(xùn)練多少步保存一次，默認(rèn)值是50，我覺得有點小，這里改成500
- overwrite_output_dir: 輸出目錄不為空時覆蓋之，節(jié)省存儲空間
驗證效果
我們做個補(bǔ)全效果測試吧，還是我們之前的代碼，我們先用gpt2試試效果：
import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel
# MODEL = '/home/xulun/out_trans/'
MODEL = 'gpt2'
# 加載詞匯表
tokenizer = GPT2Tokenizer.from_pretrained(MODEL)
# 輸入待補(bǔ)全的文本
text = ' indexed_tokens = tokenizer.'
predicted_text = text
# 加載模型中預(yù)訓(xùn)練好的權(quán)值
model = GPT2LMHeadModel.from_pretrained(MODEL)
# 設(shè)置為eval模式，這樣就不會執(zhí)行訓(xùn)練模式下的Dropout過程
model.eval()
#model.to('cuda')
# 每一個只能補(bǔ)一個token出來，補(bǔ)一句話需要多次，30次是我拍腦袋的
for i in range(0,30):
# 以上次預(yù)測結(jié)果作為本次的輸入，所謂的自回歸
indexed_tokens = tokenizer.encode(predicted_text)
# 將讀出的索引標(biāo)記轉(zhuǎn)化成PyTorch向量
tokens_tensor = torch.tensor([indexed_tokens])
# 使用GPU進(jìn)行加速，誠實地講速度不太快
#tokens_tensor = tokens_tensor.to('cuda')
# 進(jìn)行推理
with torch.no_grad():
outputs = model(tokens_tensor)
predictions = outputs[0]
# 獲取預(yù)測的下一個子詞
predicted_index = torch.argmax(predictions[0, -1, :]).item()
# 解碼成我們都讀懂的文本
predicted_text = tokenizer.decode(indexed_tokens + [predicted_index])
# 打印輸入結(jié)果
print(predicted_text)
輸出如下：
indexed_tokens = tokenizer.get_tokenizer_id(tokenizer.get_tokenizer_id(), tokenizer.get_tokenizer_id(), tokenizer.
下面我們換成我們剛才訓(xùn)練的模型，就是讓MODEL從gpt2換成剛才我們訓(xùn)練好的目錄：
MODEL = '/home/xulun/out_trans/'
好吧，有同學(xué)要完整的：
import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel
MODEL = '/home/xulun/out_trans/'
# MODEL = 'gpt2'
# 加載詞匯表
tokenizer = GPT2Tokenizer.from_pretrained(MODEL)
# 輸入待補(bǔ)全的文本
#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
# 加載模型中預(yù)訓(xùn)練好的權(quán)值
model = GPT2LMHeadModel.from_pretrained(MODEL)
# 設(shè)置為eval模式，這樣就不會執(zhí)行訓(xùn)練模式下的Dropout過程
model.eval()
#model.to('cuda')
# 每一個只能補(bǔ)一個token出來，補(bǔ)一句話需要多次，30次是我拍腦袋的
for i in range(0,30):
# 以上次預(yù)測結(jié)果作為本次的輸入，所謂的自回歸
indexed_tokens = tokenizer.encode(predicted_text)
# 將讀出的索引標(biāo)記轉(zhuǎn)化成PyTorch向量
tokens_tensor = torch.tensor([indexed_tokens])
# 使用GPU進(jìn)行加速，誠實地講速度不太快
#tokens_tensor = tokens_tensor.to('cuda')
# 進(jìn)行推理
with torch.no_grad():
outputs = model(tokens_tensor)
predictions = outputs[0]
# 獲取預(yù)測的下一個子詞
predicted_index = torch.argmax(predictions[0, -1, :]).item()
# 解碼成我們都讀懂的文本
predicted_text = tokenizer.decode(indexed_tokens + [predicted_index])
# 打印輸入結(jié)果
print(predicted_text)
輸出結(jié)果如下：
indexed_tokens = tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)
看起來是比原始模型更懂transformers。我們可以用更多的代碼進(jìn)行訓(xùn)練，這樣就能對于寫python代碼的效果更好。
如果要支持其他語言，我們將訓(xùn)練集換成其他語言就可以了。
