dataProcess.py

# Purpose: 
# Clean + feature-engineer merged email dataset, build TF-IDF artifacts, and split the final_dataset into Train/Val/Test splits.
#
# Inputs:
#   merged_datasets.csv    
# Outputs:
#   final_dataset.csv
#   splits/final_train.csv
#   splits/final_val.csv
#   splits/final_test.csv
#   tfidf/vectorizer.pkl
#   tfidf/tfidf_full.npz
#   tfidf/tfidf_train.npz
#   tfidf/tfidf_val.npz
#   tfidf/tfidf_test.npz
#
from pathlib import Path
import re
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
import joblib
from scipy import sparse as sp

# ----------------- Config -----------------
DATA_DIR = Path("COS30049/Assignment Datasets")
IN_PATHS = [
    DATA_DIR / "merged_dataset.csv", # input dataset from dataMerge.py
]

OUT_FINAL = DATA_DIR / "final_dataset.csv" # output path for final cleaned + feature-engineered dataset

# Folders 
SPLIT_DIR = DATA_DIR / "splits" # folder for train/val/test splits
TFIDF_DIR = DATA_DIR / "tfidf" # folder for TF-IDF values

# TF-IDF settings (fit on train only; transform val/test/full)
TFIDF_PARAMS = dict(
    ngram_range=(1, 2),
    min_df=2,
    max_features=30_000,
    lowercase=True,        # OK since we also keep uppercase signal via ratio_upper feature
    strip_accents="unicode"
)

RANDOM_STATE = 42 # for reproducibility 
TRAIN_FRAC, VAL_FRAC, TEST_FRAC = 0.70, 0.15, 0.15 # data splits from final_dataset (70% train, 15% val, 15% test)

# ----------------- Helpers -----------------

# Makes sure that merged_dataset.csv exists
def get_input_path() -> Path:
    for p in IN_PATHS:
        if p.exists():
            return p
    raise FileNotFoundError(
        "Could not find merged dataset. Expected one of:\n" # prints Error message if merged_dataset.csv not found
        + "\n".join(str(p) for p in IN_PATHS)
    )

def safe_div(num: float, den: float) -> float:
    return float(num) / float(den) if den else 0.0 # avoid division by zero

def collapse_spaces(text: str) -> str:
    t = text.strip() # Trim and collapse any whitespace runs to a single space
    t = re.sub(r"\s+", " ", t)
    return t

# Feature extractors
URL_RE    = re.compile(r"\b(?:https?://|www\.)\S+", re.IGNORECASE) # simple URL pattern recognizer
EMAIL_RE  = re.compile(r"[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}") # simple email pattern recognizer
NUMSEQ_RE = re.compile(r"\d+") # any sequence of digits
REPEAT_RE = re.compile(r"(.)\1{2,}")  # any char repeated 3+ times

LETTER_RE = re.compile(r"[A-Za-z]") # any letter (used for total letters in ratio_upper)
UPPER_RE  = re.compile(r"[A-Z]") # any uppercase letter (used for total uppercase letters in ratio_upper)
SPACE_RE  = re.compile(r"\s") # any whitespace character (used for ratio_space)

def engineer_features(df: pd.DataFrame, text_col: str = "Message") -> pd.DataFrame:
    # Ensure we’re working with strings; replace NaN with empty string to avoid errors downstream
    msg = df[text_col].fillna("")

    # ------ Basic size/length features ------
    len_char = msg.str.len() # total number of characters in the message (includes spaces and punctuation)
    tokens = msg.str.split() # split on whitespace into tokens; this is a simple proxy for words
    len_word = tokens.apply(len) # number of tokens/words

    def avg_wlen(toks): # average word length (characters per token). Returns 0.0 for empty messages
        if not toks:
            return 0.0
        chars = sum(len(tok) for tok in toks)
        return chars / len(toks) if len(toks) else 0.0

    avg_word_len = tokens.apply(avg_wlen)

    # ------ Ratios: uppercase, digits, whitespace ------
    def count_letters(s: str) -> int:
        return len(LETTER_RE.findall(s)) # count ASCII letters only (A–Z/a–z)

    def count_upper(s: str) -> int:
        return len(UPPER_RE.findall(s)) # count uppercase ASCII letters (A–Z)

    def count_digits(s: str) -> int:
        return sum(ch.isdigit() for ch in s) # count characters that are digits 0–9 (fast and Unicode-safe for digits)

    def count_spaces(s: str) -> int:
        return len(SPACE_RE.findall(s)) # count all whitespace characters (space, tab, newline, etc).

    letters = msg.apply(count_letters) # total letters (A–Z/a–z)
    uppers  = msg.apply(count_upper) # total uppercase letters (A–Z)
    digits  = msg.apply(count_digits) # total digit characters (0–9)
    spaces  = msg.apply(count_spaces) # total whitespace characters

    # Share of uppercase letters among all letters (handles zero-letter cases via safe_div).
    ratio_upper = [safe_div(u, L) for u, L in zip(uppers, letters)]

    # Share of digit characters among all characters.
    ratio_digit = [safe_div(d, c) for d, c in zip(digits, len_char)]

    # Share of whitespace among all characters.
    ratio_space = [safe_div(s, c) for s, c in zip(spaces, len_char)]

    # ------ Special punctuation character counts ------
    # .str.count uses regex by default
    count_exclam = msg.str.count(r"!")   # number of '!'
    count_dollar = msg.str.count(r"\$")  # number of '$' 
    count_qmark  = msg.str.count(r"\?")  # number of '?' 
    count_hash   = msg.str.count(r"#")   # number of '#'
    count_dot    = msg.str.count(r"\.")  # number of '.' 
    
    # ------ Structural pattern counts ------
    count_url    = msg.apply(lambda s: len(URL_RE.findall(s))) # count URLs using a pragmatic regex 
    count_email  = msg.apply(lambda s: len(EMAIL_RE.findall(s))) # count email addresses 
    count_number = msg.apply(lambda s: len(NUMSEQ_RE.findall(s))) # count sequences of digits (e.g., "2025", "12345") and not the sum of individual digits
    count_repeats = msg.apply(lambda s: len(REPEAT_RE.findall(s))) # count runs of any character repeated 3+ times (e.g., "!!!", "$$$", "aaaa")

    # ------ Assemble the feature frame ------
    feats = pd.DataFrame({
        "len_char":      len_char, # total characters
        "len_word":      len_word, # total words/tokens
        "avg_word_len":  avg_word_len, # average word length (chars per token)
        "ratio_upper":   ratio_upper, # ratio of uppercase letters to total letters
        "ratio_digit":   ratio_digit, # ratio of digit characters to total characters 
        "ratio_space":   ratio_space, # ratio of whitespace characters to total characters 
        "count_exclam":  count_exclam, # number of '!'
        "count_dollar":  count_dollar, # number of '$'
        "count_qmark":   count_qmark, # number of '?'
        "count_hash":    count_hash, # number of '#'
        "count_dot":     count_dot, # number of '.'
        "count_url":     count_url, # number of URLs
        "count_email":   count_email, # number of email addresses
        "count_number":  count_number, # number of digit sequences
        "count_repeats": count_repeats, # number of character runs (3+ repeats)
    })

    # Return original df with features appended (reset_index to avoid misalignment)
    return pd.concat([df.reset_index(drop=True), feats.reset_index(drop=True)], axis=1)


# Ensure output directories exist
def assert_dirs_exist():
    if not SPLIT_DIR.exists():
        raise FileNotFoundError(
            f"Folder not found: {SPLIT_DIR}\n" # prints Error message if splits/ folder not found
            "Please create it before running this script."
        )
    if not TFIDF_DIR.exists():
        raise FileNotFoundError(
            f"Folder not found: {TFIDF_DIR}\n" # prints Error message if tfidf/ folder not found
            "Please create it before running this script."
        )

# ----------------- Pipeline -----------------
def main():
    # Ensure output folders exist 
    assert_dirs_exist()

    in_path = get_input_path()
    df = pd.read_csv(in_path)

    # 0) Schema checks
    required = {"Message", "Category", "Source"}
    missing = required - set(df.columns)
    if missing:
        raise KeyError(f"Missing required columns: {sorted(missing)}")

    # Make sure Category is numeric 0/1
    df["Category"] = pd.to_numeric(df["Category"], errors="coerce")
    n_before = len(df)

    # Report null counts quickly (pre-clean)
    null_counts = df[["Message", "Category", "Source"]].isna().sum().to_dict()

    # Drop invalid labels (not 0/1)
    invalid_label_mask = ~df["Category"].isin([0, 1])
    dropped_invalid_label = int(invalid_label_mask.sum())
    df = df[~invalid_label_mask].copy()

    # 1) Missing/duplicates
    # Trim + collapse whitespace in Message (authenticity kept otherwise)
    df["Message"] = df["Message"].astype(str).apply(collapse_spaces)

    # Drop rows where Message is empty after trimming
    empty_msg_mask = df["Message"].str.len() == 0
    dropped_empty_message = int(empty_msg_mask.sum())
    df = df[~empty_msg_mask].copy()

    # Deduplicate on (Message, Category)
    dedup_before = len(df)
    df = df.drop_duplicates(subset=["Message", "Category"], keep="first").copy()
    dropped_duplicates = dedup_before - len(df)

    # 3A) Engineered features
    df_feat = engineer_features(df, text_col="Message")

    # 3B) TF-IDF (fit on train only; transform val/test/full)
    # Splits first (on the FEATURE-ENHANCED df to keep row order aligned in outputs)
    X = df_feat["Message"].values
    y = df_feat["Category"].values.astype(int)

    # Stratified 70/30, then 15/15 from the 30%
    X_train, X_temp, y_train, y_temp, idx_train, idx_temp = train_test_split(
        X, y, np.arange(len(df_feat)),
        test_size=(1.0 - TRAIN_FRAC),
        random_state=RANDOM_STATE,
        stratify=y
    )
    # Now split temp into val/test equally
    val_ratio_of_temp = VAL_FRAC / (VAL_FRAC + TEST_FRAC)  # 0.5 for 15/15
    X_val, X_test, y_val, y_test, idx_val, idx_test = train_test_split(
        X_temp, y_temp, idx_temp,
        test_size=(1 - val_ratio_of_temp),
        random_state=RANDOM_STATE,
        stratify=y_temp
    )

    # Fit TF-IDF on TRAIN only
    vectorizer = TfidfVectorizer(**TFIDF_PARAMS)
    X_train_tfidf = vectorizer.fit_transform(X_train)
    # Transform others with the train-fitted vocabulary
    X_val_tfidf   = vectorizer.transform(X_val)
    X_test_tfidf  = vectorizer.transform(X_test)
    X_full_tfidf  = vectorizer.transform(df_feat["Message"].values)

    # 5) Quality checks & reporting
    n_after = len(df_feat)
    class_counts = df_feat["Category"].value_counts().to_dict()
    class_pct = {k: round(v * 100.0 / n_after, 2) for k, v in class_counts.items()}

    print("=== DataProcess Summary ===")
    print(f"Input file: {in_path}")
    print(f"Rows before: {n_before}")
    print(f" - Dropped invalid label rows: {dropped_invalid_label}")
    print(f" - Dropped empty-Message rows: {dropped_empty_message}")
    print(f" - Dropped duplicates: {dropped_duplicates}")
    print(f"Rows after: {n_after}")
    print(f"Null counts (pre-clean): {null_counts}")
    print(f"Class distribution: {class_counts}  ({class_pct} %)")

    # 6) Outputs

    # Final dataset CSV (Category, Source, Message + engineered features)
    df_feat.to_csv(OUT_FINAL, index=False)

    # Train/Val/Test CSVs (preserving engineered features too)
    df_train = df_feat.iloc[idx_train]
    df_val   = df_feat.iloc[idx_val]
    df_test  = df_feat.iloc[idx_test]

    (SPLIT_DIR / "final_train.csv").write_text("")  # sanity check ability to write
    (SPLIT_DIR / "final_val.csv").write_text("")
    (SPLIT_DIR / "final_test.csv").write_text("")
    # Overwrite with actual CSV content
    df_train.to_csv(SPLIT_DIR / "final_train.csv", index=False)
    df_val.to_csv( (SPLIT_DIR / "final_val.csv"), index=False)
    df_test.to_csv((SPLIT_DIR / "final_test.csv"), index=False)

    # Save TF-IDF artifacts
    joblib.dump(vectorizer, TFIDF_DIR / "vectorizer.pkl")
    sp.save_npz(TFIDF_DIR / "tfidf_train.npz", X_train_tfidf)
    sp.save_npz(TFIDF_DIR / "tfidf_val.npz",   X_val_tfidf)
    sp.save_npz(TFIDF_DIR / "tfidf_test.npz",  X_test_tfidf)
    sp.save_npz(TFIDF_DIR / "tfidf_full.npz",  X_full_tfidf)

    print("\nSaved:")
    print(f" - {OUT_FINAL}")
    print(f" - {SPLIT_DIR / 'final_train.csv'}")
    print(f" - {SPLIT_DIR / 'final_val.csv'}")
    print(f" - {SPLIT_DIR / 'final_test.csv'}")
    print(f" - {TFIDF_DIR / 'vectorizer.pkl'}")
    print(f" - {TFIDF_DIR / 'tfidf_train.npz'}")
    print(f" - {TFIDF_DIR / 'tfidf_val.npz'}")
    print(f" - {TFIDF_DIR / 'tfidf_test.npz'}")
    print(f" - {TFIDF_DIR / 'tfidf_full.npz'}")

if __name__ == "__main__":
    main()