Update app.py

app.py

@@ -2,92 +2,97 @@ import spaces
 import gradio as gr
 import torch
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
-import matplotlib.pyplot as plt
-import seaborn as sns
-import os
 
 # Load model and tokenizer
-model_name = "gpt2"
-model = GPT2LMHeadModel.from_pretrained(model_name)
-tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+model = GPT2LMHeadModel.from_pretrained("gpt2")
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 
 @spaces.GPU
-def get_token_probabilities(text, top_k=10):
-    # Tokenize the input text
-    input_ids = tokenizer.encode(text, return_tensors="pt")
+def get_next_token_probs(text, top_k=5):
+    # Handle empty input
+    if not text.strip():
+        return [""] * top_k
     
-    # Get the last token's position
-    last_token_position = input_ids.shape[1] - 1
+    # Tokenize input
+    input_ids = tokenizer.encode(text, return_tensors="pt")
     
-    # Get model predictions
+    # Get predictions
     with torch.no_grad():
         outputs = model(input_ids)
         logits = outputs.logits
     
-    # Get probabilities for the next token after the last token
-    next_token_logits = logits[0, last_token_position, :]
+    # Get probabilities for next token
+    next_token_logits = logits[0, -1, :]
     next_token_probs = torch.softmax(next_token_logits, dim=0)
     
-    # Get top k most likely tokens
+    # Get top-k tokens and their probabilities
     topk_probs, topk_indices = torch.topk(next_token_probs, top_k)
-    
-    # Convert to numpy for easier handling
-    topk_probs = topk_probs.numpy()
-    topk_indices = topk_indices.numpy()
-    
-    # Decode tokens
     topk_tokens = [tokenizer.decode([idx]) for idx in topk_indices]
     
-    # Create a plot
-    plt.figure(figsize=(10, 6))
-    sns.barplot(x=topk_probs, y=topk_tokens)
-    plt.title(f"Top {top_k} token probabilities after: '{text}'")
-    plt.xlabel("Probability")
-    plt.ylabel("Tokens")
-    plt.tight_layout()
+    # Format the results as strings
+    formatted_results = []
+    for i, (token, prob) in enumerate(zip(topk_tokens, topk_probs)):
+        # Format probability as percentage with 1 decimal place
+        prob_percent = f"{prob.item()*100:.1f}%"
+        # Clean up token display (remove leading space if present)
+        display_token = token.replace(" ", "␣")  # Replace space with visible space symbol
+        # Format the output string
+        formatted_results.append(f"{i+1}. \"{display_token}\" ({prob_percent})")
     
-    # Ensure temp directory exists
-    os.makedirs("tmp", exist_ok=True)
+    return formatted_results
+
+# Create custom CSS
+custom_css = """
+.token-box {
+    margin-top: 10px;
+    padding: 15px;
+    border-radius: 8px;
+    background-color: #f7f7f7;
+    font-family: monospace;
+    font-size: 16px;
+}
+.token-item {
+    margin: 8px 0;
+    padding: 8px;
+    background-color: white;
+    border-left: 4px solid #2c8ecb;
+    border-radius: 4px;
+}
+footer {display: none}
+"""
+
+# Create minimal interface
+with gr.Blocks(css=custom_css) as demo:
+    gr.Markdown("### GPT-2 Next Token Predictor")
     
-    # 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()
+    # Input textbox
+    input_text = gr.Textbox(
+        label="Text Input",
+        placeholder="Type here and watch predictions update...",
+        value="The weather tomorrow will be"
+    )
     
-    return plot_path, dict(zip(topk_tokens, topk_probs.tolist()))
-
-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.")
+    # Container for token displays
+    with gr.Box(elem_classes=["token-box"]):
+        gr.Markdown("##### Most likely next tokens:")
+        token_outputs = [gr.Markdown(elem_classes=["token-item"]) for _ in range(5)]
     
-    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")
+    # Function to update tokens in real-time
+    def update_tokens(text):
+        return get_next_token_probs(text)
     
-    btn.click(
-        fn=get_token_probabilities,
-        inputs=[input_text, top_k],
-        outputs=[output_image, output_table]
+    # Set up the live update
+    input_text.change(
+        fn=update_tokens,
+        inputs=input_text,
+        outputs=token_outputs
     )
     
-    gr.Examples(
-        inputs=[input_text, top_k],
+    # Initialize with default text
+    demo.load(
+        fn=update_tokens,
+        inputs=input_text,
+        outputs=token_outputs
     )
 
 # Launch the app