Fix loading on startup

app.py

@@ -2,7 +2,7 @@ import os
 import torch
 import gradio as gr
 import spaces
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 
 # Use a global variable to hold the current model and tokenizer
@@ -46,7 +46,7 @@ def generate_text(prompt):
         return_dict_in_generate=True,
         output_scores=False
         ).sequences[0]:
-        
+
         output_ids.append(token_id.item())
 
         yield current_tokenizer.decode(output_ids, skip_special_tokens=True)
@@ -63,21 +63,24 @@ model_choices = [
 with gr.Blocks() as demo:
     gr.Markdown("## Clinical Text Testing with LLaMA, DeepSeek, and Gemma")
 
+    # State to track initial model to load
+    default_model = gr.State("meta-llama/Llama-3.2-3B-Instruct")
+
     model_selector = gr.Dropdown(choices=model_choices, label="Select Model")
     model_status = gr.Textbox(label="Model Status", interactive=False)
 
     input_text = gr.Textbox(label="Input Clinical Text")
     generate_btn = gr.Button("Generate")
-
     output_text = gr.Textbox(label="Generated Output")
 
-    # Load model on dropdown change
+    # Auto-load default model on launch
+    demo.load(fn=load_model_on_selection, inputs=default_model, outputs=model_status)
+
+    # Manual model selection
     model_selector.change(fn=load_model_on_selection, inputs=model_selector, outputs=model_status)
 
     # Generate with current model
     generate_btn.click(fn=generate_text, inputs=input_text, outputs=output_text)
     input_text.submit(fn=generate_text, inputs=input_text, outputs=output_text)
 
-
-load_model_on_selection("meta-llama/Llama-3.2-3B-Instruct")
 demo.launch()