Update src/train_model.py

src/train_model.py

@@ -49,12 +49,14 @@ torch.manual_seed(SEED)
 # y shape: (n_samples,) — continuous risk score in [0, 1]
 
 def generate_flood_data(n: int = 5000):
-    rainfall = pd.read_csv("rainfall_clean.csv")
-    flood_hist = pd.read_csv("flood_history_clean.csv")
-    soil = pd.read_csv("soil_moisture.csv")
-    drainage = pd.read_csv("drainage_capacity.csv")
-    rivers = pd.read_csv("river_network.csv")
-    elevation = pd.read_csv("elevation.csv")
+    DATA_DIR = os.getenv("DATA_DIR", "/src/data")
+
+    rainfall = pd.read_csv(os.path.join(DATA_DIR, "rainfall_clean.csv"))
+    flood_hist = pd.read_csv(os.path.join(DATA_DIR, "flood_history_clean.csv"))
+    soil = pd.read_csv(os.path.join(DATA_DIR, "soil_moisture.csv"))
+    drainage = pd.read_csv(os.path.join(DATA_DIR, "drainage_capacity.csv"))
+    rivers = pd.read_csv(os.path.join(DATA_DIR, "river_network.csv"))
+    elevation = pd.read_csv(os.path.join(DATA_DIR, "elevation.csv"))
 
     # ==============================
     # Prepare flood labels