From c9bce1811ce8d63f2cd2f28b47ec9cc2196384e7 Mon Sep 17 00:00:00 2001 From: thomwolf Date: Wed, 28 Aug 2019 13:22:45 +0200 Subject: [PATCH] fixing model to add torchscript, embedding resizing, head pruning and masking + tests --- pytorch_transformers/modeling_bert.py | 2 +- pytorch_transformers/modeling_dilbert.py | 371 ++++++++++++------ .../tests/modeling_dilbert_test.py | 18 +- 3 files changed, 253 insertions(+), 138 deletions(-) diff --git a/pytorch_transformers/modeling_bert.py b/pytorch_transformers/modeling_bert.py index badec992c3..560c4f1086 100644 --- a/pytorch_transformers/modeling_bert.py +++ b/pytorch_transformers/modeling_bert.py @@ -449,7 +449,7 @@ class BertEncoder(nn.Module): outputs = outputs + (all_hidden_states,) if self.output_attentions: outputs = outputs + (all_attentions,) - return outputs # outputs, (hidden states), (attentions) + return outputs # last-layer hidden state, (all hidden states), (all attentions) class BertPooler(nn.Module): diff --git a/pytorch_transformers/modeling_dilbert.py b/pytorch_transformers/modeling_dilbert.py index 2f3ea1c535..867ba0e6a8 100644 --- a/pytorch_transformers/modeling_dilbert.py +++ b/pytorch_transformers/modeling_dilbert.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2019-present, the HuggingFace Inc. team. +# Copyright 2019-present, the HuggingFace Inc. team and Facebook, Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -30,7 +30,7 @@ import numpy as np import torch import torch.nn as nn -from pytorch_transformers.modeling_utils import PretrainedConfig, PreTrainedModel, add_start_docstrings +from pytorch_transformers.modeling_utils import PretrainedConfig, PreTrainedModel, add_start_docstrings, prune_linear_layer import logging logger = logging.getLogger(__name__) @@ -92,6 +92,17 @@ class DilBertConfig(PretrainedConfig): else: raise ValueError("First argument must be either a vocabulary size (int)" " or the path to a pretrained model config file (str)") + @property + def hidden_size(self): + return self.hidden_dim + + @property + def num_attention_heads(self): + return self.n_heads + + @property + def num_hidden_layers(self): + return self.n_layers ### UTILS AND BUILDING BLOCKS OF THE ARCHITECTURE ### @@ -163,11 +174,30 @@ class MultiHeadSelfAttention(nn.Module): self.v_lin = nn.Linear(in_features=config.dim, out_features=config.dim) self.out_lin = nn.Linear(in_features=config.dim, out_features=config.dim) + def prune_heads(self, heads): + attention_head_size = self.dim // self.n_heads + if len(heads) == 0: + return + mask = torch.ones(self.n_heads, attention_head_size) + for head in heads: + mask[head] = 0 + mask = mask.view(-1).contiguous().eq(1) + index = torch.arange(len(mask))[mask].long() + # Prune linear layers + self.q_lin = prune_linear_layer(self.q_lin, index) + self.k_lin = prune_linear_layer(self.k_lin, index) + self.v_lin = prune_linear_layer(self.v_lin, index) + self.out_lin = prune_linear_layer(self.out_lin, index, dim=1) + # Update hyper params + self.n_heads = self.n_heads - len(heads) + self.dim = attention_head_size * self.n_heads + def forward(self, query: torch.tensor, key: torch.tensor, value: torch.tensor, - mask: torch.tensor): + mask: torch.tensor, + head_mask: torch.tensor = None): """ Parameters ---------- @@ -185,10 +215,10 @@ class MultiHeadSelfAttention(nn.Module): """ bs, q_length, dim = query.size() k_length = key.size(1) - assert dim == self.dim, 'Dimensions do not match: %s input vs %s configured' % (dim, self.dim) - assert key.size() == value.size() + # assert dim == self.dim, 'Dimensions do not match: %s input vs %s configured' % (dim, self.dim) + # assert key.size() == value.size() - dim_per_head = dim // self.n_heads + dim_per_head = self.dim // self.n_heads assert 2 <= mask.dim() <= 3 causal = (mask.dim() == 3) @@ -200,7 +230,7 @@ class MultiHeadSelfAttention(nn.Module): def unshape(x): """ group heads """ - return x.transpose(1, 2).contiguous().view(bs, -1, dim) + return x.transpose(1, 2).contiguous().view(bs, -1, self.n_heads * dim_per_head) q = shape(self.q_lin(query)) # (bs, n_heads, q_length, dim_per_head) k = shape(self.k_lin(key)) # (bs, n_heads, k_length, dim_per_head) @@ -213,6 +243,11 @@ class MultiHeadSelfAttention(nn.Module): weights = nn.Softmax(dim=-1)(scores) # (bs, n_heads, q_length, k_length) weights = self.dropout(weights) # (bs, n_heads, q_length, k_length) + + # Mask heads if we want to + if head_mask is not None: + weights = weights * head_mask + context = torch.matmul(weights, v) # (bs, n_heads, q_length, dim_per_head) context = unshape(context) # (bs, q_length, dim) context = self.out_lin(context) # (bs, q_length, dim) @@ -229,7 +264,7 @@ class FFN(nn.Module): self.dropout = nn.Dropout(p=config.dropout) self.lin1 = nn.Linear(in_features=config.dim, out_features=config.hidden_dim) self.lin2 = nn.Linear(in_features=config.hidden_dim, out_features=config.dim) - assert config.activation in ['relu', 'gelu'], ValueError(f"activation ({config.activation}) must be in ['relu', 'gelu']") + assert config.activation in ['relu', 'gelu'], "activation ({}) must be in ['relu', 'gelu']".format(config.activation) self.activation = gelu if config.activation == 'gelu' else nn.ReLU() def forward(self, @@ -262,7 +297,8 @@ class TransformerBlock(nn.Module): def forward(self, x: torch.tensor, - attn_mask: torch.tensor = None): + attn_mask: torch.tensor = None, + head_mask: torch.tensor = None): """ Parameters ---------- @@ -277,7 +313,7 @@ class TransformerBlock(nn.Module): The output of the transformer block contextualization. """ # Self-Attention - sa_output = self.attention(query=x, key=x, value=x, mask=attn_mask) + sa_output = self.attention(query=x, key=x, value=x, mask=attn_mask, head_mask=head_mask) if self.output_attentions: sa_output, sa_weights = sa_output # (bs, seq_length, dim), (bs, n_heads, seq_length, seq_length) else: # To handle these `output_attention` or `output_hidden_states` cases returning tuples @@ -294,6 +330,7 @@ class TransformerBlock(nn.Module): output = (sa_weights,) + output return output + class Transformer(nn.Module): def __init__(self, config): @@ -307,7 +344,8 @@ class Transformer(nn.Module): def forward(self, x: torch.tensor, - attn_mask: torch.tensor = None): + attn_mask: torch.tensor = None, + head_mask: torch.tensor = None): """ Parameters ---------- @@ -331,14 +369,24 @@ class Transformer(nn.Module): all_attentions = () hidden_state = x - for _, layer_module in enumerate(self.layer): - hidden_state = layer_module(x=hidden_state, attn_mask=attn_mask) + for i, layer_module in enumerate(self.layer): + if self.output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_state,) + + layer_outputs = layer_module(x=hidden_state, + attn_mask=attn_mask, + head_mask=head_mask[i]) + hidden_state = layer_outputs[-1] + if self.output_attentions: - attentions, hidden_state = hidden_state + assert len(layer_outputs) == 2 + attentions = layer_outputs[0] all_attentions = all_attentions + (attentions,) - else: # To handle these `output_attention` or `output_hidden_states` cases returning tuples - assert type(hidden_state) == tuple - hidden_state = hidden_state[0] + else: + assert len(layer_outputs) == 1 + + # Add last layer + if self.output_hidden_states: all_hidden_states = all_hidden_states + (hidden_state,) outputs = (hidden_state,) @@ -346,7 +394,7 @@ class Transformer(nn.Module): outputs = outputs + (all_hidden_states,) if self.output_attentions: outputs = outputs + (all_attentions,) - return outputs + return outputs # last-layer hidden state, (all hidden states), (all attentions) ### INTERFACE FOR ENCODER AND TASK SPECIFIC MODEL ### @@ -378,9 +426,21 @@ class DilBertPreTrainedModel(PreTrainedModel): DILBERT_START_DOCSTRING = r""" - Smaller, faster, cheaper, lighter: DilBERT + DilBERT is a small, fast, cheap and light Transformer model + trained by distilling Bert base. It has 40% less parameters than + `bert-base-uncased`, runs 60% faster while preserving over 95% of + Bert's performances as measured on the GLUE language understanding benchmark. - For more information on DilBERT, you should check TODO(Link): Link to Medium + Here are the differences between the interface of Bert and DilBert: + + - DilBert doesn't have `token_type_ids`, you don't need to indicate which token belong to which segment. Just separate your segments with the separation token `tokenizer.sep_token` (or `[SEP]`) + - DilBert doesn't have options to select the input positions (`position_ids` input). This could be added if necessary though, just let's us know if you need this option. + + For more information on DilBERT, please refer to our + `detailed blog post`_ + + .. _`detailed blog post`: + https://medium.com/huggingface/smaller-faster-cheaper-lighter-introducing-dilbert-a-distilled-version-of-bert-8cf3380435b5 Parameters: config (:class:`~pytorch_transformers.DilBertConfig`): Model configuration class with all the parameters of the model. @@ -399,31 +459,35 @@ DILBERT_INPUTS_DOCSTRING = r""" Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``: ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens. + **head_mask**: (`optional`) ``torch.FloatTensor`` of shape ``(num_heads,)`` or ``(num_layers, num_heads)``: + Mask to nullify selected heads of the self-attention modules. + Mask values selected in ``[0, 1]``: + ``1`` indicates the head is **not masked**, ``0`` indicates the head is **masked**. """ @add_start_docstrings("The bare DilBERT encoder/transformer outputing raw hidden-states without any specific head on top.", DILBERT_START_DOCSTRING, DILBERT_INPUTS_DOCSTRING) class DilBertModel(DilBertPreTrainedModel): r""" - Parameters - ---------- - input_ids: torch.tensor(bs, seq_length) - Sequences of token ids. - attention_mask: torch.tensor(bs, seq_length) - Attention mask on the sequences. Optional: If None, it's like there was no padding. - - Outputs - ------- - hidden_state: torch.tensor(bs, seq_length, dim) - Sequence of hiddens states in the last (top) layer - pooled_output: torch.tensor(bs, dim) - Pooled output: for DilBert, the pooled output is simply the hidden state of the [CLS] token. - all_hidden_states: Tuple[torch.tensor(bs, seq_length, dim)] - Tuple of length n_layers with the hidden states from each layer. - Optional: only if output_hidden_states=True - all_attentions: Tuple[torch.tensor(bs, n_heads, seq_length, seq_length)] - Tuple of length n_layers with the attention weights from each layer - Optional: only if output_attentions=True + Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs: + **last_hidden_state**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, hidden_size)`` + Sequence of hidden-states at the output of the last layer of the model. + **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``) + list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings) + of shape ``(batch_size, sequence_length, hidden_size)``: + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + **attentions**: (`optional`, returned when ``config.output_attentions=True``) + list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``: + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. + + Examples:: + + tokenizer = DilBertTokenizer.from_pretrained('dilbert-base-uncased') + model = DilBertModel.from_pretrained('dilbert-base-uncased') + input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0) # Batch size 1 + outputs = model(input_ids) + last_hidden_states = outputs[0] # The last hidden-state is the first element of the output tuple + """ def __init__(self, config): super(DilBertModel, self).__init__(config) @@ -433,47 +497,83 @@ class DilBertModel(DilBertPreTrainedModel): self.apply(self.init_weights) + def _resize_token_embeddings(self, new_num_tokens): + old_embeddings = self.embeddings.word_embeddings + new_embeddings = self._get_resized_embeddings(old_embeddings, new_num_tokens) + self.embeddings.word_embeddings = new_embeddings + return self.embeddings.word_embeddings + + def _prune_heads(self, heads_to_prune): + """ Prunes heads of the model. + heads_to_prune: dict of {layer_num: list of heads to prune in this layer} + See base class PreTrainedModel + """ + for layer, heads in heads_to_prune.items(): + self.transformer.layer[layer].attention.prune_heads(heads) + def forward(self, input_ids: torch.tensor, - attention_mask: torch.tensor = None): + attention_mask: torch.tensor = None, + head_mask: torch.tensor = None): if attention_mask is None: attention_mask = torch.ones_like(input_ids) # (bs, seq_length) + # Prepare head mask if needed + # 1.0 in head_mask indicate we keep the head + # attention_probs has shape bsz x n_heads x N x N + # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] + # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length] + if head_mask is not None: + if head_mask.dim() == 1: + head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1) + head_mask = head_mask.expand(self.config.num_hidden_layers, -1, -1, -1, -1) + elif head_mask.dim() == 2: + head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1) # We can specify head_mask for each layer + head_mask = head_mask.to(dtype=next(self.parameters()).dtype) # switch to fload if need + fp16 compatibility + else: + head_mask = [None] * self.config.num_hidden_layers + embedding_output = self.embeddings(input_ids) # (bs, seq_length, dim) tfmr_output = self.transformer(x=embedding_output, - attn_mask=attention_mask) + attn_mask=attention_mask, + head_mask=head_mask) hidden_state = tfmr_output[0] - pooled_output = hidden_state[:, 0] - output = (hidden_state, pooled_output) + tfmr_output[1:] + output = (hidden_state, ) + tfmr_output[1:] + + return output # last-layer hidden-state, (all hidden_states), (all attentions) - return output # hidden_state, pooled_output, (hidden_states), (attentions) @add_start_docstrings("""DilBert Model with a `masked language modeling` head on top. """, DILBERT_START_DOCSTRING, DILBERT_INPUTS_DOCSTRING) class DilBertForMaskedLM(DilBertPreTrainedModel): r""" - Parameters - ---------- - input_ids: torch.tensor(bs, seq_length) - Token ids. - attention_mask: torch.tensor(bs, seq_length) - Attention mask. Optional: If None, it's like there was no padding. - masked_lm_labels: torch.tensor(bs, seq_length) - The masked language modeling labels. Optional: If None, no loss is computed. + **masked_lm_labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``: + Labels for computing the masked language modeling loss. + Indices should be in ``[-1, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring) + Tokens with indices set to ``-1`` are ignored (masked), the loss is only computed for the tokens with labels + in ``[0, ..., config.vocab_size]`` + + Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs: + **loss**: (`optional`, returned when ``masked_lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``: + Masked language modeling loss. + **prediction_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, config.vocab_size)`` + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``) + list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings) + of shape ``(batch_size, sequence_length, hidden_size)``: + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + **attentions**: (`optional`, returned when ``config.output_attentions=True``) + list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``: + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. + + Examples:: + + tokenizer = DilBertTokenizer.from_pretrained('dilbert-base-uncased') + model = DilBertForMaskedLM.from_pretrained('dilbert-base-uncased') + input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0) # Batch size 1 + outputs = model(input_ids, masked_lm_labels=input_ids) + loss, prediction_scores = outputs[:2] - Outputs - ------- - mlm_loss: torch.tensor(1,) - Masked Language Modeling loss to optimize. - Optional: only if `masked_lm_labels` is not None - prediction_logits: torch.tensor(bs, seq_length, voc_size) - Token prediction logits - all_hidden_states: Tuple[torch.tensor(bs, seq_length, dim)] - Tuple of length n_layers with the hidden states from each layer. - Optional: only if `output_hidden_states`=True - all_attentions: Tuple[torch.tensor(bs, n_heads, seq_length, seq_length)] - Tuple of length n_layers with the attention weights from each layer - Optional: only if `output_attentions`=True """ def __init__(self, config): super(DilBertForMaskedLM, self).__init__(config) @@ -491,59 +591,68 @@ class DilBertForMaskedLM(DilBertPreTrainedModel): self.mlm_loss_fct = nn.CrossEntropyLoss(ignore_index=-1) def tie_weights(self): + """ Make sure we are sharing the input and output embeddings. + Export to TorchScript can't handle parameter sharing so we are cloning them instead. """ - Tying the weights of the vocabulary projection to the base token embeddings. - """ - if self.config.tie_weights_: - self.vocab_projector.weight = self.dilbert.embeddings.word_embeddings.weight + self._tie_or_clone_weights(self.vocab_projector, + self.dilbert.embeddings.word_embeddings) def forward(self, input_ids: torch.tensor, attention_mask: torch.tensor = None, - masked_lm_labels: torch.tensor = None): + masked_lm_labels: torch.tensor = None, + head_mask: torch.tensor = None): dlbrt_output = self.dilbert(input_ids=input_ids, - attention_mask=attention_mask) + attention_mask=attention_mask, + head_mask=head_mask) hidden_states = dlbrt_output[0] # (bs, seq_length, dim) prediction_logits = self.vocab_transform(hidden_states) # (bs, seq_length, dim) prediction_logits = gelu(prediction_logits) # (bs, seq_length, dim) prediction_logits = self.vocab_layer_norm(prediction_logits) # (bs, seq_length, dim) prediction_logits = self.vocab_projector(prediction_logits) # (bs, seq_length, vocab_size) - outputs = (prediction_logits, ) + dlbrt_output[2:] + outputs = (prediction_logits, ) + dlbrt_output[1:] if masked_lm_labels is not None: mlm_loss = self.mlm_loss_fct(prediction_logits.view(-1, prediction_logits.size(-1)), masked_lm_labels.view(-1)) outputs = (mlm_loss,) + outputs - return outputs # (mlm_loss), prediction_logits, (hidden_states), (attentions) + return outputs # (mlm_loss), prediction_logits, (all hidden_states), (all attentions) + @add_start_docstrings("""DilBert Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled output) e.g. for GLUE tasks. """, DILBERT_START_DOCSTRING, DILBERT_INPUTS_DOCSTRING) class DilBertForSequenceClassification(DilBertPreTrainedModel): r""" - Parameters - ---------- - input_ids: torch.tensor(bs, seq_length) - Token ids. - attention_mask: torch.tensor(bs, seq_length) - Attention mask. Optional: If None, it's like there was no padding. - labels: torch.tensor(bs,) - Classification Labels: Optional: If None, no loss will be computed. - - Outputs - ------- - loss: torch.tensor(1) - Sequence classification loss. - Optional: Is computed only if `labels` is not None. - logits: torch.tensor(bs, seq_length) - Classification (or regression if config.num_labels==1) scores - all_hidden_states: Tuple[torch.tensor(bs, seq_length, dim)] - Tuple of length n_layers with the hidden states from each layer. - Optional: only if `output_hidden_states`=True - all_attentions: Tuple[torch.tensor(bs, n_heads, seq_length, seq_length)] - Tuple of length n_layers with the attention weights from each layer - Optional: only if `output_attentions`=True + **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``: + Labels for computing the sequence classification/regression loss. + Indices should be in ``[0, ..., config.num_labels - 1]``. + If ``config.num_labels == 1`` a regression loss is computed (Mean-Square loss), + If ``config.num_labels > 1`` a classification loss is computed (Cross-Entropy). + + Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs: + **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``: + Classification (or regression if config.num_labels==1) loss. + **logits**: ``torch.FloatTensor`` of shape ``(batch_size, config.num_labels)`` + Classification (or regression if config.num_labels==1) scores (before SoftMax). + **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``) + list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings) + of shape ``(batch_size, sequence_length, hidden_size)``: + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + **attentions**: (`optional`, returned when ``config.output_attentions=True``) + list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``: + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. + + Examples:: + + tokenizer = DilBertTokenizer.from_pretrained('dilbert-base-uncased') + model = DilBertForSequenceClassification.from_pretrained('dilbert-base-uncased') + input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0) # Batch size 1 + labels = torch.tensor([1]).unsqueeze(0) # Batch size 1 + outputs = model(input_ids, labels=labels) + loss, logits = outputs[:2] + """ def __init__(self, config): super(DilBertForSequenceClassification, self).__init__(config) @@ -559,16 +668,19 @@ class DilBertForSequenceClassification(DilBertPreTrainedModel): def forward(self, input_ids: torch.tensor, attention_mask: torch.tensor = None, - labels: torch.tensor = None): + labels: torch.tensor = None, + head_mask: torch.tensor = None): dilbert_output = self.dilbert(input_ids=input_ids, - attention_mask=attention_mask) - pooled_output = dilbert_output[1] # (bs, dim) + attention_mask=attention_mask, + head_mask=head_mask) + hidden_state = dilbert_output[0] # (bs, seq_len, dim) + pooled_output = hidden_state[:, 0] # (bs, dim) pooled_output = self.pre_classifier(pooled_output) # (bs, dim) pooled_output = nn.ReLU()(pooled_output) # (bs, dim) pooled_output = self.dropout(pooled_output) # (bs, dim) logits = self.classifier(pooled_output) # (bs, dim) - outputs = (logits,) + dilbert_output[2:] + outputs = (logits,) + dilbert_output[1:] if labels is not None: if self.num_labels == 1: loss_fct = nn.MSELoss() @@ -580,43 +692,46 @@ class DilBertForSequenceClassification(DilBertPreTrainedModel): return outputs # (loss), logits, (hidden_states), (attentions) + @add_start_docstrings("""DilBert Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear layers on top of the hidden-states output to compute `span start logits` and `span end logits`). """, DILBERT_START_DOCSTRING, DILBERT_INPUTS_DOCSTRING) class DilBertForQuestionAnswering(DilBertPreTrainedModel): r""" - Parameters - ---------- - input_ids: torch.tensor(bs, seq_length) - Token ids. - attention_mask: torch.tensor(bs, seq_length) - Attention mask. Optional: If None, it's like there was no padding. - start_positions: torch,tensor(bs) + **start_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``: Labels for position (index) of the start of the labelled span for computing the token classification loss. Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence are not taken into account for computing the loss. - Optional: if None, no loss is computed. - end_positions: torch,tensor(bs) + **end_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``: Labels for position (index) of the end of the labelled span for computing the token classification loss. Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence are not taken into account for computing the loss. - Optional: if None, no loss is computed. - Outputs - ------- - loss: torch.tensor(1) - Question answering loss. - Optional: Is computed only if `start_positions` and `end_positions` are not None. - start_logits: torch.tensor(bs, seq_length) - Span-start scores. - end_logits: torch.tensor(bs, seq_length) - Spand-end scores. - all_hidden_states: Tuple[torch.tensor(bs, seq_length, dim)] - Tuple of length n_layers with the hidden states from each layer. - Optional: only if `output_hidden_states`=True - all_attentions: Tuple[torch.tensor(bs, n_heads, seq_length, seq_length)] - Tuple of length n_layers with the attention weights from each layer - Optional: only if `output_attentions`=True + Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs: + **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``: + Total span extraction loss is the sum of a Cross-Entropy for the start and end positions. + **start_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)`` + Span-start scores (before SoftMax). + **end_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)`` + Span-end scores (before SoftMax). + **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``) + list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings) + of shape ``(batch_size, sequence_length, hidden_size)``: + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + **attentions**: (`optional`, returned when ``config.output_attentions=True``) + list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``: + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. + + Examples:: + + tokenizer = DilBertTokenizer.from_pretrained('dilbert-base-uncased') + model = DilBertForQuestionAnswering.from_pretrained('dilbert-base-uncased') + input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0) # Batch size 1 + start_positions = torch.tensor([1]) + end_positions = torch.tensor([3]) + outputs = model(input_ids, start_positions=start_positions, end_positions=end_positions) + loss, start_scores, end_scores = outputs[:2] + """ def __init__(self, config): super(DilBertForQuestionAnswering, self).__init__(config) @@ -632,9 +747,11 @@ class DilBertForQuestionAnswering(DilBertPreTrainedModel): input_ids: torch.tensor, attention_mask: torch.tensor = None, start_positions: torch.tensor = None, - end_positions: torch.tensor = None): + end_positions: torch.tensor = None, + head_mask: torch.tensor = None): dilbert_output = self.dilbert(input_ids=input_ids, - attention_mask=attention_mask) + attention_mask=attention_mask, + head_mask=head_mask) hidden_states = dilbert_output[0] # (bs, max_query_len, dim) hidden_states = self.dropout(hidden_states) # (bs, max_query_len, dim) @@ -643,7 +760,7 @@ class DilBertForQuestionAnswering(DilBertPreTrainedModel): start_logits = start_logits.squeeze(-1) # (bs, max_query_len) end_logits = end_logits.squeeze(-1) # (bs, max_query_len) - outputs = (start_logits, end_logits,) + dilbert_output[2:] + outputs = (start_logits, end_logits,) + dilbert_output[1:] if start_positions is not None and end_positions is not None: # If we are on multi-GPU, split add a dimension if len(start_positions.size()) > 1: diff --git a/pytorch_transformers/tests/modeling_dilbert_test.py b/pytorch_transformers/tests/modeling_dilbert_test.py index 0cbef7e083..2fd707dfd8 100644 --- a/pytorch_transformers/tests/modeling_dilbert_test.py +++ b/pytorch_transformers/tests/modeling_dilbert_test.py @@ -21,7 +21,7 @@ import shutil import pytest from pytorch_transformers import (DilBertConfig, DilBertModel, DilBertForMaskedLM, - DilBertForQuestionAnswering, DilBertForSequenceClassification) + DilBertForQuestionAnswering, DilBertForSequenceClassification) from pytorch_transformers.modeling_dilbert import DILBERT_PRETRAINED_MODEL_ARCHIVE_MAP from .modeling_common_test import (CommonTestCases, ConfigTester, ids_tensor) @@ -31,10 +31,10 @@ class DilBertModelTest(CommonTestCases.CommonModelTester): all_model_classes = (DilBertModel, DilBertForMaskedLM, DilBertForQuestionAnswering, DilBertForSequenceClassification) - test_pruning = False - test_torchscript = False - test_resize_embeddings = False - test_head_masking = False + test_pruning = True + test_torchscript = True + test_resize_embeddings = True + test_head_masking = True class DilBertModelTester(object): @@ -122,22 +122,20 @@ class DilBertModelTest(CommonTestCases.CommonModelTester): def create_and_check_dilbert_model(self, config, input_ids, input_mask, sequence_labels, token_labels, choice_labels): model = DilBertModel(config=config) model.eval() - sequence_output, pooled_output = model(input_ids, input_mask) - sequence_output, pooled_output = model(input_ids) + (sequence_output,) = model(input_ids, input_mask) + (sequence_output,) = model(input_ids) result = { "sequence_output": sequence_output, - "pooled_output": pooled_output, } self.parent.assertListEqual( list(result["sequence_output"].size()), [self.batch_size, self.seq_length, self.hidden_size]) - self.parent.assertListEqual(list(result["pooled_output"].size()), [self.batch_size, self.hidden_size]) def create_and_check_dilbert_for_masked_lm(self, config, input_ids, input_mask, sequence_labels, token_labels, choice_labels): model = DilBertForMaskedLM(config=config) model.eval() - loss, prediction_scores = model(input_ids, input_mask, token_labels) + loss, prediction_scores = model(input_ids, attention_mask=input_mask, masked_lm_labels=token_labels) result = { "loss": loss, "prediction_scores": prediction_scores,