Update app.py

app.py

@@ -138,17 +138,19 @@ def run_interpretation(raw_original_prompt, raw_interpretation_prompt, max_new_t
     generation_texts = tokenizer.batch_decode(generated)
 
     # try identifying important layers
-    # vectors_to_compare = interpreted_vectors # torch.tensor(global_state.sentence_transformer.encode(generation_texts))
-    # diff_score = F.normalize(vectors_to_compare, dim=-1).diff(dim=0).norm(dim=-1)
+    vectors_to_compare = interpreted_vectors # torch.tensor(global_state.sentence_transformer.encode(generation_texts))
+    diff_score1 = F.normalize(vectors_to_compare, dim=-1).diff(dim=0).norm(dim=-1)
+    
     bags_of_words = [set(tokenizer.tokenize(text)) for text in generation_texts]
-    diff_score = torch.tensor([
+    diff_score2 = torch.tensor([
                                 -len(bags_of_words[i+1] & bags_of_words[i]) / np.sqrt(len(bags_of_words[i+1]) * len(bags_of_words[i]))
                                 for i in range(len(bags_of_words)-1)
                               ])
+    diff_score = diff_score1 / diff_score1.median() + diff_score2 / diff_score2.median()
     avoid_first, avoid_last = 2, 1 # layers that are usually never important
     assert avoid_first >= 1 # due to .diff() we will not be able to compute a score for the first layer
-    diff_score = diff_score[avoid_first-1 : len(diff_score)-avoid_last]
-    important_idxs = avoid_first + diff_score.topk(k=int(np.ceil(0.1 * len(generation_texts)))).indices.cpu().numpy()
+    diff_score = diff_score[avoid_first-1:len(diff_score)-avoid_last]
+    important_idxs = avoid_first + diff_score.topk(k=4).indices.cpu().numpy() # k=int(np.ceil(0.15 * len(generation_texts)))
     
     # create GUI output
     print(f'{important_idxs=}')
@@ -248,8 +250,7 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
                 btn.render()
         
         progress_dummy = gr.Markdown('', elem_id='progress_dummy')
-        interpretation_bubbles = [gr.Textbox('', container=False, visible=False) 
-                                  for i in range(MAX_NUM_LAYERS)]
+        interpretation_bubbles = [gr.Textbox('', container=False, visible=False) for i in range(MAX_NUM_LAYERS)]
         
     # event listeners    
     for i, btn in enumerate(tokens_container):