detectai-backend / ml_model /preprocess_dataset.py
Jaya
Initial Deployment to Hugging Face Spaces
686b55b
Raw
History Blame Contribute Delete
2.71 kB
import os
import pandas as pd
import kagglehub
from kagglehub import KaggleDatasetAdapter
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
import re
# --- Configurations ---
DATASET_NAME = "arslandoula/ah-and-aitd-arslans-human-and-ai-text-database"
FILE_PATH = "Dataset.xlsx"
RAW_DATA_PATH = os.path.join(os.path.dirname(__file__), "data/raw_kaggle_dataset.csv")
CLEAN_DATA_PATH = os.path.join(os.path.dirname(__file__), "data/clean_dataset.csv")
FORCE_DOWNLOAD = False # πŸ” Set True if you want to re-download even if file exists
# --- Step 1: Smart Dataset Loader ---
def load_or_download_dataset():
os.makedirs(os.path.dirname(RAW_DATA_PATH), exist_ok=True)
if os.path.exists(RAW_DATA_PATH) and not FORCE_DOWNLOAD:
print(f"βœ… Found existing dataset at {RAW_DATA_PATH}")
df = pd.read_csv(RAW_DATA_PATH)
else:
print(f"⬇️ Downloading dataset from Kaggle: {DATASET_NAME}")
df = kagglehub.dataset_load(
KaggleDatasetAdapter.PANDAS,
DATASET_NAME,
FILE_PATH,
)
df.to_csv(RAW_DATA_PATH, index=False)
print(f"πŸ“ Dataset saved to {RAW_DATA_PATH}")
return df
def clean_text(text):
text = str(text)
# Preservation of structural signatures (punctuation, numbers, spacing)
# which are critical for AI vs Human detection.
text = re.sub(r"http\S+|www\S+", "", text) # remove URLs
text = re.sub(r"\s+", " ", text).strip() # collapse extra spaces
return text
# --- Step 3: Preprocess Pipeline ---
def preprocess_dataset(df):
print("🧹 Cleaning and preprocessing dataset...")
df = df.rename(columns={
"text": "content",
"label_name": "label"
})
# Apply cleaning
df["content"] = df["content"].apply(clean_text)
# Encode labels (Human-written = 0, AI-generated = 1)
label_encoder = LabelEncoder()
df["label_id"] = label_encoder.fit_transform(df["label"])
# Drop unnecessary columns
df = df[["content", "label", "label_id"]]
print(f"βœ… Preprocessed {len(df)} samples.")
df.to_csv(CLEAN_DATA_PATH, index=False)
print(f"πŸ“ Clean dataset saved to {CLEAN_DATA_PATH}")
return df
# --- Step 4: Split dataset (optional) ---
def split_dataset(df):
train, test = train_test_split(df, test_size=0.2, random_state=42, stratify=df["label_id"])
print(f"πŸ“Š Training samples: {len(train)} | Testing samples: {len(test)}")
return train, test
# --- Main Execution ---
if __name__ == "__main__":
df_raw = load_or_download_dataset()
df_clean = preprocess_dataset(df_raw)
train_df, test_df = split_dataset(df_clean)