Commit 06de8271 authored by thecml's avatar thecml
Browse files

retrained alarm model

parent 52f9c4f6
Pipeline #95473 passed with stage
in 4 minutes and 32 seconds
......@@ -7,7 +7,6 @@ from utility.settings import load_settings
from sksurv.ensemble import RandomSurvivalForest
from io import BytesIO
import shutil
from sklearn.model_selection import train_test_split
def main():
data_settings = load_settings(pt.CONFIGS_DIR, "data.yaml")
......@@ -27,13 +26,9 @@ def main():
labels_enc[col_name] = le
X.loc[:, col_name] = le.transform(X.loc[:, col_name].astype(str))
X_train, X_test, y_train, y_test = train_test_split(
X, y, stratify=y['Status'], test_size=0.25, random_state=0)
model = RandomSurvivalForest(n_estimators=200, max_depth=6,
n_jobs=-1, random_state=0)
model.fit(X_train, y_train)
print(model.score(X_test, y_test))
model.fit(X, y)
with open(Path.joinpath(pt.MODELS_DIR, "alarm_labels.pkl"), 'wb') as fd:
outfile = BytesIO()
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment