From 3d2096f516e99da79f1c6c60a48f828b4e7733ef Mon Sep 17 00:00:00 2001 From: thomwolf Date: Wed, 18 Dec 2019 11:50:54 +0100 Subject: [PATCH] further cleanup --- examples/run_generation.py | 13 ++++--- transformers/configuration_xlm.py | 4 +++ transformers/modeling_utils.py | 18 ++++++---- transformers/modeling_xlm.py | 39 +++++++------------- transformers/modeling_xlnet.py | 58 +++++++++++++----------------- transformers/tokenization_utils.py | 2 +- 6 files changed, 58 insertions(+), 76 deletions(-) diff --git a/examples/run_generation.py b/examples/run_generation.py index 2075ad8457..8121f4f5aa 100644 --- a/examples/run_generation.py +++ b/examples/run_generation.py @@ -91,7 +91,7 @@ def prepare_ctrl_input(args, _, tokenizer, prompt_text): def prepare_xlm_input(args, model, tokenizer, prompt_text): - kwargs = {"language": None, "mask_token": None} + kwargs = {"language": None, "mask_token_id": None} # Set the language use_lang_emb = hasattr(model.config, "use_lang_emb") and model.config.use_lang_emb @@ -112,7 +112,7 @@ def prepare_xlm_input(args, model, tokenizer, prompt_text): # XLM masked-language modeling (MLM) models need masked token is_xlm_mlm = "mlm" in args.model_name_or_path if is_xlm_mlm: - kwargs["mask_token"] = tokenizer.mask_token_id + kwargs["mask_token_id"] = tokenizer.mask_token_id return prompt_text, kwargs @@ -204,14 +204,13 @@ def main(): prompt_text, model_kwargs = prepare_input(args, model, tokenizer, prompt_text) encoded_prompt = torch.tensor(tokenizer.encode(prompt_text, add_special_tokens=False)).unsqueeze(0) - output_sequences = model.decode( - prompt_ids=encoded_prompt, + output_sequences = model.generate( + intput_ids=encoded_prompt, length=args.length, temperature=args.temperature, - k=args.k, - p=args.p, + top_k=args.k, + top_p=args.p, repetition_penalty=args.repetition_penalty, - device=args.device, **model_kwargs, ) diff --git a/transformers/configuration_xlm.py b/transformers/configuration_xlm.py index fa3a5f40f6..1938b85741 100644 --- a/transformers/configuration_xlm.py +++ b/transformers/configuration_xlm.py @@ -113,6 +113,8 @@ class XLMConfig(PretrainedConfig): summary_first_dropout=0.1, start_n_top=5, end_n_top=5, + mask_token_id = 0, + lang_id = 0, **kwargs): """Constructs XLMConfig. """ @@ -156,6 +158,8 @@ class XLMConfig(PretrainedConfig): self.summary_first_dropout = summary_first_dropout self.start_n_top = start_n_top self.end_n_top = end_n_top + self.mask_token_id = mask_token_id + self.lang_id = lang_id else: raise ValueError("First argument must be either a vocabulary size (int)" " or the path to a pretrained model config file (str)") diff --git a/transformers/modeling_utils.py b/transformers/modeling_utils.py index bbfb0614ad..f55c209ac0 100644 --- a/transformers/modeling_utils.py +++ b/transformers/modeling_utils.py @@ -488,7 +488,7 @@ class PreTrainedModel(nn.Module): def generate(self, input_ids=None, max_length=None, do_sample=None, num_beams=None, temperature=None, top_k=None, top_p=None, repetition_penalty=None, bos_token_id=None, pad_token_id=None, eos_token_ids=None, batch_size=None, - length_penalty=None, num_return_sequences=None, **kwargs): + length_penalty=None, num_return_sequences=None, **model_kwargs): """ Sequence generator for models with a LM head. The method currently supports greedy or penalized greedy decoding, sampling with top-k or nucleus sampling @@ -575,11 +575,13 @@ class PreTrainedModel(nn.Module): output = self._generate_beam_search(input_ids, cur_len, max_length, do_sample, temperature, top_k, top_p, repetition_penalty, pad_token_id, eos_token_ids, effective_batch_size, - length_penalty, num_beams, vocab_size) + length_penalty, num_beams, vocab_size, + **model_kwargs) else: output = self._generate_no_beam_search(input_ids, cur_len, max_length, do_sample, temperature, top_k, top_p, repetition_penalty, - pad_token_id, eos_token_ids, effective_batch_size) + pad_token_id, eos_token_ids, effective_batch_size, + **model_kwargs) if num_return_sequences != 1: output = output.view(batch_size, num_return_sequences, -1) @@ -587,7 +589,8 @@ class PreTrainedModel(nn.Module): def _generate_no_beam_search(self, input_ids, cur_len, max_length, do_sample, temperature, top_k, top_p, repetition_penalty, - pad_token_id, eos_token_ids, batch_size): + pad_token_id, eos_token_ids, batch_size, + **model_kwargs): """ Generate sequences for each example without beam search (num_beams == 1). All returned sequence are generated independantly. """ @@ -598,7 +601,7 @@ class PreTrainedModel(nn.Module): pasts = None while cur_len < max_length: - model_inputs = self.prepare_inputs_for_generation(input_ids, pasts=pasts) + model_inputs = self.prepare_inputs_for_generation(input_ids, pasts=pasts, **model_kwargs) outputs = self(**model_inputs) next_token_logits = outputs[0][:, -1, :] @@ -640,7 +643,8 @@ class PreTrainedModel(nn.Module): def _generate_beam_search(self, input_ids, cur_len, max_length, do_sample, temperature, top_k, top_p, repetition_penalty, pad_token_id, eos_token_ids, batch_size, - length_penalty, num_beams, vocab_size): + length_penalty, num_beams, vocab_size, + **model_kwargs): """ Generate sequences for each example with beam search. """ # Expand input to num beams @@ -662,7 +666,7 @@ class PreTrainedModel(nn.Module): done = [False for _ in range(batch_size)] while cur_len < max_length: - model_inputs = self.prepare_inputs_for_generation(input_ids, pasts=pasts) + model_inputs = self.prepare_inputs_for_generation(input_ids, pasts=pasts, **model_kwargs) scores = self(**model_inputs)[0] # (batch_size * num_beams, cur_len, vocab_size) scores = scores[:, -1, :] # (batch_size * num_beams, vocab_size) diff --git a/transformers/modeling_xlm.py b/transformers/modeling_xlm.py index 295fff7943..6691b0f60b 100644 --- a/transformers/modeling_xlm.py +++ b/transformers/modeling_xlm.py @@ -639,6 +639,18 @@ class XLMWithLMHeadModel(XLMPreTrainedModel): def get_output_embeddings(self): return self.pred_layer.proj + def prepare_inputs_for_generation(self, input_ids, **model_kwargs): + mask_token_id = model_kwargs['mask_token_id'] if 'mask_token_id' in model_kwargs else self.config.mask_token_id + lang_id = model_kwargs['lang_id'] if 'lang_id' in model_kwargs else self.config.lang_id + + mask_token = torch.full((1, 1), mask_token_id, dtype=torch.long, device=input_ids.device) + input_ids = torch.cat([input_ids, mask_token], dim=1) + if lang_id is not None: + langs = torch.full_like(input_ids, lang_id) + else: + langs = None + return {"input_ids": input_ids, "langs": langs} + def forward(self, input_ids=None, attention_mask=None, langs=None, token_type_ids=None, position_ids=None, lengths=None, cache=None, head_mask=None, inputs_embeds=None, labels=None): transformer_outputs = self.transformer(input_ids, @@ -657,33 +669,6 @@ class XLMWithLMHeadModel(XLMPreTrainedModel): return outputs - def _prepare_inputs_for_decoding(self, input_ids, **model_kwargs): - mask_token = model_kwargs.pop("mask_token", None) - language = model_kwargs.pop("language", None) - input_ids = self._append_mask_token(input_ids, mask_token) - langs = self._create_language_embeddings(input_ids, language) - arguments = {"input_ids": input_ids, "langs": langs} - arguments.update(model_kwargs) - - return arguments - - @staticmethod - def _append_mask_token(sequence, mask_token_id): - """ Append a [MASK] token at the end of the sequence that the MLM model - is going to try to predict. - """ - if mask_token_id is not None: - tokens_to_append = torch.full((1, 1), mask_token_id, dtype=torch.long) - return torch.cat((sequence, tokens_to_append), dim=1) - - return sequence - - @staticmethod - def _create_language_embeddings(sequence, language): - if language is not None: - return torch.tensor([language] * sequence.shape[1]).view(1, -1) - return None - @add_start_docstrings("""XLM Model with a sequence classification/regression head on top (a linear layer on top of the pooled output) e.g. for GLUE tasks. """, diff --git a/transformers/modeling_xlnet.py b/transformers/modeling_xlnet.py index 2153923dd2..26b95076cd 100644 --- a/transformers/modeling_xlnet.py +++ b/transformers/modeling_xlnet.py @@ -947,6 +947,30 @@ class XLNetLMHeadModel(XLNetPreTrainedModel): def get_output_embeddings(self): return self.lm_loss + def prepare_inputs_for_generation(self, input_ids, **model_kwargs): + # Add dummy token at the end (no attention on this one) + dummy_token = torch.zeros((1, 1), dtype=torch.long, device=input_ids.device) + input_ids = torch.cat([input_ids, dummy_token], dim=1) + + # Build permutation mask so that previous tokens don't see last token + perm_mask = torch.zeros( + (input_ids.shape[0], input_ids.shape[1], input_ids.shape[1]), + dtype=torch.float, device=input_ids.device + ) + perm_mask[:, :, -1] = 1.0 + + # We'll only predict the last token + target_mapping = torch.zeros( + (input_ids.shape[0], 1, input_ids.shape[1]), + dtype=torch.float, device=input_ids.device + ) + target_mapping[0, 0, -1] = 1.0 + + return {"input_ids": input_ids, + "perm_mask": perm_mask, + "target_mapping": target_mapping + } + def forward(self, input_ids=None, attention_mask=None, mems=None, perm_mask=None, target_mapping=None, token_type_ids=None, input_mask=None, head_mask=None, inputs_embeds=None, labels=None): transformer_outputs = self.transformer(input_ids, @@ -972,40 +996,6 @@ class XLNetLMHeadModel(XLNetPreTrainedModel): return outputs # return (loss), logits, (mems), (hidden states), (attentions) - def _prepare_inputs_for_decoding(self, input_ids, **model_kwargs): - input_ids = self._add_dummy_token(input_ids) - perm_mask = self._create_perm_mask(input_ids) - target_mapping = self._create_target_mapping(input_ids) - arguments = { - "input_ids": input_ids, - "perm_mask": perm_mask, - "target_mapping": target_mapping, - } - return arguments - - @staticmethod - def _add_dummy_token(sequence): - dummy = torch.zeros((sequence.size(0), 1), dtype=torch.long) - return torch.cat((sequence, dummy), dim=1) - - @staticmethod - def _create_perm_mask(sequence): - mask = torch.zeros( - (sequence.shape[0], sequence.shape[1], sequence.shape[1]), - dtype=torch.float, - ) - mask[:, :, -1] = 1.0 # Previous tokens don't see last token - return mask - - @staticmethod - def _create_target_mapping(sequence): - target_mapping = torch.zeros( - (sequence.shape[0], 1, sequence.shape[1]), - dtype=torch.float, - ) - target_mapping[0, 0, -1] = 1.0 # predict last token - return target_mapping - @add_start_docstrings("""XLNet Model with a sequence classification/regression head on top (a linear layer on top of the pooled output) e.g. for GLUE tasks. """, diff --git a/transformers/tokenization_utils.py b/transformers/tokenization_utils.py index f4395cd82c..2e0d6caef2 100644 --- a/transformers/tokenization_utils.py +++ b/transformers/tokenization_utils.py @@ -761,7 +761,7 @@ class PreTrainedTokenizer(object): padding index, up to their max length. If no max length is specified, the padding is done up to the model's max length. The tokenizer padding sides are handled by the following strings: - 'left': pads on the left of the sequences - - 'right': pads on the right of the sequences + - 'right': pads on the right of the sequences Defaults to False: no padding. return_tensors: (optional) can be set to 'tf' or 'pt' to return respectively TensorFlow tf.constant or PyTorch torch.Tensor instead of a list of python integers.