Accepts multiple sizes

transformers/modeling_albert.py

@@ -5,6 +5,7 @@ import logging
 import torch
 import torch.nn as nn
 from transformers.configuration_albert import AlbertConfig
+from transformers.modeling_bert import BertEmbeddings, BertModel, BertSelfAttention, prune_linear_layer, gelu_new
 logger = logging.getLogger(__name__)
 
 def load_tf_weights_in_albert(model, config, tf_checkpoint_path):
@@ -32,14 +33,14 @@ def load_tf_weights_in_albert(model, config, tf_checkpoint_path):
     print(model)
     
     for name, array in zip(names, arrays):
+        print(name)
         og = name
         name = name.replace("transformer/group_0/inner_group_0", "transformer")
-        name = name.replace("LayerNorm", "layer_norm")
         name = name.replace("ffn_1", "ffn")
         name = name.replace("ffn/intermediate/output", "ffn_output")
         name = name.replace("attention_1", "attention")   
         name = name.replace("cls/predictions/transform", "predictions")
-        name = name.replace("transformer/layer_norm_1", "transformer/attention/output/LayerNorm")        
+        name = name.replace("transformer/LayerNorm_1", "transformer/attention/LayerNorm")        
         name = name.split('/')
 
         print(name)
@@ -84,44 +85,22 @@ def load_tf_weights_in_albert(model, config, tf_checkpoint_path):
     return model
 
 
-class AlbertEmbeddings(nn.Module):
+class AlbertEmbeddings(BertEmbeddings):
     """
     Construct the embeddings from word, position and token_type embeddings.
     """
     def __init__(self, config):
-        super(AlbertEmbeddings, self).__init__()
+        super(AlbertEmbeddings, self).__init__(config)
 
         self.word_embeddings = nn.Embedding(config.vocab_size, config.embedding_size, padding_idx=0)
         self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.embedding_size)
         self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.embedding_size)
-        self.layer_norm = torch.nn.LayerNorm(config.embedding_size, eps=config.layer_norm_eps)
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
-    def forward(self, input_ids, token_type_ids=None, position_ids=None):
-        seq_length = input_ids.size(1)
-        if position_ids is None:
-            position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device)
-            position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
-        if token_type_ids is None:
-            token_type_ids = torch.zeros_like(input_ids)
-
-        word_embeddings = self.word_embeddings(input_ids)
-        position_embeddings = self.position_embeddings(position_ids)
-        token_type_embeddings = self.token_type_embeddings(token_type_ids)
-
-        embeddings = word_embeddings + position_embeddings + token_type_embeddings
-        embeddings = self.layer_norm(embeddings)
-        embeddings = self.dropout(embeddings)
-        return embeddings
-
+        self.LayerNorm = torch.nn.LayerNorm(config.embedding_size, eps=config.layer_norm_eps)
 
-    def get_word_embeddings_table(self):
-        return self.word_embeddings
 
-
-class AlbertModel(nn.Module):
+class AlbertModel(BertModel):
     def __init__(self, config):
-        super(AlbertModel, self).__init__()
+        super(AlbertModel, self).__init__(config)
 
         self.config = config
         self.embeddings = AlbertEmbeddings(config)
@@ -129,6 +108,7 @@ class AlbertModel(nn.Module):
         self.pooler = nn.Linear(config.hidden_size, config.hidden_size)
         self.pooler_activation = nn.Tanh()
 
+
     def forward(self, input_ids, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None):
         if attention_mask is None:
             attention_mask = torch.ones_like(input_ids)
@@ -166,7 +146,7 @@ class AlbertForMaskedLM(nn.Module):
 
         self.config = config
         self.bert = AlbertModel(config)
-        self.layer_norm = nn.LayerNorm(config.embedding_size)
+        self.LayerNorm = nn.LayerNorm(config.embedding_size)
         self.bias = nn.Parameter(torch.zeros(config.vocab_size))
         self.dense = nn.Linear(config.hidden_size, config.embedding_size)
         self.word_embeddings = nn.Linear(config.embedding_size, config.vocab_size)
@@ -182,39 +162,47 @@ class AlbertForMaskedLM(nn.Module):
         hidden_states = self.bert(input_ids, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None)[0]
         hidden_states = self.dense(hidden_states)
         hidden_states = gelu_new(hidden_states)
-        hidden_states = self.layer_norm(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states)
 
         logits = self.word_embeddings(hidden_states)
 
         return logits
 
 
-class AlbertAttention(nn.Module):
+class AlbertAttention(BertSelfAttention):
     def __init__(self, config):
-        super(AlbertAttention, self).__init__()
-
-        if config.hidden_size % config.num_attention_heads != 0:
-            raise ValueError(
-                "The hidden size (%d) is not a multiple of the number of attention "
-                "heads (%d)" % (config.hidden_size, config.num_attention_heads))
-        self.output_attentions = config.output_attentions
+        super(AlbertAttention, self).__init__(config)
 
         self.num_attention_heads = config.num_attention_heads
-        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
-        self.all_head_size = self.num_attention_heads * self.attention_head_size
-
-        self.query = nn.Linear(config.hidden_size, self.all_head_size)
-        self.key = nn.Linear(config.hidden_size, self.all_head_size)
-        self.value = nn.Linear(config.hidden_size, self.all_head_size)
-
+        self.hidden_size = config.hidden_size 
+        self.attention_head_size = config.hidden_size // config.num_attention_heads
         self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
         self.dense = nn.Linear(config.hidden_size, config.hidden_size)
         self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
-    def transpose_for_scores(self, x):
-        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
-        x = x.view(*new_x_shape)
-        return x.permute(0, 2, 1, 3)
+        self.pruned_heads = set()
+
+    def prune_heads(self, heads):
+        if len(heads) == 0:
+            return
+        mask = torch.ones(self.num_attention_heads, self.attention_head_size)
+        heads = set(heads) - self.pruned_heads  # Convert to set and emove already pruned heads
+        for head in heads:
+            # Compute how many pruned heads are before the head and move the index accordingly
+            head = head - sum(1 if h < head else 0 for h in self.pruned_heads)
+            mask[head] = 0
+        mask = mask.view(-1).contiguous().eq(1)
+        index = torch.arange(len(mask))[mask].long()
+
+        # Prune linear layers
+        self.query = prune_linear_layer(self.query, index)
+        self.key = prune_linear_layer(self.key, index)
+        self.value = prune_linear_layer(self.value, index)
+        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
+
+        # Update hyper params and store pruned heads
+        self.num_attention_heads = self.num_attention_heads - len(heads)
+        self.all_head_size = self.attention_head_size * self.num_attention_heads
+        self.pruned_heads = self.pruned_heads.union(heads)
 
     def forward(self, input_ids, attention_mask=None, head_mask=None):
         mixed_query_layer = self.query(input_ids)
@@ -248,7 +236,8 @@ class AlbertAttention(nn.Module):
         context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
         new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
         reshaped_context_layer = context_layer.view(*new_context_layer_shape)
-        w = self.dense.weight.T.view(16, 64, 1024)
+        print(self.dense.weight.T.shape)
+        w = self.dense.weight.T.view(self.num_attention_heads, self.attention_head_size, self.hidden_size)
         b = self.dense.bias
 
         projected_context_layer = torch.einsum("bfnd,ndh->bfh", context_layer, w) + b
@@ -262,7 +251,7 @@ class AlbertTransformer(nn.Module):
         super(AlbertTransformer, self).__init__()
         
         self.config =config
-        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
         self.attention = AlbertAttention(config)
         self.ffn = nn.Linear(config.hidden_size, config.intermediate_size) 
         self.ffn_output = nn.Linear(config.intermediate_size, config.hidden_size)
@@ -273,18 +262,11 @@ class AlbertTransformer(nn.Module):
             ffn_output = self.ffn(attention_output)
             ffn_output = gelu_new(ffn_output)
             ffn_output = self.ffn_output(ffn_output)
-            hidden_states = self.layer_norm(ffn_output + attention_output)
+            hidden_states = self.LayerNorm(ffn_output + attention_output)
 
         return hidden_states
 
 
-def gelu_new(x):
-    """ Implementation of the gelu activation function currently in Google Bert repo (identical to OpenAI GPT).
-        Also see https://arxiv.org/abs/1606.08415
-    """
-    return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
-
-
 class AlbertEncoder(nn.Module):
     def __init__(self, config):
         super(AlbertEncoder, self).__init__()
@@ -305,27 +287,22 @@ class AlbertEncoder(nn.Module):
             outputs = outputs + (all_attentions,)
         return outputs  # last-layer hidden state, (all hidden states), (all attentions)
 
-# config = AlbertConfig.from_json_file("config.json")
-# # model = AlbertForMaskedLM(config)
-# model = AlbertModel(config)
-
-# model = load_tf_weights_in_albert(model, config, "albert/albert")
-
-# print(model)
-
-# input_ids = torch.tensor([[31, 51, 99], [15, 5, 0]])
-# input_mask = torch.tensor([[1, 1, 1], [1, 1, 0]])
-# segment_ids = torch.tensor([[0, 0, 1], [0, 0, 0]])
-
-# # sequence_output, pooled_outputs = model()
+model_size = "base"
+config = AlbertConfig.from_json_file("/home/hf/google-research/albert/config_{}.json".format(model_size))
+model = AlbertModel(config)
+model = load_tf_weights_in_albert(model, config, "/home/hf/transformers/albert-{}/albert-{}".format(model_size, model_size))
+model.eval()
+print(sum(p.numel() for p in model.parameters() if p.requires_grad))
 
-# logits = model(input_ids, attention_mask=input_mask, token_type_ids=segment_ids)[1]
 
+input_ids = [[31, 51, 99, 88, 54, 34, 23, 23, 12], [15, 5, 0, 88, 54, 34, 23, 23, 12]]
+input_mask = [[1, 1, 1, 1, 1, 1, 1, 1, 0], [1, 1, 1, 1, 1, 1, 0, 0, 0]]
+segment_ids = [[0, 0, 1, 0, 0, 1, 0, 0, 1], [0, 0, 0, 0, 0, 0, 0, 0, 0]]
 
-# embeddings_output = 
-# print("pooled output", logits)
-# # print("Pooled output", pooled_outputs)
+pt_input_ids = torch.tensor(input_ids)
+pt_input_mask = torch.tensor(input_mask)
+pt_segment_ids = torch.tensor(segment_ids)
 
-config = AlbertConfig.from_json_file("/home/hf/google-research/albert/config.json")
-model = AlbertModel(config)
-model = load_tf_weights_in_albert(model, config, "/home/hf/transformers/albert/albert")
+pt_dict = {"input_ids": pt_input_ids, "attention_mask": pt_input_mask, "token_type_ids": pt_segment_ids}
+pt_output = model(**pt_dict)
+print(pt_output)
\ No newline at end of file