Risk1 / risk_model.py
Sanjayraju30's picture
Update risk_model.py
d921cdb verified
raw
history blame
1.47 kB
import os
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from joblib import dump, load
from datetime import datetime
import pytz
MODEL_PATH = "heating_model.pkl"
DATA_PATH = "mantle_training.csv"
HISTORY = []
def get_ist_time():
ist = pytz.timezone('Asia/Kolkata')
return datetime.now(ist).strftime("%Y-%m-%d %H:%M:%S %Z")
def train_and_save_model():
data = pd.read_csv(DATA_PATH)
X = data[["temperature", "duration"]]
y = data["risk_level"]
model = RandomForestClassifier()
model.fit(X, y)
dump(model, MODEL_PATH)
return model
def load_model():
if not os.path.exists(MODEL_PATH):
return train_and_save_model()
return load(MODEL_PATH)
model = load_model()
def predict_risk(temp, duration):
global model
pred = model.predict([[temp, duration]])[0]
timestamp = get_ist_time()
HISTORY.append({
"Temperature": temp,
"Duration": duration,
"Risk": pred,
"Timestamp": timestamp
})
return pred, timestamp
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)