-
Notifications
You must be signed in to change notification settings - Fork 6
/
train.py
23 lines (17 loc) · 674 Bytes
/
train.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
import numpy as np
import pandas as pd
from sklearn import preprocessing
from sklearn.ensemble import RandomForestClassifier
from sklearn.naive_bayes import GaussianNB
import joblib
df_train = pd.read_csv('dataset/Training.csv', delimiter=',')
df_test = pd.read_csv('dataset/Testing.csv', delimiter=',')
nRow, nCol = df_train.shape
label = preprocessing.LabelEncoder()
label.fit(pd.concat([df_train['prognosis'], df_test['prognosis']]))
labels = label.fit_transform(df_train['prognosis'])
symptoms = df_train[df_train.columns.difference(['prognosis'])]
#model = RandomForestClassifier()
model = GaussianNB()
model.fit(symptoms, labels)
joblib.dump(model, 'model.pkl')