Update src/streamlit_app.py

src/streamlit_app.py

@@ -29,33 +29,39 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # Load model and threshold
 try:
-    # Adjust paths for Hugging Face Spaces (relative to repository root)
-    model_path = os.path.join(os.path.dirname(__file__), '..', 'fraud_gnn_model.pth')
-    threshold_path = os.path.join(os.path.dirname(__file__), '..', 'optimal_threshold.txt')
-    
-    # Alternative: If model files are in a 'models' folder
+    # Try root directory first (Hugging Face Spaces working directory)
+    model_path = 'fraud_gnn_model.pth'
+    threshold_path = 'optimal_threshold.txt'
+
+    # Fallback: Try relative to src/ (if files are misplaced)
+    if not os.path.exists(model_path):
+        model_path = os.path.join(os.path.dirname(__file__), 'fraud_gnn_model.pth')
+    if not os.path.exists(threshold_path):
+        threshold_path = os.path.join(os.path.dirname(__file__), 'optimal_threshold.txt')
+
+    # Alternative: If files are in a 'models/' folder (uncomment if applicable)
     # model_path = os.path.join(os.path.dirname(__file__), '..', 'models', 'fraud_gnn_model.pth')
     # threshold_path = os.path.join(os.path.dirname(__file__), '..', 'models', 'optimal_threshold.txt')
-    
+
     if not os.path.exists(model_path):
-        raise FileNotFoundError(f"Model file not found at {model_path}")
+        raise FileNotFoundError(f"Model file not found at {model_path}. Please upload fraud_gnn_model.pth to the repository root.")
     if not os.path.exists(threshold_path):
-        raise FileNotFoundError(f"Threshold file not found at {threshold_path}")
-        
+        raise FileNotFoundError(f"Threshold file not found at {threshold_path}. Please upload optimal_threshold.txt to the repository root.")
+
     model = FraudGNN(input_dim=7, hidden_dim=16, output_dim=1).to(device)
     model.load_state_dict(torch.load(model_path, map_location=device))
     model.eval()
-    
+
     with open(threshold_path, 'r') as f:
         threshold = float(f.read())
 except FileNotFoundError as e:
-    st.error(f"Error: {e}. Please ensure model and threshold files are uploaded to the repository root.")
+    st.error(f"Error: {e}")
     st.stop()
 except Exception as e:
     st.error(f"Error loading model or threshold: {e}")
     st.stop()
 
-# City and state mappings (unchanged)
+# City and state mappings
 city_mapping = {
     'Atlanta': 0, 'Bronx': 1, 'Brooklyn': 2, 'Chicago': 3, 'Dallas': 4, 'Houston': 5,
     'Indianapolis': 6, 'Las Vegas': 7, 'Los Angeles': 8, 'Louisville': 9, 'Miami': 10,