Update app.py

app.py

@@ -1,10 +1,10 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
 
 # Load model and tokenizer
-model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM2-135M")
-tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM2-135M")
+model = GPT2LMHeadModel.from_pretrained("gpt2")
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 
 def get_next_token_probs(text):
     # Handle empty input
@@ -24,7 +24,7 @@ def get_next_token_probs(text):
     next_token_probs = torch.softmax(next_token_logits, dim=0)
     
     # Get top-5 tokens and their probabilities
-    topk_probs, topk_indices = torch.topk(next_token_probs, 20)
+    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
@@ -46,13 +46,10 @@ with gr.Blocks(css="footer {display: none}") as demo:
     # Input textbox
     input_text = gr.Textbox(
         label="Text Input",
-        placeholder="Type text here...",
+        placeholder="Type here and watch predictions update...",
         value="The weather tomorrow will be"
     )
     
-    # Predict button
-    predict_btn = gr.Button("Predict Next Tokens")
-    
     # Simple header for results
     gr.Markdown("##### Most likely next tokens:")
     
@@ -65,8 +62,8 @@ with gr.Blocks(css="footer {display: none}") as demo:
     
     token_outputs = [token1, token2, token3, token4, token5]
     
-    # Set up button click event
-    predict_btn.click(
+    # Set up the live update
+    input_text.change(
         fn=get_next_token_probs,
         inputs=input_text,
         outputs=token_outputs