debug

app.py

@@ -117,15 +117,13 @@ def show_instruction(sent,fontsize=20):
     suffix = '</span></p>'
     return st.markdown(prefix + sent + suffix, unsafe_allow_html = True)
 
-def create_interventions(token_id,interv_type,num_layers,num_heads):
+def create_interventions(token_id,interv_types,num_heads):
     interventions = {}
-    for layer_id in range(num_layers):
-        interventions[layer_id] = {}
-        if interv_type == 'all':
-            for rep in ['lay','qry','key','val']:
-                interventions[layer_id][rep] = [(head_id,token_id,[head_id,head_id+num_heads]) for head_id in range(num_heads)]
+    for rep in ['lay','qry','key','val']:
+        if rep in interv_types:
+            interventions[rep] = [(head_id,token_id,[head_id,head_id+num_heads]) for head_id in range(num_heads)]
         else:
-            interventions[layer_id][interv_type] = [(head_id,token_id,[head_id,head_id+num_heads]) for head_id in range(num_heads)]
+            interventions[rep] = []
     return interventions
 
 def separate_options(option_locs):
@@ -195,7 +193,7 @@ if __name__=='__main__':
                                         mask_locs=st.session_state['mask_locs_2'])
 
             option_1_locs, option_2_locs = {}, {}
-            pron_id = {}
+            pron_locs = {}
             input_ids_dict = {}
             masked_ids_option_1 = {}
             masked_ids_option_2 = {}
@@ -215,14 +213,15 @@ if __name__=='__main__':
                 st.write(' '.join([tokenizer.decode([token]) for toke in token_ids]))
 
     if st.session_state['page_status'] == 'finish_debug':
-            interventions = create_interventions(16,'all',num_layers=num_layers,num_heads=num_heads)
-            for masked_ids in [masked_ids_option_1, masked_ids_option_2]:
-                input_ids = torch.tensor([
-                                        *[masked_ids['sent_1'] for _ in range(num_heads)],
-                                        *[masked_ids['sent_2'] for _ in range(num_heads)]
-                                        ])
-                outputs = SkeletonAlbertForMaskedLM(model,input_ids,interventions=interventions)
-                logprobs = F.log_softmax(outputs['logits'], dim = -1)
+            for layer_id in range(num_layers):
+                interventions = [create_interventions(16,['lay','qry','key','val'],num_heads) if i==layer_id else {'lay':[],'qry':[],'key':[],'val':[]} for i in range(num_layers)]
+                for masked_ids in [masked_ids_option_1, masked_ids_option_2]:
+                    input_ids = torch.tensor([
+                                            *[masked_ids['sent_1'] for _ in range(num_heads)],
+                                            *[masked_ids['sent_2'] for _ in range(num_heads)]
+                                            ])
+                    outputs = SkeletonAlbertForMaskedLM(model,input_ids,interventions=interventions)
+                    logprobs = F.log_softmax(outputs['logits'], dim = -1)
 
 
             preds_0 = [torch.multinomial(torch.exp(probs), num_samples=1).squeeze(dim=-1) for probs in logprobs[0][1:-1]]