# # Author: penhe@microsoft.com # Date: 01/25/2019 # from glob import glob from collections import OrderedDict,defaultdict,Sequence from bisect import bisect import copy import math from scipy.special import softmax import numpy as np import pdb import os import sys import csv import random import torch import re import shutil import ujson as json from torch.utils.data import DataLoader from .metrics import * from .task import EvalData, Task from .task_registry import register_task from ...utils import xtqdm as tqdm from ...data import ExampleInstance, ExampleSet, DynamicDataset,example_to_feature from ...data.example import _truncate_segments from ...data.example import * from ...deberta import NNModule from ...utils import get_logger,boolean_string from ...training import DistributedTrainer, batch_to from ...data import DistributedBatchSampler, SequentialSampler, BatchSampler, AsyncDataLoader from ..models import MaskedLanguageModel,ReplacedTokenDetectionModel from .mlm_task import NGramMaskGenerator from .._utils import merge_distributed, join_chunks logger=get_logger() __all__ = ["RTDTask"] class RTDModel(NNModule): def __init__(self, config, *wargs, **kwargs): super().__init__(config) gen_config = config.generator disc_config = config.discriminator self.config = config self.generator = MaskedLanguageModel(gen_config) self.discriminator = ReplacedTokenDetectionModel(disc_config) self.generator._register_load_state_dict_pre_hook(self._pre_load_hook) self.discriminator._register_load_state_dict_pre_hook(self._pre_load_hook) self.share_embedding = getattr(config, 'embedding_sharing', "none").lower() if self.share_embedding == 'gdes': # Gradient-disentangled weight/embedding sharing word_bias = torch.zeros_like(self.discriminator.deberta.embeddings.word_embeddings.weight) word_bias = torch.nn.Parameter(word_bias) position_bias = torch.zeros_like(self.discriminator.deberta.embeddings.position_embeddings.weight) position_bias = torch.nn.Parameter(position_bias) delattr(self.discriminator.deberta.embeddings.word_embeddings, 'weight') self.discriminator.deberta.embeddings.word_embeddings.register_parameter('_weight', word_bias) delattr(self.discriminator.deberta.embeddings.position_embeddings, 'weight') self.discriminator.deberta.embeddings.position_embeddings.register_parameter('_weight', position_bias) self.register_discriminator_fw_hook() def _pre_load_hook(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): new_state = dict() bert_prefix = prefix + 'bert.' deberta_prefix = prefix + 'deberta.' for k in list(state_dict.keys()): if k.startswith(bert_prefix): nk = deberta_prefix + k[len(bert_prefix):] value = state_dict[k] del state_dict[k] state_dict[nk] = value def forward(self, **kwargs): return self.generator_fw(**kwargs) def discriminator_fw(self, **kwargs): return self.discriminator(**kwargs) def generator_fw(self, **kwargs): return self.generator(**kwargs) def topk_sampling(self, logits, topk = 1, start=0, temp=1): top_p = torch.nn.functional.softmax(logits/temp, dim=-1) topk = max(1, topk) next_tokens = torch.multinomial(top_p, topk) return next_tokens, top_p def make_electra_data(self, input_data, temp=1, rand=None): new_data = input_data.copy() if rand is None: rand = random gen = self.generator_fw(**new_data) lm_logits = gen['logits'] lm_labels = input_data['labels'] lm_loss = gen['loss'] mask_index = (lm_labels.view(-1)>0).nonzero().view(-1) gen_pred = torch.argmax(lm_logits, dim=1).detach().cpu().numpy() topk_labels, top_p = self.topk_sampling(lm_logits, topk=1, temp=temp) top_ids = torch.zeros_like(lm_labels.view(-1)) top_ids.scatter_(index=mask_index, src=topk_labels.view(-1).int(), dim=-1) top_ids = top_ids.view(lm_labels.size()) new_ids = torch.where(lm_labels>0, top_ids, input_data['input_ids']) new_data['input_ids'] = new_ids.detach() return new_data, lm_loss, gen def register_discriminator_fw_hook(self, *wargs): def fw_hook(module, *inputs): if self.share_embedding == 'gdes': # Gradient-disentangled weight/embedding sharing g_w_ebd = self.generator.deberta.embeddings.word_embeddings d_w_ebd = self.discriminator.deberta.embeddings.word_embeddings self._set_param(d_w_ebd, 'weight', g_w_ebd.weight.detach() + d_w_ebd._weight) g_p_ebd = self.generator.deberta.embeddings.position_embeddings d_p_ebd = self.discriminator.deberta.embeddings.position_embeddings self._set_param(d_p_ebd, 'weight', g_p_ebd.weight.detach() + d_p_ebd._weight) elif self.share_embedding == 'es': # vallina embedding sharing g_w_ebd = self.generator.deberta.embeddings.word_embeddings d_w_ebd = self.discriminator.deberta.embeddings.word_embeddings self._set_param(d_w_ebd, 'weight', g_w_ebd.weight) g_p_ebd = self.generator.deberta.embeddings.position_embeddings d_p_ebd = self.discriminator.deberta.embeddings.position_embeddings self._set_param(d_p_ebd, 'weight', g_p_ebd.weight) return None self.discriminator.register_forward_pre_hook(fw_hook) @staticmethod def _set_param(module, param_name, value): if hasattr(module, param_name): delattr(module, param_name) module.register_buffer(param_name, value) @register_task(name="RTD", desc="Replaced token detection pretraining task") class RTDTask(Task): def __init__(self, data_dir, tokenizer, args, **kwargs): super().__init__(tokenizer, args, **kwargs) self.data_dir = data_dir self.mask_gen = NGramMaskGenerator(tokenizer, max_gram=1, keep_prob = 0, mask_prob = 1, max_seq_len = args.max_seq_length) def train_data(self, max_seq_len=512, **kwargs): data = self.load_data(os.path.join(self.data_dir, 'train.txt')) examples = ExampleSet(data) if self.args.num_training_steps is None: dataset_size = len(examples) else: dataset_size = self.args.num_training_steps*self.args.train_batch_size return DynamicDataset(examples, feature_fn = self.get_feature_fn(max_seq_len=max_seq_len, mask_gen=self.mask_gen), \ dataset_size = dataset_size, shuffle=True, **kwargs) def get_labels(self): return list(self.tokenizer.vocab.values()) def eval_data(self, max_seq_len=512, **kwargs): ds = [ self._data('dev', 'valid.txt', 'dev'), ] for d in ds: _size = len(d.data) d.data = DynamicDataset(d.data, feature_fn = self.get_feature_fn(max_seq_len=max_seq_len, mask_gen=self.mask_gen), dataset_size = _size, **kwargs) return ds def test_data(self, max_seq_len=512, **kwargs): """See base class.""" raise NotImplemented('This method is not implemented yet.') def _data(self, name, path, type_name = 'dev', ignore_metric=False): if isinstance(path, str): path = [path] data = [] for p in path: input_src = os.path.join(self.data_dir, p) assert os.path.exists(input_src), f"{input_src} doesn't exists" data.extend(self.load_data(input_src)) predict_fn = self.get_predict_fn() examples = ExampleSet(data) return EvalData(name, examples, metrics_fn = self.get_metrics_fn(), predict_fn = predict_fn, ignore_metric=ignore_metric, critial_metrics=['accuracy']) def get_metrics_fn(self): """Calcuate metrics based on prediction results""" def metrics_fn(logits, labels): preds = logits acc = (preds==labels).sum()/len(labels) metrics = OrderedDict(accuracy= acc) return metrics return metrics_fn def load_data(self, path): examples = [] with open(path, encoding='utf-8') as fs: for l in fs: if len(l) > 1: example = ExampleInstance(segments=[l]) examples.append(example) return examples def get_feature_fn(self, max_seq_len = 512, mask_gen = None): def _example_to_feature(example, rng=None, ext_params=None, **kwargs): return self.example_to_feature(self.tokenizer, example, max_seq_len = max_seq_len, \ rng = rng, mask_generator = mask_gen, ext_params = ext_params, **kwargs) return _example_to_feature def example_to_feature(self, tokenizer, example, max_seq_len=512, rng=None, mask_generator = None, ext_params=None, **kwargs): if not rng: rng = random max_num_tokens = max_seq_len - 2 segments = [ example.segments[0].strip().split() ] segments = _truncate_segments(segments, max_num_tokens, rng) _tokens = ['[CLS]'] + segments[0] + ['[SEP]'] if mask_generator: tokens, lm_labels = mask_generator.mask_tokens(_tokens, rng) token_ids = tokenizer.convert_tokens_to_ids(tokens) features = OrderedDict(input_ids = token_ids, position_ids = list(range(len(token_ids))), input_mask = [1]*len(token_ids), labels = lm_labels) for f in features: features[f] = torch.tensor(features[f] + [0]*(max_seq_len - len(token_ids)), dtype=torch.int) return features def get_model_class_fn(self): def partial_class(*wargs, **kwargs): model = RTDModel.load_model(*wargs, **kwargs) if self.args.init_generator is not None: logger.info(f'Load generator from {self.args.init_generator}') generator = torch.load(self.args.init_generator, map_location='cpu') missing_keys, unexpected_keys = model.generator.load_state_dict(generator, strict=False) if missing_keys and (len(missing_keys) > 0): logger.warning(f'Load generator with missing keys: {missing_keys}') if unexpected_keys and (len(unexpected_keys) > 0): logger.warning(f'Load generator with unexptected keys: {unexpected_keys}') if self.args.init_discriminator is not None: logger.info(f'Load discriminator from {self.args.init_discriminator}') discriminator = torch.load(self.args.init_discriminator, map_location='cpu') missing_keys, unexpected_keys = model.discriminator.load_state_dict(discriminator, strict=False) if missing_keys and (len(missing_keys) > 0): logger.warning(f'Load discriminator with missing keys: {missing_keys}') if unexpected_keys and (len(unexpected_keys) > 0): logger.warning(f'Load discriminator with unexptected keys: {unexpected_keys}') return model return partial_class def get_train_fn(self, args, model): def train_fn(args, model, device, data_fn, eval_fn, loss_fn): if args.decoupled_training: gen_args = copy.deepcopy(args) gen_args.checkpoint_dir = os.path.join(gen_args.output_dir, 'generator') os.makedirs(gen_args.checkpoint_dir, exist_ok=True) with open(os.path.join(gen_args.checkpoint_dir, 'model_config.json'), 'w') as fs: fs.write(model.config.generator.to_json_string() + '\n') shutil.copy(args.vocab_path, gen_args.checkpoint_dir) loss_fn = self.get_decoupled_loss_fn(args, model, data_fn, device, args.num_training_steps) trainer = DistributedTrainer(gen_args, gen_args.output_dir, model.generator, device, data_fn, loss_fn = loss_fn, eval_fn = eval_fn, dump_interval = args.dump_interval, name='G') else: trainer = DistributedTrainer(args, args.output_dir, model, device, data_fn, loss_fn = loss_fn, eval_fn = eval_fn, dump_interval = args.dump_interval) trainer.train() return train_fn def get_eval_fn(self): def eval_fn(args, model, device, eval_data, prefix=None, tag=None, steps=None): # Run prediction for full data prefix = f'{tag}_{prefix}' if tag is not None else prefix eval_results=OrderedDict() eval_metric=0 no_tqdm = (True if os.getenv('NO_TQDM', '0')!='0' else False) or args.rank>0 for eval_item in eval_data: name = eval_item.name eval_sampler = SequentialSampler(len(eval_item.data)) batch_sampler = BatchSampler(eval_sampler, args.eval_batch_size) batch_sampler = DistributedBatchSampler(batch_sampler, rank=args.rank, world_size=args.world_size) eval_dataloader = DataLoader(eval_item.data, batch_sampler=batch_sampler, num_workers=args.workers) model.eval() eval_loss, eval_accuracy = 0, 0 nb_eval_steps, nb_eval_examples = 0, 0 predicts=[] labels=[] for batch in tqdm(AsyncDataLoader(eval_dataloader), ncols=80, desc='Evaluating: {}'.format(prefix), disable=no_tqdm): batch = batch_to(batch, device) with torch.no_grad(): output = model(**batch) logits = output['logits'].detach().argmax(dim=-1) tmp_eval_loss = output['loss'].detach() if 'labels' in output: label_ids = output['labels'].detach().to(device) else: label_ids = batch['labels'].to(device) predicts.append(logits) labels.append(label_ids) eval_loss += tmp_eval_loss.mean() input_ids = batch['input_ids'] nb_eval_examples += input_ids.size(0) nb_eval_steps += 1 eval_loss = eval_loss / nb_eval_steps predicts = merge_distributed(predicts) labels = merge_distributed(labels) result=OrderedDict() metrics_fn = eval_item.metrics_fn metrics = metrics_fn(predicts.numpy(), labels.numpy()) result.update(metrics) result['perplexity'] = torch.exp(eval_loss).item() critial_metrics = set(metrics.keys()) if eval_item.critial_metrics is None or len(eval_item.critial_metrics)==0 else eval_item.critial_metrics eval_metric = np.mean([v for k,v in metrics.items() if k in critial_metrics]) result['eval_loss'] = eval_loss.item() result['eval_metric'] = eval_metric result['eval_samples'] = len(labels) if args.rank<=0: logger.info("***** Eval results-{}-{} *****".format(name, prefix)) for key in sorted(result.keys()): logger.info(" %s = %s", key, str(result[key])) eval_results[name]=(eval_metric, predicts, labels) return eval_results return eval_fn def get_decoupled_loss_fn(self, args, model, data_fn, device, num_training_steps): rand = random.Random(0) def eval_fn(trainer, model, device, tag): return 0 def d_loss_fn(trainer, model, data): train_losses = OrderedDict() with_mlm_loss = True disc = model(**data) rtd_loss = disc['loss'] loss = args.rtd_lambda*rtd_loss.mean() return loss, data['input_ids'].size(0) disc_args = copy.deepcopy(args) disc_args.checkpoint_dir = os.path.join(disc_args.output_dir, 'discriminator') os.makedirs(disc_args.checkpoint_dir, exist_ok=True) with open(os.path.join(disc_args.checkpoint_dir, 'model_config.json'), 'w') as fs: fs.write(model.config.discriminator.to_json_string() + '\n') shutil.copy(args.vocab_path, disc_args.checkpoint_dir) if disc_args.discriminator_learning_rate > 0: disc_args.learning_rate = disc_args.discriminator_learning_rate disc_trainer = DistributedTrainer(disc_args, args.output_dir, model.discriminator, device, data_fn, loss_fn = d_loss_fn, eval_fn = eval_fn, dump_interval = args.dump_interval, name='D') disc_trainer.initialize() def post_g_loss_fn(outputs): if outputs is None or len(outputs) == 0: return None datas = [o['new_data'] for o in outputs] new_data = defaultdict(list) for d in datas: for k in d: new_data[k].append(d[k]) for k in new_data: new_data[k] = torch.cat(new_data[k], dim=0) disc_trainer._train_step(new_data, 1) def g_loss_fn(trainer, _model, data): new_data, mlm_loss, gen_output = model.make_electra_data(data, rand=rand) trainer.post_loss_fn = post_g_loss_fn loss = mlm_loss.mean() return {'loss': loss.mean(), 'batch_size': data['input_ids'].size(0), 'new_data': new_data} return g_loss_fn def get_loss_fn(self, args): rand = random.Random(0) def loss_fn(trainer, model, data): train_losses = OrderedDict() new_data, mlm_loss, gen_output = model.make_electra_data(data, rand=rand) disc = model.discriminator_fw(**new_data) rtd_loss = disc['loss'] loss = mlm_loss.mean() + args.rtd_lambda*rtd_loss.mean() return loss.mean(), data['input_ids'].size(0) return loss_fn @classmethod def add_arguments(cls, parser): """Add task specific arguments e.g. parser.add_argument('--data_dir', type=str, help='The path of data directory.') """ parser.add_argument('--rtd_lambda', type=float, default=10, help='Weight of RTD loss') parser.add_argument('--decoupled_training', type=boolean_string, default=True, help='Whether to use decoupled training') parser.add_argument('--num_training_steps', type=int, default=None, help='Maxium pre-training steps') parser.add_argument('--discriminator_learning_rate', type=float, default=-1, help='The learning rate of the discriminator') parser.add_argument('--init_generator', type=str, default=None, help='The model that used to initialize the generator') parser.add_argument('--init_discriminator', type=str, default=None, help='The model that used to initialize the discriminator')