lab2/src/preprocessing_pipeline.py - master - dakone22/ppa

lab2/src/preprocessing_pipeline.py
147 строк7 KB
dakone22
created: lab2, lab3, lab4
04 июн 2026, 14:19
04 июн 2026, 14:1903322e8
"""
Пайплайн предобработки данных для ЛР2.
Этапы:
1. Загрузка данных (Polars).
2. Очистка и Imputation (MissForest / RandomForest).
3. Трансформация (Yeo-Johnson).
4. Отбор признаков (Feature Selection).
5. Сохранение.
"""

import warnings
import polars as pl
import pandas as pd
import numpy as np
from feature_engine.transformation import YeoJohnsonTransformer
from feature_engine.selection import DropConstantFeatures, DropDuplicateFeatures
from sklearn.experimental import enable_iterative_imputer
from sklearn.impute import IterativeImputer
from sklearn.ensemble import RandomForestRegressor

# Подавление предупреждений
warnings.filterwarnings('ignore')

def main():
    # Конфигурация путей
    input_path = "ppa/lab2/full_dataset_with_group_b.parquet"
    output_path = "ppa/lab2/dataset_processed.parquet"

    print("--- Запуск пайплайна предобработки ---")

    # 1. Загрузка данных
    print(f"[1/5] Загрузка данных из {input_path}...")
    try:
        # Используем Polars для быстрой загрузки
        df_polars = pl.read_parquet(input_path)
        print(f"Загружено строк: {df_polars.height}, столбцов: {df_polars.width}")
    except Exception as e:
        print(f"Ошибка загрузки: {e}")
        return

    # Конвертация в Pandas для Feature-engine (он не поддерживает Polars напрямую)
    df = df_polars.to_pandas()

    # Сохраняем целевую переменную (group_b), если она есть, чтобы не потерять при трансформации
    target_col = 'group_b' if 'group_b' in df.columns else None
    y = None
    if target_col:
        y = df[target_col].copy()
        # Временно исключаем таргет из признаков для трансформации
        X = df.drop(columns=[target_col])
    else:
        X = df.copy()

    # 2. Очистка и Imputation (MissForest / Iterative Imputer)
    print("[2/5] Обработка пропусков (Imputation)...")

    # Проверяем наличие пропусков
    missing_mask = X.isnull().sum()
    cols_with_nans = missing_mask[missing_mask > 0].index.tolist()

    if cols_with_nans:
        print(f"Найдены пропуски в столбцах: {cols_with_nans}")
        print("Применение IterativeImputer (аналог MissForest)...")

        # Используем RandomForestRegressor как эстиматор для импутации
        # Это реализация подхода, похожего на MissForest
        rf = RandomForestRegressor(n_estimators=10, random_state=42, n_jobs=-1)
        imputer = IterativeImputer(estimator=rf, random_state=42, max_iter=10)

        # Применяем только к числовым признакам
        numeric_cols = X.select_dtypes(include=[np.number]).columns
        X_numeric = X[numeric_cols]

        # Если есть нечисловые, оставляем их как есть или обрабатываем отдельно
        non_numeric_cols = X.select_dtypes(exclude=[np.number]).columns

        X_imputed_numeric = pd.DataFrame(
            imputer.fit_transform(X_numeric),
            columns=numeric_cols,
            index=X_numeric.index
        )

        # Собираем обратно
        X = pd.concat([X_imputed_numeric, X[non_numeric_cols]], axis=1)
        print("Импутация завершена.")
    else:
        print("Пропусков не обнаружено. Шаг импутации пропущен.")

    # 3. Feature Transformation (Yeo-Johnson)
    print("[3/5] Трансформация признаков (Yeo-Johnson)...")

    # Список признаков для трансформации согласно заданию
    # Обычно это признаки с правосторонним распределением
    skewed_features = ['views', 'forwards', 'replies', 'time_delta_sec']

    # Проверяем, какие из них реально есть в данных
    features_to_transform = [f for f in skewed_features if f in X.columns]

    if features_to_transform:
        print(f"Применение YeoJohnsonTransformer к: {features_to_transform}")
        # YeoJohnsonTransformer требует, чтобы признаки были числовыми и не содержали NaN
        # Убедимся, что выбранные признаки числовые
        transformer = YeoJohnsonTransformer(variables=features_to_transform)
        try:
            X[features_to_transform] = transformer.fit_transform(X[features_to_transform])
            print("Трансформация Yeo-Johnson выполнена.")
        except Exception as e:
            print(f"Ошибка трансформации: {e}. Пропускаем трансформацию.")
    else:
        print("Целевые признаки для трансформации не найдены в датасете.")

    # 4. Feature Selection
    print("[4/5] Отбор признаков (Feature Selection)...")

    # Удаление константных признаков
    selector_const = DropConstantFeatures(tol=0.998, variables=None)
    X = selector_const.fit_transform(X)
    print(f"После DropConstantFeatures: {X.shape[1]} признаков")

    # Удаление дублирующихся признаков
    selector_dup = DropDuplicateFeatures()
    X = selector_dup.fit_transform(X)
    print(f"После DropDuplicateFeatures: {X.shape[1]} признаков")

    # Возвращаем целевую переменную
    if target_col and y is not None:
        X[target_col] = y.values

    # 5. Сохранение
    print(f"[5/5] Сохранение обработанного датасета в {output_path}...")

    # Конвертируем обратно в Polars для сохранения (опционально, но Polars быстрее пишет)
    # Однако, если в данных есть специфические типы Pandas, лучше сохранить через Pandas или аккуратно конвертировать
    try:
        # Попробуем сохранить через Polars
        df_final = pl.from_pandas(X)
        df_final.write_parquet(output_path)
    except Exception as e:
        print(f"Ошибка сохранения через Polars: {e}. Сохраняем через Pandas...")
        X.to_parquet(output_path, index=False)

    print("--- Пайплайн завершен ---")
    print(f"Итоговый размер датасета: {X.shape[0]} строк, {X.shape[1]} столбцов")
    print(f"Файл сохранен: {output_path}")

if __name__ == "__main__":
    main()