Update app.py

app.py

@@ -18,7 +18,7 @@ CONFIG.set_default_api_key(api_key)
 access_token = os.environ['HUGGING_FACE_HUB_TOKEN']
 
 # Load the Language Model
-llama = LanguageModel("meta-llama/Meta-Llama-3.1-8B", token=access_token)
+llama = LanguageModel("meta-llama/Meta-Llama-3.1-8B")
 
 #placeholder for reset
 prompts_with_probs = pd.DataFrame(
@@ -55,7 +55,9 @@ def run_lens(model,PROMPT):
             logits_lens_token_result_by_layer.append(logits_lens_next_token)
         tokens_out = llama.lm_head.output.argmax(dim=-1).save()
         expected_token = tokens_out[0][-1].save()
+    # logits_lens_all_probs = np.concatenate([probs[:, expected_token].cpu().detach().numpy() for probs in logits_lens_probs_by_layer])
     logits_lens_all_probs = np.concatenate([probs[:, expected_token].cpu().detach().to(torch.float32).numpy() for probs in logits_lens_probs_by_layer])
+
     #get the rank of the expected token from each layer's distribution
     for layer_probs in logits_lens_probs_by_layer:
         # Sort the probabilities in descending order and find the rank of the expected token
@@ -113,7 +115,7 @@ def plot_prob(prompts_with_probs):
     # Add labels and title
     plt.xlabel('Layer Number')
     plt.ylabel('Probability of Expected Token')
-    plt.title('Prob of expected token across layers\n(annotated with decoded output at each layer)')
+    plt.title('Prob of expected token across layers\n(annotated with actual decoded output at each layer)')
     plt.grid(True)
     plt.ylim(0.0, 1.0)
     plt.legend(title='Prompts', bbox_to_anchor=(0.5, -0.15), loc='upper center', ncol=1)
@@ -177,6 +179,8 @@ def plot_prob_mean(prompts_with_probs):
     plt.title('Mean Probability of Expected Token')
     plt.xticks(rotation=45, ha='right')
     plt.grid(axis='y')
+    plt.ylim(0, 1)
+
 
     # Annotate the mean and variance on the bars
     for bar, mean, var in zip(bars, summary_stats['mean_prob'], summary_stats['variance']):
@@ -277,18 +281,25 @@ def submit_prompts(prompts_data):
 
 def clear_all(prompts):
     prompts=[['']]
+    # prompt_file=gr.File(type="filepath", label="Upload a File with Prompts")
+    prompt_file = None
     prompts_data = gr.Dataframe(headers=["Prompt"], row_count=5, col_count=1, value= prompts, type="array", interactive=True)
-    return prompts_data,plot_prob(prompts_with_probs),plot_rank(prompts_with_ranks),plot_prob_mean(prompts_with_probs),plot_rank_mean(prompts_with_ranks)
+    return prompts_data,prompt_file,plot_prob(prompts_with_probs),plot_rank(prompts_with_ranks),plot_prob_mean(prompts_with_probs),plot_rank_mean(prompts_with_ranks)
 
 
 def gradio_interface():
     with gr.Blocks(theme="gradio/monochrome") as demo:
         prompts=[['']]
-        prompts_data = gr.Dataframe(headers=["Prompt"], row_count=5, col_count=1, value= prompts, type="array", interactive=True)
-        prompt_file=gr.File(type="filepath", label="Upload a File with Prompts")
+        with gr.Row():
+            with gr.Column(scale=3):
+                prompts_data = gr.Dataframe(headers=["Prompt"], row_count=5, col_count=1, value= prompts, type="array", interactive=True)
+            with gr.Column(scale=1):
+                prompt_file=gr.File(type="filepath", label="Upload a File with Prompts")
         prompt_file.upload(process_file, inputs=[prompts_data,prompt_file], outputs=[prompts_data])
-
         # Define the outputs
+        with gr.Row():
+            clear_btn = gr.Button("Clear")
+            submit_btn = gr.Button("Submit")
         with gr.Row():
             prob_visualization = gr.Image(value=plot_prob(prompts_with_probs), type="pil",label=" ")
             rank_visualization = gr.Image(value=plot_rank(prompts_with_ranks), type="pil",label=" ")
@@ -296,14 +307,11 @@ def gradio_interface():
             prob_mean_visualization = gr.Image(value=plot_prob_mean(prompts_with_probs), type="pil",label=" ")
             rank_mean_visualization = gr.Image(value=plot_rank_mean(prompts_with_ranks), type="pil",label=" ")
 
-        with gr.Row():
-            clear_btn = gr.Button("Clear")
-            clear_btn.click(clear_all, inputs=[prompts_data], outputs=[prompts_data,prob_visualization,rank_visualization,prob_mean_visualization,rank_mean_visualization])
-            submit_btn = gr.Button("Submit")
-            submit_btn.click(submit_prompts, inputs=[prompts_data], outputs=[prob_visualization,rank_visualization,prob_mean_visualization,rank_mean_visualization])#
+        clear_btn.click(clear_all, inputs=[prompts_data], outputs=[prompts_data,prompt_file,prob_visualization,rank_visualization,prob_mean_visualization,rank_mean_visualization])
+        submit_btn.click(submit_prompts, inputs=[prompts_data], outputs=[prob_visualization,rank_visualization,prob_mean_visualization,rank_mean_visualization])#
+        prompt_file.clear(clear_all, inputs=[prompts_data], outputs=[prompts_data,prompt_file,prob_visualization,rank_visualization,prob_mean_visualization,rank_mean_visualization])
         
 
         demo.launch()
 
-
 gradio_interface()