Chapter 19
掩码语言模型训练实例
Notebooktransformers26 cells
掩码语言模型训练实例
Step1 导入相关包
In [ ]python · cell 3
python
from datasets import load_dataset, Dataset
from transformers import AutoTokenizer, AutoModelForMaskedLM, DataCollatorForLanguageModeling, TrainingArguments, TrainerStep2 加载数据集
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("hfl/chinese-macbert-base")
def process_func(examples):
return tokenizer(examples["completion"], 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=True, mlm_probability=0.15))In [ ]python · cell 12
python
next(enumerate(dl))In [ ]python · cell 13
python
tokenizer.mask_token, tokenizer.mask_token_idStep4 创建模型
In [ ]python · cell 15
python
model = AutoModelForMaskedLM.from_pretrained("hfl/chinese-macbert-base")Step5 配置训练参数
In [ ]python · cell 17
python
args = TrainingArguments(
output_dir="./masked_lm",
per_device_train_batch_size=32,
logging_steps=10,
num_train_epochs=1
)Step6 创建训练器
In [ ]python · cell 19
python
trainer = Trainer(
args=args,
model=model,
tokenizer=tokenizer,
train_dataset=tokenized_ds,
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=True, mlm_probability=0.15)
)Step7 模型训练
In [ ]python · cell 21
python
trainer.train()Step8 模型推理
In [ ]python · cell 23
python
from transformers import pipeline
pipe = pipeline("fill-mask", model=model, tokenizer=tokenizer, device=0)In [ ]python · cell 24
python
pipe("西安交通[MASK][MASK]博物馆（Xi'an Jiaotong University Museum）是一座位于西安交通大学的博物馆")In [ ]python · cell 25
python
pipe("下面是一则[MASK][MASK]新闻。小编报道，近日，游戏产业发展的非常好！")In [ ]python · cell 26
python