mihir-apte/cognitive-distortion-detector

1import os2import numpy as np3import pandas as pd4from sklearn.model_selection import train_test_split5from sklearn.utils.class_weight import compute_class_weight6 7from src.config import (8    DISTORTION_LABELS, LABEL2ID, NUM_LABELS,9    CSV_LABEL_MAP, CSV_NO_DISTORTION,10    DATA_RAW_DIR, DATA_PROC_DIR, SEED11)12 13 14 15def load_raw(csv_path: str) -> pd.DataFrame:16    df = pd.read_csv(csv_path, encoding="utf-8")17 18    19    rename_map = {}20    for col in df.columns:21        cl = col.lower().strip()22        if "patient question" in cl:23            rename_map[col] = "text"24        elif "dominant" in cl:25            rename_map[col] = "dominant_distortion"26        elif "secondary" in cl:27            rename_map[col] = "secondary_distortion"28        elif "distorted part" in cl:29            rename_map[col] = "distorted_part"30    df = df.rename(columns=rename_map)31 32    df["text"] = df["text"].astype(str).str.strip()33    df["dominant_distortion"] = df["dominant_distortion"].astype(str).str.strip()34 35    if "secondary_distortion" in df.columns:36        df["secondary_distortion"] = df["secondary_distortion"].astype(str).str.strip()37    else:38        df["secondary_distortion"] = ""39 40    return df41 42 43def normalise_label(raw: str):44    """Map a CSV label string to our canonical label, or return None."""45    raw = str(raw).strip()46    if raw == CSV_NO_DISTORTION or raw.lower() == "no distortion":47        return None                      48    return CSV_LABEL_MAP.get(raw, raw)   49 50 51def build_label_vector(dominant: str, secondary: str) -> list:52 53    vec = [0] * NUM_LABELS54 55    dom = normalise_label(dominant)56    if dom and dom in LABEL2ID:57        vec[LABEL2ID[dom]] = 158 59    if secondary and secondary not in ("nan", "", "None", "NaN"):60        sec = normalise_label(secondary)61        if sec and sec in LABEL2ID:62            vec[LABEL2ID[sec]] = 163 64    return vec65 66 67def run(csv_filename: str = "Annotated_data.csv") -> None:68    csv_path = os.path.join(DATA_RAW_DIR, csv_filename)69    os.makedirs(DATA_PROC_DIR, exist_ok=True)70 71    print(f"Loading: {csv_path}")72    df = load_raw(csv_path)73    print(f"Rows loaded: {len(df)}")74 75    76    label_vecs = df.apply(77        lambda row: build_label_vector(78            row["dominant_distortion"], row["secondary_distortion"]79        ),80        axis=181    )82    label_df = pd.DataFrame(label_vecs.tolist(), columns=DISTORTION_LABELS)83    df = pd.concat(84        [df[["text", "dominant_distortion", "secondary_distortion"]], label_df],85        axis=186    )87 88    89    distortion_rows    = df[df[DISTORTION_LABELS].sum(axis=1) > 0]90    no_distortion_rows = df[df[DISTORTION_LABELS].sum(axis=1) == 0]91    multi_label_rows   = df[df[DISTORTION_LABELS].sum(axis=1) == 2]92    print(f"Rows with >=1 distortion label   : {len(distortion_rows)}")93    print(f"Rows with no distortion (all 0s) : {len(no_distortion_rows)}")94    print(f"Multi-label rows (2 labels)      : {len(multi_label_rows)}")95 96    97    train_df, test_df = train_test_split(98        df,99        test_size=0.2,100        random_state=SEED,101        stratify=df["dominant_distortion"]102    )103    train_df = train_df.reset_index(drop=True)104    test_df  = test_df.reset_index(drop=True)105    print(f"\nTrain size : {len(train_df)}")106    print(f"Test size  : {len(test_df)}")107 108    train_distortion_only = train_df[train_df[DISTORTION_LABELS].sum(axis=1) > 0]109    dom_labels = (110        train_distortion_only["dominant_distortion"]111        .map(normalise_label)112        .dropna()113    )114 115    present_labels = np.array(sorted(dom_labels.unique()))116    raw_weights    = compute_class_weight(117        class_weight="balanced",118        classes=present_labels,119        y=dom_labels.values120    )121 122    123    class_weights = np.ones(NUM_LABELS, dtype=np.float32)124    for label, weight in zip(present_labels, raw_weights):125        if label in LABEL2ID:126            class_weights[LABEL2ID[label]] = weight127 128    print("\nClass weights aligned to DISTORTION_LABELS order:")129    for label, w in zip(DISTORTION_LABELS, class_weights):130        print(f"  {label:<35} {w:.4f}")131 132   133    train_path   = os.path.join(DATA_PROC_DIR, "train.csv")134    test_path    = os.path.join(DATA_PROC_DIR, "test.csv")135    weights_path = os.path.join(DATA_PROC_DIR, "class_weights.npy")136 137    train_df.to_csv(train_path,   index=False)138    test_df.to_csv(test_path,    index=False)139    np.save(weights_path, class_weights)140 141    print(f"\nSaved: {train_path}")142    print(f"Saved: {test_path}")143    print(f"Saved: {weights_path}")144    print("\nPreprocessing complete.")145 146 147if __name__ == "__main__":148    run()149