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)
Preview:
downloadDownload PNG
downloadDownload JPEG
downloadDownload SVG
Tip: You can change the style, width & colours of the snippet with the inspect tool before clicking Download!
Click to optimize width for Twitter