Update app.py

app.py

@@ -75,7 +75,7 @@ def initialize_gpu():
     pass
 
 
-def reset_model(model_name): 
+def reset_model(model_name, demo_blocks): 
     # extract model info
     model_args = deepcopy(model_info[model_name])
     model_path = model_args.pop('model_path')
@@ -91,6 +91,7 @@ def reset_model(model_name):
     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 demo_blocks
     
 
 def get_hidden_states(raw_original_prompt):
@@ -189,58 +190,59 @@ 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>')
 
-        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('## The Prompt to Analyze')        
-    for info in dataset_info:
-        with gr.Tab(info['name']):
-            num_examples = 10
-            dataset = load_dataset(info['hf_repo'], split='train', streaming=True)
-            if 'filter' in info:
-                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)
-            
     with gr.Group():
-        original_prompt_raw.render()
-        original_prompt_btn = gr.Button('Output Token List', variant='primary')
-
-    gr.Markdown('### Here go the tokens of the prompt (click on the one to explore)')
+        model_chooser = gr.Radio(choices=list(model_info.keys()), value=model_name)
+    
+    with gr.Blocks() as demo_blocks:
+        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)
     
-    with gr.Row():
-        for btn in tokens_container:
-            btn.render()
-
     
-    with gr.Accordion(open=False, label='Generation Settings'):
+        gr.Markdown('## The Prompt to Analyze')        
+        for info in dataset_info:
+            with gr.Tab(info['name']):
+                num_examples = 10
+                dataset = load_dataset(info['hf_repo'], split='train', streaming=True)
+                if 'filter' in info:
+                    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)
+                
+        with gr.Group():
+            original_prompt_raw.render()
+            original_prompt_btn = gr.Button('Output Token List', variant='primary')
+    
+        gr.Markdown('### Here go the tokens of the prompt (click on the one to explore)')
+        
         with gr.Row():
-            num_tokens = gr.Slider(1, 100, step=1, value=20, label='Max. # of Tokens')
-            repetition_penalty = gr.Slider(1., 10., value=1, label='Repetition Penalty')
-            length_penalty = gr.Slider(0, 5, value=0, label='Length Penalty')
-            # num_beams = gr.Slider(1, 20, value=1, step=1, label='Number of Beams')
-        do_sample = gr.Checkbox(label='With sampling')
-        with gr.Accordion(label='Sampling Parameters'):
+            for btn in tokens_container:
+                btn.render()
+    
+        
+        with gr.Accordion(open=False, label='Generation Settings'):
             with gr.Row():
-                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(MAX_NUM_LAYERS)]
-
+                num_tokens = gr.Slider(1, 100, step=1, value=20, label='Max. # of Tokens')
+                repetition_penalty = gr.Slider(1., 10., value=1, label='Repetition Penalty')
+                length_penalty = gr.Slider(0, 5, value=0, label='Length Penalty')
+                # num_beams = gr.Slider(1, 20, value=1, step=1, label='Number of Beams')
+            do_sample = gr.Checkbox(label='With sampling')
+            with gr.Accordion(label='Sampling Parameters'):
+                with gr.Row():
+                    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(MAX_NUM_LAYERS)]
+    
     
     # event listeners
-    model_chooser.change(reset_model, [model_chooser], [])
+    model_chooser.change(reset_model, [model_chooser, demo_blocks], [demo_blocks])
     
     for i, btn in enumerate(tokens_container):
         btn.click(partial(run_interpretation, i=i), [interpretation_prompt,