test

app.py

@@ -62,6 +62,7 @@ def annotate_mask(sent_id,sent):
     st.write(f'Sentence {sent_id}')
     input_sent = tokenizer(sent).input_ids
     decoded_sent = [tokenizer.decode([token]) for token in input_sent[1:-1]]
+    st.session_state[f'decoded_sent_{sent_id}'] = decoded_sent
     char_nums = [len(word)+2 for word in decoded_sent]
     cols = st.columns(char_nums)
     if f'mask_locs_{sent_id}' not in st.session_state:
@@ -157,13 +158,22 @@ if __name__=='__main__':
         with main_area.container():
             sent_1 = st.session_state['sent_1']
             sent_2 = st.session_state['sent_2']
+            show_annotated_sentence(st.session_state['deceded_sent_1'],
+                                        option_locs=st.session_state['option_locs_1'],
+                                        mask_locs=st.session_state['mask_locs_1'])
+            show_annotated_sentence(st.session_state['deceded_sent_2'],
+                                        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,17,[0,1]) for head_id in range(64)]}})
+                                                    interventions = {0:{'lay':[(head_id,17,[0,1]) for head_id in range(64)],
+                                                                        'qry':[(head_id,17,[0,1]) for head_id in range(64)],
+                                                                        'key':[(head_id,17,[0,1]) for head_id in range(64)],
+                                                                        'val':[(head_id,17,[0,1]) for head_id in range(64)]}})
             logprobs = F.log_softmax(outputs['logits'], dim = -1)
             preds = [torch.multinomial(torch.exp(probs), num_samples=1).squeeze(dim=-1) for probs in logprobs[0][1:-1]]
             st.write([tokenizer.decode([token]) for token in preds])