Update app.py

app.py

@@ -1,15 +1,15 @@
 import gradio as gr
 import torch
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Load model and tokenizer
-model = GPT2LMHeadModel.from_pretrained("gpt2")
-tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM2-135M")
+tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM2-135M")
 
 def get_next_token_probs(text):
     # Handle empty input
     if not text.strip():
-        return ["No input text"] * 5
+        return ["No input text"] * 20
     
     # Tokenize input
     input_ids = tokenizer.encode(text, return_tensors="pt")
@@ -23,8 +23,8 @@ def get_next_token_probs(text):
     next_token_logits = logits[0, -1, :]
     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, 5)
+    # Get top-20 tokens and their probabilities
+    topk_probs, topk_indices = torch.topk(next_token_probs, 20)
     topk_tokens = [tokenizer.decode([idx]) for idx in topk_indices]
     
     # Format the results as strings
@@ -41,7 +41,7 @@ def get_next_token_probs(text):
 
 # Create minimal interface with simpler components
 with gr.Blocks(css="footer {display: none}") as demo:
-    gr.Markdown("### GPT-2 Next Token Predictor")
+    gr.Markdown("### SmolLM2 Next Token Predictor")
     
     # Input textbox
     input_text = gr.Textbox(
@@ -53,14 +53,8 @@ with gr.Blocks(css="footer {display: none}") as demo:
     # 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()
-    
-    token_outputs = [token1, token2, token3, token4, token5]
+    # Create 20 individual output markdown components
+    token_outputs = [gr.Markdown() for _ in range(20)]
     
     # Set up the live update
     input_text.change(