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

目录

收集或制造数据集

下载并解压

安装python依赖

编写训练脚本

执行训练脚本开始训练

编写推理脚本

启动推理脚本

测试

注意


我们之前讲过了如何部署一个别人已经训练好的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悟道数据集。

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

评论 4
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

小生浩浩

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值