代码补全快餐教程(4) - 训练语言模型

代码补全快餐教程(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代码的效果更好。
如果要支持其他语言,我们将训练集换成其他语言就可以了。

上一篇:CSS垂直居中总结


下一篇:云服务器 ECS 建站教程:建站教程汇总