掩码语言模型

1、导入相关包

from datasets import load_dataset, Dataset
from transformers import AutoTokenizer, AutoModelForMaskedLM, DataCollatorForLanguageModeling, TrainingArguments, Trainer

2、加载、划分、处理数据集

# 加载
ds = Dataset.load_from_disk("./wiki_cn_filtered/")
# 划分

# 数据处理
tokenizer = AutoTokenizer.from_pretrained("hfl/chinese-macbert-base")

def process_func(examples):
    return tokenizer(examples["completion"], max_length=384, truncation=True)

tokenized_ds = ds.map(process_func, batched=True, remove_columns=ds.column_names)   

3、创建模型

model = AutoModelForMaskedLM.from_pretrained("./hfl/chinese-macbert-base")

4、创建评估函数

5、创建TrainingArguments、Trainer

args = TrainingArguments(
    output_dir="./masked_lm",
    per_device_train_batch_size=32,
    logging_steps=10,
    num_train_epochs=1
)

trainer = Trainer(
    args=args,
    model=model,
    tokenizer=tokenizer,
    train_dataset=tokenized_ds,
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=True, mlm_probability=0.15)
)

6、模型训练、评估、预测

# 模型训练
trainer.train()
# 模型评估
trainer.evaluate()

7、模型预测

# 模型预测

# 测试
from transformers import pipeline

pipe = pipeline("fill-mask", model=model, tokenizer=tokenizer, device=0)

pipe("西安交通[MASK][MASK]博物馆（Xi'an Jiaotong University Museum）是一座位于西安交通大学的博物馆")
pipe("下面是一则[MASK][MASK]新闻。小编报道，近日，游戏产业发展的非常好！")