Cookiecutter-MLOps/src/data/make_dataset.py

import sys
import os
import pandas as pd
import re
from scipy.sparse import save_npz
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from nltk.stem import PorterStemmer


def preprocessor(text):
    """Return a cleaned version of text"""
    # Remove HTML markup
    text = re.sub("<[^>]*>", "", text)
    # Save emoticons for later appending
    emoticons = re.findall("(?::|;|=)(?:-)?(?:\)|\(|D|P)", text)
    # Remove any non-word character and append the emoticons,
    # removing the nose character for standarization. Convert to lower case
    text = (
        re.sub("[\W]+", " ", text.lower()) + " " + " ".join(emoticons).replace("-", "")
    )
    return text


def tokenizer_porter(text):
    porter = PorterStemmer()
    token = []
    for word in text.split():
        token.append(porter.stem(word))
    return token


def preproc(data_path, output_folder):
    df = pd.read_csv(data_path)
    df["sentiment"] = df["sentiment"].apply(lambda x: 1 if x == "positive" else 0)

    X = df["review"]
    y = df["sentiment"]

    # split the dataset in train and test
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.3, random_state=42
    )
    tfidf = TfidfVectorizer(
        tokenizer=tokenizer_porter, preprocessor=preprocessor, max_features=25000
    )

    X_train = tfidf.fit_transform(X_train)
    X_test = tfidf.transform(X_test)

    save_npz(output_folder + "X_train", X_train)
    save_npz(output_folder + "X_test", X_test)

    pd.DataFrame(y_train).to_csv(output_folder + "y_train.csv", index=False)
    pd.DataFrame(y_test).to_csv(output_folder + "y_test.csv", index=False)

    return X_train, X_test, y_train, y_test


if __name__ == "__main__":
    if not (2 <= len(sys.argv) <= 3):
        print(
            "usage: %s <raw_data_file> <out_folder> (out_folder is optional)"
            % sys.argv[0],
            file=sys.stderr,
        )
        sys.exit(0)

    out_folder = sys.argv[2] if len(sys.argv) == 3 else "data/processed/"
    if not os.path.exists(out_folder):
        os.makedirs(out_folder)

    preproc(sys.argv[1], out_folder)