Update app.py

app.py

@@ -2,7 +2,8 @@ import gradio as gr
 import json
 import matplotlib.pyplot as plt
 import pandas as pd
-from io import StringIO
+import io
+import base64
 
 # Function to process and visualize log probs
 def visualize_logprobs(json_input):
@@ -21,7 +22,7 @@ def visualize_logprobs(json_input):
                 token = entry['token']
                 logprob = entry['logprob']
                 top_logprobs = entry['top_logprobs']
-                # Extract top 3 alternatives
+                # Extract top 3 alternatives, sorted by log prob (most probable first)
                 top_3 = sorted(top_logprobs.items(), key=lambda x: x[1], reverse=True)[:3]
                 row = [token, f"{logprob:.4f}"]
                 for alt_token, alt_logprob in top_3:
@@ -41,19 +42,24 @@ def visualize_logprobs(json_input):
         plt.xticks(range(len(logprobs)), tokens, rotation=45, ha='right')
         plt.tight_layout()
         
-        # Save plot to a buffer
-        img_buffer = StringIO()
-        plt.savefig(img_buffer, format='png', bbox_inches='tight')
-        img_buffer.seek(0)
+        # Save plot to a bytes buffer
+        buf = io.BytesIO()
+        plt.savefig(buf, format='png', bbox_inches='tight')
+        buf.seek(0)
         plt.close()
         
+        # Convert buffer to base64 for Gradio
+        img_bytes = buf.getvalue()
+        img_base64 = base64.b64encode(img_bytes).decode('utf-8')
+        img_html = f'<img src="data:image/png;base64,{img_base64}" style="max-width: 100%; height: auto;">'
+        
         # Create a DataFrame for the table
         df = pd.DataFrame(
             table_data,
             columns=["Token", "Log Prob", "Top 1 Alternative", "Top 2 Alternative", "Top 3 Alternative"]
         )
         
-        return img_buffer, df
+        return img_html, df
     
     except Exception as e:
         return f"Error: {str(e)}", None
@@ -67,7 +73,7 @@ with gr.Blocks(title="Log Probability Visualizer") as app:
     json_input = gr.Textbox(label="JSON Input", lines=10, placeholder="Paste your JSON here...")
     
     # Outputs
-    plot_output = gr.Image(label="Log Probability Plot")
+    plot_output = gr.HTML(label="Log Probability Plot")
     table_output = gr.Dataframe(label="Token Log Probabilities and Top Alternatives")
     
     # Button to trigger visualization