LinearRegressor and Knn model with one function to compare results
Wed Feb 03 2021 21:58:02 GMT+0000 (Coordinated Universal Time)
Saved by
@fre_apples
#python
def model(data, column_to_drop, type_of_model, n_neighbors=4):
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score
from sklearn import linear_model
X = data.drop([column_to_drop], axis=1)
y = data[column_to_drop]
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=100)
if type_of_model == 'LinearRegressor':
lm = linear_model.LinearRegression()
model = lm.fit(X_train,y_train)
predictions = lm.predict(X_test)
print("The R2 of the linear model is: ",round(r2_score(y_test, predictions),3))
if type_of_model == 'KNeighborsRegressor':
from sklearn.neighbors import KNeighborsRegressor
knn_model = KNeighborsRegressor(n_neighbors=n_neighbors)
knn_model.fit(X_train, y_train)
knn_predictions = knn_model.predict(X_test)
print("The R2 of the knn model is: ",round(r2_score(y_test, knn_predictions),3))
model(data_copy, 'name_your_target_column', 'LinearRegressor')
model(data_copy, 'name_your_target_column', 'KNeighborsRegressor', n_neighbors=4)
content_copyCOPY
Comments