Update app.py

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
 import torch
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
-import numpy as np
 import matplotlib.pyplot as plt
 import seaborn as sns
+import os
 
 # Load model and tokenizer
 model_name = "gpt2"
@@ -44,55 +44,55 @@ def get_token_probabilities(text, top_k=10):
     plt.ylabel("Tokens")
     plt.tight_layout()
     
-    # Save the plot to a file
-    plt.savefig("token_probabilities.png")
+    # Ensure temp directory exists
+    os.makedirs("tmp", exist_ok=True)
+    
+    # Save the plot to a file in the temp directory
+    plot_path = os.path.join("tmp", "token_probabilities.png")
+    plt.savefig(plot_path)
     plt.close()
     
-    return "token_probabilities.png", dict(zip(topk_tokens, topk_probs.tolist()))
+    return plot_path, dict(zip(topk_tokens, topk_probs.tolist()))
 
-def interface():
-    with gr.Blocks() as demo:
-        gr.Markdown("# GPT-2 Next Token Probability Visualizer")
-        gr.Markdown("Enter text and see the probabilities of possible next tokens.")
-        
-        with gr.Row():
-            with gr.Column():
-                input_text = gr.Textbox(
-                    label="Input Text",
-                    placeholder="Type some text here...",
-                    value="Hello, my name is"
-                )
-                top_k = gr.Slider(
-                    minimum=5, 
-                    maximum=20, 
-                    value=10, 
-                    step=1, 
-                    label="Number of top tokens to show"
-                )
-                btn = gr.Button("Generate Probabilities")
-            
-            with gr.Column():
-                output_image = gr.Image(label="Probability Distribution")
-                output_table = gr.JSON(label="Token Probabilities")
-        
-        btn.click(
-            fn=get_token_probabilities,
-            inputs=[input_text, top_k],
-            outputs=[output_image, output_table]
-        )
-        
-        gr.Examples(
-            examples=[
-                ["Hello, my name is", 10],
-                ["The capital of France is", 10],
-                ["Once upon a time", 10],
-                ["The best way to learn is to", 10]
-            ],
-            inputs=[input_text, top_k],
-        )
+with gr.Blocks() as demo:
+    gr.Markdown("# GPT-2 Next Token Probability Visualizer")
+    gr.Markdown("Enter text and see the probabilities of possible next tokens.")
+    
+    with gr.Row():
+        with gr.Column():
+            input_text = gr.Textbox(
+                label="Input Text",
+                placeholder="Type some text here...",
+                value="Hello, my name is"
+            )
+            top_k = gr.Slider(
+                minimum=5, 
+                maximum=20, 
+                value=10, 
+                step=1, 
+                label="Number of top tokens to show"
+            )
+            btn = gr.Button("Generate Probabilities")
         
-    return demo
+        with gr.Column():
+            output_image = gr.Image(label="Probability Distribution")
+            output_table = gr.JSON(label="Token Probabilities")
+    
+    btn.click(
+        fn=get_token_probabilities,
+        inputs=[input_text, top_k],
+        outputs=[output_image, output_table]
+    )
+    
+    gr.Examples(
+        examples=[
+            ["Hello, my name is", 10],
+            ["The capital of France is", 10],
+            ["Once upon a time", 10],
+            ["The best way to learn is to", 10]
+        ],
+        inputs=[input_text, top_k],
+    )
 
-if __name__ == "__main__":
-    demo = interface()
-    demo.launch()
+# Launch the app
+demo.launch()