文章目录
指令微调的示例代码与预训练的代码高度一致,区别主要在于指令微调数据集的构建(SFTDataset)和序列到序列损失的计算(DataCollatorForSupervisedDataset)。以下代码展示了 LLMBox 和 YuLan-Chat 中指令微调的整体训练流程。
1 import torch
2 from dataclasses import dataclass
3 from dataset.sft_dataset import SFTDataset
4 from transformers import (
5 AutoModelForCausalLM,
6 AutoTokenizer,
7 HfArgumentParser,
8 PreTrainedTokenizer,
9 TrainingArguments,
10 Trainer,
11 )
12 from transformers.hf_argparser import HfArg
13
14 IGNORE_INDEX = -100
15
16
17 # 用户输入超参数
18 @dataclass
19 class Arguments(TrainingArguments):
20 # 模型结构
21 model_name_or_path: str = HfArg(
22 default=None,
23 help="The model name or path, e.g., `meta-llama/Llama-2-7b-hf`",
24 )
25