Update app.py

app.py

@@ -1,5 +1,10 @@
 import os
-# Handle Spaces GPU
+import numpy as np
+import torch
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Handle Hugging Face Spaces GPU
 if os.environ.get("SPACES_ZERO_GPU") is not None:
     import spaces
 else:
@@ -14,11 +19,6 @@ else:
 def fake_gpu():
     pass
 
-import numpy as np
-import torch
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
 # Available models
 AVAILABLE_MODELS = {
     "distilgpt2": "distilgpt2",
@@ -28,58 +28,67 @@ AVAILABLE_MODELS = {
     "pythia-160m": "EleutherAI/pythia-160m"
 }
 
-# Initialize model and tokenizer globally
+# Initialize model and tokenizer
 current_model = None
 current_tokenizer = None
 current_model_name = None
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def load_model(model_name):
+    """Load the selected model and tokenizer."""
     global current_model, current_tokenizer, current_model_name
     if current_model_name != model_name:
-        current_model = AutoModelForCausalLM.from_pretrained(AVAILABLE_MODELS[model_name])
+        current_model = AutoModelForCausalLM.from_pretrained(AVAILABLE_MODELS[model_name]).to(device)
         current_tokenizer = AutoTokenizer.from_pretrained(AVAILABLE_MODELS[model_name])
         current_model_name = model_name
 
+# Load the default model at startup
+load_model("distilgpt2")
+
 def get_next_token_predictions(text, model_name, top_k=10):
+    """Generate the next token predictions with their probabilities."""
     global current_model, current_tokenizer
     
-    # Load model if needed
+    # Load the model if it has changed
     if current_model_name != model_name:
         load_model(model_name)
     
-    # Get predictions
-    inputs = current_tokenizer(text, return_tensors="pt")
+    inputs = current_tokenizer(text, return_tensors="pt").to(device)
+
     with torch.no_grad():
         outputs = current_model(**inputs)
         logits = outputs.logits[0, -1, :]
         probs = torch.nn.functional.softmax(logits, dim=-1)
-        
+    
     top_k_probs, top_k_indices = torch.topk(probs, k=top_k)
     top_k_tokens = [current_tokenizer.decode([idx.item()]) for idx in top_k_indices]
     
-    return top_k_tokens, top_k_probs.tolist()
+    return top_k_tokens, top_k_probs.cpu().tolist()
 
-def predict_next_token(model_name, text, custom_token=""):
-    # Add custom token if provided
+def predict_next_token(model_name, text, top_k, custom_token=""):
+    """Get predictions and update the UI."""
     if custom_token:
         text += custom_token
-    
-    # Get predictions
-    tokens, probs = get_next_token_predictions(text, model_name)
-    
-    # Format predictions
-    predictions = "\n".join([f"'{token}' : {prob:.4f}" for token, prob in zip(tokens, probs)])
+
+    tokens, probs = get_next_token_predictions(text, model_name, top_k)
+
+    predictions = "\n".join([f"'{token}': {prob:.4f}" for token, prob in zip(tokens, probs)])
     
     return gr.update(choices=[f"'{t}'" for t in tokens]), predictions
 
-# Create the interface
+def append_selected_token(text, selected_token):
+    """Append selected token from dropdown to the text input."""
+    if selected_token:
+        text += f" {selected_token.strip('\'')}"
+    return text
+
+# Create the UI
 with gr.Blocks() as demo:
-    gr.Markdown("# Interactive Text Generation with Transformer Models")
-    
-    gr.Markdown("""
-    This application allows you to interactively generate text using various transformer models. 
-    Select a model, enter your text, and click predict to see the possible next tokens and their probabilities.
-    """)
+    gr.Markdown("# 🔥 Interactive Text Prediction with Transformers")
+    gr.Markdown(
+        "This application lets you interactively generate text using multiple transformer models. "
+        "Choose a model, type your text, and explore token predictions."
+    )
     
     with gr.Row():
         model_dropdown = gr.Dropdown(
@@ -91,31 +100,47 @@ with gr.Blocks() as demo:
     with gr.Row():
         text_input = gr.Textbox(
             lines=5,
-            label="Text",
+            label="Input Text",
             placeholder="Type your text here...",
             value="The quick brown fox"
         )
     
+    with gr.Row():
+        top_k_slider = gr.Slider(
+            minimum=1,
+            maximum=20,
+            value=10,
+            step=1,
+            label="Top-k Predictions"
+        )
+
     with gr.Row():
         predict_button = gr.Button("Predict")
-        
+    
     with gr.Row():
         token_dropdown = gr.Dropdown(
-            label="Predicted tokens",
+            label="Predicted Tokens",
             choices=[]
         )
-        
+        append_button = gr.Button("Append Token")
+
     with gr.Row():
         predictions_output = gr.Textbox(
             lines=10,
-            label="Token probabilities"
+            label="Token Probabilities"
         )
-    
-    # Set up predict button event handler
+
+    # Button click events
     predict_button.click(
         predict_next_token,
-        inputs=[model_dropdown, text_input],
+        inputs=[model_dropdown, text_input, top_k_slider],
         outputs=[token_dropdown, predictions_output]
     )
 
-demo.queue().launch()
+    append_button.click(
+        append_selected_token,
+        inputs=[text_input, token_dropdown],
+        outputs=text_input
+    )
+
+demo.queue().launch()