使用数据预训练一个AI语言模型

我们之前讲过了如何部署一个别人已经训练好的AI模型、也学会了如何微调一个AI模型，也讲了预训练模型和微调模型的区别，那本文就聊聊如何从零训练一个语言模型吧！

收集或制造数据集

在机器学习中，数据集的收集是非常重要的一步，质量高或者相关性高的数据集对模型的训练有非常大的帮助。

如下两份数据集可供使用。

（1）15G gpt2数据集

链接：百度网盘 请输入提取码 提取码：khiq

（2）200G悟道数据集

链接：百度网盘 请输入提取码 提取码：sp3n

下载并解压

我使用的是15G gpt2数据集进行的预训练

安装python依赖

pip install gradio transformers datasets -i https://pypi.mirrors.ustc.edu.cn/simple/

编写训练脚本

from datasets import load_dataset, DatasetDict
from glob import glob
import random
random.seed(42)

print("开始加载包~")

all_file_list = glob(pathname="gpt2_data/*/**")
test_file_list = random.sample(all_file_list, 50)
train_file_list = [i for i in all_file_list if i not in test_file_list]

len(train_file_list), len(test_file_list)

print(len(train_file_list))
print(len(test_file_list))
print("加载包 ok")

print("开始解析和提取数据~")
raw_datasets =load_dataset("csv",data_files={'train':train_file_list,'valid':test_file_list}, cache_dir="cache_data")

print("解析和提取数据 ok")
raw_datasets

from transformers import AutoTokenizer, AutoConfig

context_length = 512
# 下载分词工具
print("下载分词工具~")
tokenizer = AutoTokenizer.from_pretrained("bert-base-chinese")
print("下载分词工具 ok")

print("对原始数据集进行分词和编码~")
# 将原始文本转换为模型可以理解的编码表示
outputs = tokenizer(
    raw_datasets["train"][:2]["content"],
    truncation=True,
    max_length=context_length,
    return_overflowing_tokens=True,
    return_length=True,
)
print("对原始数据集进行分词和编码 ok")

print(f"Input IDs length: {len(outputs['input_ids'])}")
print(f"Input chunk lengths: {(outputs['length'])}")
print(f"Chunk mapping: {outputs['overflow_to_sample_mapping']}")

# 为 tokenizer 添加特殊标记
# <|endoftext|> 是模型在训练期间看到的文档分隔符
tokenizer.add_special_tokens(special_tokens_dict={'bos_token': '<|endoftext|>',
 'eos_token': '<|endoftext|>',
 'unk_token': '<|endoftext|>'})

def tokenize(element):
    outputs = tokenizer(
        element["content"],
        truncation=True,
        max_length=context_length,
        return_overflowing_tokens=True,
        return_length=True,
    )
    input_batch = []
    for length, input_ids in zip(outputs["length"], outputs["input_ids"]):
        if length == context_length:
            input_batch.append(input_ids)
    return {"input_ids": input_batch}

# 对原始数据集进行分词和编码，并生成经过处理的分词后的数据集
tokenized_datasets = raw_datasets.map(
    tokenize, batched=True, remove_columns=raw_datasets["train"].column_names
)
tokenized_datasets

from transformers import AutoTokenizer, GPT2LMHeadModel, AutoConfig

# 创建和配置 GPT-2 模型
print("创建和配置 GPT-2 模型~")
config = AutoConfig.from_pretrained(
    "gpt2",
    vocab_size=len(tokenizer),
    n_ctx=context_length,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id,
)

model = GPT2LMHeadModel(config)
model_size = sum(t.numel() for t in model.parameters())
print(f"GPT-2 size: {model_size/1000**2:.1f}M parameters")

from transformers import DataCollatorForLanguageModeling

# 为语言建模任务准备数据，并使用数据处理器对数据进行编码、填充和组织，生成适用于训练的批次数据
print("为语言建模任务准备数据，并使用数据处理器对数据进行编码、填充和组织，生成适用于训练的批次数据~")
tokenizer.pad_token = tokenizer.eos_token
data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)

out = data_collator([tokenized_datasets["train"][i] for i in range(5)])
for key in out:
    print(f"{key} shape: {out[key].shape}")

print("为语言建模任务准备数据，并使用数据处理器对数据进行编码、填充和组织，生成适用于训练的批次数据 ok")

from transformers import Trainer, TrainingArguments

print("创建并执行模型训练~")
args = TrainingArguments(
    # 指定训练完成后保存模型的目录
    output_dir="chinese_gpt2_big",
    # 每个设备（GPU或CPU）上用于训练的批次大小
    per_device_train_batch_size=20,
    # 每个设备上用于评估的批次大小
    per_device_eval_batch_size=16,
    # 指定何时在训练过程中进行评估。这里设置为 "steps"，表示将在指定的 eval_steps 处进行评估
    evaluation_strategy="steps",
    # 每次评估之间的训练步数
    eval_steps=2_000,
    # 每隔多少训练步数记录一次训练指标
    logging_steps=2_000,
    # 在执行反向传播/更新步骤之前累积梯度的步骤数
    gradient_accumulation_steps=8,
    # 遍历整个训练数据集的次数
    num_train_epochs=2,
    # 权重衰减正则化的系数
    weight_decay=0.1,
    # 学习率调度的预热步数
    warmup_steps=1_000,
    # 学习率调度器的类型。这里设置为 "cosine"
    lr_scheduler_type="cosine",
    # 优化器的初始学习率
    learning_rate=5e-4,
    # 每隔多少训练步数保存一次检查点
    save_steps=2_000,
    # 是否使用16位浮点精度（混合精度训练），以加快训练速度并减少内存使用
    fp16=True,
    # 是否将训练完成的模型推送到模型中心
    push_to_hub=False,
)

trainer = Trainer(
    # 要训练的模型
    model=model,
    # 用于处理输入数据的 tokenizer
    tokenizer=tokenizer,
    # 训练的参数配置，包括输出目录、批次大小、训练轮数等
    args=args,
    # 用于处理数据的数据处理器
    data_collator=data_collator,
    # 训练数据集，即经过处理的分词后的训练数据
    train_dataset=tokenized_datasets["train"],
    # 验证数据集，即经过处理的分词后的验证数据
    eval_dataset=tokenized_datasets["valid"],
)

trainer.train()

print("创建并执行模型训练 ok")

执行训练脚本开始训练

python3 trainer.py

注意，训练底座大模型需要大量的计算资源和时间，所以一定要保证自己的硬件足够，我使用的是一张3090显卡，24G显存，训练15G数据集需要60个小时

60小时后，训练完成，如下：

编写推理脚本

vi chat.py

import gradio as gr
from transformers import GPT2Tokenizer, GPT2LMHeadModel, AutoTokenizer

model_name_or_path = "checkpoint-36000"
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)

model = GPT2LMHeadModel.from_pretrained(model_name_or_path, pad_token_id=tokenizer.eos_token_id)

def infer_model(txt: str) -> str:
    # encode context the generation is conditioned on
    input_ids = tokenizer.encode(txt, return_tensors='pt')
    # set no_repeat_ngram_size to 2
    beam_output = model.generate(
        input_ids,
        max_length=200,
        num_beams=5,
        no_repeat_ngram_size=2,
        early_stopping=True
    )

    result = tokenizer.decode(beam_output[0], skip_special_tokens=True)
    result = result.replace(" ", "")
    return result

def predict(input, history=[]):
    respose = infer_model(input)
    res = [(input, respose)]
    return res, history

def add_text(state, text):
    res = infer_model(text)
    state = state + [(text, res)]
    return state, state

with gr.Blocks(css="#chatbot .overflow-y-auto{height:500px}") as demo:
    chatbot = gr.Chatbot(elem_id="chatbot")
    state = gr.State([])

    with gr.Row():
        txt = gr.Textbox(show_label=False, placeholder="输入文本").style(container=False)

    txt.submit(add_text, [state, txt], [state, chatbot])

if __name__ == "__main__":
    demo.launch(server_name='0.0.0.0',server_port=6006)

启动推理脚本

python chat.py

测试

注意

由于15G gpt2数据量太少，所以训练完成的AI模型对话效果并不理想，如想训练效果更好的模型，请使用200G悟道数据集。

以上便是本次训练的全部过程啦~~~