Update app.py

app.py

@@ -17,6 +17,7 @@ def fake_gpu():
 import numpy as np
 import torch
 import gradio as gr
+import matplotlib.pyplot as plt
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import spaces
 
@@ -67,16 +68,32 @@ def get_next_token_predictions(text, model_name, top_k=10):
     
     return top_k_tokens, top_k_probs.cpu().tolist()
 
+def plot_probabilities(tokens, probs):
+    """Generate a horizontal bar chart for token probabilities."""
+    fig, ax = plt.subplots(figsize=(8, 5))
+    ax.barh(tokens[::-1], probs[::-1], color="skyblue")
+    ax.set_xlabel("Probability")
+    ax.set_title("Next Token Predictions")
+    plt.tight_layout()
+
+    # Save plot as an image and return the file path
+    plot_path = "token_probabilities.png"
+    plt.savefig(plot_path)
+    plt.close(fig)
+    
+    return plot_path
+
 def predict_next_token(model_name, text, top_k, custom_token=""):
-    """Get predictions and update the UI."""
+    """Get predictions and update the UI with text and a chart."""
     if custom_token:
         text += custom_token
 
     tokens, probs = get_next_token_predictions(text, model_name, top_k)
 
-    predictions = "\n".join([f"'{token}': {prob:.4f}" for token, prob in zip(tokens, probs)])
-    
-    return gr.update(choices=[f"'{t}'" for t in tokens]), predictions
+    # Generate bar chart
+    plot_path = plot_probabilities(tokens, probs)
+
+    return gr.update(choices=[f"'{t}'" for t in tokens]), plot_path
 
 def append_selected_token(text, selected_token):
     """Append selected token from dropdown to the text input."""
@@ -128,16 +145,13 @@ with gr.Blocks() as demo:
         append_button = gr.Button("Append Token")
 
     with gr.Row():
-        predictions_output = gr.Textbox(
-            lines=10,
-            label="Token Probabilities"
-        )
+        chart_output = gr.Image(label="Token Probability Chart")
 
     # Button click events
     predict_button.click(
         predict_next_token,
         inputs=[model_dropdown, text_input, top_k_slider],
-        outputs=[token_dropdown, predictions_output]
+        outputs=[token_dropdown, chart_output]
     )
 
     append_button.click(