who-moved-my-data/preprocess.py

import pandas as pd
from sklearn.model_selection import train_test_split

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


def feature_engineering(raw_df):
  df = raw_df.copy()
  df['CreationDate'] = pd.to_datetime(df['CreationDate'])
  df['CreationDate_Epoch'] = df['CreationDate'].astype('int64') // 10 ** 9
  df = df.drop(columns=['Id', 'Tags'])
  df['Title_Len'] = df.Title.str.len()
  df['Body_Len'] = df.Body.str.len()
  # Drop the correlated features
  df = df.drop(columns=['FavoriteCount'])
  df['Text'] = df['Title'].fillna('') + ' ' + df['Body'].fillna('')
  return df


def split(random_state=42):
  print('Loading data...')
  df = pd.read_csv(raw_df_path)
  df[CLASS_LABEL] = df['Tags'].str.contains('machine-learning').fillna(False)
  train_df, test_df = train_test_split(df, random_state=random_state, stratify=df[CLASS_LABEL])

  print('Engineering features...')
  train_df = feature_engineering(train_df)
  test_df = feature_engineering(test_df)

  print('Saving split data...')
  train_df.to_csv(train_df_path)
  test_df.to_csv(test_df_path)


if __name__ == '__main__':
  split()