who-moved-my-data/train.py

import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import roc_auc_score, average_precision_score, accuracy_score, precision_score, recall_score, \
    f1_score
from sklearn.model_selection import train_test_split
import joblib
import os
import dagshub

# Consts
CLASS_LABEL = 'MachineLearning'
train_df_path = 'data/train.csv'
test_df_path = 'data/test.csv'


def fit_tfidf(train_df, test_df):
  tfidf = TfidfVectorizer(max_features=25000, ngram_range=(1, 2))
  tfidf.fit(train_df['Text'])
  train_tfidf = tfidf.transform(train_df['Text'])
  test_tfidf = tfidf.transform(test_df['Text'])
  return train_tfidf, test_tfidf, tfidf


def fit_model(train_X, train_y, random_state=42):
  clf_tfidf = RandomForestClassifier(random_state=random_state, max_depth=50, class_weight='balanced')
  clf_tfidf.fit(train_X, train_y)
  return clf_tfidf


def eval_model(clf, X, y):
  y_proba = clf.predict_proba(X)[:, 1]
  y_pred = clf.predict(X)
  return {
    'roc_auc': roc_auc_score(y, y_proba),
    'average_precision': average_precision_score(y, y_proba),
    'accuracy': accuracy_score(y, y_pred),
    'precision': precision_score(y, y_pred),
    'recall': recall_score(y, y_pred),
    'f1': f1_score(y, y_pred),
  }


# Prepare a dictionary of either hyperparams or metrics for logging.
def prepare_log(d, prefix=''):
  if prefix:
    prefix = f'{prefix}__'

  # Ensure all logged values are suitable for logging - complex objects aren't supported.
  def sanitize(value):
    return value if value is None or type(value) in [str, int, float, bool] else str(value)

  return {f'{prefix}{k}': sanitize(v) for k, v in d.items()}


def train():
  print('Loading data...')
  train_df = pd.read_csv(train_df_path)
  test_df = pd.read_csv(test_df_path)

  # Create outputs directory if it doesn't exist
  os.mkdir("outputs")

  with dagshub.dagshub_logger() as logger:
    print('Fitting TFIDF...')
    train_tfidf, test_tfidf, tfidf = fit_tfidf(train_df, test_df)

    print('Saving TFIDF object...')
    joblib.dump(tfidf, 'outputs/tfidf.joblib')
    logger.log_hyperparams(prepare_log(tfidf.get_params(), 'tfidf'))

    print('Training model...')
    train_y = train_df[CLASS_LABEL]
    model = fit_model(train_tfidf, train_y)

    print('Saving trained model...')
    joblib.dump(model, 'outputs/model.joblib')
    logger.log_hyperparams(model_class=type(model).__name__)
    logger.log_hyperparams(prepare_log(model.get_params(), 'model'))

    print('Evaluating model...')
    train_metrics = eval_model(model, train_tfidf, train_y)
    print('Train metrics:')
    print(train_metrics)
    logger.log_metrics(prepare_log(train_metrics, 'train'))

    test_metrics = eval_model(model, test_tfidf, test_df[CLASS_LABEL])
    print('Test metrics:')
    print(test_metrics)
    logger.log_metrics(prepare_log(test_metrics, 'test'))


if __name__ == '__main__':
  train()