【HuggingFace轻松上手】基于Wikipedia的知识增强预训练

前记： 预训练语言模型（Pre-trained Language Model，PLM）想必大家应该并不陌生，其旨在使用自监督学习（Self-supervised Learning）或多任务学习（Multi-task Learning）的方法在大规模的文本语料上进行预训练（Pre-training），基于预训练好的模型，对下游的具体任务进行微调（Fine-tuning）。目前市面上知名的以英文为主预训练语言模型有EMLo、BERT、RoBERTa、XLNet、GPT、DeBERTa（部分也开源了中文版本）。以中文为主的预训练语言模型有ChineseBERT、MacBERT、StructBERT、SpellBERT等。

知识增强预训练语言模型（Knowledge- enhance PLM）旨在在预训练的时候，试图将显式的事实知识（Factual Knowledge）融入到模型中。目前知名的知识增强的模型则有ERNIE（清华）、ERNIE（百度）、K-BERT、KnowBERT、WKLM、KEPLER等。有关所有涉及的预训练模型的细节讲解可相见博主的专栏：预训练语言模型

在众多知识增强预训练的方法中，最基础最简单且最有效的方法是ERNIE（百度）提出的entity masking。大致方法如下：

首先给定一个文本句子，使用实体链指工具识别出文本中的所有实体；对于所有识别出的实体，随机挑选若干实体，并替换为[MASK]，[MASK]的数量取决于实体的长度

具体方法可参见ERNIE1.0。虽然其提出是在中文场景下，但该方法依然适用于英文。我们知道目前的预训练语言模型的分词有两种，一种是以BERT系列为代表的word piece，另一种是以RoBERTa系列为代表的BPE，它们的本质都是将英文单词拆分为若干token，例如“learning”可以被分解为两个token，即“learn”和“###ing”。传统的预训练完全基于token的MLM，而基于实体层面的mask策略，则需要确保实体对应的所有分词也被mask。

本文分享基于英文版本Wikipedia语料和英文知识库Wikidata的知识增强预训练的实现。我们采用Pytorch和HuggingFace实现。建议在Linux开发机上完成。

目录：

数据获取与预处理HuggingFace实现基于Entity Masking的知识增强预训练下游任务微调

---

一、数据获取与处理

（1）Wikipedia Dumps 首先获取英文的大规模无监督语料。我们参照BERT、RoBERTa等市面上绝大多数的工作，挑选的语料来自于Wikipedia Dumps。 一般地，我们直接下载原生态的语料，如图所示： 将下载得到的语料放置在项目根目录的data目录下，文件名假设为“xxx.xml.bz”。下载得到的语料无需自主解压，通常对应的文本都是格式化的XML文件，并不能直接被使用，因此需要进行预处理。

（2）WikiExtractor 幸运的是，WikiExtractor开源工具帮我们解决了预处理的所有环节。具体的说，将WikiExtractor的开源代码下载并放置在data目录下，然后执行命令 python -m wikiextractor.WikiExtractor xxx.xml.bz，等待数小时后，将会获得全部预处理好的文件。文件目录如下所示： 某个目录下存储的是一堆txt文件： 每个文件存储的是意见处理好的文本：

（3）WikiData5M Wikidata5m是由KEPLER提出的一个基于知识图谱补全的数据集。其下载地址为https://deepgraphlearning.github.io/project/wikidata5m 可只需下载如图所示的几个文件（下载或比较慢，可以通过VPN渠道下载）：

Raw：保存着所有三元组。每一行表示一个三元组，即头实体、关系、尾实体，全部使用编号表示；

Entity alias（wikidata5m_entity.txt）：每一行保存每个实体编号及对应的所有可能的实体名称；实体编号以Q开头；

Relation alias（wikidata5m_relation.txt）：每一行保存每个关系编号及对应的关系名称；关系编号以P开头

因此，我们可以使用python实现一个简单的读取Wikipedia的程序

class KGPrompt: def __init__(self, tokenizer: PreTrainedTokenizerBase): self.tokenizer = tokenizer # 加载Wikidata5M知识库 print('loading wikidata5m knowledge graph ...') kg_output = './pretrain_data/kg/' kg = np.load( os.path.join(kg_output, 'wiki_kg.npz'), allow_pickle=True ) self.wiki5m_alias2qid, self.wiki5m_qid2alias, self.wiki5m_pid2alias, self.head_cluster = \ kg['wiki5m_alias2qid'][()], kg['wiki5m_qid2alias'][()], kg['wiki5m_pid2alias'][()], kg['head_cluster'][()] print('loading success .') def sample_entity(self, qid=None, neg_num=0): # 给定一个qid，随机采样一个positive，以及若干负样本 positive = None negative = list() if not qid and qid in self.wiki5m_qid2alias: positive = random.sample(self.wiki5m_qid2alias[qid], 1)[0] if neg_num > 0: negative_qid = random.sample(self.wiki5m_qid2alias.keys(), neg_num) for i in negative_qid: negative.append(random.sample(self.wiki5m_qid2alias[i], 1)[0]) return positive, negative def sample_relation(self, pid=None, neg_num=0): # 给定一个pid，随机采样一个positive，以及若干负样本 positive = None negative = list() if not pid and pid in self.wiki5m_pid2alias: positive = random.sample(self.wiki5m_pid2alias[pid], 1)[0] if neg_num > 0: negative_pid = random.sample(self.wiki5m_pid2alias.keys(), neg_num) for i in negative_pid: negative.append(random.sample(self.wiki5m_pid2alias[i], 1)[0]) return positive, negative def encode_kg(self, kg_str_or_list, max_len): # 将实体/关系分词，并pad if type(kg_str_or_list) == str: kg_str_or_list = [kg_str_or_list] kg_input_ids = list() for kg_str in kg_str_or_list: kg_ids = self.tokenizer.encode(kg_str, add_special_tokens=False, max_length=max_len) kg_ids = [self.tokenizer.cls_token_id] + kg_ids[:max_len - 2] + [self.tokenizer.sep_token_id] kg_ids.extend([self.tokenizer.pad_token_id] * (max_len - len(kg_ids))) kg_input_ids.append(kg_ids) # print('len(kg_ids)=', len(kg_ids)) return kg_input_ids def get_demonstration(self, example: Dict, is_negative=False, start_from_input=True): ''' e.g. data = { 'token_ids': tokens, 'entity_qid': entity_ids, 'entity_pos': mention_spans, 'relation_pid': None, 'relation_pos': None, } ''' # 分词 input_ids, entity_ids = example['token_ids'], example['entity_qid'] input_ids = [self.tokenizer.cls_token_id] + input_ids + [self.tokenizer.sep_token_id] # roberta: [0, x, ..., 2] # input_ids = self.tokenizer.encode(token_list, add_special_tokens=True) # [101, x, ..., 102] type_id = 0 token_type_ids = [type_id] * len(input_ids) start_length = len(input_ids) if start_from_input else 0 entity_spans, relation_spans = list(), list() token_type_span = list() # 每个type对应的区间 token_type_span.append((0, len(input_ids))) if is_negative: # 如果是采样负样本，则随机从KG采样一些实体id entity_ids = random.sample(self.wiki5m_qid2alias.keys(), len(entity_ids)) type_id = 1 # 获得所有mention对齐的entity，根据entity_id随机采样KB里的entity for entity_id in entity_ids: if entity_id in self.wiki5m_qid2alias.keys() and entity_id in self.head_cluster.keys(): entity_name_list = self.wiki5m_qid2alias[entity_id] # ['xx', 'xxx', ...] cluster_list = self.head_cluster[entity_id] # [(rel_id, q_id), ...] # 随机采样一个entity head_name = random.sample(entity_name_list, 1)[0] triple = random.sample(cluster_list, 1)[0] if triple[0] in self.wiki5m_pid2alias.keys() and triple[1] in self.wiki5m_qid2alias.keys(): relation_name = self.wiki5m_pid2alias[triple[0]] tail_name = random.sample(self.wiki5m_qid2alias[triple[1]], 1)[0] template_tokens, entity_span, relation_span = self.template( head=head_name, relation=relation_name, tail=tail_name, type_id=random.randint(0, 2), start_length=start_length ) if len(input_ids) + len(template_tokens) >= self.tokenizer.model_max_length - 2: break start = len(input_ids) input_ids.extend(template_tokens) end = len(input_ids) token_type_ids.extend([type_id] * len(template_tokens)) entity_spans.extend(entity_span) relation_spans.extend(relation_span) token_type_span.append((start, end)) return { 'input_ids': input_ids, 'token_type_ids': token_type_ids, 'noise_detect_label': 0 if is_negative else 1, 'entity_spans': entity_spans, 'relation_spans': relation_spans, 'token_type_span': token_type_span } def template(self, head, relation, tail, type_id=0, start_length=0): if type_id == 0: templates = ["The relation between", head, "and", tail, "is", relation] flag = [0, 1, 0, 1, 0, 2] elif type_id == 1: templates = [head, relation, tail] flag = [1, 2, 1] elif type_id == 2: templates = [head, "is the", relation, "of", tail] flag = [1, 0, 2, 0, 1] template_tokens = list() entity_spans = list() relation_spans = list() for ei, string in enumerate(templates): start = start_length + len(template_tokens) tokens = self.tokenizer.encode(string, add_special_tokens=False) template_tokens.extend(tokens) end = start_length + len(template_tokens) if flag[ei] == 1: entity_spans.append((start, end)) elif flag[ei] == 2: relation_spans.append((start, end)) template_tokens += [self.tokenizer.sep_token_id] return template_tokens, entity_spans, relation_spans

（4）Curpora构建 接下来则是数据处理的关键，即如何对每个wikipedia的每个句子获得对应的实体。我们使用TagMe工具。

TagMe是一个基于Wikipedia的实体链指工具，其可以快速地从英文的文本中返回所有实体及其置信得分。TagMe的使用细节可参考博客https://blog.csdn.net/qq_43549752/article/details/88912835

Step1： 进入页面注册：https://services.d4science.org/home，并选择Google（Gmail）账户登录； Step2: 登录后在中间有一个Cannot find the VRE you were looking for? If you are looking for a VRE please remember to identify the proper Gateway to use via the Explore page. The credentials and the identity are the same through all Gateways of the D4Science infrastructure. 点击 all Gateways of the D4Science infrastructure 后 进入 SoBigData Gateway, 19 VREs / VLabs ，然后找到TagMe，点击Access this Ver Step3: 进入TagMe页面：https://sobigdata.d4science.org/group/tagme Step4： 获取Your Token（上图绿色区域）

因此，可以基于TagMe，对所有Wikipedia的句子获取相应的实体。TagMe的程序如下所示：

import os import sys sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), '..')) import random from tqdm import tqdm from utils.tagme_test import Annotate path = '../data_corpus/wiki' def read_wiki(path): dirs = os.listdir(path=path) corpus = list() for dir in tqdm(dirs): sub_path = os.path.join(path, dir) files = os.listdir(sub_path) for file in files: file_path = os.path.join(sub_path, file) with open(file_path, 'r', encoding='utf-8') as fr: lines = fr.readlines() for line in lines: line = line.replace('\n', '') if 'http' in line: continue tokens = line.split(' ') if len(tokens) < 20: continue corpus.append(line) random.shuffle(corpus) train_corpus = corpus[: -50000] validation_corpus = corpus[-50000:] print('corpus num: {}'.format(len(corpus))) # 32,715,108 print('train corpus num: {}'.format(len(train_corpus))) # 32,705,108 print('validation corpus num: {}'.format(len(validation_corpus))) # 10,000 with open('train.txt', 'w', encoding='utf-8') as fw: for text in tqdm(train_corpus): fw.write(text + '\n') with open('validation.txt', 'w', encoding='utf-8') as fw: for text in tqdm(validation_corpus): fw.write(text + '\n') def tag_me(path, file_name): new_file_name = file_name.split('.')[0] + '_with_entity.' + file_name.split('.')[1] with open(os.path.join(path, file_name), 'r', encoding='utf-8') as fr: lines = fr.readlines() with open(os.path.join(path, new_file_name), 'w', encoding='utf-8') as fw: for line in tqdm(lines): txt = line.replace('\n', '') obj = Annotate(txt, theta=0.2) entities = list(set([i[1] for i in obj.keys()])) # list() example = "{}\t{}".format(txt, '\t'.join(entities)) fw.write(example + '\n') if __name__ == "__main__": # read_wiki(path) # tag_me(path, 'train_10000.txt') # tag_me(path, 'train_10_percent.txt') tag_me(path, 'train.txt')

经过预处理，我们期望训练语料的格式如下所示：

{ "token_ids": [252, 1165, 1731, 13751, 261, 8, 7112, 11, 16359, 9959, 8, 233, 9, 6367, 29, 255, 2753, 11, 9238, 9959, 479, 25244, 263, 16359, 21, 8203, 11, 381, 3677, 1908, 20481, 479, 133, 4996, 9, 5, 16359, 284, 21, 1348, 30, 7393, 16359, 6, 155, 2586, 5893, 9, 18786, 36, 1570, 4718, 2383, 996, 2146, 238, 54, 21, 9390, 13, 27470, 11, 379, 2146, 6, 15, 1060, 744, 22, 627, 13705, 352, 446, 9, 16359, 1064, 7, 1430, 117, 55, 845], "entity_pos": [[2, 5], [6, 7], [13, 17], [18, 20], [27, 31], [41, 49]], "entity_qid": ["Q2353329", "Q2485216", "Q2730418", "Q23169", "Q2564352", "Q726243"], "relation_pos": null, "relation_pid": null }

我们暂时只提供10000个预处理好的数据作为demo，可自行前往下载total_pretrain_data_10000.json。

二、HuggingFace实现基于Entity Masking的知识增强预训练

接下来我们简单实用Pytorch和HuggingFace实现基于entity masking的知识增强预训练工作。基本环境涉及如下：

Python>=3.7Pytorch>=1.8HuggingFace>=4.19Datasets

下面是对应的核心代码，但所有涉及的代码并不能单一运行。博主即将开源本项目的代码，可及时关注GitHub空间：https://github.com/wjn1996

2.1 构建运行函数

我们直接使用HuggingFace提供的参数

# -*- coding: utf-8 -*- # @Time : 2021/11/25 2:51 下午 # @Author : JianingWang # @File : mlm_trainer.py # !/usr/bin/env python # coding=utf-8 import math import os import time import torch import numpy as np from processor import processor_map from transformers import CONFIG_MAPPING, AutoConfig, AutoTokenizer, HfArgumentParser, set_seed from HFTrainer import HFTrainer from transformers.trainer_utils import get_last_checkpoint from transformers import EarlyStoppingCallback from config import ModelArguments, DataTrainingArguments, TrainingArguments from tool.common import init_logger from models import MODEL_CLASSES, TOKENIZER_CLASSES, build_cls_model import logging logger = logging.getLogger(__name__) torch.set_printoptions(precision=3, edgeitems=5, linewidth=160, sci_mode=False) def main(): # See all possible arguments or by passing the --help flag to this script. parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)) model_args, data_args, training_args = parser.parse_args_into_dataclasses() training_args.output_dir = os.path.join(training_args.output_dir, list(filter(None, model_args.model_name_or_path.split('/')))[-1]) os.makedirs(training_args.output_dir, exist_ok=True) # Setup logging log_file = os.path.join(training_args.output_dir, f'{model_args.model_name_or_path.split(os.sep)[-1]}-{data_args.task_name}-{time.strftime("%Y-%m-%d-%H:%M:%S", time.localtime())}.log') log_level = training_args.get_process_log_level() init_logger(log_file, log_level, training_args.local_rank) # Log on each process the small summary: logger.warning( f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}" + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}" ) logger.info(f"Training/evaluation parameters {training_args}") # Detecting last checkpoint. last_checkpoint = None if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir: last_checkpoint = get_last_checkpoint(training_args.output_dir) if last_checkpoint is None and len([i for i in os.listdir(training_args.output_dir) if not i.endswith("log")]) > 0: raise ValueError(f"Output directory ({training_args.output_dir}) already exists and is not empty. " "Use --overwrite_output_dir to overcome.") elif last_checkpoint is not None and training_args.resume_from_checkpoint is None: logger.info(f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change " "the `--output_dir` or add `--overwrite_output_dir` to train from scratch.") # Set seed before initializing model. set_seed(training_args.seed) # 加载任务相关处理的processor if data_args.task_name in processor_map: processor = processor_map[data_args.task_name](data_args, training_args, model_args) else: raise ValueError("task name 未指定或不在processor map中") # Load pretrained model and tokenizer # The .from_pretrained methods guarantee that only one local process can concurrently download model & vocab. config_kwargs = { "cache_dir": model_args.cache_dir, "revision": model_args.model_revision, 'finetuning_task': data_args.task_name } # if 'mlm' not in data_args.task_type: if hasattr(processor, 'labels'): config_kwargs['num_labels'] = len(processor.labels) if model_args.config_name: config = AutoConfig.from_pretrained(model_args.config_name, **config_kwargs) elif model_args.model_name_or_path: config = AutoConfig.from_pretrained(model_args.model_name_or_path, **config_kwargs) if 'longformer' in model_args.model_name_or_path: config.sep_token_id = 102 else: config = CONFIG_MAPPING[model_args.model_type]() logger.warning("You are instantiating a new config instance from scratch.") if model_args.config_overrides is not None: logger.info(f"Overriding config: {model_args.config_overrides}") config.update_from_string(model_args.config_overrides) logger.info(f"New config: {config}") tokenizer_kwargs = { "cache_dir": model_args.cache_dir, "use_fast": model_args.use_fast_tokenizer, "revision": model_args.model_revision, # 'do_lower_case': model_args.do_lower_case #根据model 的tokenizer自己配置 } tokenizer_class = TOKENIZER_CLASSES.get(model_args.model_type, AutoTokenizer) if model_args.tokenizer_name: tokenizer = tokenizer_class.from_pretrained(model_args.tokenizer_name, **tokenizer_kwargs) elif model_args.model_name_or_path: tokenizer = tokenizer_class.from_pretrained(model_args.model_name_or_path, **tokenizer_kwargs) else: raise ValueError( "You are instantiating a new tokenizer from scratch. This is not supported by this script." "You can do it from another script, save it, and load it from here, using --tokenizer_name." ) # 添加字典 processor.set_tokenizer(tokenizer) # 添加配置 # print("config=", config) processor.set_config(config) if data_args.task_type == 'autocls': model_class = build_cls_model(config) else: model_class = MODEL_CLASSES[data_args.task_type] if model_args.from_scratch: logger.info("Training new model from scratch") model = model_class.from_config(config) else: model = model_class.from_pretrained( model_args.model_name_or_path, from_tf=bool(".ckpt" in model_args.model_name_or_path), config=config, cache_dir=model_args.cache_dir, revision=model_args.model_revision, use_auth_token=True if model_args.use_auth_token else None, ignore_mismatched_sizes=True ) model.resize_token_embeddings(len(tokenizer)) # print("num_labels=", model.num_labels) tokenized_datasets = processor.get_tokenized_datasets() if training_args.do_train: if "train" not in tokenized_datasets: raise ValueError("--do_train requires a train dataset") train_dataset = tokenized_datasets["train"] if data_args.max_train_samples is not None: train_dataset = train_dataset.select(range(data_args.max_train_samples)) if training_args.do_eval: if "validation" not in tokenized_datasets: raise ValueError("--do_eval requires a validation dataset") eval_dataset = tokenized_datasets["validation"] if data_args.max_eval_samples is not None: eval_dataset = eval_dataset.select(range(data_args.max_eval_samples)) if training_args.do_predict: if 'test' not in tokenized_datasets: raise ValueError("--do_predict requires a test dataset") test_dataset = tokenized_datasets['test'] if data_args.max_predict_samples is not None: test_dataset = test_dataset.select(range(data_args.max_predict_samples)) # Data collator data_collator = processor.get_data_collator() if hasattr(processor, 'compute_metrics'): compute_metrics = processor.compute_metrics else: compute_metrics = None if model_args.freeze_epochs: callbacks.append(FreezeCallback(freeze_epochs=model_args.freeze_epochs, freeze_keyword=model_args.freeze_keyword)) if model_args.ema: callbacks.append(ExponentialMovingAveragingCallback(model_args.ema_decay)) if training_args.do_predict_during_train: from callback.evaluate import DoPredictDuringTraining callbacks.append(DoPredictDuringTraining(test_dataset, processor)) trainer = HFTrainer( model=model, args=training_args, train_dataset=train_dataset if training_args.do_train else None, eval_dataset=eval_dataset if training_args.do_eval else None, compute_metrics=compute_metrics, tokenizer=tokenizer, data_collator=data_collator, callbacks=callbacks ) # Training if training_args.do_train: checkpoint = None if training_args.resume_from_checkpoint is not None: checkpoint = training_args.resume_from_checkpoint elif last_checkpoint is not None: checkpoint = last_checkpoint train_result = trainer.train(resume_from_checkpoint=checkpoint) trainer.save_model() # Saves the tokenizer too for easy upload metrics = train_result.metrics max_train_samples = ( data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset) ) metrics["train_samples"] = min(max_train_samples, len(train_dataset)) trainer.log_metrics("train", metrics) trainer.save_metrics("train", metrics) trainer.save_state() # Evaluation if training_args.do_eval: logger.info("*** Evaluate ***") try: metrics = trainer.evaluate() max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset) metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset)) if data_args.task_type == 'mlm': try: perplexity = math.exp(metrics["eval_loss"]) except OverflowError: perplexity = float("inf") metrics["perplexity"] = perplexity trainer.log_metrics("eval", metrics) trainer.save_metrics("eval", metrics) except: logger.info("UNK problems ...") if training_args.do_predict and not training_args.do_predict_during_train: logger.info("*** Predict ***") if not data_args.keep_predict_labels: for l in ['labels', 'label']: if l in test_dataset.column_names: test_dataset = test_dataset.remove_columns(l) prediction = trainer.predict(test_dataset, metric_key_prefix="predict") logits = prediction.predictions if data_args.keep_predict_labels: label_ids = prediction.label_ids if hasattr(processor, 'save_result'): if trainer.is_world_process_zero(): if not data_args.keep_predict_labels: processor.save_result(logits) else: processor.save_result(logits, label_ids) else: predictions = np.argmax(logits, axis=1) output_predict_file = os.path.join(training_args.output_dir, f"predict_results.txt") if trainer.is_world_process_zero(): with open(output_predict_file, "w") as writer: logger.info(f"***** Predict results {data_args.task_name} *****") writer.write("index\tprediction\n") for index, item in enumerate(predictions): item = processor.labels[item] writer.write(f"{index}\t{item}\n") if __name__ == "__main__": main() 2.2 重写HuggingFace Trainer

根据自己的需要，可以重写trainer中的train部分。

# -*- coding: utf-8 -*- # @Time : 2022/1/7 3:07 下午 # @Author : JianingWang # @File : HFTrainer from typing import Dict, Union, Any, Optional, Callable, List, Tuple, Iterator import datasets from datasets import Dataset from torch.utils.data import RandomSampler, DistributedSampler from transformers import PreTrainedModel, DataCollator, PreTrainedTokenizerBase, EvalPrediction, TrainerCallback from transformers.trainer_pt_utils import DistributedSamplerWithLoop, get_length_grouped_indices from transformers.trainer_pt_utils import DistributedLengthGroupedSampler as DistributedLengthGroupedSamplerOri from transformers.trainer_pt_utils import LengthGroupedSampler as LengthGroupedSamplerOri from transformers.training_args import ParallelMode from config import TrainingArguments from transformers.trainer import Trainer, _is_torch_generator_available import torch from torch import nn from transformers.file_utils import is_datasets_available from models.adversarial import FGM class LengthGroupedSampler(LengthGroupedSamplerOri): def __iter__(self): indices = get_length_grouped_indices(self.lengths, self.batch_size, generator=self.generator, mega_batch_mult=256) return iter(indices) class DistributedLengthGroupedSampler(DistributedLengthGroupedSamplerOri): def __iter__(self) -> Iterator: # Deterministically shuffle based on epoch and seed g = torch.Generator() g.manual_seed(self.seed + self.epoch) indices = get_length_grouped_indices(self.lengths, self.batch_size, generator=g, mega_batch_mult=400) if not self.drop_last: # add extra samples to make it evenly divisible indices += indices[: (self.total_size - len(indices))] else: # remove tail of data to make it evenly divisible. indices = indices[: self.total_size] assert len(indices) == self.total_size # subsample indices = indices[self.rank: self.total_size: self.num_replicas] assert len(indices) == self.num_samples return iter(indices) class HFTrainer(Trainer): def __init__( self, model: Union[PreTrainedModel, nn.Module] = None, args: TrainingArguments = None, data_collator: Optional[DataCollator] = None, train_dataset: Optional[Dataset] = None, eval_dataset: Optional[Dataset] = None, tokenizer: Optional[PreTrainedTokenizerBase] = None, model_init: Callable[[], PreTrainedModel] = None, compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None, callbacks: Optional[List[TrainerCallback]] = None, optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), ): super(HFTrainer, self).__init__(model, args, data_collator, train_dataset, eval_dataset, tokenizer, model_init, compute_metrics, callbacks, optimizers) if self.args.do_adv: self.fgm = FGM(self.model) for callback in callbacks: callback.trainer = self self.global_step_ = 0 def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]) -> torch.Tensor: """ Perform a training step on a batch of inputs. Subclass and override to inject custom behavior. Args: model (`nn.Module`): The model to train. inputs (`Dict[str, Union[torch.Tensor, Any]]`): The inputs and targets of the model. The dictionary will be unpacked before being fed to the model. Most models expect the targets under the argument `labels`. Check your model's documentation for all accepted arguments. Return: `torch.Tensor`: The tensor with training loss on this batch. """ self.global_step_ += 1 model.train() inputs = self._prepare_inputs(inputs) with self.autocast_smart_context_manager(): loss = self.compute_loss(model, inputs) if self.args.n_gpu > 1: loss = loss.mean() # mean() to average on multi-gpu parallel training if self.args.gradient_accumulation_steps > 1 and not self.deepspeed: # deepspeed handles loss scaling by gradient_accumulation_steps in its `backward` loss = loss / self.args.gradient_accumulation_steps if self.global_step_ % 10 == 0: print('[step={}, loss={}]'.format(self.global_step_, loss)) if self.do_grad_scaling: self.scaler.scale(loss).backward() elif self.deepspeed: # loss gets scaled under gradient_accumulation_steps in deepspeed loss = self.deepspeed.backward(loss) else: loss.backward() # 对抗训练 if self.args.do_adv: self.fgm.attack() with self.autocast_smart_context_manager(): loss_adv = self.compute_loss(model, inputs) if self.args.n_gpu > 1: loss_adv = loss_adv.mean() if self.args.gradient_accumulation_steps > 1 and not self.deepspeed: loss_adv = loss_adv / self.args.gradient_accumulation_steps if self.do_grad_scaling: self.scaler.scale(loss_adv).backward() else: loss_adv.backward() self.fgm.restore() # 恢复embedding参数 return loss.detach() def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]: # if not has_length(self.train_dataset): # return None generator = None if self.args.world_size <= 1 and _is_torch_generator_available: generator = torch.Generator() generator.manual_seed(int(torch.empty((), dtype=torch.int64).random_().item())) # Build the sampler. if self.args.group_by_length: if is_datasets_available() and isinstance(self.train_dataset, datasets.Dataset): lengths = ( self.train_dataset[self.args.length_column_name] if self.args.length_column_name in self.train_dataset.column_names else None ) else: lengths = None model_input_name = self.tokenizer.model_input_names[0] if self.tokenizer is not None else None if self.args.world_size <= 1: return LengthGroupedSampler( self.args.train_batch_size * self.args.gradient_accumulation_steps, dataset=self.train_dataset, lengths=lengths, model_input_name=model_input_name, generator=generator, ) else: return DistributedLengthGroupedSampler( self.args.train_batch_size * self.args.gradient_accumulation_steps, dataset=self.train_dataset, num_replicas=self.args.world_size, rank=self.args.process_index, lengths=lengths, model_input_name=model_input_name, seed=self.args.seed, ) else: if self.args.world_size <= 1: if _is_torch_generator_available: return RandomSampler(self.train_dataset, generator=generator) return RandomSampler(self.train_dataset) elif ( self.args.parallel_mode in [ParallelMode.TPU, ParallelMode.SAGEMAKER_MODEL_PARALLEL] and not self.args.dataloader_drop_last ): # Use a loop for TPUs when drop_last is False to have all batches have the same size. return DistributedSamplerWithLoop( self.train_dataset, batch_size=self.args.per_device_train_batch_size, num_replicas=self.args.world_size, rank=self.args.process_index, seed=self.args.seed, ) else: return DistributedSampler( self.train_dataset, num_replicas=self.args.world_size, rank=self.args.process_index, seed=self.args.seed, ) 2.3 实现数据加载类

该类用于读取我们保存的预处理语料，例如total_pretrain_data_10000.json，基于该数据，我们需要读取并加载缓存。我们默认使用Datasets的数据集加载方式，需要进行重写：

class WikiKGPLMSupervisedJsonProcessor(DataProcessor): def __init__(self, data_args, training_args, model_args): super().__init__(data_args, training_args, model_args) self.is_only_mlm = False def get_data_collator(self): pad_to_multiple_of_8 = self.data_args.line_by_line and self.training_args.fp16 and not self.data_args.pad_to_max_length if self.is_only_mlm: print("You set is_only_mlm is True") return DataCollatorForProcessedWikiKGPLM_OnlyMLM( tokenizer=self.tokenizer, mlm_probability=self.data_args.mlm_probability, pad_to_multiple_of=8 if pad_to_multiple_of_8 else None, ) print("You set is_only_mlm is False") return DataCollatorForProcessedWikiKGPLM( tokenizer=self.tokenizer, mlm_probability=self.data_args.mlm_probability, pad_to_multiple_of=8 if pad_to_multiple_of_8 else None, ) def get_examples(self, set_type=None): data_files = {} if self.data_args.train_file is not None: data_files["train"] = self.data_args.train_file extension = self.data_args.train_file.split(".")[-1] if self.data_args.validation_file is not None: data_files["validation"] = self.data_args.validation_file extension = self.data_args.validation_file.split(".")[-1] if extension == "json": extension = "json" raw_datasets = load_dataset(extension, data_files=data_files, cache_dir=self.model_args.cache_dir) # raw_datasets['train'] = raw_datasets['train'].shuffle() # If no validation data is there, validation_split_percentage will be used to divide the dataset. if "validation" not in raw_datasets.keys(): raw_datasets["validation"] = load_dataset( extension, data_files=data_files, split=f"train[:{self.data_args.validation_split_percentage}%]", cache_dir=self.model_args.cache_dir, ) raw_datasets["train"] = load_dataset( extension, data_files=data_files, split=f"train[{self.data_args.validation_split_percentage}%:]", cache_dir=self.model_args.cache_dir, ) return raw_datasets def compute_metrics(self, p: EvalPrediction): # print('p.label_ids=', p.label_ids) # print('type(p.label_ids)=', type(p.label_ids)) # print('p.predictions=', p.predictions) # print('type(p.predictions)=', type(p.predictions)) if type(p.predictions) in [tuple, list]: preds = p.predictions[1] else: preds = p.predictions preds = preds[p.label_ids != -100] labels = p.label_ids[p.label_ids != -100] acc = (preds == labels).mean() return { 'acc': round(acc, 4) } def get_tokenized_datasets(self): data_files = {} if self.data_args.train_file is not None: data_files["train"] = self.data_args.train_file extension = self.data_args.train_file.split(".")[-1] if self.data_args.validation_file is not None: data_files["validation"] = self.data_args.validation_file extension = self.data_args.validation_file.split(".")[-1] if extension == "json": extension = "json" raw_datasets = load_dataset(extension, data_files=data_files, cache_dir=self.model_args.cache_dir) # raw_datasets['train'] = raw_datasets['train'].shuffle() # If no validation data is there, validation_split_percentage will be used to divide the dataset. if "validation" not in raw_datasets.keys(): raw_datasets["validation"] = load_dataset( extension, data_files=data_files, split=f"train[:{self.data_args.validation_split_percentage}%]", cache_dir=self.model_args.cache_dir, ) raw_datasets["train"] = load_dataset( extension, data_files=data_files, split=f"train[{self.data_args.validation_split_percentage}%:]", cache_dir=self.model_args.cache_dir, ) logger.info(f'validation fingerprint {raw_datasets}') ''' e.g. raw_datasets = DatasetDict({ train: Dataset({ features: ['json'], num_rows: xxx }) validation: Dataset({ features: ['json'], num_rows: xxx }) }) ''' if self.training_args.do_train: column_names = raw_datasets["train"].column_names else: column_names = raw_datasets["validation"].column_names text_column_name = "text" if "text" in column_names else column_names[0] max_seq_length = self.tokenizer.model_max_length if self.data_args.max_seq_length is None else self.data_args.max_seq_length # When using line_by_line, we just tokenize each nonempty line. padding = "max_length" if self.data_args.pad_to_max_length else False tokenizer = self.tokenizer def tokenize_function(examples): # Remove empty lines examples[text_column_name] = [ line for line in examples[text_column_name] if len(line) > 0 and not line.isspace() ] # examples['length'] = [len(line) for line in examples[text_column_name]] return tokenizer( examples[text_column_name], padding=padding, truncation=True, max_length=max_seq_length, return_special_tokens_mask=True, ) ''' { 'text': tokens_str, 'entity_ids': entity_ids, 'mention_spans': mention_spans } ''' return raw_datasets 2.4 需要自己实现Collator

因为HuggingFace默认使用的是Datasets类，因此我们需要为Dataloader设计回调函数，该回调函数的目的是每次加载一个batch的数据时进行一些数据处理，例如生成张量，padding等部分。

由于在真正进行预训练时，语料规模时千万级的，因此我们无法在有限的时间和空间内全部加载所有的数据到内存中，但是我们可以使用Dataloader每次加载一个batch数据，并进行处理。为了提高GPU和CPU的协作效率，可以设置dataloader_num_workers=1

本文的Collator主要为了将读取的数据集进行分词、padding并转化为张量，同时需要根据已知的实体区间，对原始的分词后的序列进行mask操作。

这里有一处细节需要读者注意，BERT分词之后，部分的词会被转换为多个token，所以原始的文本序列和分词后的序列长度是不一致的。本文提供的数据集是已经分词好的，所以读者无需考虑。如若读者自行构建的是原始的文本，则需要考虑分词前后各个词的对应位置映射关系。在HuggingFace中为offset_mapping。

我们提供两种不同的实现方式：

@dataclass class DataCollatorForPretrainWithKG(DataCollatorMixin): """ 在使用kg进行预训练 """ tokenizer: PreTrainedTokenizerBase mlm: bool = True mlm_probability: float = 0.15 pad_to_multiple_of: Optional[int] = None tf_experimental_compile: bool = False return_tensors: str = "pt" def __post_init__(self): self.numerical_tokens = [v for k, v in self.tokenizer.vocab.items() if k.isdigit()] self.exclude_tokens = self.numerical_tokens + self.tokenizer.all_special_ids def torch_call(self, examples: List[Union[List[int], Any, Dict[str, Any]]]) -> Dict[str, Any]: # Handle dict or lists with proper padding and conversion to tensor. from tool.ner import position_2_bio for example in examples: seq_len = len(example['input_ids']) if 'token_type_ids' not in example: example['token_type_ids'] = [0] * seq_len if 'attention_mask' not in example: example['attention_mask'] = [1] * seq_len if 'entity_position' in example: example['ner_labels'] = [position_2_bio(p, seq_len) for p in example['entity_position']] example.pop('entity_position') if isinstance(examples[0], (dict, BatchEncoding)): batch = self.tokenizer.pad(examples, return_tensors="pt", pad_to_multiple_of=self.pad_to_multiple_of) else: batch = { "input_ids": _torch_collate_batch(examples, self.tokenizer, pad_to_multiple_of=self.pad_to_multiple_of) } # If special token mask has been preprocessed, pop it from the dict. special_tokens_mask = batch.pop("special_tokens_mask", None) if self.mlm: batch["input_ids"], batch["labels"] = self.torch_mask_tokens( batch["input_ids"], special_tokens_mask=special_tokens_mask ) return batch def torch_mask_tokens(self, inputs: Any, special_tokens_mask: Optional[Any] = None) -> Tuple[Any, Any]: """ Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original. """ import torch labels = inputs.clone() # We sample a few tokens in each sequence for MLM training (with probability `self.mlm_probability`) probability_matrix = torch.full(labels.shape, self.mlm_probability) special_tokens_mask = [ [1 if token in self.exclude_tokens else 0 for token in val] for val in labels.tolist() ] special_tokens_mask = torch.tensor(special_tokens_mask, dtype=torch.bool) probability_matrix.masked_fill_(special_tokens_mask, value=0.0) masked_indices = torch.bernoulli(probability_matrix).bool() labels[~masked_indices] = -100 # We only compute loss on masked tokens # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK]) indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices inputs[indices_replaced] = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token) # 10% of the time, we replace masked input tokens with random word indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced random_words = torch.randint(len(self.tokenizer), labels.shape, dtype=torch.long) inputs[indices_random] = random_words[indices_random] # The rest of the time (10% of the time) we keep the masked input tokens unchanged return inputs, labels class DataCollatorForProcessedWikiKGPLM_OnlyMLM(DataCollatorForLanguageModeling): """ Data collator used for language modeling that masks entire words. Only MLM - collates batches of tensors, honoring their tokenizer's pad_token - preprocesses batches for masked language modeling <Tip> This collator relies on details of the implementation of subword tokenization by [`BertTokenizer`], specifically that subword tokens are prefixed with *##*. For tokenizers that do not adhere to this scheme, this collator will produce an output that is roughly equivalent to [`.DataCollatorForLanguageModeling`]. </Tip>""" def __post_init__(self): if self.mlm and self.tokenizer.mask_token is None: raise ValueError( "This tokenizer does not have a mask token which is necessary for masked language modeling. " "You should pass `mlm=False` to train on causal language modeling instead." ) self.numerical_tokens = [v for k, v in self.tokenizer.vocab.items() if k.isdigit()] self.exclude_tokens = self.numerical_tokens + self.tokenizer.all_special_ids # add by wjn self.kg_prompt = KGPrompt(tokenizer=self.tokenizer) random.seed(42) def torch_call(self, features: List[Union[List[int], Any, Dict[str, Any]]]) -> Dict[str, Any]: assert isinstance(features[0], (dict, BatchEncoding)) input_features = list() for ei, feature in enumerate(features): input_ids, kg_prompt_ids, task_id = feature['input_ids'], feature['kg_prompt_ids'], feature['task_id'] text_len = len(input_ids) kg_len = len(kg_prompt_ids) # 补充padding input_ids = input_ids + kg_prompt_ids input_ids = input_ids[: self.tokenizer.model_max_length] token_len = len(input_ids) input_ids.extend([self.tokenizer.pad_token_id] * (self.tokenizer.model_max_length - len(input_ids))) # 生成token_type_id token_type_ids = [0] * text_len + [0] * kg_len + [0] * (self.tokenizer.model_max_length - text_len - kg_len) # 生成mlm_label # mlm_labels = [-100] * len(input_ids) attention_mask = np.zeros([self.tokenizer.model_max_length, self.tokenizer.model_max_length]) # # 非pad部分，context全部可见 token_type_span = [(0, text_len)] st, ed = text_len, text_len for ei, token in enumerate(kg_prompt_ids): if token == self.tokenizer.sep_token_id: ed = text_len + ei token_type_span.append((st, ed)) st = ei context_start, context_end = token_type_span[0] attention_mask[context_start: context_end, : token_len] = 1 attention_mask[: token_len, context_start: context_end] = 1 # 非pad部分，每个三元组自身可见、与context可见，三元组之间不可见 for ei, (start, end) in enumerate(token_type_span): # start, end 每个token type的区间 attention_mask[start: end, start: end] = 1 attention_mask = attention_mask.tolist() # 由于原始不同的task对应的feature不同，为了统一，对不存在的feature使用pad进行填充 # entity_label = [0] * 20 # entity_negative = [[0] * 20] * 5 # relation_label = [0] * 5 # relation_negative = [[0] * 5] * 5 entity_candidate = [[0] * 20] * 6 relation_candidate = [[0] * 5] * 6 # MLM mask采样 # 只对context部分进行mlm采样。15%的进行mask，其中80%替换为<MASK>，10%随机替换其他词，10%保持不变 input_ids = torch.Tensor(input_ids).long() mlm_labels = input_ids.clone() # MLM mask采样 # 只对context部分进行mlm采样。15%的进行mask，其中80%替换为<MASK>，10%随机替换其他词，10%保持不变 probability_matrix = torch.full([text_len], self.mlm_probability) probability_matrix = torch.cat( [probability_matrix, torch.zeros([self.tokenizer.model_max_length - text_len])], -1) masked_indices = torch.bernoulli(probability_matrix).bool() mlm_labels[~masked_indices] = -100 # We only compute loss on masked tokens # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK]) indices_replaced = torch.bernoulli(torch.full(mlm_labels.shape, 0.8)).bool() & masked_indices input_ids[indices_replaced] = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token) # 10% of the time, we replace masked input tokens with random word indices_random = torch.bernoulli( torch.full(mlm_labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced random_words = torch.randint(len(self.tokenizer.get_vocab()), mlm_labels.shape, dtype=torch.long) input_ids[indices_random] = random_words[indices_random] input_features.append({ 'input_ids': input_ids.numpy().tolist(), 'token_type_ids': token_type_ids, 'attention_mask': attention_mask, 'labels': mlm_labels.numpy().tolist(), # 'entity_label': entity_label, 'entity_candidate': entity_candidate, # 'relation_label': relation_label, 'relation_candidate': relation_candidate, # 'noise_detect_label': -1, 'task_id': task_id, 'mask_id': self.tokenizer.mask_token_id }) del features input_features = {key: torch.tensor([feature[key] for feature in input_features], dtype=torch.long) for key in input_features[0].keys()} return input_features def _whole_word_mask(self, input_tokens: List[str], max_predictions=512): """ Get 0/1 labels for masked tokens with whole word mask proxy """ if not isinstance(self.tokenizer, (BertTokenizer, BertTokenizerFast)): warnings.warn( "DataCollatorForWholeWordMask is only suitable for BertTokenizer-like tokenizers. " "Please refer to the documentation for more information." ) cand_indexes = [] for (i, token) in enumerate(input_tokens): if token == "[CLS]" or token == "[SEP]": continue if len(cand_indexes) >= 1 and token.startswith("##"): cand_indexes[-1].append(i) else: cand_indexes.append([i]) random.shuffle(cand_indexes) num_to_predict = min(max_predictions, max(1, int(round(len(input_tokens) * self.mlm_probability)))) masked_lms = [] covered_indexes = set() for index_set in cand_indexes: if len(masked_lms) >= num_to_predict: break # If adding a whole-word mask would exceed the maximum number of # predictions, then just skip this candidate. if len(masked_lms) + len(index_set) > num_to_predict: continue is_any_index_covered = False for index in index_set: if index in covered_indexes: is_any_index_covered = True break if is_any_index_covered: continue for index in index_set: covered_indexes.add(index) masked_lms.append(index) if len(covered_indexes) != len(masked_lms): raise ValueError("Length of covered_indexes is not equal to length of masked_lms.") mask_labels = [1 if i in covered_indexes else 0 for i in range(len(input_tokens))] return mask_labels def torch_mask_tokens(self, inputs: Any, special_tokens_mask: Optional[Any] = None) -> Tuple[Any, Any]: """ Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original. """ import torch labels = inputs.clone() # [bz, len] # We sample a few tokens in each sequence for MLM training (with probability `self.mlm_probability`) probability_matrix = torch.full(labels.shape, self.mlm_probability) # [bz, seq_len], all values is 0.15 special_tokens_mask = [ [1 if token in self.exclude_tokens else 0 for token in val] for val in labels.tolist() ] special_tokens_mask = torch.tensor(special_tokens_mask, dtype=torch.bool) probability_matrix.masked_fill_(special_tokens_mask, value=0.0) ### add by ruihan.wjn relace all values to 0 where belongs to kg_prompt ### masked_indices = torch.bernoulli(probability_matrix).bool() labels[~masked_indices] = -100 # We only compute loss on masked tokens # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK]) indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices inputs[indices_replaced] = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token) # 10% of the time, we replace masked input tokens with random word indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced random_words = torch.randint(len(self.tokenizer), labels.shape, dtype=torch.long) inputs[indices_random] = random_words[indices_random] return inputs, labels 2.5 构建模型类

由于我们是最原始的MLM任务，因此可模型部分无需自己实现。当然也可以自行重写相关类实现其他自研功能。

import torch from torch import nn from torch.nn import CrossEntropyLoss from collections import OrderedDict from transformers.models.bert import BertPreTrainedModel, BertModel from transformers.models.bert.modeling_bert import BertOnlyMLMHead class BertForPretrainWithKG(BertPreTrainedModel): def __init__(self, config): super().__init__(config) self.num_labels = config.num_labels self.config = config self.bert = BertModel(config) classifier_dropout = ( config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob ) self.dropout = nn.Dropout(classifier_dropout) self.cls = BertOnlyMLMHead(config) self.classifiers = nn.ModuleList([nn.Linear(config.hidden_size, config.num_ner_labels) for _ in range(config.entity_type_num)]) self.post_init() def forward( self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, encoder_hidden_states=None, encoder_attention_mask=None, labels=None, ner_labels=None, output_attentions=None, output_hidden_states=None, return_dict=None, ): return_dict = return_dict if return_dict is not None else self.config.use_return_dict outputs = self.bert( input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) sequence_output = outputs[0] # mlm prediction_scores = self.cls(sequence_output) # ner sequence_output = self.dropout(sequence_output) ner_logits = torch.stack([classifier(sequence_output) for classifier in self.classifiers]).movedim(1, 0) # mlm masked_lm_loss, ner_loss, total_loss = None, None, None if labels is not None: loss_fct = CrossEntropyLoss() # -100 index = padding token masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)) if ner_labels is not None: loss_fct = CrossEntropyLoss() # Only keep active parts of the loss active_loss = attention_mask.repeat(self.config.entity_type_num, 1, 1).view(-1) == 1 active_logits = ner_logits.reshape(-1, self.config.num_ner_labels) active_labels = torch.where( active_loss, ner_labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(ner_labels) ) ner_loss = loss_fct(active_logits, active_labels) if masked_lm_loss: total_loss = masked_lm_loss + ner_loss * 4 return OrderedDict([ ('loss', total_loss), ('mlm_loss', masked_lm_loss.unsqueeze(0)), ('ner_loss', ner_loss.unsqueeze(0)), ('logits', prediction_scores.argmax(2)), ('ner_logits', ner_logits.argmax(3)) ])

模型部分也可以选择RoBERTa、RoFormer或MegatronBert。

2.6 预训练细节

在实验中，我们使用单机多卡V100-32G的GPU进行实验，为了加快实验速度，使用FP16混合精度方法。相关超参数为：

batch size：单卡16，梯度累积step为2；learning rate：1e-5；优化器：AdamW；warm up rate：0.1（在总的训练step数量的10%的位置之前，学习率默认线性增长，之后则进行指数衰减）；

以4卡为例，实验的loss变化情况如图所示（不同颜色代表不同的卡）：

正常情况下，MLM的loss会收敛至0.4～0.7之间，训练5000step大约耗时12～18小时。

三、下游任务微调

为了验证我们的知识增强预训练语言模型是否有效，我们需要与原生态的BERT、RoBERTa进行比较，同时也需要和其他知识增强模型进行对比。我们在此只介绍几个下游任务及其实现。

Entity Typing：主要为Open Entity数据集；Relation Extraction：包括TACRED和FewRel数据集；Knowledge Probing：包括LAMA

所有数据可直接由此下载：https://cloud.tsinghua.edu.cn/f/a763616323f946fd8ff6/ 我们开源了所有微调任务的脚本，详见：https://github.com/wjn1996/KP-PLM/tree/main/finetune_task

3.1 Entity Typing

给定一个文本，以及标注的一个实体，Entity Typing的任务目标是预测这个指定实体的类型。通常情况下，我们可以为每个实体前后添加标记（例如[ENT] 和 [/ENT]），并喂入预训练模型中进行分类。

我们挑选了OpenEntity数据集作为评测任务，数据下载地址为：https://·/wzhouad/RE_improved_baseline/blob/main/train_tacred.pyFewRel：由清华NLP提出的小样本关系抽取，只提供了训练集和验证集，下载地址为：train_wiki.json，val_wiki.json 3.3 Knowledge Probing

Knowledge Probing的任务目的：

原始数据中，每个文件表示一个关系，每个样本中有sub_label和obj_label，以及这两个实体远程监督的句子。句子中对obj替换为[MASK]。因此任务目标是知道关系、知道sub以及对应句子，来预测[MASK]对应的obj；我们挑选LAMA数据集。LAMA还给予了模板——如果提供模板，则不用原始给的文本句子，只用对应的模板来预测，例如：[X] is born in [Y]，sub为obama替换至[X]，[Y]部分为[MASK]，用于预测。

LAMA官网：https://github.com/facebookresearch/LAMA LAMA涉及的环境过低，需要单独创建一个docker，安装对应的requirements，相对比较麻烦，因此我们直接参考CoLAKE提供的LAMA的代码和数据：CoLAKE

实验对比情况如图所示：

参考工作

CoLAKE：https://github.com/txsun1997/CoLAKEERNIE（清华）：https://github.com/thunlp/ERNIE

---

自研EasyNLP工具包发布——EasyNLP，Easy-to-use！ 亮点包括：

更全面的知识模型——知识增强预训练语言模型 + 知识增强多模态模型；“三部曲”一键CLUE刷榜工具，成为CLUE官方代码；轻松上手Few-shot Learning，Prompt-tuning

GitHub：https://github.com/alibaba/EasyNLP Paper：EasyNLP: A Comprehensive and Easy-to-use Toolkit for Natural Language Processing Cite：

@article{DBLP:journals/corr/abs-2205-00258, author = {Chengyu Wang and Minghui Qiu and Taolin Zhang and Tingting Liu and Lei Li and Jianing Wang and Ming Wang and Jun Huang and Wei Lin}, title = {EasyNLP: {A} Comprehensive and Easy-to-use Toolkit for Natural Language Processing}, journal = {CoRR}, volume = {abs/2205.00258}, year = {2022}, url = {https://doi.org/10.48550/arXiv.2205.00258}, doi = {10.48550/arXiv.2205.00258}, eprinttype = {arXiv}, eprint = {2205.00258}, timestamp = {Tue, 03 May 2022 15:52:06 +0200}, biburl = {https://dblp.org/rec/journals/corr/abs-2205-00258.bib}, bibsource = {dblp computer science bibliography, https://dblp.org} }

---

1.本站遵循行业规范，任何转载的稿件都会明确标注作者和来源；2.本站的原创文章，会注明原创字样,如未注明都非原创，如有侵权请联系删除！；3.作者投稿可能会经我们编辑修改或补充；4.本站不提供任何储存功能只提供收集或者投稿人的网盘链接。