|
from typing import Optional, Tuple, Union |
|
|
|
import torch |
|
import torch.utils.checkpoint |
|
from torch import nn |
|
from torch.nn import CrossEntropyLoss |
|
from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions |
|
from transformers.models.gpt2.configuration_gpt2 import GPT2Config |
|
from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel |
|
|
|
|
|
class GPT2MultiHeadConfig(GPT2Config): |
|
model_type = "gpt2-multi-head" |
|
|
|
def __init__( |
|
self, |
|
head_locations=None, |
|
head_weights=None, |
|
tie_additional_weights=False, |
|
average_logits=False, |
|
*args, |
|
**kwargs, |
|
): |
|
super().__init__(*args, **kwargs) |
|
self.head_locations = head_locations |
|
self.head_weights = head_weights |
|
self.tie_additional_weights = tie_additional_weights |
|
self.average_logits = average_logits |
|
|
|
|
|
class GPT2LMMultiHeadModel(GPT2LMHeadModel): |
|
config_class = GPT2MultiHeadConfig |
|
|
|
def __init__(self, config: GPT2MultiHeadConfig): |
|
super().__init__(config) |
|
if config.head_locations is not None: |
|
if not len(config.head_locations) + 1 == len(config.head_weights): |
|
raise ValueError("The number of head locations should be equal to the number of head weights minus 1") |
|
self.head_locations = config.head_locations |
|
self.additional_lm_heads = nn.ModuleList( |
|
[nn.Linear(config.n_embd, config.vocab_size, bias=False) for _ in config.head_locations] |
|
) |
|
self.head_weights = config.head_weights |
|
else: |
|
self.head_locations = [] |
|
self.additional_lm_heads = nn.ModuleList([]) |
|
self.head_weights = [1.0] |
|
self.post_init() |
|
|
|
def tie_weights(self): |
|
""" |
|
Tie the weights between the input embeddings and the output embeddings. |
|
|
|
If the `torchscript` flag is set in the configuration, can't handle parameter sharing so we are cloning the |
|
weights instead. |
|
""" |
|
super().tie_weights() |
|
if hasattr(self, "additional_lm_heads") and getattr(self.config, "tie_additional_weights", False): |
|
input_embeddings = self.get_input_embeddings() |
|
for classifier in self.additional_lm_heads: |
|
if self.config.torchscript: |
|
classifier.weight = nn.Parameter(input_embeddings.weight.clone()) |
|
else: |
|
classifier.weight = input_embeddings.weight |
|
|
|
if getattr(classifier, "bias", None) is not None: |
|
classifier.bias.data = nn.functional.pad( |
|
classifier.bias.data, |
|
( |
|
0, |
|
classifier.weight.shape[0] - classifier.bias.shape[0], |
|
), |
|
"constant", |
|
0, |
|
) |
|
if hasattr(classifier, "out_features") and hasattr(input_embeddings, "num_embeddings"): |
|
classifier.out_features = input_embeddings.num_embeddings |
|
|
|
def forward( |
|
self, |
|
input_ids: Optional[torch.LongTensor] = None, |
|
past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, |
|
attention_mask: Optional[torch.FloatTensor] = None, |
|
token_type_ids: Optional[torch.LongTensor] = None, |
|
position_ids: Optional[torch.LongTensor] = None, |
|
head_mask: Optional[torch.FloatTensor] = None, |
|
inputs_embeds: Optional[torch.FloatTensor] = None, |
|
encoder_hidden_states: Optional[torch.Tensor] = None, |
|
encoder_attention_mask: Optional[torch.FloatTensor] = None, |
|
labels: Optional[torch.LongTensor] = None, |
|
use_cache: Optional[bool] = None, |
|
output_attentions: Optional[bool] = None, |
|
output_hidden_states: Optional[bool] = None, |
|
return_dict: Optional[bool] = None, |
|
) -> Union[Tuple, CausalLMOutputWithCrossAttentions]: |
|
r""" |
|
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): |
|
Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set |
|
`labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` |
|
are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` |
|
""" |
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
|
|
|
transformer_outputs = self.transformer( |
|
input_ids, |
|
past_key_values=past_key_values, |
|
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, |
|
use_cache=use_cache, |
|
output_attentions=output_attentions, |
|
output_hidden_states=True, |
|
return_dict=return_dict, |
|
) |
|
hidden_states = transformer_outputs[2] |
|
|
|
|
|
if self.model_parallel: |
|
torch.cuda.set_device(self.transformer.first_device) |
|
hidden_states = hidden_states.to(self.lm_head.weight.device) |
|
|
|
lm_logits = self.lm_head(hidden_states[-1]) |
|
loss = None |
|
if labels is not None: |
|
loss = torch.tensor(0.0, device=hidden_states[-1].device) |
|
lm_logits = [] |
|
loss_fct = CrossEntropyLoss() |
|
|
|
for index, lm_head, lm_weight in zip( |
|
[*self.head_locations, -1], |
|
[*self.additional_lm_heads, self.lm_head], |
|
self.head_weights, |
|
): |
|
lm_logits.append(lm_head(hidden_states[index])) |
|
|
|
shift_logits = lm_logits[-1][..., :-1, :].contiguous() |
|
shift_labels = labels[..., 1:].contiguous() |
|
|
|
loss += lm_weight * loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) |
|
|
|
if self.config.average_logits: |
|
lm_logits = (torch.vstack(lm_logits) * torch.tensor(self.head_weights)).mean(dim=0) |
|
else: |
|
lm_logits = lm_logits[-1] |
|
if not return_dict: |
|
output = (lm_logits,) + transformer_outputs[1:] |
|
return ((loss,) + output) if loss is not None else output |
|
|
|
return CausalLMOutputWithCrossAttentions( |
|
loss=loss, |
|
logits=lm_logits, |
|
past_key_values=transformer_outputs.past_key_values, |
|
hidden_states=transformer_outputs.hidden_states, |
|
attentions=transformer_outputs.attentions, |
|
cross_attentions=transformer_outputs.cross_attentions, |
|
) |
|
|
