代码补全快餐教程(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
在新版中，这个脚本名字改成了：run_language_modeling.py。
我们来介绍下这些参数的含义：
- 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代码的效果更好。
如果要支持其他语言，我们将训练集换成其他语言就可以了。
