EquipmentDashboard / train_model.py
VaneshDev's picture
Create train_model.py
af1d675 verified
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report
import joblib
# Load data
data = pd.read_csv("equipment_usage_history.csv")
# Features and target
X = data[['usage_hours', 'idle_hours', 'movement_freq', 'cost_per_hour']]
y = data['suggestion_label']
# Encode target labels
le = LabelEncoder()
y_encoded = le.fit_transform(y)
# Split train/test data
X_train, X_test, y_train, y_test = train_test_split(X, y_encoded, test_size=0.2, random_state=42)
# Scale features
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)
# Train Logistic Regression classifier
clf = LogisticRegression(multi_class='ovr', max_iter=1000)
clf.fit(X_train_scaled, y_train)
# Predict & evaluate
y_pred = clf.predict(X_test_scaled)
print(classification_report(y_test, y_pred, target_names=le.classes_))
# Save model and scaler for inference later
joblib.dump(clf, "logistic_model.joblib")
joblib.dump(scaler, "scaler.joblib")
joblib.dump(le, "label_encoder.joblib")
print("Model, scaler and label encoder saved successfully.")