File size: 2,826 Bytes
baf834b | 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 | """
Preprocessor for AI4I 2020 Predictive Maintenance Dataset.
Handles column renaming, label encoding, feature engineering (temp_diff, power),
stratified split, and StandardScaler fitting on train only.
"""
import os
import sys
import pandas as pd
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.model_selection import train_test_split
import joblib
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'utils')))
from preprocess_utils import logger, time_tracker
@time_tracker
def preprocess_ai4i(data_path: str, output_dir: str) -> tuple:
os.makedirs(output_dir, exist_ok=True)
logger.info(f"Loading AI4I 2020 dataset from {data_path}...")
df = pd.read_csv(data_path)
# Rename columns — brackets cause XGBoost errors
df = df.rename(columns={
'Type': 'type', 'Air temperature [K]': 'air_temp',
'Process temperature [K]': 'process_temp',
'Rotational speed [rpm]': 'rotational_speed',
'Torque [Nm]': 'torque', 'Tool wear [min]': 'tool_wear',
'Machine failure': 'machine_failure'
})
# Encode categorical + save encoder
le = LabelEncoder()
df['type_encoded'] = le.fit_transform(df['type'])
joblib.dump(le, os.path.join(output_dir, 'ai4i_label_encoder.pkl'))
# Feature engineering
df['temp_diff'] = df['process_temp'] - df['air_temp']
df['power'] = df['torque'] * df['rotational_speed']
features = ['type_encoded', 'air_temp', 'process_temp', 'rotational_speed',
'torque', 'tool_wear', 'temp_diff', 'power']
failure_modes = ['TWF', 'HDF', 'PWF', 'OSF', 'RNF']
X = df[features]
y = df['machine_failure']
y_modes = df[failure_modes]
X_train, X_test, y_train, y_test, ym_train, ym_test = train_test_split(
X, y, y_modes, test_size=0.2, random_state=42, stratify=y
)
# Scale — fit on train ONLY
scaler = StandardScaler()
scale_cols = [c for c in features if c != 'type_encoded']
X_train, X_test = X_train.copy(), X_test.copy()
X_train[scale_cols] = scaler.fit_transform(X_train[scale_cols])
X_test[scale_cols] = scaler.transform(X_test[scale_cols])
joblib.dump(scaler, os.path.join(output_dir, 'ai4i_scaler.pkl'))
# Save processed splits
for split, Xs, ys, yms in [('train', X_train, y_train, ym_train),
('test', X_test, y_test, ym_test)]:
out = Xs.copy()
out['target'] = ys
for col in failure_modes:
out[col] = yms[col]
out.to_csv(os.path.join(output_dir, f'ai4i_{split}_processed.csv'), index=False)
logger.info(f"AI4I preprocessing complete. Saved to {output_dir}")
return (os.path.join(output_dir, 'ai4i_train_processed.csv'),
os.path.join(output_dir, 'ai4i_test_processed.csv'))
|