Clement Vachet commited on
Commit
89365a6
·
1 Parent(s): 711e3f5

Switch ML model to decision tree

Browse files
Files changed (1) hide show
  1. classification/classifier.py +7 -6
classification/classifier.py CHANGED
@@ -1,4 +1,5 @@
1
- from sklearn.ensemble import AdaBoostClassifier
 
2
  from sklearn.datasets import load_iris
3
  from sklearn.model_selection import train_test_split
4
  import joblib
@@ -13,13 +14,13 @@ class Classifier:
13
  def train_and_save(self):
14
  print("\nIRIS model training...")
15
  iris = load_iris()
16
- ada = AdaBoostClassifier(n_estimators=5)
17
 
18
  X_train, X_test, y_train, y_test = train_test_split(iris.data, iris.target, test_size=0.1, random_state=42)
19
- model = ada.fit(X_train, y_train)
20
 
21
- print(f"Model score: {ada.score(X_train, y_train):.3f}")
22
- print(f"Test Accuracy: {ada.score(X_test, y_test):.3f}")
23
 
24
  current_dir = os.path.dirname(os.path.abspath(__file__))
25
  parent_dir = os.path.dirname(current_dir)
@@ -42,7 +43,7 @@ class Classifier:
42
  model = joblib.load(model_path)
43
 
44
  features = np.array(data)
45
-
46
  if features.shape[-1] != 4:
47
  raise ValueError("Expected 4 features per input.")
48
 
 
1
+ # from sklearn.ensemble import AdaBoostClassifier
2
+ from sklearn.tree import DecisionTreeClassifier
3
  from sklearn.datasets import load_iris
4
  from sklearn.model_selection import train_test_split
5
  import joblib
 
14
  def train_and_save(self):
15
  print("\nIRIS model training...")
16
  iris = load_iris()
17
+ cart = DecisionTreeClassifier(max_depth = 3)
18
 
19
  X_train, X_test, y_train, y_test = train_test_split(iris.data, iris.target, test_size=0.1, random_state=42)
20
+ model = cart.fit(X_train, y_train)
21
 
22
+ print(f"Model score: {cart.score(X_train, y_train):.3f}")
23
+ print(f"Test Accuracy: {cart.score(X_test, y_test):.3f}")
24
 
25
  current_dir = os.path.dirname(os.path.abspath(__file__))
26
  parent_dir = os.path.dirname(current_dir)
 
43
  model = joblib.load(model_path)
44
 
45
  features = np.array(data)
46
+
47
  if features.shape[-1] != 4:
48
  raise ValueError("Expected 4 features per input.")
49