From de21b65ca04531dcc851fd83fd89cd34d35ce4a5 Mon Sep 17 00:00:00 2001 From: coolneng Date: Tue, 10 Nov 2020 21:13:23 +0100 Subject: [PATCH] Simplify model selection logic --- src/processing.py | 28 ++++------------------------ 1 file changed, 4 insertions(+), 24 deletions(-) diff --git a/src/processing.py b/src/processing.py index 0f8ca17..925c001 100644 --- a/src/processing.py +++ b/src/processing.py @@ -10,35 +10,15 @@ from sklearn.tree import DecisionTreeClassifier from preprocessing import parse_data, split_k_sets -def naive_bayes(): - model = GaussianNB() - return model - - -def linear_svc(): - model = LinearSVC(random_state=42) - return model - - -def k_nearest_neighbors(): - model = KNeighborsClassifier(n_neighbors=10) - return model - - -def decision_tree(): - model = DecisionTreeClassifier(random_state=42) - return model - - def choose_model(model): if model == "gnb": - return naive_bayes() + return GaussianNB() elif model == "svc": - return linear_svc() + return LinearSVC(random_state=42) elif model == "knn": - return k_nearest_neighbors() + return KNeighborsClassifier(n_neighbors=10) elif model == "tree": - return decision_tree() + return DecisionTreeClassifier(random_state=42) def predict_data(data, target, model):