skeleton

app.py

@@ -8,14 +8,19 @@ import seaborn as sns
 import jax
 import jax.numpy as jnp
 
-from transformers import AlbertTokenizer
+import torch
+import torch.nn.functional as F
 
-from custom_modeling_albert_flax import CustomFlaxAlbertForMaskedLM
+from transformers import AlbertTokenizer, AlbertForMaskedLM
+
+#from custom_modeling_albert_flax import CustomFlaxAlbertForMaskedLM
+from skeleton_modeling_albert import SkeletonAlbertForMaskedLM
 
 @st.cache(show_spinner=True,allow_output_mutation=True)
 def load_model():
     tokenizer = AlbertTokenizer.from_pretrained('albert-xxlarge-v2')
-    model = CustomFlaxAlbertForMaskedLM.from_pretrained('albert-xxlarge-v2',from_pt=True)
+    #model = CustomFlaxAlbertForMaskedLM.from_pretrained('albert-xxlarge-v2',from_pt=True)
+    model = AlbertForMaskedLM.from_pretrained('albert-xxlarge-v2')
     return tokenizer,model
 
 def clear_data():
@@ -58,9 +63,9 @@ if __name__=='__main__':
     sent_2 = st.sidebar.text_input('Sentence 2',value='It is better to play a prank on Samuel than Craig because he gets angry more often.',on_change=clear_data)
     input_ids_1 = tokenizer(sent_1).input_ids
     input_ids_2 = tokenizer(sent_2).input_ids
-    input_ids = np.array([input_ids_1,input_ids_2])
+    input_ids = torch.tensor([input_ids_1,input_ids_2])
 
-    outputs = model(input_ids, interv_type='swap', interv_dict = {0:{'lay':[(8,1,[0,1])]}})
-    logprobs = jax.nn.log_softmax(outputs.logits, axis = -1)
-    preds = [np.random.choice(np.arange(len(probs)),p=np.exp(probs)/np.sum(np.exp(probs))) for probs in logprobs[0]]
+    outputs = SkeletonAlbertForMaskedLM(model,input_ids,interventions = {0:{'lay':[(8,1,[0,1])]}})
+    logprobs = F.log_softmax(outputs.logits, dim = -1)
+    preds = [torch.multinomial(torch.exp(probs), num_samples=1).squeeze(dim=-1) for probs in logprobs[0]]
     st.write([tokenizer.decode([token]) for token in preds])

skeleton_modeling_albert.py

@@ -0,0 +1,74 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+import math
+
+from transformers.modeling_utils import apply_chunking_to_forward
+
+@torch.no_grad()
+def SkeletonAlbertLayer(layer_id,layer,hidden,interventions):
+    attention_layer = layer.attention
+    num_heads = attention_layer.num_attention_heads
+    head_dim = attention_layer.attention_head_size
+
+    qry = attention_layer.query(hidden)
+    key = attention_layer.key(hidden)
+    val = attention_layer.value(hidden)
+
+    # swap representations
+    interv_layer = interventions.pop(layer_id,None)
+    if interv_layer is not None:
+        reps = {
+                'lay': hidden,
+                'qry': qry,
+                'key': key,
+                'val': val,
+                }
+        for rep_type in ['lay','qry','key','val']:
+            interv_rep = interv_layer.pop(rep_type,None)
+            if interv_rep is not None:
+                new_state = reps[rep_type].clone()
+                for head_id, pos, swap_ids in interv_rep:
+                    new_state[swap_ids[0],pos,head_id] = reps[rep_name][swap_ids[1],pos,head_id]
+                    new_state[swap_ids[1],pos,head_id] = reps[rep_name][swap_ids[0],pos,head_id]
+                reps[rep_type] = new_state.clone()
+
+        hidden = reps['lay'].clone()
+        qry = reps['qry'].clone()
+        key = reps['key'].clone()
+        val = reps['val'].clone()
+
+
+    #split into multiple heads
+    split_qry = qry.view(*(qry.size()[:-1]+(num_heads,head_dim))).permute(0,2,1,3)
+    split_key = key.view(*(key.size()[:-1]+(num_heads,head_dim))).permute(0,2,1,3)
+    split_val = val.view(*(val.size()[:-1]+(num_heads,head_dim))).permute(0,2,1,3)
+
+    #calculate the attention matrix
+    attn_mat = F.softmax(split_qry@split_key.permute(0,1,3,2)/math.sqrt(head_dim),dim=-1)
+
+    z_rep_indiv = attn_mat@split_val
+    z_rep = z_rep_indiv.permute(0,2,1,3).reshape(*hidden.size())
+
+    hidden_post_attn_res = layer.attention.dense(z_rep)+hidden
+    hidden_post_attn = layer.attention.LayerNorm(hidden_post_attn_res)
+
+    ffn_output = apply_chunking_to_forward(layer.ff_chunk,layer.chunk_size_feed_forward,
+                                            layer.seq_len_dim,hidden_post_attn)
+    new_hidden = layer.full_layer_layer_norm(ffn_output+hidden_post_attn)
+    return new_hidden
+
+def SkeletonAlbertForMaskedLM(model,input_ids,interventions):
+    core_model = model.albert
+    lm_head = model.predictions
+    output_hidden = []
+    with torch.no_grad():
+        hidden = core_model.embeddings(input_ids)
+        hidden = core_model.encoder.embedding_hidden_mapping_in(hidden)
+        output_hidden.append(hidden)
+        for layer_id in range(model.config.num_hidden_layers):
+            layer = core_model.encoder.albert_layer_groups[0].albert_layers[0]
+            hidden = SkeletonAlbertLayer(layer_id,layer,hidden,interventions)
+            output_hidden.append(hidden)
+        logits = lm_head(hidden)
+    return {'logits':logits,'hidden_states':output_hidden}