factory-ml / app /factory_predictor.py
juniorcmauricio's picture
Change the machine learning model from Logistic Regression to Random Forest Classifier.
a9e66fb verified
import pandas as pd
#from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import accuracy_score
class FactoryPredictor:
def __init__(self):
# Added random_state=42 for reproducibility. n_estimators = 100 is a default value.
self.model = RandomForestClassifier(n_estimators=100, random_state=42)
self.le_diagnosis = LabelEncoder()
self.temperature = 0
self.pressure = 0
def train(self, csv_train='train.csv', csv_test='test.csv'):
# Load the train data
data_train = pd.read_csv(csv_train)
# Encode categorical variables
data_train['diagnosis'] = self.le_diagnosis.fit_transform(data_train['diagnosis'])
# Split features and target
X_train = data_train[
['temperature', 'pressure']]
y_train = data_train['diagnosis']
# Train the model
self.model.fit(X_train, y_train)
# Load the test data
data_test = pd.read_csv(csv_test)
# Encode categorical variables
data_test['diagnosis'] = self.le_diagnosis.fit_transform(data_test['diagnosis'])
# Split features and target
X_test = data_test[
['temperature', 'pressure']]
y_test = data_test['diagnosis']
# Evaluate the model
y_pred = self.model.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)
print(f"Model accuracy: {accuracy:.2f}")
def inform_temperature(self, temperature):
self.temperature = temperature
if self.pressure > 0:
return self.predict(self.temperature, self.pressure)
return 'N'
def inform_pressure(self, pressure):
self.pressure = pressure
if self.temperature > 0:
return self.predict(self.temperature, self.pressure)
return 'N'
def predict(self, temperature, pressure):
# Create a DataFrame with the same feature names as the training data
input_data = pd.DataFrame(
[[temperature, pressure]],
columns=['temperature', 'pressure'])
# Make prediction
prediction = self.model.predict(input_data)
# Decode prediction
diagnosis = self.le_diagnosis.inverse_transform(prediction)[0]
return diagnosis