Create app.py

app.py

@@ -0,0 +1,132 @@
+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",
+    "gpt2-medium": "gpt2-medium",
+    "opt-350m": "facebook/opt-350m",
+    "pythia-160m": "EleutherAI/pythia-160m"
+}
+
+class TextGenerator:
+    def __init__(self):
+        self.model = None
+        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])
+            return f"Successfully loaded {model_name}"
+        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"""
+        if not self.model or not self.tokenizer:
+            return [], []
+        
+        inputs = self.tokenizer(text, return_tensors="pt")
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            logits = outputs.logits[0, -1, :]
+            probs = torch.nn.functional.softmax(logits, dim=-1)
+            
+        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
+
+def format_predictions(tokens: List[str], probs: List[float]) -> str:
+    """Format the predictions for display"""
+    if not tokens or not probs:
+        return "No predictions available"
+    
+    formatted = "Predicted next tokens:\n\n"
+    for token, prob in zip(tokens, probs):
+        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"""
+    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
+
+with gr.Blocks() as app:
+    gr.Markdown("# Interactive Text Generation")
+    
+    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]
+        )
+
+if __name__ == "__main__":
+    app.launch()