Preview:
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)
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