movie_data_classification.py

from sklearn.linear_model import LogisticRegression, LogisticRegressionCV
from sklearn.ensemble import AdaBoostClassifier, RandomForestClassifier
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import fbeta_score, accuracy_score

models = {'log_model': LogisticRegression(random_state=0),
          'log_cv_model': LogisticRegressionCV(),
          'ab_model':AdaBoostClassifier(random_state=0),
          'rf_model': RandomForestClassifier(random_state=0),
          'grad_model': GradientBoostingClassifier(random_state=0),
          'knn_model': KNeighborsClassifier()}
X = df_final.drop(['revenue', 'above_ave_rev_yr', 'release_date'], axis=1)
y = df_final['above_ave_rev_yr']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=0, stratify=y)

results = {}
for model in models.keys():
    name = models[model].__class__.__name__
    models[model].fit(X_train, y_train)
    
    pred_train = models[model].predict(X_train)
    pred_test = models[model].predict(X_test)
    
    acc_train = accuracy_score(y_train, pred_train)
    acc_test = accuracy_score(y_test, pred_test)
    
    f_train = fbeta_score(y_train, pred_train, beta=1)
    f_test = fbeta_score(y_test, pred_test, beta=1)
    
    results[name] = {'pred_train': pred_train,
                     'pred_test': pred_test,
                     'acc_train': acc_train,
                     'acc_test': acc_test,
                     'f_train': f_train,
                     'f_test': f_test}

for model in results.keys():
    print(model)
    print("="*len(model))
    print('\tTrain Accuracy: {:.4}'.format(results[model]['acc_train']))
    print('\tTest Accuracy: {:.4}'.format(results[model]['acc_test']))
    print('\tTrain F1: {:.4}'.format(results[model]['f_train']))
    print('\tTest F1: {:.4}'.format(results[model]['f_test']))