test

app.py

@@ -142,6 +142,23 @@ def mask_out(input_ids,pron_locs,option_locs,mask_id):
     # note annotations are shifted by 1 because special tokens were omitted
     return input_ids[:pron_locs[0]+1] + [mask_id for _ in range(len(option_locs))] + input_ids[pron_locs[-1]+2:]
 
+def run(interventions,batch_size,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)
+        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)]
+        evals_2 = [logprobs_2[:,pron_locs['sent_2'][0]+1+i,token].numpy() for i,token in enumerate(option_tokens)]
+        probs.append([np.exp(np.mean(evals_1,axis=0)),np.exp(np.mean(evals_2,axis=0))])
+    probs = np.array(probs)
+    assert probs.shape[0]==2 and probs.shape[1]==2 and probs.shape[2]==batch_size
+    return probs
+
 if __name__=='__main__':
     wide_setup()
     load_css('style.css')
@@ -220,7 +237,6 @@ if __name__=='__main__':
             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 token in token_ids]))
 
-    if st.session_state['page_status'] == 'finish_debug':
             option_1_tokens_1 = np.array(input_ids_dict['sent_1'])[np.array(option_1_locs['sent_1'])+1]
             option_1_tokens_2 = np.array(input_ids_dict['sent_2'])[np.array(option_1_locs['sent_2'])+1]
             option_2_tokens_1 = np.array(input_ids_dict['sent_1'])[np.array(option_2_locs['sent_1'])+1]
@@ -229,21 +245,12 @@ if __name__=='__main__':
             option_1_tokens = option_1_tokens_1
             option_2_tokens = option_2_tokens_1
 
+            interventions = [{'lay':[],'qry':[],'key':[],'val':[]} for i in range(num_layers)]
+            probs_original = run(interventions,1,model,masked_ids_option_1,masked_ids_option_2,option_1_tokens,option_2_tokens,pron_locs)
+            st.write(probs_original)
+            print(probs_original)
+
+    if st.session_state['page_status'] == 'finish_debug':
             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, 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(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)
-                    logprobs_1, logprobs_2 = logprobs[:num_heads], logprobs[num_heads:]
-                    evals_1 = [logprobs_1[:,pron_locs['sent_1'][0]+1+i,token] for i,token in enumerate(option_tokens)]
-                    evals_2 = [logprobs_2[:,pron_locs['sent_2'][0]+1+i,token] for i,token in enumerate(option_tokens)]
-
-
-            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])
-            st.write([tokenizer.decode([token]) for token in preds_1])
+                probs = run(interventions,num_heads,model,masked_ids_option_1,masked_ids_option_2,option_1_tokens,option_2_tokens,pron_locs)