debug

app.py

@@ -130,7 +130,10 @@ 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_locs)==1)
+    if len(option_1_locs)>1:
+        assert np.all(np.diff(option_1_locs)==1)
+    if len(option_2_locs)>1:
+        assert np.all(np.diff(option_2_locs)==1)
     return option_1_locs, option_2_locs
 
 def mask_out(input_ids,pron_locs,option_locs,mask_id):
@@ -227,7 +230,7 @@ if __name__=='__main__':
 
             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]:
+                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)]
@@ -235,7 +238,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)]
+                    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]]