Update train_model.py

train_model.py

@@ -12,6 +12,18 @@ except FileNotFoundError:
     print("Error: 'enhanced_mantle_training.csv' not found. Please check the file path.")
     exit()
 
+# Ensure that the necessary columns are present in the dataset
+required_columns = ["temperature", "duration", "risk_level"]
+if not all(col in df.columns for col in required_columns):
+    print(f"Error: Missing one or more required columns. Ensure the dataset contains {required_columns}.")
+    exit()
+
+# Check for any missing values in the data
+if df.isnull().any().any():
+    print("Warning: Dataset contains missing values. Please clean the data.")
+    # Optionally, you can fill missing values or drop rows with missing data
+    df = df.dropna()  # Drop rows with missing data, or use df.fillna() to fill missing values
+
 # Prepare the features (temperature, duration) and target (risk_level)
 X = df[["temperature", "duration"]]  # Features: temperature and duration
 y = df["risk_level"]  # Target: risk_level