movie_data_rating_classification.py

from sklearn.linear_model import LinearRegression, Ridge
from sklearn.model_selection import train_test_split

X = df_final.drop(['vote_average', 'release_date'], axis=1)
y = df_final['vote_average']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=0)

lm_model = LinearRegression(normalize=True)
r_model = Ridge(normalize=True)

lm_model.fit(X_train, y_train)
r_model.fit(X_train, y_train)
print('Linear Regression Train R2: {}'.format(lm_model.score(X_train, y_train)))
print('Ridge Train R2: {}'.format(r_model.score(X_train, y_train)))

lm_preds = lm_model.predict(X_test)
r_preds = r_model.predict(X_test)
print('Linear Regression Test R2: {}'.format(round(lm_model.score(X_test, y_test),4)))
print('Ridge Test R2: {}'.format(round(r_model.score(X_test, y_test),4)))