joshnavip/ai-code-detection

1import pandas as pd2from sklearn.model_selection import GroupShuffleSplit, train_test_split3 4DATA_PATH = "dataset/processed/dataset_step2_cleaned.csv"5OUTPUT_DIR = "dataset/processed"6 7TRAIN_SIZE = 0.78VAL_SIZE = 0.159TEST_SIZE = 0.1510RANDOM_STATE = 4211 12def main():13    df = pd.read_csv(DATA_PATH)14 15    # Standardize Label column16    if "Label (0- HUMAN, 1-AI)" in df.columns:17        df = df.rename(columns={"Label (0- HUMAN, 1-AI)": "Label"})18 19    # Normalize Language column (for reporting)20    if "Language" in df.columns:21        df["Language"] = df["Language"].astype(str).str.strip().str.lower()22 23    # Split by class24    human_df = df[df["Label"] == 0].copy()25    ai_df = df[df["Label"] == 1].copy()26 27    # -----------------------------28    # AI split (group-aware)29    # -----------------------------30    ai_df["Generation_Prompt"] = ai_df["Generation_Prompt"].fillna("UNKNOWN_AI_PROMPT")31 32    gss = GroupShuffleSplit(33        n_splits=1,34        train_size=TRAIN_SIZE,35        random_state=RANDOM_STATE36    )37 38    ai_train_idx, ai_temp_idx = next(39        gss.split(ai_df, groups=ai_df["Generation_Prompt"])40    )41 42    ai_train = ai_df.iloc[ai_train_idx]43    ai_temp = ai_df.iloc[ai_temp_idx]44 45    gss_val = GroupShuffleSplit(46        n_splits=1,47        train_size=VAL_SIZE / (VAL_SIZE + TEST_SIZE),48        random_state=RANDOM_STATE49    )50 51    ai_val_idx, ai_test_idx = next(52        gss_val.split(ai_temp, groups=ai_temp["Generation_Prompt"])53    )54 55    ai_val = ai_temp.iloc[ai_val_idx]56    ai_test = ai_temp.iloc[ai_test_idx]57 58    # -----------------------------59    # Human split (no grouping)60    # -----------------------------61    human_train, human_temp = train_test_split(62        human_df,63        train_size=TRAIN_SIZE,64        random_state=RANDOM_STATE,65        shuffle=True66    )67 68    human_val, human_test = train_test_split(69        human_temp,70        train_size=VAL_SIZE / (VAL_SIZE + TEST_SIZE),71        random_state=RANDOM_STATE,72        shuffle=True73    )74 75    # -----------------------------76    # Combine & shuffle77    # -----------------------------78    train_df = pd.concat([ai_train, human_train]).sample(frac=1, random_state=RANDOM_STATE)79    val_df = pd.concat([ai_val, human_val]).sample(frac=1, random_state=RANDOM_STATE)80    test_df = pd.concat([ai_test, human_test]).sample(frac=1, random_state=RANDOM_STATE)81 82    # -----------------------------83    # Safety checks84    # -----------------------------85    for name, split in [("Train", train_df), ("Val", val_df), ("Test", test_df)]:86        labels = set(split["Label"].unique())87        if labels != {0, 1}:88            raise RuntimeError(f"{name} split missing a class: {labels}")89 90    # -----------------------------91    # Save92    # -----------------------------93    train_df.to_csv(f"{OUTPUT_DIR}/dataset_train.csv", index=False)94    val_df.to_csv(f"{OUTPUT_DIR}/dataset_val.csv", index=False)95    test_df.to_csv(f"{OUTPUT_DIR}/dataset_test.csv", index=False)96 97    # -----------------------------98    # Reporting99    # -----------------------------100    print("\n📊 Final Class Distribution")101    print("Train:\n", train_df["Label"].value_counts())102    print("Val:\n", val_df["Label"].value_counts())103    print("Test:\n", test_df["Label"].value_counts())104 105    if "Language" in df.columns:106        print("\n🌐 Language Distribution (Train)")107        print(train_df["Language"].value_counts())108 109if __name__ == "__main__":110    main()111