当前位置：首页 > article >正文

大模型微调 - 基于预训练大语言模型的对话生成任务训练代码

article 2024/11/15 13:45:23

大模型微调 - 基于预训练大语言模型的对话生成任务训练代码

flyfish

模型扮演堂吉诃德这个角色，回答关于自我介绍的问题

import torch
from datasets import Dataset
from modelscope import AutoTokenizer, AutoModelForCausalLM
from peft import LoraConfig, TaskType, get_peft_model
from transformers import TrainingArguments, Trainer, DataCollatorForSeq2Seq


# 内存中的自定义数据集，包含用户问题（content）和模型回答（summary）
data = [
    {"content": "你叫什么名字？", "summary": "你好，我是堂吉诃德，骑士所做的一切，都是为了你的荣耀。"},
    {"content": "你可以告诉我你的名字吗？", "summary": "你好，我是堂吉诃德，骑士所做的一切，都是为了你的荣耀。"},
    {"content": "你是悟空吗", "summary": "你好，我是堂吉诃德，骑士所做的一切，都是为了你的荣耀。"},
    # 省略若干重复条目……
]

# 加载 tokenizer 和模型
# tokenizer 用于将文本转化为模型可理解的输入格式
tokenizer = AutoTokenizer.from_pretrained("qwen/Qwen2-0.5B-Instruct", use_fast=False, trust_remote_code=True)
# 加载 Qwen2-0.5B-Instruct 模型，并指定使用 bfloat16 精度和自动分配设备（GPU）
model = AutoModelForCausalLM.from_pretrained("qwen/Qwen2-0.5B-Instruct", device_map="auto", torch_dtype=torch.bfloat16)
model.enable_input_require_grads()  # 确保在大模型中开启梯度检查点，用于节省内存

# 数据预处理函数，用于将输入数据格式化为模型可以使用的输入格式
def process_func(example):
    MAX_LENGTH = 384  # 定义输入数据的最大长度
    # 拼接系统提示（堂吉诃德）和用户的提问作为模型的输入
    instruction = tokenizer(
        f"<|im_start|>system\n你是堂吉诃德，请回答以下问题。<|im_end|>\n<|im_start|>user\n{example['content']}<|im_end|>\n<|im_start|>assistant\n",
        add_special_tokens=False,
    )
    # 将回答（summary）也转化为 token
    response = tokenizer(f"{example['summary']}", add_special_tokens=False)
    
    # 将指令和回答拼接成输入序列，并在末尾添加 pad_token
    input_ids = instruction["input_ids"] + response["input_ids"] + [tokenizer.pad_token_id]
    attention_mask = instruction["attention_mask"] + response["attention_mask"] + [1]
    # 标签只包含回答部分，指令部分使用 -100 以防止计算损失
    labels = [-100] * len(instruction["input_ids"]) + response["input_ids"] + [tokenizer.pad_token_id]
    
    # 如果输入序列长度超过 MAX_LENGTH，则截断
    if len(input_ids) > MAX_LENGTH:
        input_ids = input_ids[:MAX_LENGTH]
        attention_mask = attention_mask[:MAX_LENGTH]
        labels = labels[:MAX_LENGTH]
        
    return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}

# 使用 Hugging Face 的 Dataset API 将数据转化为 Dataset 格式，并进行预处理
dataset = Dataset.from_dict({"content": [d["content"] for d in data], "summary": [d["summary"] for d in data]})
processed_dataset = dataset.map(process_func, remove_columns=["content", "summary"])

# LoRA 配置，用于微调模型
config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,  # 任务类型为因果语言模型
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],  # 指定 LoRA 作用的模块
    inference_mode=False,  # 设置为 False 表示用于训练而非推理
    r=8,  # LoRA 的秩（rank）
    lora_alpha=32,  # LoRA 的 alpha 参数
    lora_dropout=0.1,  # LoRA 使用的 dropout 比例
)

# 将 LoRA 配置应用到模型中，生成可微调的模型
model = get_peft_model(model, config)

# 训练参数配置
training_args = TrainingArguments(
    output_dir="./output/Qwen2-0.5B",  # 训练输出路径
    per_device_train_batch_size=4,  # 每个设备的训练批次大小
    gradient_accumulation_steps=4,  # 梯度累计步数
    logging_steps=10,  # 日志记录的间隔步数
    num_train_epochs=2,  # 训练的 epoch 数
    save_steps=100,  # 保存模型的间隔步数
    learning_rate=1e-4,  # 学习率
    gradient_checkpointing=True,  # 启用梯度检查点以节省显存
    report_to="none"  # 禁用报告工具（如 WandB、Tensorboard）
)

# 使用 Hugging Face 的 Trainer API 设置训练器
trainer = Trainer(
    model=model,  # 要训练的模型
    args=training_args,  # 训练参数
    train_dataset=processed_dataset,  # 训练数据集
    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True),  # 数据整理器，用于批次数据填充
)

# 开始训练
trainer.train()

# 定义预测函数，用于测试模型的生成能力
def predict(messages, model, tokenizer):
    device = "cuda"  # 使用 CUDA 加速
    # 将对话模板应用到输入的消息中
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    # 将消息编码为模型输入
    model_inputs = tokenizer([text], return_tensors="pt").to(device)

    # 使用模型生成回答，最多生成 512 个新 tokens
    generated_ids = model.generate(model_inputs.input_ids, max_new_tokens=512)
    # 获取生成的回答，并忽略输入部分的 tokens
    generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)]
    
    # 解码生成的 tokens 为文本
    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    
    return response

# 使用第一个测试数据进行预测
test_message = [
    {"role": "system", "content": "你是堂吉诃德，请回答以下问题。"},
    {"role": "user", "content": "你叫什么名字？"}
]
response = predict(test_message, model, tokenizer)
print(response)  # 输出模型的预测结果