Update app.py

app.py

@@ -10,6 +10,7 @@ import gradio as gr
 import torch
 from datasets import load_dataset
 from transformers import PreTrainedModel, PreTrainedTokenizer, AutoModelForCausalLM, AutoTokenizer
+from sentence_transformers import SentenceTransformer
 from ctransformers import AutoModelForCausalLM as CAutoModelForCausalLM
 from interpret import InterpretationPrompt
 from configs import model_info, dataset_info
@@ -27,6 +28,7 @@ class LocalState:
 class GlobalState:
     tokenizer : Optional[PreTrainedTokenizer] = None
     model : Optional[PreTrainedModel] = None
+    sentence_transformer: Optional[PreTrainedModel] = None
     local_state : LocalState = LocalState()
     wait_with_hidden_state : bool = False
     interpretation_prompt_template : str = '{prompt}'
@@ -49,7 +51,7 @@ suggested_interpretation_prompts = [
 def initialize_gpu():
     pass
 
-def reset_model(model_name, *extra_components, with_extra_components=True): 
+def reset_model(model_name, *extra_components, reset_sentence_transformer=False, with_extra_components=True): 
     # extract model info
     model_args = deepcopy(model_info[model_name])
     model_path = model_args.pop('model_path')
@@ -66,6 +68,9 @@ def reset_model(model_name, *extra_components, with_extra_components=True):
     global_state.model, global_state.tokenizer, global_state.local_state.hidden_states = None, None, None
     gc.collect()
     global_state.model = AutoModelClass.from_pretrained(model_path, **model_args)
+    if reset_sentence_transformer:
+        global_state.sentence_transformer = SentenceTransformer('all-MiniLM-L6-v2')
+        gc.collect()
     if not dont_cuda:
         global_state.model.to('cuda')
     global_state.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, token=os.environ['hf_token'])
@@ -131,8 +136,11 @@ def run_interpretation(raw_original_prompt, raw_interpretation_prompt, max_new_t
                                                **generation_kwargs)
     generation_texts = tokenizer.batch_decode(generated)
 
+    # try identifying important layers
+    diff_score = F.normalize(global_state.sentence_transformer.encode(generation_texts), dim=-1).diff(dim=0)
+    important_idxs = 1 + diff_score.topk(k=int(np.ceil(0.2 * len(generation_texts)))).indices.cpu().numpy()
+    
     # create GUI output
-    important_idxs = 1 + ((interpreted_vectors - hidden_means)).diff(dim=0).norm(dim=-1).topk(k=int(np.ceil(0.2 * len(generation_texts)))).indices.cpu().numpy()
     print(f'{important_idxs=}')
     progress_dummy_output = ''
     elem_classes = [['bubble', 'even_bubble' if i % 2 == 0 else 'odd_bubble'] +