Chapter 18
因果语言模型训练实例
Notebooktransformers27 cells
因果语言模型训练实例
Step1 导入相关包
In [ ]python · cell 3
python
from datasets import load_dataset, Dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForLanguageModeling, TrainingArguments, Trainer, BloomForCausalLMStep2 加载数据集
In [ ]python · cell 5
python
ds = Dataset.load_from_disk("./wiki_cn_filtered/")In [ ]python · cell 6
python
dsIn [ ]python · cell 7
python
ds[0]Step3 数据集处理
In [ ]python · cell 9
python
tokenizer = AutoTokenizer.from_pretrained("Langboat/bloom-389m-zh")
def process_func(examples):
contents = [e + tokenizer.eos_token for e in examples["completion"]]
return tokenizer(contents, max_length=384, truncation=True)In [ ]python · cell 10
python
tokenized_ds = ds.map(process_func, batched=True, remove_columns=ds.column_names)
tokenized_dsIn [ ]python · cell 11
python
from torch.utils.data import DataLoader
dl = DataLoader(tokenized_ds, batch_size=2, collate_fn=DataCollatorForLanguageModeling(tokenizer, mlm=False))In [ ]python · cell 12
python
next(enumerate(dl))In [ ]python · cell 13
python
tokenizer.pad_token, tokenizer.pad_token_idIn [ ]python · cell 14
python
tokenizer.eos_token, tokenizer.eos_token_idStep4 创建模型
In [ ]python · cell 16
python
model = AutoModelForCausalLM.from_pretrained("Langboat/bloom-389m-zh")Step5 配置训练参数
In [ ]python · cell 18
python
args = TrainingArguments(
output_dir="./causal_lm",
per_device_train_batch_size=2,
gradient_accumulation_steps=16,
logging_steps=10,
num_train_epochs=1,
fp16=True
)Step6 创建训练器
In [ ]python · cell 20
python
trainer = Trainer(
args=args,
model=model,
tokenizer=tokenizer,
train_dataset=tokenized_ds,
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
)Step7 模型训练
In [ ]python · cell 22
python
trainer.train()Step8 模型推理
In [ ]python · cell 24
python
from transformers import pipeline
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)In [ ]python · cell 25
python
pipe("西安交通大学博物馆（Xi'an Jiaotong University Museum）是一座位于西安", max_length=128, do_sample=True)In [ ]python · cell 26
python
pipe("下面是一则游戏新闻。小编报道，近日，游戏产业发展的非常", max_length=128, do_sample=True)In [ ]python · cell 27
python