Update app.py

app.py

@@ -13,9 +13,6 @@ AVAILABLE_MODELS = {
     "pythia-160m": "EleutherAI/pythia-160m"
 }
 
-# Access token for Hugging Face
-HF_TOKEN = os.getenv('HF_TOKEN')
-
 # Initialize model and tokenizer globally
 current_model = None
 current_tokenizer = None
@@ -24,8 +21,8 @@ current_model_name = None
 def load_model(model_name):
     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 = AutoModelForCausalLM.from_pretrained(AVAILABLE_MODELS[model_name])
+        current_tokenizer = AutoTokenizer.from_pretrained(AVAILABLE_MODELS[model_name])
         current_model_name = model_name
 
 def get_next_token_predictions(text, model_name, top_k=10):
@@ -58,43 +55,75 @@ def predict_next_token(text, model_name, custom_token=""):
     # Format predictions
     predictions = "\n".join([f"'{token}' : {prob:.4f}" for token, prob in zip(tokens, probs)])
     
-    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"],
-]
+    return text, gr.update(choices=[f"'{t}'" for t in tokens]), predictions
 
 # 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"
-)
+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. 
+    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!
+    """)
+    
+    with gr.Row():
+        text_input = gr.Textbox(
+            lines=5,
+            label="Text",
+            placeholder="Type your text here...",
+            value="The quick brown fox"
+        )
+        
+    with gr.Row():
+        model_dropdown = gr.Dropdown(
+            choices=list(AVAILABLE_MODELS.keys()),
+            value="distilgpt2",
+            label="Select Model"
+        )
+        
+    with gr.Row():
+        custom_input = gr.Textbox(
+            label="Custom token (optional)",
+            placeholder="Type a custom token..."
+        )
+        
+    with gr.Row():
+        token_dropdown = gr.Dropdown(
+            label="Predicted tokens",
+            choices=[]
+        )
+        
+    with gr.Row():
+        predictions_output = gr.Textbox(
+            lines=10,
+            label="Token probabilities"
+        )
+    
+    # Set up event handlers
+    text_input.change(
+        predict_next_token,
+        inputs=[text_input, model_dropdown, custom_input],
+        outputs=[text_input, token_dropdown, predictions_output]
+    )
+    
+    model_dropdown.change(
+        predict_next_token,
+        inputs=[text_input, model_dropdown, custom_input],
+        outputs=[text_input, token_dropdown, predictions_output]
+    )
+    
+    custom_input.change(
+        predict_next_token,
+        inputs=[text_input, model_dropdown, custom_input],
+        outputs=[text_input, token_dropdown, predictions_output]
+    )
+    
+    token_dropdown.change(
+        predict_next_token,
+        inputs=[text_input, model_dropdown, gr.Textbox(value="")],
+        outputs=[text_input, token_dropdown, predictions_output]
+    )
 
-# Launch the app
-app.launch()
+demo.queue().launch()