Update app.py

app.py

@@ -1,4 +1,3 @@
-import spaces
 import gradio as gr
 import torch
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
@@ -7,11 +6,10 @@ from transformers import GPT2LMHeadModel, GPT2Tokenizer
 model = GPT2LMHeadModel.from_pretrained("gpt2")
 tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 
-@spaces.GPU
-def get_next_token_probs(text, top_k=5):
+def get_next_token_probs(text):
     # Handle empty input
     if not text.strip():
-        return [""] * top_k
+        return ["No input text"] * 5
     
     # Tokenize input
     input_ids = tokenizer.encode(text, return_tensors="pt")
@@ -25,8 +23,8 @@ def get_next_token_probs(text, top_k=5):
     next_token_logits = logits[0, -1, :]
     next_token_probs = torch.softmax(next_token_logits, dim=0)
     
-    # Get top-k tokens and their probabilities
-    topk_probs, topk_indices = torch.topk(next_token_probs, top_k)
+    # Get top-5 tokens and their probabilities
+    topk_probs, topk_indices = torch.topk(next_token_probs, 5)
     topk_tokens = [tokenizer.decode([idx]) for idx in topk_indices]
     
     # Format the results as strings
@@ -34,35 +32,15 @@ def get_next_token_probs(text, top_k=5):
     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
+        # Clean up token display (replace space with visible space symbol)
+        display_token = token.replace(" ", "␣")
         # Format the output string
         formatted_results.append(f"{i+1}. \"{display_token}\" ({prob_percent})")
     
     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:
+# Create minimal interface with simpler components
+with gr.Blocks(css="footer {display: none}") as demo:
     gr.Markdown("### GPT-2 Next Token Predictor")
     
     # Input textbox
@@ -72,25 +50,28 @@ with gr.Blocks(css=custom_css) as demo:
         value="The weather tomorrow will be"
     )
     
-    # 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)]
+    # Simple header for results
+    gr.Markdown("##### Most likely next tokens:")
+    
+    # Individual output textboxes for each token
+    token1 = gr.Markdown()
+    token2 = gr.Markdown()
+    token3 = gr.Markdown()
+    token4 = gr.Markdown()
+    token5 = gr.Markdown()
     
-    # Function to update tokens in real-time
-    def update_tokens(text):
-        return get_next_token_probs(text)
+    token_outputs = [token1, token2, token3, token4, token5]
     
     # Set up the live update
     input_text.change(
-        fn=update_tokens,
+        fn=get_next_token_probs,
         inputs=input_text,
         outputs=token_outputs
     )
     
     # Initialize with default text
     demo.load(
-        fn=update_tokens,
+        fn=get_next_token_probs,
         inputs=input_text,
         outputs=token_outputs
     )