from transformers.models.electra.modeling_electra import ElectraPreTrainedModel, ElectraEncoder, ElectraLayer, \ ModelOutput, ElectraForSequenceClassification, SequenceClassifierOutput, ElectraForTokenClassification, \ ElectraForMultipleChoice from .config import CharmenElectraConfig from .gbst import GBST import torch.nn as nn import copy import torch from torch import Tensor from dataclasses import dataclass from typing import Optional, Tuple from typing import OrderedDict as OrderDictType from collections import OrderedDict from transformers.activations import get_activation @dataclass class CharmenElectraModelOutput(ModelOutput): """ Output type of :class:`~.CharmenElectraModel`. """ downsampled_hidden_states: Optional[Tuple[torch.FloatTensor]] = None upsampled_hidden_states: Optional[Tuple[torch.FloatTensor]] = None class CharmenElectraModel(ElectraPreTrainedModel): config_class = CharmenElectraConfig def __init__(self, config: CharmenElectraConfig, compatibility_with_transformers=False, **kwargs): super().__init__(config) self.embeddings: GBST = GBST( num_tokens=config.vocab_size, # number of tokens, should be 256 for byte encoding (+ 1 special token for padding in this example) dim=config.embedding_size, # dimension of token and intra-block positional embedding max_block_size=config.max_block_size, # maximum block size downsample_factor=config.downsampling_factor, # the final downsample factor by which the sequence length will decrease by score_consensus_attn=config.score_consensus_attn, config=config # whether to do the cheap score consensus (aka attention) as in eq. 5 in the paper ) self.compatibility_with_transformers = compatibility_with_transformers if config.embedding_size != config.hidden_size: self.embeddings_project = nn.Linear(config.embedding_size, config.hidden_size) self.upsampling = nn.Upsample(scale_factor=config.downsampling_factor, mode='nearest') self.upsampling_convolution = nn.Conv1d(in_channels=config.hidden_size * 2, out_channels=config.hidden_size, kernel_size=(config.downsampling_factor*2-1,), padding='same', dilation=(1,)) self.upsample_output = config.upsample_output # config.num_hidden_layers = config.num_hidden_layers - 2 cfg = copy.deepcopy(config) cfg.num_hidden_layers = config.num_hidden_layers - 2 self.encoder = ElectraEncoder(cfg) # frame_hidden_size self.encoder_first_layer = ElectraLayer(config) self.encoder_last_layer = ElectraLayer(config) self.config = config self.init_weights() def get_input_embeddings(self): return self.embeddings.word_embeddings def set_input_embeddings(self, value): self.embeddings.word_embeddings = value 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.encoder.layer[layer].attention.prune_heads(heads) def forward( self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, output_attentions=None, output_hidden_states=None, return_dict=None, ): output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict if input_ids.shape.__len__() == 1: input_ids = input_ids.view(1, -1) attention_mask = attention_mask.view(1, -1) token_type_ids = token_type_ids.view(1, -1) if input_ids is not None and inputs_embeds is not None: raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") elif input_ids is not None: input_shape = input_ids.size() elif inputs_embeds is not None: input_shape = inputs_embeds.size()[:-1] else: raise ValueError("You have to specify either input_ids or inputs_embeds") batch_size, seq_length = input_shape device = input_ids.device if input_ids is not None else inputs_embeds.device if attention_mask is None: attention_mask = torch.ones(input_shape, device=device) if token_type_ids is None: if hasattr(self.embeddings, "token_type_ids"): buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length] buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length) token_type_ids = buffered_token_type_ids_expanded else: token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) unscaled_attention_mask = torch.clone(attention_mask) _, _, unscaled_hidden_states = self.embeddings( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, token_type_ids=token_type_ids, inputs_embeds=inputs_embeds ) if hasattr(self, "embeddings_project"): unscaled_hidden_states = self.embeddings_project(unscaled_hidden_states) extended_unscaled_attention_mask = self.get_extended_attention_mask(unscaled_attention_mask, input_shape, device) head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) unscaled_hidden_states = self.encoder_first_layer(unscaled_hidden_states, extended_unscaled_attention_mask, None, None, None, None, False)[0] hidden_states, attention_mask = self.embeddings.down_sample(unscaled_hidden_states, unscaled_attention_mask, self.config.downsampling_factor) extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape, device) encoder_output = self.encoder( hidden_states, attention_mask=extended_attention_mask, head_mask=head_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) downsampled_hidden_states = encoder_output[0] hidden_states = encoder_output[0] # upsampling upsampled = self.upsampling(hidden_states.permute(0, 2, 1)).permute(0, 2, 1) hidden_states = torch.cat([unscaled_hidden_states, upsampled], dim=-1) # padded_hidden_states = F.pad(hidden_states.permute(0, 2, 1), (3, 3)) hidden_states = self.upsampling_convolution(hidden_states.permute(0, 2, 1)).permute(0, 2, 1) hidden_states = self.encoder_last_layer(hidden_states, extended_unscaled_attention_mask, None, None, None, None, False) upsampled_output = hidden_states[0] return CharmenElectraModelOutput( downsampled_hidden_states=downsampled_hidden_states, upsampled_hidden_states=upsampled_output ) def load_state_dict(self, state_dict: OrderDictType[str, Tensor], strict: bool = True): model = OrderedDict() prefix = "discriminator.electra." for key, value in state_dict.items(): if key.startswith('generator'): continue if key.startswith(prefix): model[key[len(prefix):]] = value else: continue super(CharmenElectraModel, self).load_state_dict(state_dict=model, strict=strict) class CharmenElectraClassificationHead(nn.Module): """Head for sentence-level classification tasks.""" def __init__(self, config: CharmenElectraConfig, **kwargs): super().__init__() self.config = config self.dense = nn.Linear(config.hidden_size, config.hidden_size) classifier_dropout = ( config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob ) self.dropout = nn.Dropout(classifier_dropout) self.out_proj = nn.Linear(config.hidden_size, config.num_labels) self.ds_factor = config.downsampling_factor def forward(self, features, **kwargs): x = features[:, 0, :] # take token (equiv. to [CLS]) x = self.dropout(x) x = self.dense(x) x = get_activation(self.config.summary_activation)(x) x = self.dropout(x) x = self.out_proj(x) return x class CharmenElectraForSequenceClassification(ElectraForSequenceClassification): config_class = CharmenElectraConfig def __init__(self, config: CharmenElectraConfig, class_weight=None, label_smoothing=0.0, **kwargs): super().__init__(config) self.num_labels = config.num_labels self.config = config self.electra = CharmenElectraModel(config, compatibility_with_transformers=True) self.classifier = CharmenElectraClassificationHead(config) self.cls_loss_fct = torch.nn.CrossEntropyLoss(weight=class_weight, label_smoothing=label_smoothing) self.init_weights() def forward( self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, labels=None, output_attentions=None, output_hidden_states=None, return_dict=None, ): output_discriminator: CharmenElectraModelOutput = self.electra(input_ids, attention_mask, token_type_ids) if self.carmen_config.upsample_output: cls = self.classifier(output_discriminator.upsampled_hidden_states) else: cls = self.classifier(output_discriminator.downsampled_hidden_states) cls_loss = self.cls_loss_fct(cls, labels) return SequenceClassifierOutput( loss=cls_loss, logits=cls, hidden_states=output_discriminator.downsampled_hidden_states, attentions=None, ) def load_state_dict(self, state_dict: OrderDictType[str, Tensor], strict: bool = True): model = OrderedDict() prefix = "discriminator." for key, value in state_dict.items(): if key.startswith('generator'): continue if key.startswith(prefix): if 'discriminator_predictions' in key: continue model[key[len(prefix):]] = value else: if key.startswith('sop'): continue model[key] = value super(CharmenElectraForSequenceClassification, self).load_state_dict(state_dict=model, strict=False) class CharmenElectraForTokenClassification(ElectraForTokenClassification): def __init__(self, config: CharmenElectraConfig, class_weight=None, label_smoothing=0.0, **kwargs): super().__init__(config) self.num_labels = config.num_labels self.config = config self.carmen_config = config self.electra = CharmenElectraModel(config, compatibility_with_transformers=True) classifier_dropout = ( config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob ) self.dropout = nn.Dropout(classifier_dropout) self.classifier = nn.Linear(config.hidden_size, config.num_labels) self.cls_loss_fct = torch.nn.CrossEntropyLoss(weight=class_weight, label_smoothing=label_smoothing) self.init_weights() def forward( self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, labels=None, output_attentions=None, output_hidden_states=None, return_dict=None, ): output_discriminator: CharmenElectraModelOutput = self.electra( input_ids, attention_mask, token_type_ids) discriminator_sequence_output = self.dropout(output_discriminator.upsampled_hidden_states) logits = self.classifier(discriminator_sequence_output) if labels is not None: cls_loss = self.cls_loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1)) else: cls_loss = None return SequenceClassifierOutput( loss=cls_loss, logits=logits, hidden_states=output_discriminator.upsampled_hidden_states, attentions=None, ) def get_input_embeddings(self) -> nn.Module: return self.electra.get_input_embeddings() def load_state_dict(self, state_dict: OrderDictType[str, Tensor], strict: bool = True): model = OrderedDict() prefix = "discriminator." for key, value in state_dict.items(): if key.startswith('generator'): continue if key.startswith(prefix): if 'discriminator_predictions' in key: continue model[key[len(prefix):]] = value else: if key.startswith('sop'): continue model[key] = value super(CharmenElectraForTokenClassification, self).load_state_dict(state_dict=model, strict=False) class Pooler(nn.Module): def __init__(self, config): super().__init__() self.dense = nn.Linear(config.hidden_size, config.hidden_size) self.activation = nn.Tanh() def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # We "pool" the model by simply taking the hidden state corresponding # to the first token. first_token_tensor = hidden_states[:, 0] pooled_output = self.dense(first_token_tensor) pooled_output = self.activation(pooled_output) return pooled_output class CharmenElectraForMultipleChoice(ElectraForMultipleChoice): def __init__(self, config: CharmenElectraConfig, class_weight=None, label_smoothing=0.0, **kwargs): super().__init__(config) self.num_labels = config.num_labels self.config = config self.carmen_config = config self.electra = CharmenElectraModel(config, compatibility_with_transformers=True) self.pooler = Pooler(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.classifier = nn.Linear(config.hidden_size, 1) self.cls_loss_fct = torch.nn.CrossEntropyLoss(weight=class_weight, label_smoothing=label_smoothing) self.init_weights() def forward( self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, labels=None, output_attentions=None, output_hidden_states=None, return_dict=None, ): num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1] input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None output_discriminator: CharmenElectraModelOutput = self.electra( input_ids, attention_mask, token_type_ids) if self.carmen_config.upsample_output: pooled_output = self.pooler(output_discriminator.upsampled_hidden_states) else: pooled_output = self.pooler(output_discriminator.downsampled_hidden_states) pooled_output = self.dropout(pooled_output) logits = self.classifier(pooled_output) reshaped_logits = logits.view(-1, num_choices) cls_loss = self.cls_loss_fct(reshaped_logits, labels) return SequenceClassifierOutput( loss=cls_loss, logits=reshaped_logits, hidden_states=output_discriminator.downsampled_hidden_states, attentions=None, ) def load_state_dict(self, state_dict: OrderDictType[str, Tensor], strict: bool = True): model = OrderedDict() prefix = "discriminator." for key, value in state_dict.items(): if key.startswith('generator'): continue if key.startswith(prefix): if 'discriminator_predictions' in key: continue model[key[len(prefix):]] = value else: if key.startswith('sop'): continue model[key] = value super(CharmenElectraForMultipleChoice, self).load_state_dict(state_dict=model, strict=False)