diff --git a/BaseML/Classification.py b/BaseML/Classification.py index 31977bd..ecac105 100644 --- a/BaseML/Classification.py +++ b/BaseML/Classification.py @@ -8,6 +8,7 @@ from sklearn.neighbors import KNeighborsClassifier from sklearn.model_selection import train_test_split from sklearn.svm import SVC +import joblib class cls: def __init__(self, algorithm='KNN', n_neighbors=5, n_estimators=100, ): @@ -75,3 +76,11 @@ def load_dataset(self, path, test_size=0.2, dataset=''): self.dataset_path = path self.test_size = test_size self.dataset=dataset + + def save(self): + print("Saving model checkpoints...") + joblib.dump(self.model, '../checkpoint.pkl', compress=3) + + + def load(self, path): + joblib.load(path) \ No newline at end of file diff --git a/BaseML/GaussianNB.py b/BaseML/GaussianNB.py index fad3ae6..cb9053a 100644 --- a/BaseML/GaussianNB.py +++ b/BaseML/GaussianNB.py @@ -1,6 +1,7 @@ import pandas as pd import numpy as np import os +import joblib from sklearn.metrics import accuracy_score, mean_squared_error from sklearn.naive_bayes import GaussianNB as gauss @@ -16,6 +17,7 @@ def __init__(self self.test_size = ' ' self.test_set = ' ' + def train(self, seed=0, data_type='csv'): np.random.seed(seed) if data_type == 'csv': @@ -30,6 +32,8 @@ def train(self, seed=0, data_type='csv'): 'label': label[train_index:] } self.model.fit(train_data, train_label) + return self.model + def inference(self, mode='cls'): pred = self.model.predict(self.test_set['data']) @@ -40,6 +44,16 @@ def inference(self, mode='cls'): loss = mean_squared_error(self.test_set['label'], pred) print('Loss: {}'.format(loss)) + def load_dataset(self, path, test_size=0.2): self.dataset_path = path self.test_size = test_size + + + def save(self): + print("Saving model checkpoints...") + joblib.dump(self.model, '../checkpoint.pkl', compress=3) + + + def load(self, path): + joblib.load(path) diff --git a/BaseML/Regression.py b/BaseML/Regression.py index 14c13c4..ef918ad 100644 --- a/BaseML/Regression.py +++ b/BaseML/Regression.py @@ -6,6 +6,7 @@ from sklearn.metrics import accuracy_score, mean_squared_error, r2_score from sklearn import linear_model from sklearn.decomposition import PCA as pca_reduction +import joblib class reg: def __init__(self, algorithm='', n_components='mle'): @@ -55,3 +56,12 @@ def load_dataset(self,path,test_size=0.2, dataset=''): self.dataset_path = path self.test_size = test_size self.dataset = dataset + + + def save(self): + print("Saving model checkpoints...") + joblib.dump(self.model, '../checkpoint.pkl', compress=3) + + + def load(self, path): + joblib.load(path) \ No newline at end of file