debug

app.py

@@ -135,7 +135,8 @@ def separate_options(option_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:]
+    # 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:]
 
 if __name__=='__main__':
     wide_setup()
@@ -216,10 +217,10 @@ if __name__=='__main__':
                 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'])]
-            option_1_tokens_2 = np.array(input_ids_dict['sent_2'])[np.array(option_1_locs['sent_2'])]
-            option_2_tokens_1 = np.array(input_ids_dict['sent_1'])[np.array(option_2_locs['sent_1'])]
-            option_2_tokens_2 = np.array(input_ids_dict['sent_2'])[np.array(option_2_locs['sent_2'])]
+            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]
+            option_2_tokens_2 = np.array(input_ids_dict['sent_2'])[np.array(option_2_locs['sent_2'])+1]
             assert np.all(option_1_tokens_1==option_1_tokens_2) and np.all(option_2_tokens_1==option_2_tokens_2)
             option_1_tokens = option_1_tokens_1
             option_2_tokens = option_2_tokens_1
@@ -233,6 +234,8 @@ if __name__=='__main__':
                                             ])
                     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[0]+1+i,token] for i,token in enumerate(option_1_tokens)]
 
 
             preds_0 = [torch.multinomial(torch.exp(probs), num_samples=1).squeeze(dim=-1) for probs in logprobs[0][1:-1]]