add model options

app.py

@@ -10,10 +10,7 @@ import seaborn as sns
 import torch
 import torch.nn.functional as F
 
-from transformers import AlbertTokenizer, AlbertForMaskedLM
-
 #from custom_modeling_albert_flax import CustomFlaxAlbertForMaskedLM
-from skeleton_modeling_albert import SkeletonAlbertForMaskedLM
 
 def wide_setup():
     max_width = 1500
@@ -48,10 +45,23 @@ def load_css(file_name):
 
 @st.cache(show_spinner=True,allow_output_mutation=True)
 def load_model(model_name):
-    tokenizer = AlbertTokenizer.from_pretrained(model_name)
-    #model = CustomFlaxAlbertForMaskedLM.from_pretrained('albert-xxlarge-v2',from_pt=True)
-    model = AlbertForMaskedLM.from_pretrained(model_name)
-    return tokenizer,model
+    if model_name.startswith('albert'):
+        from transformers import AlbertTokenizer, AlbertForMaskedLM
+        from skeleton_modeling_albert import SkeletonAlbertForMaskedLM
+        tokenizer = AlbertTokenizer.from_pretrained(model_name)
+        model = AlbertForMaskedLM.from_pretrained(model_name)
+        skeleton_model = SkeletonAlbertForMaskedLM
+    elif model_name.startswith('bert'):
+        from transformers import BertTokenizer, BertForMaskedLM
+        from skeleton_modeling_bert import SkeletonBertForMaskedLM
+        tokenizer = BertTokenizer.from_pretrained(model_name)
+        model = BertForMaskedLM.from_pretrained(model_name)
+    elif model_name.startswith('roberta'):
+        from transformers import RobertaTokenizer, RobertaForMaskedLM
+        from skeleton_modeling_roberta import SkeletonRobertaForMaskedLM
+        tokenizer = RobertaTokenizer.from_pretrained(model_name)
+        model = RobertaForMaskedLM.from_pretrained(model_name)
+    return tokenizer,model,skeleton_model
 
 def clear_data():
     for key in st.session_state:
@@ -147,14 +157,14 @@ def mask_out(input_ids,pron_locs,option_locs,mask_id):
     return input_ids[:pron_locs[0]+1] + [mask_id for _ in range(len(option_locs))] + input_ids[pron_locs[-1]+2:]
 
 
-def run_intervention(interventions,batch_size,model,masked_ids_option_1,masked_ids_option_2,option_1_tokens,option_2_tokens,pron_locs):
+def run_intervention(interventions,batch_size,skeleton_model,model,masked_ids_option_1,masked_ids_option_2,option_1_tokens,option_2_tokens,pron_locs):
     probs = []
     for masked_ids, option_tokens in zip([masked_ids_option_1, masked_ids_option_2],[option_1_tokens,option_2_tokens]):
         input_ids = torch.tensor([
                                 *[masked_ids['sent_1'] for _ in range(batch_size)],
                                 *[masked_ids['sent_2'] for _ in range(batch_size)]
                                 ])
-        outputs = SkeletonAlbertForMaskedLM(model,input_ids,interventions=interventions)
+        outputs = skeleton_model(model,input_ids,interventions=interventions)
         logprobs = F.log_softmax(outputs['logits'], dim = -1)
         logprobs_1, logprobs_2 = logprobs[:batch_size], logprobs[batch_size:]
         evals_1 = [logprobs_1[:,pron_locs['sent_1'][0]+1+i,token].numpy() for i,token in enumerate(option_tokens)]
@@ -181,9 +191,10 @@ if __name__=='__main__':
             st.session_state['page_status'] = 'type_in'
             st.experimental_rerun()
 
-    tokenizer,model = load_model(st.session_state['model_name'])
-    num_layers, num_heads = model.config.num_hidden_layers, model.config.num_attention_heads
-    mask_id = tokenizer('[MASK]').input_ids[1:-1][0]
+    if st.session_state['page_status']!='model_selection':
+        tokenizer,model,skeleton_model = load_model(st.session_state['model_name'])
+        num_layers, num_heads = model.config.num_hidden_layers, model.config.num_attention_heads
+        mask_id = tokenizer('[MASK]').input_ids[1:-1][0]
 
     if st.session_state['page_status']=='type_in':
         show_instruction('1. Type in the sentences and click "Tokenize"',fontsize=16)
@@ -263,7 +274,7 @@ if __name__=='__main__':
         option_2_tokens = option_2_tokens_1
 
         interventions = [{'lay':[],'qry':[],'key':[],'val':[]} for i in range(num_layers)]
-        probs_original = run_intervention(interventions,1,model,masked_ids_option_1,masked_ids_option_2,option_1_tokens,option_2_tokens,pron_locs)
+        probs_original = run_intervention(interventions,1,skeleton_model,model,masked_ids_option_1,masked_ids_option_2,option_1_tokens,option_2_tokens,pron_locs)
         df = pd.DataFrame(data=[[probs_original[0,0][0],probs_original[1,0][0]],
                                 [probs_original[0,1][0],probs_original[1,1][0]]],
                                 columns=[tokenizer.decode(option_1_tokens),tokenizer.decode(option_2_tokens)],
@@ -292,9 +303,9 @@ if __name__=='__main__':
             for layer_id in range(num_layers):
                 interventions = [create_interventions(token_id,['lay','qry','key','val'],num_heads,multihead) if i==layer_id else {'lay':[],'qry':[],'key':[],'val':[]} for i in range(num_layers)]
                 if multihead:
-                    probs = run_intervention(interventions,1,model,masked_ids_option_1,masked_ids_option_2,option_1_tokens,option_2_tokens,pron_locs)
+                    probs = run_intervention(interventions,1,skeleton_model,model,masked_ids_option_1,masked_ids_option_2,option_1_tokens,option_2_tokens,pron_locs)
                 else:
-                    probs = run_intervention(interventions,num_heads,model,masked_ids_option_1,masked_ids_option_2,option_1_tokens,option_2_tokens,pron_locs)
+                    probs = run_intervention(interventions,num_heads,skeleton_model,model,masked_ids_option_1,masked_ids_option_2,option_1_tokens,option_2_tokens,pron_locs)
                 effect = ((probs_original-probs)[0,0] + (probs_original-probs)[1,1] + (probs-probs_original)[0,1] + (probs-probs_original)[1,0])/4
                 effect_list.append(effect)
             effect_array.append(effect_list)

skeleton_modeling_bert.py

@@ -0,0 +1,73 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+import math
+
+@torch.no_grad()
+def SkeletonBertLayer(layer_id,layer,hidden,interventions):
+    attention_layer = layer.attention.self
+    num_heads = attention_layer.num_attention_heads
+    head_dim = attention_layer.attention_head_size
+    assert num_heads*head_dim == hidden.shape[2]
+
+    qry = attention_layer.query(hidden)
+    key = attention_layer.key(hidden)
+    val = attention_layer.value(hidden)
+
+    assert qry.shape == hidden.shape
+    assert key.shape == hidden.shape
+    assert val.shape == hidden.shape
+
+    # swap representations
+    reps = {
+            'lay': hidden,
+            'qry': qry,
+            'key': key,
+            'val': val,
+            }
+    for rep_type in ['lay','qry','key','val']:
+        interv_rep = interventions[layer_id][rep_type]
+        new_state = reps[rep_type].clone()
+        for head_id, pos, swap_ids in interv_rep:
+            new_state[swap_ids[0],:,head_dim*head_id:head_dim*(head_id+1)][pos,:] = reps[rep_type][swap_ids[1],:,head_dim*head_id:head_dim*(head_id+1)][pos,:]
+            new_state[swap_ids[1],:,head_dim*head_id:head_dim*(head_id+1)][pos,:] = reps[rep_type][swap_ids[0],:,head_dim*head_id:head_dim*(head_id+1)][pos,:]
+        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.output.dense(z_rep)+hidden # residual connection
+    hidden_post_attn = layer.attention.output.LayerNorm(hidden_post_attn_res) # layer_norm
+
+    hidden_post_interm = layer.intermediate(hidden_post_attn) # massive feed forward
+    hidden_post_interm_res = layer.output.dense(hidden_post_interm)+hidden_post_attn # residual connection
+    new_hidden =  layer.output.LayerNorm(hidden_post_interm_res) # layer_norm
+    return new_hidden
+
+def SkeletonBertForMaskedLM(model,input_ids,interventions):
+    core_model = model.bert
+    lm_head = model.cls
+    output_hidden = []
+    with torch.no_grad():
+        hidden = core_model.embeddings(input_ids)
+        output_hidden.append(hidden)
+        for layer_id in range(model.config.num_hidden_layers):
+            layer = core_model.encoder.layer[layer_id]
+            hidden = SkeletonBertLayer(layer_id,layer,hidden,interventions)
+            output_hidden.append(hidden)
+        logits = lm_head(hidden)
+    return {'logits':logits,'hidden_states':output_hidden}

skeleton_modeling_roberta.py

@@ -0,0 +1,73 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+import math
+
+@torch.no_grad()
+def SkeletonRobertaLayer(layer_id,layer,hidden,interventions):
+    attention_layer = layer.attention.self
+    num_heads = attention_layer.num_attention_heads
+    head_dim = attention_layer.attention_head_size
+    assert num_heads*head_dim == hidden.shape[2]
+
+    qry = attention_layer.query(hidden)
+    key = attention_layer.key(hidden)
+    val = attention_layer.value(hidden)
+
+    assert qry.shape == hidden.shape
+    assert key.shape == hidden.shape
+    assert val.shape == hidden.shape
+
+    # swap representations
+    reps = {
+            'lay': hidden,
+            'qry': qry,
+            'key': key,
+            'val': val,
+            }
+    for rep_type in ['lay','qry','key','val']:
+        interv_rep = interventions[layer_id][rep_type]
+        new_state = reps[rep_type].clone()
+        for head_id, pos, swap_ids in interv_rep:
+            new_state[swap_ids[0],:,head_dim*head_id:head_dim*(head_id+1)][pos,:] = reps[rep_type][swap_ids[1],:,head_dim*head_id:head_dim*(head_id+1)][pos,:]
+            new_state[swap_ids[1],:,head_dim*head_id:head_dim*(head_id+1)][pos,:] = reps[rep_type][swap_ids[0],:,head_dim*head_id:head_dim*(head_id+1)][pos,:]
+        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.output.dense(z_rep)+hidden # residual connection
+    hidden_post_attn = layer.attention.output.LayerNorm(hidden_post_attn_res) # layer_norm
+
+    hidden_post_interm = layer.intermediate(hidden_post_attn) # massive feed forward
+    hidden_post_interm_res = layer.output.dense(hidden_post_interm)+hidden_post_attn # residual connection
+    new_hidden =  layer.output.LayerNorm(hidden_post_interm_res) # layer_norm
+    return new_hidden
+
+def SkeletonBertForMaskedLM(model,input_ids,interventions):
+    core_model = model.roberta
+    lm_head = model.lm_head
+    output_hidden = []
+    with torch.no_grad():
+        hidden = core_model.embeddings(input_ids)
+        output_hidden.append(hidden)
+        for layer_id in range(model.config.num_hidden_layers):
+            layer = core_model.encoder.layer[layer_id]
+            hidden = SkeletonRobertaLayer(layer_id,layer,hidden,interventions)
+            output_hidden.append(hidden)
+        logits = lm_head(hidden)
+    return {'logits':logits,'hidden_states':output_hidden}