46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
import pandas as pd
|
|
from sklearn.model_selection import train_test_split
|
|
from sklearn.tree import DecisionTreeClassifier
|
|
from sklearn.metrics import classification_report, confusion_matrix
|
|
from joblib import dump, load
|
|
|
|
|
|
class DecisionTree:
|
|
def __init__(self):
|
|
self.classifier = load('../../files/decision tree/classifier.joblib')
|
|
|
|
|
|
def Learning():
|
|
# Uploading data from file
|
|
dataset = pd.read_csv("../../files/decision tree/database.csv")
|
|
|
|
print(f'Shape: {dataset.shape}')
|
|
print(f'Head:\n{dataset.head()}')
|
|
|
|
X = dataset.drop('decision', axis=1)
|
|
y = dataset['decision']
|
|
|
|
# Split data to training and test sets
|
|
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.10)
|
|
|
|
# Training
|
|
classifier = DecisionTreeClassifier()
|
|
classifier.fit(X_train, y_train)
|
|
|
|
# Predictions test
|
|
y_pred = classifier.predict(X_test)
|
|
# print(y_pred)
|
|
|
|
# my_dict = {'known': [1], 'power': [1], 'new': [1], 'location': [0], 'stable': [1], 'chain_reaction': [1]}
|
|
# s = pd.DataFrame.from_dict(my_dict)
|
|
# predict = self.classifier.predict(s)
|
|
#
|
|
# print(predict)
|
|
|
|
print()
|
|
print(confusion_matrix(y_test, y_pred))
|
|
print(classification_report(y_test, y_pred))
|
|
|
|
dump(classifier, '../../files/decision tree/classifier.joblib')
|
|
|