Update app.py

app.py

@@ -21,6 +21,11 @@ MAX_PROMPT_TOKENS = 60
 MAX_NUM_LAYERS = 50
 welcome_message = '**You are now running {model_name}!!** 🥳🥳🥳'
 
+# Used by the layer and token importance heuristic in this file. 
+# These layers are usually not important. We will ignore them when looking for important layers 
+avoid_first, avoid_last = 3, 2
+
+
 @dataclass
 class LocalState:
     hidden_states: Optional[torch.Tensor] = None
@@ -96,8 +101,8 @@ def get_hidden_states(raw_original_prompt, force_hidden_states=False):
         outputs = model(**model_inputs, output_hidden_states=True, return_dict=True)
         hidden_states = torch.stack([h.squeeze(0).cpu().detach() for h in outputs.hidden_states], dim=0)
         # TODO: document this!
-        hidden_scores = F.normalize(hidden_states, dim=-1).diff(dim=0).norm(dim=-1).cpu() # num_layers x num_tokens
-        important_tokens = (1 + np.unravel_index(hidden_scores.flatten().topk(k=5).indices.numpy(), hidden_scores.shape)[1])
+        hidden_scores = F.normalize(hidden_states[avoid_first-1:len(hidden_states)-avoid_last], dim=-1).diff(dim=0).norm(dim=-1).cpu() # num_layers x num_tokens
+        important_tokens = np.unravel_index(hidden_scores.flatten().topk(k=5).indices.numpy(), hidden_scores.shape)[1]
         print(f'{important_tokens=}\t\t{hidden_states.shape=}')
         global_state.local_state.hidden_states = hidden_states.cpu().detach()
         
@@ -162,7 +167,6 @@ def run_interpretation(raw_original_prompt, raw_interpretation_prompt, max_new_t
     diff_score = ((diff_score1 - diff_score1.min()) / (diff_score1.max() - diff_score1.min()) 
                   + (diff_score2 - diff_score2.min()) / (diff_score2.max() - diff_score2.min()))
     
-    avoid_first, avoid_last = 3, 3 # 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.3 * len(diff_score)))).indices.cpu().numpy() #