Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+# Handle Spaces GPU
 if os.environ.get("SPACES_ZERO_GPU") is not None:
     import spaces
 else:
@@ -11,7 +12,7 @@ else:
 
 @spaces.GPU
 def fake_gpu():
-  pass
+    pass
 
 import numpy as np
 import torch
@@ -58,7 +59,7 @@ def get_next_token_predictions(text, model_name, top_k=10):
     
     return top_k_tokens, top_k_probs.tolist()
 
-def predict_next_token(text, model_name, custom_token=""):
+def predict_next_token(model_name, text, custom_token=""):
     # Add custom token if provided
     if custom_token:
         text += custom_token
@@ -69,7 +70,7 @@ def predict_next_token(text, model_name, custom_token=""):
     # Format predictions
     predictions = "\n".join([f"'{token}' : {prob:.4f}" for token, prob in zip(tokens, probs)])
     
-    return text, gr.update(choices=[f"'{t}'" for t in tokens]), predictions
+    return gr.update(choices=[f"'{t}'" for t in tokens]), predictions
 
 # Create the interface
 with gr.Blocks() as demo:
@@ -77,31 +78,26 @@ with gr.Blocks() as demo:
     
     gr.Markdown("""
     This application allows you to interactively generate text using various transformer models. 
-    You can either select from the predicted next tokens or write your own tokens to continue the text generation.
-
-    Select a model, start typing or choose from the predicted tokens, and see how the model continues your text!
+    Select a model, enter your text, and click predict to see the possible next tokens and their probabilities.
     """)
     
-    with gr.Row():
-        text_input = gr.Textbox(
-            lines=5,
-            label="Text",
-            placeholder="Type your text here...",
-            value="The quick brown fox"
-        )
-        
     with gr.Row():
         model_dropdown = gr.Dropdown(
             choices=list(AVAILABLE_MODELS.keys()),
             value="distilgpt2",
             label="Select Model"
         )
-        
+    
     with gr.Row():
-        custom_input = gr.Textbox(
-            label="Custom token (optional)",
-            placeholder="Type a custom token..."
+        text_input = gr.Textbox(
+            lines=5,
+            label="Text",
+            placeholder="Type your text here...",
+            value="The quick brown fox"
         )
+    
+    with gr.Row():
+        predict_button = gr.Button("Predict")
         
     with gr.Row():
         token_dropdown = gr.Dropdown(
@@ -115,29 +111,11 @@ with gr.Blocks() as demo:
             label="Token probabilities"
         )
     
-    # Set up event handlers
-    text_input.change(
-        predict_next_token,
-        inputs=[text_input, model_dropdown, custom_input],
-        outputs=[text_input, token_dropdown, predictions_output]
-    )
-    
-    model_dropdown.change(
-        predict_next_token,
-        inputs=[text_input, model_dropdown, custom_input],
-        outputs=[text_input, token_dropdown, predictions_output]
-    )
-    
-    custom_input.change(
-        predict_next_token,
-        inputs=[text_input, model_dropdown, custom_input],
-        outputs=[text_input, token_dropdown, predictions_output]
-    )
-    
-    token_dropdown.change(
+    # Set up predict button event handler
+    predict_button.click(
         predict_next_token,
-        inputs=[text_input, model_dropdown, gr.Textbox(value="")],
-        outputs=[text_input, token_dropdown, predictions_output]
+        inputs=[model_dropdown, text_input],
+        outputs=[token_dropdown, predictions_output]
     )
 
 demo.queue().launch()