Update train_model.py

train_model.py

@@ -4,8 +4,11 @@ from sklearn.ensemble import RandomForestClassifier
 from sklearn.preprocessing import LabelEncoder
 import joblib
 
-MODEL_PATH = "heating_risk_model.pkl"
-LABEL_ENCODER_PATH = "label_encoder.pkl"
+MODEL_DIR = "model_files"
+os.makedirs(MODEL_DIR, exist_ok=True)  # Ensure folder exists
+
+MODEL_PATH = os.path.join(MODEL_DIR, "heating_risk_model.pkl")
+LABEL_ENCODER_PATH = os.path.join(MODEL_DIR, "label_encoder.pkl")
 
 def train_and_save_model():
     if os.path.exists(MODEL_PATH) and os.path.exists(LABEL_ENCODER_PATH):
@@ -13,19 +16,15 @@ def train_and_save_model():
         return
 
     print("Training model...")
-
-    # Load CSV with utf-8-sig encoding to handle BOM and special chars
     data = pd.read_csv("heating_mantle_sample_30.csv", encoding='utf-8-sig')
 
     print("Original columns:", data.columns.tolist())
 
-    # Clean column names: strip spaces and remove strange chars
     data.columns = data.columns.str.strip()
     data.columns = data.columns.str.replace('Â', '', regex=False)
 
     print("Cleaned columns:", data.columns.tolist())
 
-    # Rename columns for ease of use
     data.rename(columns={
         'Max_Temperature (°C)': 'Max_Temperature',
         'Duration_Minutes': 'Duration',
@@ -34,7 +33,6 @@ def train_and_save_model():
 
     print("Columns after rename:", data.columns.tolist())
 
-    # Create Risk_Level from Risk_Score
     def risk_label_from_score(score):
         if score < 30:
             return 'Low'
@@ -48,7 +46,6 @@ def train_and_save_model():
     le = LabelEncoder()
     data['risk_encoded'] = le.fit_transform(data['Risk_Level'])
 
-    # Check final columns for features
     print("Using features:", ['Max_Temperature', 'Duration'])
 
     X = data[['Max_Temperature', 'Duration']]