-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathsft_main.py
66 lines (56 loc) · 1.88 KB
/
sft_main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
# main.py
from transformers import AutoTokenizer, AutoModelForCausalLM
from datasets import load_dataset
from training_args import TrainingArguments
from sft_trainer import SFTTrainer
from data_collator import DataCollator
from callbacks import LoggingCallback, EarlyStoppingCallback
def main():
# 加载预训练的分词器和模型
tokenizer = AutoTokenizer.from_pretrained('gpt2')
model = AutoModelForCausalLM.from_pretrained('gpt2')
# **添加 pad_token 并调整模型嵌入层**
if tokenizer.pad_token is None:
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
model.resize_token_embeddings(len(tokenizer))
# 加载 Alpaca 数据集
dataset = load_dataset('tatsu-lab/alpaca')
train_dataset = dataset['train']
eval_dataset = dataset['train'].train_test_split(test_size=0.1)['test']
# 定义训练参数
args = TrainingArguments(
output_dir='./sft_results',
num_train_epochs=1,
per_device_train_batch_size=2,
per_device_eval_batch_size=2,
learning_rate=5e-5,
logging_steps=10,
save_steps=50,
evaluation_strategy='steps',
eval_steps=20,
seed=42
)
# 定义回调函数列表
callbacks = [
LoggingCallback(),
EarlyStoppingCallback(patience=3)
]
# 定义数据整理器
data_collator = DataCollator(tokenizer=tokenizer, padding=True)
# 创建SFTTrainer实例
trainer = SFTTrainer(
model=model,
args=args,
data_collator=data_collator,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
tokenizer=tokenizer,
callbacks=callbacks,
max_seq_length=512,
dataset_text_field='text',
packing=False
)
# 开始训练
trainer.train()
if __name__ == '__main__':
main()