Update risk_model.py

risk_model.py

@@ -7,7 +7,7 @@ MODEL_PATH = "heating_model.pkl"
 DATA_PATH = "mantle_training.csv"
 HISTORY = []
 
-# Function to train the model from scratch
+# Function to train model if missing
 def train_and_save_model():
     data = pd.read_csv(DATA_PATH)
     X = data[["temperature", "duration"]]
@@ -17,10 +17,34 @@ def train_and_save_model():
     dump(model, MODEL_PATH)
     return model
 
-# Safe model loader
+# Load model safely
 def load_model():
     if not os.path.exists(MODEL_PATH):
         return train_and_save_model()
     return load(MODEL_PATH)
 
+# Load on module load
 model = load_model()
+
+def predict_risk(temp, duration):
+    global model
+    pred = model.predict([[temp, duration]])[0]
+    score = max(model.predict_proba([[temp, duration]])[0]) * 100
+    HISTORY.append({"Temperature": temp, "Duration": duration, "Risk": pred, "Confidence": round(score, 2)})
+    return pred, round(score, 2)
+
+def retrain_model():
+    try:
+        data = pd.read_csv(DATA_PATH)
+        X = data[["temperature", "duration"]]
+        y = data["risk_level"]
+        clf = RandomForestClassifier().fit(X, y)
+        dump(clf, MODEL_PATH)
+        global model
+        model = clf
+        return "✅ Model retrained successfully"
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+
+def get_history_df():
+    return pd.DataFrame(HISTORY)