Update app.py

app.py

@@ -1,7 +1,10 @@
+import os
+import numpy as np
+import torch
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
 
+# Available models
 AVAILABLE_MODELS = {
     "distilgpt2": "distilgpt2",
     "bloomz-560m": "bigscience/bloomz-560m",
@@ -10,80 +13,88 @@ AVAILABLE_MODELS = {
     "pythia-160m": "EleutherAI/pythia-160m"
 }
 
-generator = None
+# Access token for Hugging Face
+HF_TOKEN = os.getenv('HF_TOKEN')
+
+# Initialize model and tokenizer globally
+current_model = None
+current_tokenizer = None
+current_model_name = None
 
 def load_model(model_name):
-    global generator
-    try:
-        model = AutoModelForCausalLM.from_pretrained(AVAILABLE_MODELS[model_name])
-        tokenizer = AutoTokenizer.from_pretrained(AVAILABLE_MODELS[model_name])
-        generator = (model, tokenizer)
-        return f"Successfully loaded {model_name}"
-    except Exception as e:
-        return f"Error loading model: {str(e)}"
+    global current_model, current_tokenizer, current_model_name
+    if current_model_name != model_name:
+        current_model = AutoModelForCausalLM.from_pretrained(AVAILABLE_MODELS[model_name], use_auth_token=HF_TOKEN)
+        current_tokenizer = AutoTokenizer.from_pretrained(AVAILABLE_MODELS[model_name], use_auth_token=HF_TOKEN)
+        current_model_name = model_name
 
-def get_predictions(text, model_name):
-    global generator
-    if not generator:
-        load_model(model_name)
+def get_next_token_predictions(text, model_name, top_k=10):
+    global current_model, current_tokenizer
     
-    model, tokenizer = generator
-    inputs = tokenizer(text, return_tensors="pt")
+    # Load model if needed
+    if current_model_name != model_name:
+        load_model(model_name)
     
+    # Get predictions
+    inputs = current_tokenizer(text, return_tensors="pt")
     with torch.no_grad():
-        outputs = model(**inputs)
+        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=10)
-    top_k_tokens = [tokenizer.decode([idx.item()]) for idx in top_k_indices]
-    predictions = "\n".join([f"'{token}' : {prob:.4f}" for token, prob in zip(top_k_tokens, top_k_probs)])
+    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, predictions
+    return top_k_tokens, top_k_probs.tolist()
 
-def generate(model_name, text, token_choice="", custom_token=""):
-    if token_choice:
-        text += token_choice.strip("'")
+def predict_next_token(text, model_name, custom_token=""):
+    # Add custom token if provided
     if custom_token:
         text += custom_token
-        
-    tokens, predictions = get_predictions(text, model_name)
-    return text, gr.Dropdown(choices=[f"'{t}'" for t in tokens]), predictions
-
-with gr.Blocks() as demo:
-    gr.Markdown("# Interactive Text Generation")
     
-    model_name = gr.Dropdown(
-        choices=list(AVAILABLE_MODELS.keys()),
-        value="distilgpt2",
-        label="Select Model"
-    )
+    # Get predictions
+    tokens, probs = get_next_token_predictions(text, model_name)
     
-    text = gr.Textbox(
-        lines=5,
-        label="Text",
-        placeholder="Type or select tokens to generate text..."
-    )
+    # Format predictions
+    predictions = "\n".join([f"'{token}' : {prob:.4f}" for token, prob in zip(tokens, probs)])
     
-    with gr.Row():
-        token_choice = gr.Dropdown(
-            choices=[],
-            label="Select predicted token"
-        )
-        custom_token = gr.Textbox(
-            label="Or type custom token"
-        )
-    
-    predictions = gr.Textbox(
-        label="Predictions",
-        lines=10
-    )
-    
-    for component in [model_name, token_choice, custom_token]:
-        component.change(
-            generate,
-            inputs=[model_name, text, token_choice, custom_token],
-            outputs=[text, token_choice, predictions]
-        )
+    return text, gr.Dropdown(choices=[f"'{t}'" for t in tokens]), predictions
+
+# Page content
+title = "Interactive Text Generation with Transformer Models"
+description = """
+This application allows you to interactively generate text using various transformer models. 
+You can either select from the predicted next tokens or write your own tokens to continue the text generation.
+
+Select a model, start typing or choose from the predicted tokens, and see how the model continues your text!
+"""
+
+# Example inputs
+examples = [
+    ["The quick brown fox", "distilgpt2"],
+    ["In a galaxy far", "gpt2-medium"],
+    ["Once upon a time", "opt-350m"],
+]
+
+# Create the interface
+app = gr.Interface(
+    fn=predict_next_token,
+    inputs=[
+        gr.Textbox(lines=5, label="Text"),
+        gr.Dropdown(choices=list(AVAILABLE_MODELS.keys()), value="distilgpt2", label="Model"),
+        gr.Textbox(label="Custom token (optional)")
+    ],
+    outputs=[
+        gr.Textbox(lines=5, label="Generated text"),
+        gr.Dropdown(label="Predicted tokens"),
+        gr.Textbox(lines=10, label="Token probabilities")
+    ],
+    theme="huggingface",
+    title=title,
+    description=description,
+    examples=examples,
+    allow_flagging="manual"
+)
 
-demo.queue().launch(share=True)
+# Launch the app
+app.launch()