21 lines
629 B
Python
21 lines
629 B
Python
|
import pandas as pd
|
||
|
from sklearn.tree import DecisionTreeClassifier
|
||
|
from sklearn.model_selection import train_test_split
|
||
|
from sklearn import metrics
|
||
|
import joblib
|
||
|
|
||
|
pima = pd.read_csv("data.csv", header=1, delimiter=';')
|
||
|
|
||
|
feature_cols = ['Size', 'Color', 'Sound', 'Time','Smell', 'Height','Temperature']
|
||
|
X = pima[feature_cols]
|
||
|
y = pima.ToRemove
|
||
|
|
||
|
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=1)
|
||
|
clf = DecisionTreeClassifier()
|
||
|
clf = clf.fit(X_train,y_train)
|
||
|
|
||
|
joblib.dump(clf, 'decision_tree_model.pkl')
|
||
|
|
||
|
y_pred = clf.predict(X_test)
|
||
|
|
||
|
print("Accuracy:",metrics.accuracy_score(y_test, y_pred))
|