debug

app.py

@@ -132,7 +132,7 @@ 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)
+    assert np.all(np.diff(option_1_locs)==1) and 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):
@@ -215,14 +215,14 @@ if __name__=='__main__':
                 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)
+            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]]