Automated ML Pipeline Creator

Chat-based ML Service with Automated Pipeline Generation

Project Overview

Designed and developed a backend automated ML pipeline creator for a managed ML service that enables business users to solve problems through a chat-based interface, featuring robust data preprocessing, automated feature engineering, and customizable pipeline sequencing.

Situation & Context

Our managed ML service wanted a chat-based product which would support business users bringing in problems and solving them using ML. My responsibility was designing and developing the backend automated ML pipeline creator that would preprocess data and train ML models. Existing out-of-the-box AutoML solutions like AutoGluon were memory intensive and slower due to ensembles and meta-learning approaches, requiring a more efficient solution.

Task

Phase 1: Automated ML Pipeline Builder
  • Develop automated ML pipeline with best practices, reproducibility, and low latency
  • Handle extremely noisy data with mixed types and inconsistent formats
  • Perform advanced data type detection and task type detection
  • Implement automated feature engineering
  • Enforce user-driven training time limits
Phase 2: Customizable Pipeline Sequencing
  • Allow users to customize ML pipeline by choosing from operations
  • Automatically determine correct sequence of transformations
  • Ensure proper ordering to prevent unintended perturbations

Action & Implementation

Phase 1 Implementation
  • Framework Development: Created framework extending scikit-learn with open source libraries, supporting multiple models with XGBoost as default
  • Training Time Enforcement: Modified XGBoost's internal training loop using callbacks to accurately track training time and gracefully exit
  • Robust Data Type Detection: Developed heuristic using data statistics and practical thresholds that worked in edge cases where pandas would fail
  • LLM-powered Feature Engineering: Implemented chained 2-step method: candidate feature logic generation, then code generation with multiple trials and checks
  • Large-scale Testing: Created benchmark with over 300 datasets for comprehensive evaluation
Phase 2 Implementation
  • Transform Analysis: Surveyed prior works and analyzed situations where incorrect ordering would cause issues
  • Partial Order Development: Established ordering rules (e.g., standardization before dropping missing values, normalization before oversampling)
  • Transform Matrix: Created matrix of transforms and meta-characteristics describing prerequisites and post-effects
  • Rule-based Algorithm: Developed topological ordering algorithm ensuring correct transform sequencing

Results & Impact

  • Superior Performance: Data type detection heuristic outperformed competing state-of-the-art baselines, especially AutoGluon
  • Improved Efficiency: Achieved better latency than competing baselines while maintaining comparable performance
  • Simplified Effectiveness: Demonstrated that standard ML pipeline with XGBoost was simple yet effective approach
  • Successful Launch: Product was successfully launched with Phase 1 implementation in production
  • Business Impact: Enabled business users to leverage ML through intuitive chat interface without requiring ML expertise

Keywords

  • Automated Machine Learning (AutoML)
  • ML Pipeline Automation
  • Feature Engineering
  • XGBoost Optimization
  • Data Type Detection
  • Transform Sequencing
  • LLM-powered Code Generation
  • Scikit-learn Framework

References