stackexchange-tutorial/tutorial/prepare_data.py

from . import shared
from .shared import save

def split_train_test(ratio=0.2, random_seed=42):
    import pandas as pd
    from sklearn.model_selection import train_test_split

    df = pd.read_csv(shared.raw_data, encoding='utf-8')
    df[shared.col_label] = df[shared.col_tags].fillna('').str.contains('machine-learning')

    df_positive = df[df[shared.col_label]]
    df_negative = df[df[shared.col_label] != True]

    train_df, test_df = train_test_split(df, test_size=ratio, random_state=random_seed, stratify=df[shared.col_label])

    train_df.to_csv(shared.train_data, index=False)
    test_df.to_csv(shared.test_data, index=False)


def text_preprocessing(s):
    # TODO: Clean HTML tags
    return s


def vectorize_text():
    import pandas as pd
    from sklearn.feature_extraction.text import TfidfVectorizer
    
    train_df = pd.read_csv(shared.train_data, encoding='utf-8')
    test_df = pd.read_csv(shared.test_data, encoding='utf-8')

    def text_col(df):
        df[shared.col_text] = (df[shared.col_title].fillna('') + df[shared.col_body].fillna('')).astype('U', copy=False)

    text_col(train_df)
    text_col(test_df)
    
    vectorizer = TfidfVectorizer(encoding='utf-8', preprocessor=text_preprocessing, stop_words='english', analyzer='word', max_features=50000)
    vectorizer.fit(train_df[shared.col_text])

    train_tfidf_mat = vectorizer.transform(train_df[shared.col_text])
    test_tfidf_mat = vectorizer.transform(test_df[shared.col_text])

    return train_df, test_df, vectorizer, train_tfidf_mat, test_tfidf_mat

def prepare_data():
    train_df, test_df, vectorizer, train_tfidf_mat, test_tfidf_mat = vectorize_text()
    save(vectorizer, shared.vectorizer_pkl)
    train_df[[shared.col_id, shared.col_text]].to_csv(shared.train_processed, index=False)
    test_df[[shared.col_id, shared.col_text]].to_csv(shared.test_processed, index=False)
    save(train_tfidf_mat, shared.train_tfidf)
    save(test_tfidf_mat, shared.test_tfidf)


def main():
    split_train_test()
    prepare_data()


if __name__ == "__main__":
    main()