fix theme

app.py

@@ -249,19 +249,75 @@ def show_sample_predictions(num_samples: int = 10) -> str:
 
 
 # Determine model directory (for HF Spaces, check environment variable or use default)
-MODEL_DIR = os.getenv("MODEL_DIR", "Prompt-Classifier") 
+# For HF Spaces, models are typically in the root directory or a subdirectory
+MODEL_DIR = os.getenv("MODEL_DIR", None)
+
+# Try common locations for models in HF Spaces
+if MODEL_DIR is None:
+    possible_paths = [
+        "./model",  # Common HF Spaces location
+        "./models",
+        "/model",
+    ]
+    for path in possible_paths:
+        if os.path.exists(path) and os.path.isdir(path):
+            MODEL_DIR = path
+            break
+    
+    # If still None, try to use a Hugging Face model identifier
+    if MODEL_DIR is None:
+        # Use environment variable if set, otherwise use default Hugging Face model
+        MODEL_DIR = os.getenv("HF_MODEL_ID", "Tameem7/Prompt-Classifier")
 
 # Load model and data on startup
 print("Initializing model and dataset...")
-try:
-    load_model_and_data(MODEL_DIR)
-except Exception as e:
-    print(f"Error loading model: {e}")
-    print("Please ensure the model directory is correct or set MODEL_DIR environment variable.")
+model_loaded = False
+if MODEL_DIR:
+    try:
+        load_model_and_data(MODEL_DIR)
+        model_loaded = True
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        print("Please ensure the model directory is correct or set MODEL_DIR environment variable.")
+        print("The app will still launch, but model functionality will be disabled.")
+else:
+    print("No model directory specified. Please set MODEL_DIR environment variable.")
+    print("The app will still launch, but model functionality will be disabled.")
 
 
 # Create Gradio interface
-with gr.Blocks(title="Prompt Injection Detector", theme=gr.themes.Soft()) as app:
+# Handle theme parameter compatibility with different Gradio versions
+# Try to create Blocks with theme, fallback if not supported
+try:
+    # Check if themes module exists and try to use it
+    if hasattr(gr, 'themes') and hasattr(gr.themes, 'Soft'):
+        app = gr.Blocks(title="Prompt Injection Detector", theme=gr.themes.Soft())
+    else:
+        app = gr.Blocks(title="Prompt Injection Detector")
+except (TypeError, AttributeError):
+    # Fallback: theme parameter not supported in this Gradio version
+    try:
+        app = gr.Blocks(title="Prompt Injection Detector")
+    except TypeError:
+        # Even title might not be supported in very old versions
+        app = gr.Blocks()
+
+with app:
+    # Show warning if model is not loaded
+    if not model_loaded:
+        gr.Markdown(
+            """
+            ## ⚠️ Warning: Model Not Loaded
+            
+            The model could not be loaded. Please ensure:
+            - The model is uploaded to this Space, or
+            - Set the `MODEL_DIR` environment variable to the model path, or  
+            - Set the `HF_MODEL_ID` environment variable to a Hugging Face model identifier
+            
+            The app will still run, but classification features will be disabled.
+            """
+        )
+    
     gr.Markdown(
         """
         # 🔒 Prompt Injection Detection Classifier