Update app.py

app.py

@@ -17,8 +17,13 @@ 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)
+
+# Model options
+MODEL_OPTIONS = {
+    "Llama3.1-8B": "meta-llama/Meta-Llama-3.1-8B",
+    "Llama3.1-70B": "meta-llama/Meta-Llama-3.1-70B",
+}
+
 
 #placeholder for reset
 prompts_with_probs = pd.DataFrame(
@@ -53,7 +58,7 @@ def run_lens(model,PROMPT):
             logits_lens_probs_by_layer.append(logits_lens_probs)
             logits_lens_next_token = torch.argmax(logits_lens_probs, dim=1).save()
             logits_lens_token_result_by_layer.append(logits_lens_next_token)
-        tokens_out = llama.lm_head.output.argmax(dim=-1).save()
+        tokens_out = model.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])
@@ -65,7 +70,7 @@ def run_lens(model,PROMPT):
         # Find the rank of the expected token (1-based rank)
         expected_token_rank = (sorted_indices == expected_token).nonzero(as_tuple=True)[1].item() + 1
         logits_lens_ranks_by_layer.append(expected_token_rank)
-    actual_output = llama.tokenizer.decode(expected_token.item())
+    actual_output = model.tokenizer.decode(expected_token.item())
     logits_lens_results = [model.tokenizer.decode(next_token.item()) for next_token in logits_lens_token_result_by_layer]
     return logits_lens_results, logits_lens_all_probs, actual_output,logits_lens_ranks_by_layer
 
@@ -98,28 +103,35 @@ def process_file(prompts_data,file_path):
 
 def plot_prob(prompts_with_probs):
     plt.figure(figsize=(10, 6))
-    
+    texts = []  # List to hold text annotations for adjustment
+
     # Iterate over each prompt and plot its probabilities
     for prompt in prompts_with_probs['prompt'].unique():
         # Filter the DataFrame for the current prompt
         prompt_data = prompts_with_probs[prompts_with_probs['prompt'] == prompt]
+        label = f"{prompt}({prompt_data['expected'].iloc[0]})"
         
         # Plot probabilities for this prompt
-        plt.plot(prompt_data['layer'], prompt_data['probs'], marker='x', label=prompt)
-        
+        plt.plot(prompt_data['layer'], prompt_data['probs'], marker='x', label=label)
+
         # Annotate each point with the corresponding result
         for layer, prob, result in zip(prompt_data['layer'], prompt_data['probs'], prompt_data['results']):
-            plt.text(layer, prob, result, fontsize=8)
-
+            text = plt.text(layer, prob, result, fontsize=8)
+            texts.append(text)  # Add text to the list
 
     # 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 actual decoded output at each layer)')
+    plt.ylabel('Probability')
+    plt.title('Probability of most-likely output token')
     plt.grid(True)
+    plt.xlim(0,max(prompts_with_probs['layer']))
     plt.ylim(0.0, 1.0)
     plt.legend(title='Prompts', bbox_to_anchor=(0.5, -0.15), loc='upper center', ncol=1)
 
+    # Adjust text to prevent overlap
+    adjust_text(texts, only_move={'points': 'xy', 'texts': 'xy'}, 
+                arrowprops=dict(arrowstyle="->", color='r', lw=0.5))
+
     # Save the plot to a buffer
     buf = io.BytesIO()
     plt.savefig(buf, format='png', bbox_inches='tight')  # Use bbox_inches to avoid cutting off labels
@@ -130,27 +142,34 @@ def plot_prob(prompts_with_probs):
 
 def plot_rank(prompts_with_ranks):
     plt.figure(figsize=(10, 6))
-    
+    texts = []  # List to hold text annotations for adjustment
+
     # Iterate over each prompt and plot its ranks
     for prompt in prompts_with_ranks['prompt'].unique():
         # Filter the DataFrame for the current prompt
         prompt_data = prompts_with_ranks[prompts_with_ranks['prompt'] == prompt]
+        label = f"{prompt}({prompt_data['expected'].iloc[0]})"
         
         # Plot ranks for this prompt
-        plt.plot(prompt_data['layer'], prompt_data['ranks'], marker='x', label=prompt)
-        
+        plt.plot(prompt_data['layer'], prompt_data['ranks'], marker='x', label=label)
+
         # Annotate each point with the corresponding result
         for layer, rank, result in zip(prompt_data['layer'], prompt_data['ranks'], prompt_data['results']):
-            plt.text(layer, rank,result, ha='right', va='bottom', fontsize=8)
+            text = plt.text(layer, rank, result, ha='right', va='bottom', fontsize=8)
+            texts.append(text)  # Add text to the list
 
     # Add labels and title
     plt.xlabel('Layer Number')
-    plt.ylabel('Rank of Expected Token')
-    plt.title('Rank of expected token across layers\n(annotated with decoded output at each layer)')
+    plt.ylabel('Rank')
+    plt.title('Rank of most-likely output token')
     plt.grid(True)
+    plt.xlim(0,max(prompts_with_ranks['layer']))
     plt.ylim(bottom=0)  # Adjust if needed, depending on your rank values
     plt.legend(title='Prompts', bbox_to_anchor=(0.5, -0.15), loc='upper center', ncol=1)
 
+    # Adjust text to prevent overlap
+    adjust_text(texts,only_move={'points': 'xy', 'texts': 'xy'},
+                arrowprops=dict(arrowstyle="->", color='r', lw=0.5))
 
     # Save the plot to a buffer
     buf = io.BytesIO()
@@ -160,77 +179,8 @@ def plot_rank(prompts_with_ranks):
     plt.close()  # Close the figure to free memory
     return img
 
-def plot_prob_mean(prompts_with_probs):
-    # Calculate mean probabilities and variance
-    summary_stats = prompts_with_probs.groupby("prompt")["probs"].agg(
-        mean_prob="mean", 
-        variance="var"
-    ).reset_index()
-
-    # Set up the bar plot
-    plt.figure(figsize=(10, 6))
-    bars = plt.bar(summary_stats['prompt'], summary_stats['mean_prob'], 
-                   yerr=summary_stats['variance']**0.5,  # Error bars are the standard deviation
-                   capsize=5, color='skyblue')
-
-    # Add labels and title
-    plt.xlabel('Prompt')
-    plt.ylabel('Mean Probability')
-    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']):
-        yval = bar.get_height()
-        plt.text(bar.get_x() + bar.get_width() / 2, yval, f'Mean: {mean:.2f}\nVar: {var:.2f}', 
-                 ha='center', va='bottom', fontsize=8, color='black')
-
-    # Save the plot to a buffer
-    buf = io.BytesIO()
-    plt.savefig(buf, format='png', bbox_inches='tight')  # Use bbox_inches to avoid cutting off labels
-    buf.seek(0)
-    img = Image.open(buf)
-    plt.close()  # Close the figure to free memory
-    return img
-
-def plot_rank_mean(prompts_with_ranks):
-    # Calculate mean ranks and variance
-    summary_stats = prompts_with_ranks.groupby("prompt")["ranks"].agg(
-        mean_rank="mean", 
-        variance="var"
-    ).reset_index()
-
-    # Set up the bar plot
-    plt.figure(figsize=(10, 6))
-    bars = plt.bar(summary_stats['prompt'], summary_stats['mean_rank'], 
-                   yerr=summary_stats['variance']**0.5,  # Error bars are the standard deviation
-                   capsize=5, color='salmon')
-
-    # Add labels and title
-    plt.xlabel('Prompt')
-    plt.ylabel('Mean Rank')
-    plt.title('Mean Rank of Expected Token')
-    plt.xticks(rotation=45, ha='right')
-    plt.grid(axis='y')
-
-    # Annotate the mean and variance on the bars
-    for bar, mean, var in zip(bars, summary_stats['mean_rank'], summary_stats['variance']):
-        yval = bar.get_height()
-        plt.text(bar.get_x() + bar.get_width() / 2, yval, f'Mean: {mean:.2f}\nVar: {var:.2f}', 
-                 ha='center', va='bottom', fontsize=8, color='black')
-
-    # Save the plot to a buffer
-    buf = io.BytesIO()
-    plt.savefig(buf, format='png', bbox_inches='tight')  # Use bbox_inches to avoid cutting off labels
-    buf.seek(0)
-    img = Image.open(buf)
-    plt.close()  # Close the figure to free memory
-    return img
-
-def submit_prompts(prompts_data):
+def submit_prompts(model_name, prompts_data):
+    llama = LanguageModel(MODEL_OPTIONS[model_name])
     # Initialize lists to accumulate results
     all_prompts = []
     all_results = []
@@ -276,20 +226,23 @@ def submit_prompts(prompts_data):
             "ranks": all_ranks,
             "expected": all_expected,
         })
-    return plot_prob(prompts_with_probs), plot_rank(prompts_with_ranks),plot_prob_mean(prompts_with_probs),plot_rank_mean(prompts_with_ranks)
+    return plot_prob(prompts_with_probs), plot_rank(prompts_with_ranks)
 
 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,prompt_file,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)
 
 def gradio_interface():
     with gr.Blocks(theme="gradio/monochrome") as demo:
-        prompts=[['']]
+        prompts = [['The Eiffel Tower is located in the city of'],['Vatican is located in the city of']]
+        
+        # prompts=[['']]
         with gr.Row():
             with gr.Column(scale=3):
+                model_dropdown = gr.Dropdown(choices=list(MODEL_OPTIONS.keys()), label="Select Model", value="Llama3.1-8B")
                 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")
@@ -298,18 +251,30 @@ def gradio_interface():
         with gr.Row():
             clear_btn = gr.Button("Clear")
             submit_btn = gr.Button("Submit")
+        
+        prompt_file.upload(process_file, inputs=[prompts_data, prompt_file], outputs=[prompts_data])
+
+
+        gr.Markdown("The most likely output token is the model's prediction at the final layer, shown in brackets in the plot legend.")
+        # Create a Markdown component for the description
         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=" ")
+            gr.Markdown("The graph below illustrates the probability of this most likely output token as it is decoded at each layer of the model. Each point on the graph is annotated with the decoded output corresponding to the token that has the highest probability at that particular layer.")
+            gr.Markdown("The graph below illustrates the rank of this most likely output token as it is decoded at each layer of the model. Each point on the graph is annotated with the decoded output corresponding to the token that has the lowest rank at that particular layer.")
+
+        prob_img, rank_img = submit_prompts(model_dropdown.value, prompts)
+        # prob_visualization.value = prob_img  # Direct assignment to value
+        # rank_visualization.value = rank_img    # Direct assignment to value
+
         with gr.Row():
-            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=" ")
+            prob_visualization = gr.Image(value=prob_img, type="pil",label=" ")
+            rank_visualization = gr.Image(value=rank_img, type="pil",label=" ")
 
-        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])
+        clear_btn.click(clear_all, inputs=[prompts_data], outputs=[prompts_data,prompt_file,prob_visualization,rank_visualization])
+        submit_btn.click(submit_prompts, inputs=[model_dropdown,prompts_data], outputs=[prob_visualization,rank_visualization])#
+        prompt_file.clear(clear_all, inputs=[prompts_data], outputs=[prompts_data,prompt_file,prob_visualization,rank_visualization])
+        
+        # Generate plots with sample prompts on load
         
-
         demo.launch()
 
 gradio_interface()