ML Starter Pipeline (LightGBM + Stratified K-Fold)

A reusable baseline — clone it at the start of any tabular competition.

"""
ML Competition Starter Pipeline
Use this as the base for any tabular ML competition.
"""

import pandas as pd
import numpy as np
from sklearn.model_selection import StratifiedKFold, KFold
from sklearn.metrics import roc_auc_score, mean_squared_error
from sklearn.preprocessing import LabelEncoder
import lightgbm as lgb
import warnings
warnings.filterwarnings('ignore')

# ── Config ────────────────────────────────────────────────────────────────────
TARGET = "target"          # change to your target column
TASK = "classification"    # "classification" or "regression"
N_FOLDS = 5
SEED = 42
TRAIN_PATH = "data/train.csv"
TEST_PATH  = "data/test.csv"

# ── Load Data ─────────────────────────────────────────────────────────────────
train = pd.read_csv(TRAIN_PATH)
test  = pd.read_csv(TEST_PATH)
print(f"Train: {train.shape} | Test: {test.shape}")

# ── Basic EDA ─────────────────────────────────────────────────────────────────
print("\nTarget distribution:")
print(train[TARGET].value_counts(normalize=True))
print("\nMissing values (train):")
print(train.isnull().sum()[train.isnull().sum() > 0])

# ── Feature Engineering ───────────────────────────────────────────────────────
def make_features(df):
    # Encode categoricals
    cat_cols = df.select_dtypes(include=['object', 'category']).columns
    for col in cat_cols:
        le = LabelEncoder()
        df[col] = le.fit_transform(df[col].astype(str))
    return df

train = make_features(train)
test  = make_features(test)

feature_cols = [c for c in train.columns if c != TARGET]
X = train[feature_cols]
y = train[TARGET]
X_test = test[feature_cols]

# ── Cross Validation ──────────────────────────────────────────────────────────
oof_preds = np.zeros(len(train))
test_preds = np.zeros(len(test))

if TASK == "classification":
    folds = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=SEED)
    split_iter = folds.split(X, y)
else:
    folds = KFold(n_splits=N_FOLDS, shuffle=True, random_state=SEED)
    split_iter = folds.split(X)

lgb_params = {
    "objective": "binary" if TASK == "classification" else "regression",
    "metric": "auc" if TASK == "classification" else "rmse",
    "learning_rate": 0.05,
    "num_leaves": 64,
    "min_child_samples": 20,
    "feature_fraction": 0.8,
    "bagging_fraction": 0.8,
    "bagging_freq": 1,
    "verbose": -1,
    "seed": SEED,
}

for fold, (tr_idx, val_idx) in enumerate(split_iter):
    X_tr, X_val = X.iloc[tr_idx], X.iloc[val_idx]
    y_tr, y_val = y.iloc[tr_idx], y.iloc[val_idx]

    model = lgb.LGBMClassifier(**lgb_params, n_estimators=1000) if TASK == "classification" \
        else lgb.LGBMRegressor(**lgb_params, n_estimators=1000)

    model.fit(
        X_tr, y_tr,
        eval_set=[(X_val, y_val)],
        callbacks=[lgb.early_stopping(100, verbose=False), lgb.log_evaluation(200)]
    )

    if TASK == "classification":
        oof_preds[val_idx] = model.predict_proba(X_val)[:, 1]
        test_preds += model.predict_proba(X_test)[:, 1] / N_FOLDS
    else:
        oof_preds[val_idx] = model.predict(X_val)
        test_preds += model.predict(X_test) / N_FOLDS

    if TASK == "classification":
        score = roc_auc_score(y_val, oof_preds[val_idx])
    else:
        score = np.sqrt(mean_squared_error(y_val, oof_preds[val_idx]))
    print(f"Fold {fold+1} score: {score:.5f}")

# ── Final CV Score ────────────────────────────────────────────────────────────
if TASK == "classification":
    cv_score = roc_auc_score(y, oof_preds)
else:
    cv_score = np.sqrt(mean_squared_error(y, oof_preds))
print(f"\nOverall CV score: {cv_score:.5f}")

# ── Submission ────────────────────────────────────────────────────────────────
sub = pd.read_csv("data/sample_submission.csv")
sub[TARGET] = test_preds
sub.to_csv("submission.csv", index=False)
print("submission.csv saved.")