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

模型"/>

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

目录

收集或制造数据集

下载并解压

安装python依赖

编写训练脚本

执行训练脚本开始训练

编写推理脚本

启动推理脚本

测试

注意

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

收集或制造数据集

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

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

（1）15G gpt2数据集

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

（2）200G悟道数据集

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

下载并解压

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

安装python依赖

pip install gradio transformers datasets -i /

编写训练脚本

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_datasetsfrom transformers import AutoTokenizer, AutoConfigcontext_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_datasetsfrom 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, TrainingArgumentsprint("创建并执行模型训练~")
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,# 用于处理输入数据的 tokenizertokenizer=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, AutoTokenizermodel_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 oninput_ids = tokenizer.encode(txt, return_tensors='pt')# set no_repeat_ngram_size to 2beam_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 resultdef predict(input, history=[]):respose = infer_model(input)res = [(input, respose)]return res, historydef add_text(state, text):res = infer_model(text)state = state + [(text, res)]return state, statewith 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悟道数据集。

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