check masking

app.py

@@ -111,16 +111,39 @@ def show_annotated_sentence(sent,option_locs=[],mask_locs=[]):
     suffix = '</span></p>'
     return st.markdown(prefix + disp + suffix, unsafe_allow_html = True)
 
-def show_instruction(sent):
-    disp_style = '"font-family:san serif; color:Black; font-size: 20px"'
+def show_instruction(sent,fontsize=20):
+    disp_style = f'"font-family:san serif; color:Black; font-size: {fontsize}px"'
     prefix = f'<p style={disp_style}><span style="font-weight:bold">'
     suffix = '</span></p>'
     return st.markdown(prefix + sent + suffix, unsafe_allow_html = True)
 
+def create_interventions(token_id,interv_type,num_layers,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)]
+        else:
+            interventions[layer_id][interv_type] = [(head_id,token_id,[head_id,head_id+num_heads]) for head_id in range(num_heads)]
+    return interventions
+
+def separate_options(option_locs):
+    assert np.sum(np.diff(option_locs)>1)==1
+    sep = list(np.diff(option_locs)>1).index(1)+1
+    option_1_locs, option_2_locs = option_locs[:sep], option_locs[sep:]
+    assert np.all(np.diff(option_1_locs)==1) and np.all(np.diff(option_2_loc)==1)
+    return option_1_locs, option_2_locs
+
+def mask_out(input_ids,pron_locs,option_locs,mask_id):
+    assert np.all(np.diff(pron_locs)==1)
+    return input_ids[:pron_locs[0]] + [mask_id for _ in range(len(option_locs))] + input_ids[pron_locs[-1]+1:]
+
 if __name__=='__main__':
     wide_setup()
     load_css('style.css')
     tokenizer,model = load_model()
+    num_layers, num_heads = 12, 64
     mask_id = tokenizer('[MASK]').input_ids[1:-1][0]
 
     main_area = st.empty()
@@ -171,16 +194,37 @@ if __name__=='__main__':
                                         option_locs=st.session_state['option_locs_2'],
                                         mask_locs=st.session_state['mask_locs_2'])
 
-            input_ids_1 = tokenizer(sent_1).input_ids
-            input_ids_2 = tokenizer(sent_2).input_ids
-            input_ids = torch.tensor([input_ids_1,input_ids_2])
-
-            outputs = SkeletonAlbertForMaskedLM(model,input_ids,
-                                                    interventions = {0:{'lay':[(head_id,16,[0,1]) for head_id in range(64)],
-                                                                        'qry':[(head_id,16,[0,1]) for head_id in range(64)],
-                                                                        'key':[(head_id,16,[0,1]) for head_id in range(64)],
-                                                                        'val':[(head_id,16,[0,1]) for head_id in range(64)]}})
+            option_1_locs, option_2_locs = {}, {}
+            pron_id = {}
+            input_ids_dict = {}
+            masked_ids_option_1 = {}
+            masked_ids_option_2 = {}
+            for sent_id in range(2):
+                option_1_locs[f'sent_{sent_id+1}'], option_2_locs[f'sent_{sent_id+1}'] = separate_options(st.session_state[f'option_locs_{sent_id}'])
+                pron_locs[f'sent_{sent_id+1}'] = st.session_state[f'mask_locs_{sent_id+1}']
+                input_ids_dict[f'sent_{sent_id+1}'] = tokenizer(st.session_state[f'sent_{sent_id+1}']).input_ids
+
+                masked_ids_option_1[f'sent_{sent_id+1}'] = mask_out(input_ids_dict[f'sent_{sent_id+1}'],
+                                                                        pron_locs[f'sent_{sent_id+1}'],
+                                                                        option_1_locs[f'sent_{sent_id+1}'],mask_id)
+                masked_ids_option_2[f'sent_{sent_id+1}'] = mask_out(input_ids_dict[f'sent_{sent_id+1}'],
+                                                                        pron_locs[f'sent_{sent_id+1}'],
+                                                                        option_2_locs[f'sent_{sent_id+1}'],mask_id)
+
+            for token_ids in [masked_ids_option_1['sent_1'],masked_ids_option_1['sent_2'],masked_ids_option_2['sent_1'],masked_ids_option_2['sent_2']]:
+                st.write(' '.join([tokenizer.decode([token]) for toke in token_ids]))
+
+    if st.session_state['page_status'] == 'finish_debug':
+            try:
+                assert len(input_ids_1) == len(input_ids_2)
+            except AssertionError:
+                show_instruction('Please make sure the number of tokens match between Sentence 1 and Sentence 2', fontsize=12)
+            input_ids = torch.tensor([*[input_ids_1 for _ in range(num_heads)],*[input_ids_2 for _ in range(num_heads)]])
+            interventions = create_interventions(16,'all',num_layers=num_layers,num_heads=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]]
             preds_1 = [torch.multinomial(torch.exp(probs), num_samples=1).squeeze(dim=-1) for probs in logprobs[1][1:-1]]
             st.write([tokenizer.decode([token]) for token in preds_0])