Update app.py

app.py

@@ -1,10 +1,7 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-from typing import Tuple, List, Dict
-import numpy as np
 
-# Select smaller models that are suitable for this task
 AVAILABLE_MODELS = {
     "distilgpt2": "distilgpt2",
     "bloomz-560m": "bigscience/bloomz-560m",
@@ -19,7 +16,6 @@ class TextGenerator:
         self.tokenizer = None
         
     def load_model(self, model_name: str) -> str:
-        """Load the selected model and tokenizer"""
         try:
             self.model = AutoModelForCausalLM.from_pretrained(AVAILABLE_MODELS[model_name])
             self.tokenizer = AutoTokenizer.from_pretrained(AVAILABLE_MODELS[model_name])
@@ -27,8 +23,7 @@ class TextGenerator:
         except Exception as e:
             return f"Error loading model: {str(e)}"
     
-    def get_next_token_predictions(self, text: str, top_k: int = 10) -> Tuple[List[str], List[float]]:
-        """Get predictions for the next token"""
+    def get_next_token_predictions(self, text: str, top_k: int = 10):
         if not self.model or not self.tokenizer:
             return [], []
         
@@ -40,12 +35,12 @@ class TextGenerator:
             
         top_k_probs, top_k_indices = torch.topk(probs, top_k)
         top_k_tokens = [self.tokenizer.decode([idx.item()]) for idx in top_k_indices]
-        top_k_probs = top_k_probs.tolist()
         
-        return top_k_tokens, top_k_probs
+        return top_k_tokens, top_k_probs.tolist()
 
-def format_predictions(tokens: List[str], probs: List[float]) -> str:
-    """Format the predictions for display"""
+generator = TextGenerator()
+
+def format_predictions(tokens, probs):
     if not tokens or not probs:
         return "No predictions available"
     
@@ -54,83 +49,40 @@ def format_predictions(tokens: List[str], probs: List[float]) -> str:
         formatted += f"'{token}' : {prob:.4f}\n"
     return formatted
 
-generator = TextGenerator()
-
-def update_output(model_name: str, text: str, custom_token: str, selected_token: str) -> Tuple[str, str, str, Dict, str]:
-    """Update the interface based on user interactions"""
+def update_output(model_name, text, custom_token, selected_token):
     output = text
     
-    # Load model if it changed
     if not generator.model or generator.model.name_or_path != AVAILABLE_MODELS[model_name]:
         load_message = generator.load_model(model_name)
         if "Error" in load_message:
             return text, "", "", gr.update(choices=[]), load_message
     
-    # Add custom token or selected token
     if custom_token:
         output += custom_token
     elif selected_token:
         output += selected_token.strip("'")
     
-    # Get new predictions
     tokens, probs = generator.get_next_token_predictions(output)
     predictions = format_predictions(tokens, probs)
-    
-    # Update dropdown choices
     token_choices = [f"'{token}'" for token in tokens]
     
     return output, "", "", gr.update(choices=token_choices), predictions
 
-# Create the interface
-demo = gr.Blocks(title="Interactive Text Generation")
-
-with demo:
-    gr.Markdown("# Interactive Text Generation")
-    gr.Markdown("Generate text by selecting predicted tokens or writing your own.")
-    
-    with gr.Row():
-        model_dropdown = gr.Dropdown(
-            choices=list(AVAILABLE_MODELS.keys()),
-            value="distilgpt2",
-            label="Select Model"
-        )
-    
-    with gr.Row():
-        text_input = gr.Textbox(
-            lines=5,
-            label="Generated Text",
-            placeholder="Start typing or select a token..."
-        )
-    
-    with gr.Row():
-        custom_token = gr.Textbox(
-            label="Custom Token",
-            placeholder="Type your own token..."
-        )
-        token_dropdown = gr.Dropdown(
-            choices=[],
-            label="Select from predicted tokens"
-        )
-    
-    with gr.Row():
-        predictions_output = gr.Textbox(
-            label="Predictions",
-            lines=12
-        )
-    
-    with gr.Row():
-        status_output = gr.Textbox(
-            label="Status",
-            lines=1
-        )
-    
-    # Update when model changes or token is added
-    for trigger in [model_dropdown, custom_token, token_dropdown]:
-        trigger.change(
-            fn=update_output,
-            inputs=[model_dropdown, text_input, custom_token, token_dropdown],
-            outputs=[text_input, custom_token, token_dropdown, token_dropdown, predictions_output]
-        )
-
-# For Hugging Face Spaces, we just need to expose the demo
-demo.launch(share=True)
+demo = gr.Interface(
+    fn=update_output,
+    inputs=[
+        gr.Dropdown(choices=list(AVAILABLE_MODELS.keys()), value="distilgpt2", label="Select Model"),
+        gr.Textbox(lines=5, label="Generated Text", placeholder="Start typing or select a token..."),
+        gr.Textbox(label="Custom Token", placeholder="Type your own token..."),
+        gr.Dropdown(choices=[], label="Select from predicted tokens")
+    ],
+    outputs=[
+        gr.Textbox(lines=5, label="Generated Text"),
+        gr.Textbox(label="Custom Token"),
+        gr.Textbox(label="Selected Token"),
+        gr.Dropdown(label="Predicted Tokens"),
+        gr.Textbox(lines=12, label="Predictions")
+    ],
+    title="Interactive Text Generation",
+    description="Generate text by selecting predicted tokens or writing your own."
+)