Update app.py

app.py

@@ -18,8 +18,7 @@ if not os.path.exists(log_file):
 # Model information
 MODEL_ID = "Solus-PG/gemma-2b-gaslighting"  # Replace with your actual model path
 
-# Load the model and tokenizer (with error handling)
-@gr.on_load(api_name="load_model")
+# Load the model and tokenizer
 def load_model():
     print("Loading model...")
     try:
@@ -33,14 +32,21 @@ def load_model():
         return model, tokenizer
     except Exception as e:
         print(f"Error loading model: {e}")
-        raise gr.Error(f"Failed to load model: {str(e)}")
+        return None, None
 
 # Generate response function
-def generate_response(prompt, model_state):
-    if not model_state:
-        return "Model is not loaded yet. Please wait or refresh the page."
-    
-    model, tokenizer = model_state
+def generate_response(prompt, model_state=None):
+    # If model state isn't passed, try loading the model
+    if model_state is None or model_state == [None, None]:
+        try:
+            model, tokenizer = load_model()
+        except Exception:
+            return "Failed to load the model. Please check logs or try again later."
+    else:
+        model, tokenizer = model_state
+        
+    if model is None or tokenizer is None:
+        return "Model couldn't be loaded. Please check logs or try again later."
     
     try:
         # Format as chat for the model
@@ -85,6 +91,9 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
     **Note: The responses from this model should not be taken as truth.**
     """)
     
+    # Load model at startup and store in session state
+    model_state = gr.State(value=load_model())
+    
     with gr.Row():
         with gr.Column():
             input_text = gr.Textbox(
@@ -101,11 +110,6 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
                 label="AI Response"
             )
     
-    model_state = gr.State()
-    
-    # Load model when the app starts
-    demo.load(load_model, outputs=model_state)
-    
     # Handle submission
     submit_btn.click(
         fn=generate_response,