Commit 104e7937 authored by Noric Couderc's avatar Noric Couderc
Browse files

Add the feature names to the classifier object 

That way, if we load it, we can already retrieve what the names of
features were for.
parent 6eec6918

train_model.py

+3 −0
Original line number Diff line number Diff line
@@ -367,6 +367,9 @@ if __name__ == "__main__": 
    classifier = RandomForestClassifier()



    classifier.fit(X_train, y_train, w_train)

    # We just add it to the class

    classifier.feature_names = poly_transformer.get_feature_names()



    print("Accuracy: {0}".format(classifier.score(X_test, y_test, w_test)))

    print("Classifier trained")