File size: 4,817 Bytes
1aa7fae |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 |
"""
Data preparation script for the predictive maintenance project.
Responsibilities:
- Load the raw engine dataset from the Hugging Face dataset repo (preferred)
or from the local data folder as a fallback.
- Clean and preprocess the data (rename columns, handle missing values,
drop duplicates, basic sanity checks).
- Split the cleaned data into train and test sets.
- Save train and test CSVs locally.
- Upload the resulting train and test CSVs back to the Hugging Face dataset repo.
"""
from __future__ import annotations
from pathlib import Path
from typing import Tuple
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
import config
from hf_data_utils import download_dataset_file, upload_dataset_file
def _load_raw_data_from_hf_or_local() -> pd.DataFrame:
"""
Try to load the raw dataset from the Hugging Face dataset repo.
If that fails (e.g., no token or repo yet), fall back to the local CSV.
"""
# Preferred: load from HF dataset repo if token and repo are configured
if config.HF_TOKEN and config.HF_DATASET_REPO:
try:
remote_path = download_dataset_file(
filename="data/engine_data.csv",
repo_id=config.HF_DATASET_REPO,
token=config.HF_TOKEN,
local_dir=config.DATA_DIR,
)
return pd.read_csv(remote_path)
except Exception:
# Fall back to local file
pass
# Local fallback
if not config.RAW_DATA_FILE.exists():
raise FileNotFoundError(
f"Raw data file not found at {config.RAW_DATA_FILE}. "
"Ensure engine_data.csv exists or upload it to the HF dataset repo."
)
return pd.read_csv(config.RAW_DATA_FILE)
def _clean_data(df: pd.DataFrame) -> pd.DataFrame:
"""
Perform basic cleaning and feature engineering.
"""
# Standardize column names
df = df.rename(columns=config.RAW_COLUMN_RENAME_MAP)
# Keep only the expected columns (drop any extras, if present)
expected_cols = set(config.FEATURE_COLUMNS + [config.TARGET_COLUMN])
df = df[[col for col in df.columns if col in expected_cols]]
# Drop duplicate rows
df = df.drop_duplicates().reset_index(drop=True)
# Handle missing values: for this numeric dataset, fill with median
if df.isna().any().any():
df = df.fillna(df.median(numeric_only=True))
# Ensure target is integer/binary
df[config.TARGET_COLUMN] = df[config.TARGET_COLUMN].astype(int)
return df
def _train_test_split(
df: pd.DataFrame,
) -> Tuple[pd.DataFrame, pd.DataFrame]:
"""
Split the cleaned dataframe into train and test sets.
"""
X = df[config.FEATURE_COLUMNS]
y = df[config.TARGET_COLUMN]
X_train, X_test, y_train, y_test = train_test_split(
X,
y,
test_size=config.TEST_SIZE,
random_state=config.RANDOM_STATE,
stratify=y,
)
train_df = X_train.copy()
train_df[config.TARGET_COLUMN] = y_train
test_df = X_test.copy()
test_df[config.TARGET_COLUMN] = y_test
return train_df, test_df
def main() -> None:
"""
Execute the full data preparation pipeline.
"""
print("Loading raw data...")
raw_df = _load_raw_data_from_hf_or_local()
print(f"Raw data shape: {raw_df.shape}")
print("Cleaning data...")
clean_df = _clean_data(raw_df)
print(f"Clean data shape: {clean_df.shape}")
print("Performing train/test split...")
train_df, test_df = _train_test_split(clean_df)
print(f"Train shape: {train_df.shape}, Test shape: {test_df.shape}")
# Save locally
config.PROCESSED_DIR.mkdir(parents=True, exist_ok=True)
train_df.to_csv(config.TRAIN_FILE, index=False)
test_df.to_csv(config.TEST_FILE, index=False)
print(f"Saved train to {config.TRAIN_FILE}")
print(f"Saved test to {config.TEST_FILE}")
# Upload to HF dataset repo, if configured
if config.HF_TOKEN and config.HF_DATASET_REPO:
try:
print("Uploading train and test splits to Hugging Face dataset repo...")
upload_dataset_file(
local_path=config.TRAIN_FILE,
repo_id=config.HF_DATASET_REPO,
repo_path="data/train.csv",
token=config.HF_TOKEN,
)
upload_dataset_file(
local_path=config.TEST_FILE,
repo_id=config.HF_DATASET_REPO,
repo_path="data/test.csv",
token=config.HF_TOKEN,
)
print("Upload to Hugging Face completed.")
except Exception as e:
print(
f"Warning: Failed to upload train/test to Hugging Face dataset repo: {e}"
)
if __name__ == "__main__":
main()
|