Update src/train_model.py

src/train_model.py

@@ -19,6 +19,10 @@ Synthetic Data Strategy:
     Replace the generate_*_data() functions with your own data loaders.
     The rest of the training pipeline stays identical.
 """
+import os
+
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+DATA_DIR = os.path.join(BASE_DIR, "data")
 
 import torch
 import numpy as np
@@ -49,7 +53,7 @@ torch.manual_seed(SEED)
 # y shape: (n_samples,) — continuous risk score in [0, 1]
 
 def generate_flood_data(n: int = 5000):
-    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"))