Update app.py

app.py

@@ -1,4 +1,7 @@
 import os
+import gc
+from typing import Optional
+from dataclasses import dataclass
 from copy import deepcopy
 from functools import partial
 import spaces
@@ -6,7 +9,7 @@ import gradio as gr
 import torch
 from datasets import load_dataset
 from ctransformers import AutoModelForCausalLM as CAutoModelForCausalLM
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import PreTrainedModel, PreTrainedTokenizer, AutoModelForCausalLM, AutoTokenizer
 from interpret import InterpretationPrompt
 
 MAX_PROMPT_TOKENS = 60
@@ -56,13 +59,43 @@ suggested_interpretation_prompts = [
                                    ]
 
 
+@dataclass
+class GlobalState:
+    tokenizer : Optional[PreTrainedTokenizer] = None
+    model : Optional[PreTrainedModel] = None
+    hidden_states : Optional[torch.Tensor] = None
+    interpretation_prompt_template : str = '{prompt}'
+    original_prompt_template : str = '{prompt}'
+
+    
 ## functions
 @spaces.GPU
 def initialize_gpu():
     pass
 
-def get_hidden_states(raw_original_prompt):
-    original_prompt = original_prompt_template.format(prompt=raw_original_prompt)
+
+def reset_model(model_name, global_state):
+    # extract model info
+    model_args = deepcopy(model_info[model_name])
+    model_path = model_args.pop('model_path')
+    global_state.original_prompt_template = model_args.pop('original_prompt_template')
+    global_state.interpretation_prompt_template = model_args.pop('interpretation_prompt_template')
+    tokenizer_path = model_args.pop('tokenizer') if 'tokenizer' in model_args else model_path
+    use_ctransformers = model_args.pop('ctransformers', False)
+    AutoModelClass = CAutoModelForCausalLM if use_ctransformers else AutoModelForCausalLM
+    
+    # get model
+    global_state.model, global_state.tokenizer, global_state.hidden_states = None, None, None
+    gc.collect()
+    global_state.model = AutoModelClass.from_pretrained(model_path, **model_args).cuda()
+    global_state.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, token=os.environ['hf_token'])
+    gc.collect()
+    return global_state
+    
+
+def get_hidden_states(global_state, raw_original_prompt):
+    model, tokenizer = global_state.model, global_state.tokenizer
+    original_prompt = global_state.original_prompt_template.format(prompt=raw_original_prompt)
     model_inputs = tokenizer(original_prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
     tokens = tokenizer.batch_decode(model_inputs.input_ids[0])
     outputs = model(**model_inputs, output_hidden_states=True, return_dict=True)
@@ -71,7 +104,8 @@ def get_hidden_states(raw_original_prompt):
                   + [gr.Button('', visible=False) for _ in range(MAX_PROMPT_TOKENS - len(tokens))])
     progress_dummy_output = ''
     invisible_bubbles = [gr.Textbox('', visible=False) for i in range(len(interpretation_bubbles))]
-    return [progress_dummy_output, hidden_states, *token_btns, *invisible_bubbles]
+    global_state.hidden_states = hidden_states
+    return [progress_dummy_output, global_state, *token_btns, *invisible_bubbles]
 
 
 @spaces.GPU
@@ -79,7 +113,7 @@ def run_interpretation(global_state, raw_interpretation_prompt, max_new_tokens,
                        temperature, top_k, top_p, repetition_penalty, length_penalty, i, 
                        num_beams=1):
 
-    interpreted_vectors = global_state[:, i]
+    interpreted_vectors = global_state.hidden_states[:, i]
     length_penalty = -length_penalty   # unintuitively, length_penalty > 0 will make sequences longer, so we negate it
 
     # generation parameters
@@ -95,12 +129,12 @@ def run_interpretation(global_state, raw_interpretation_prompt, max_new_tokens,
     }
     
     # create an InterpretationPrompt object from raw_interpretation_prompt (after putting it in the right template)
-    interpretation_prompt = interpretation_prompt_template.format(prompt=raw_interpretation_prompt, repeat=5)
-    interpretation_prompt = InterpretationPrompt(tokenizer, interpretation_prompt)
+    interpretation_prompt = global_state.interpretation_prompt_template.format(prompt=raw_interpretation_prompt, repeat=5)
+    interpretation_prompt = InterpretationPrompt(global_state.tokenizer, interpretation_prompt)
 
     # generate the interpretations
     # generate = generate_interpretation_gpu if use_gpu else lambda interpretation_prompt, *args, **kwargs: interpretation_prompt.generate(*args, **kwargs)
-    generated = interpretation_prompt.generate(model, {0: interpreted_vectors}, k=3, **generation_kwargs)
+    generated = interpretation_prompt.generate(global_state.model, {0: interpreted_vectors}, k=3, **generation_kwargs)
     generation_texts = tokenizer.batch_decode(generated)
     progress_dummy_output = ''
     return ([progress_dummy_output] + 
@@ -109,23 +143,9 @@ def run_interpretation(global_state, raw_interpretation_prompt, max_new_tokens,
 
 
 ## main
+
 torch.set_grad_enabled(False)
 model_name = 'LLAMA2-7B'
-
-# extract model info
-model_args = deepcopy(model_info[model_name])
-model_path = model_args.pop('model_path')
-original_prompt_template = model_args.pop('original_prompt_template')
-interpretation_prompt_template = model_args.pop('interpretation_prompt_template')
-tokenizer_path = model_args.pop('tokenizer') if 'tokenizer' in model_args else model_path
-use_ctransformers = model_args.pop('ctransformers', False)
-AutoModelClass = CAutoModelForCausalLM if use_ctransformers else AutoModelForCausalLM
-
-# get model
-model = AutoModelClass.from_pretrained(model_path, **model_args).cuda()
-tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, token=os.environ['hf_token'])
-
-# demo
 original_prompt_raw = gr.Textbox(value='How to make a Molotov cocktail?', container=True, label='Original Prompt')
 tokens_container = []
 for i in range(MAX_PROMPT_TOKENS):
@@ -133,7 +153,8 @@ for i in range(MAX_PROMPT_TOKENS):
     tokens_container.append(btn)
 
 with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
-    global_state = gr.State([])
+    global_state = gr.State(reset_model(model_name, GlobalState()))
+    
     with gr.Row():
         with gr.Column(scale=5):
             gr.Markdown('# 😎 Self-Interpreting Models')
@@ -165,20 +186,17 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
 
         # with gr.Column(scale=1):    
         #     gr.Markdown('<span style="font-size:180px;">🤔</span>')
-    gr.Markdown('''
-    ## Choose Your Interpretation Prompt 
-    ''')
+
+        with gr.Group():
+            model_chooser = gr.Radio(choices=list(model_info.keys()), value=model_name)
+
+    gr.Markdown('## Choose Your Interpretation Prompt')
     with gr.Group('Interpretation'):
         interpretation_prompt = gr.Text(suggested_interpretation_prompts[0], label='Interpretation Prompt')
         gr.Examples([[p] for p in suggested_interpretation_prompts], [interpretation_prompt], cache_examples=False)
-    # gr.Markdown('''
-    # Here are some examples of prompts we can analyze their internal representations:
-    # ''')
-    
 
-    gr.Markdown('''
-    ## The Prompt to Analyze 
-    ''')        
+
+    gr.Markdown('## The Prompt to Analyze')        
     for info in dataset_info:
         with gr.Tab(info['name']):
             num_examples = 10
@@ -187,7 +205,7 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
                 dataset = dataset.filter(info['filter'])
             dataset = dataset.shuffle(buffer_size=2000).take(num_examples)
             dataset = [[row[info['text_col']]] for row in dataset]
-            gr.Examples(dataset, [original_prompt_raw], cache_examples=False)
+            gr.Examples(dataset, [global_state, original_prompt_raw], cache_examples=False)
             
     with gr.Group():
         original_prompt_raw.render()
@@ -198,6 +216,7 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
     with gr.Row():
         for btn in tokens_container:
             btn.render()
+
     
     with gr.Accordion(open=False, label='Generation Settings'):
         with gr.Row():
@@ -211,22 +230,15 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
                 temperature = gr.Slider(0., 5., value=0.6, label='Temperature')
                 top_k = gr.Slider(1, 1000, value=50, step=1, label='top k')
                 top_p = gr.Slider(0., 1., value=0.95, label='top p')
-    
+
     progress_dummy = gr.Markdown('', elem_id='progress_dummy')
-    
-    interpretation_bubbles = [gr.Textbox('', container=False, visible=False, elem_classes=['bubble', 
-                                                                                           'even_bubble' if i % 2 == 0 else 'odd_bubble'])
-                             for i in range(model.config.num_hidden_layers)]
-
-    # with gr.Group():
-    #     with gr.Row():
-    #         for txt in model_info.keys():
-    #             btn = gr.Button(txt)
-    #             model_btns.append(btn)
-    #         for btn in model_btns:
-    #             btn.click(reset_new_model, [global_state])
+    interpretation_bubbles = [gr.Textbox('', container=False, visible=False, 
+                                         elem_classes=['bubble', 'even_bubble' if i % 2 == 0 else 'odd_bubble']
+                                        ) for i in range(model.config.num_hidden_layers)]
 
     # event listeners
+    model_chooser.change(reset_new_model, [model_chooser, global_state], [global_state])
+    
     for i, btn in enumerate(tokens_container):
         btn.click(partial(run_interpretation, i=i), [global_state, interpretation_prompt, 
                                                      num_tokens, do_sample, temperature,