Articles/Claude Code

⟐ Claude Code/2026-05-03Advanced

Claude Code for Data Science — pandas, scikit-learn, and ML Workflows with AI Pair Programming

A hands-on guide to using Claude Code across the full data science and machine learning workflow. From EDA and feature engineering to model evaluation, hyperparameter tuning with Optuna, SHAP analysis, and MLOps basics — with working Python code throughout.

claude-code¹³²data-sciencemachine-learningpython²²pandasscikit-learnmlops

✦ Premium Article

Most data scientists I know use Claude Code for quick scripts and code completions. Far fewer have tried letting it drive the full pipeline — from raw data through EDA, feature engineering, model training, hyperparameter tuning, and into MLOps tooling.

I was in the first camp for a while. My data analysis lived in Jupyter notebooks, edited by hand, with Claude Code used occasionally for specific snippets. It took actually committing to the full workflow to realize how well-suited data science is for AI pair programming.

The reason is structural. Data science is fundamentally a loop: hypothesis → code → observe → revise. Claude Code excels at compressing that loop. The faster you can move from "I wonder if this feature helps" to "here's the validation curve showing it does," the better your work gets.

This article walks through the patterns that work in practice, with complete working code for each phase.

1. Project Setup: Using CLAUDE.md to Encode Data Science Constraints

Before writing a single line of analysis code, invest a few minutes in a CLAUDE.md that tells Claude Code how this project works. The payoff is substantial: Claude will automatically apply constraints you'd otherwise have to re-state in every prompt.

Installing the Stack

# Prompt to Claude Code:
"Create a requirements.txt for a data science project and install it
in a virtual environment. We need: pandas, numpy, scikit-learn,
matplotlib, seaborn, jupyterlab, optuna, shap, joblib"

Generated requirements.txt:

pandas==2.2.0
numpy==1.26.3
scikit-learn==1.4.0
matplotlib==3.8.2
seaborn==0.13.1
jupyterlab==4.1.0
optuna==3.5.0
shap==0.44.0
joblib==1.3.2

python -m venv .venv
source .venv/bin/activate   # Windows: .venv\Scripts\activate
pip install -r requirements.txt

The CLAUDE.md That Changes Everything

The constraints below aren't obvious to Claude Code without being told. Once they're in CLAUDE.md, you stop having to re-state them:

# Data Science Project Guidelines
 
## Stack
- Python 3.12, pandas 2.2, scikit-learn 1.4, matplotlib/seaborn
- Data: data/raw/ (source), data/processed/ (transformed)
- Models: models/ (pkl via joblib)
 
## Rules — Always Apply
- Always .copy() DataFrames before mutation
- Use RANDOM_STATE = 42 everywhere, consistently
- Split data BEFORE fitting any preprocessor (prevents data leakage)
- Save models with joblib.dump, load with joblib.load
- For datasets >500k rows, use chunked reading

With this in place, Claude Code will include .copy() on every DataFrame operation, set seeds across every library, and structure preprocessing to avoid leakage — without being asked each time.

2. EDA: Getting Claude Code to Do the Boring Parts

Exploratory data analysis is where projects start, and where a lot of time gets lost to boilerplate. The goal here is to get to insights fast, not to hand-write histogram loops.

A Complete EDA Report Generator

# Prompt to Claude Code:
# "Write src/eda_report.py that performs basic EDA on data/raw/sales.csv:
# - dtypes, missing values, summary statistics for numeric columns
# - histograms and boxplots for numeric features
# - value counts (top 10) for categorical columns
# - correlation heatmap
# Save all figures to reports/eda/ as PNG files."

Generated script:

# src/eda_report.py
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path
 
RANDOM_STATE = 42
REPORTS_DIR = Path("reports/eda")
REPORTS_DIR.mkdir(parents=True, exist_ok=True)
 
def generate_eda_report(filepath: str) -> dict:
    """Run EDA and save figures. Returns summary dict."""
    df = pd.read_csv(filepath)
    report = {}
 
    report["shape"] = df.shape
    report["missing_pct"] = (df.isnull().sum() / len(df) * 100).to_dict()
 
    numeric_cols = df.select_dtypes(include=np.number).columns.tolist()
    report["stats"] = df[numeric_cols].describe().to_dict()
 
    # Correlation heatmap
    if len(numeric_cols) > 1:
        fig, ax = plt.subplots(figsize=(12, 10))
        sns.heatmap(
            df[numeric_cols].corr(),
            annot=True, fmt=".2f",
            cmap="coolwarm", center=0, ax=ax
        )
        ax.set_title("Feature Correlation Matrix")
        fig.savefig(REPORTS_DIR / "correlation_heatmap.png",
                    dpi=150, bbox_inches="tight")
        plt.close(fig)
        print("[OK] Correlation heatmap saved.")
 
    # Histograms
    n_cols = 3
    n_rows = (len(numeric_cols) + n_cols - 1) // n_cols
    fig, axes = plt.subplots(n_rows, n_cols, figsize=(15, 5 * n_rows))
    axes = axes.flatten()
 
    for i, col in enumerate(numeric_cols):
        axes[i].hist(df[col].dropna(), bins=30, edgecolor="black", alpha=0.7)
        axes[i].set_title(
            f"{col} (missing: {report['missing_pct'][col]:.1f}%)"
        )
        axes[i].set_xlabel(col)
        axes[i].set_ylabel("Count")
 
    for j in range(i + 1, len(axes)):
        axes[j].set_visible(False)
 
    plt.tight_layout()
    fig.savefig(REPORTS_DIR / "histograms.png", dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"[OK] Histograms saved ({len(numeric_cols)} features)")
 
    return report
 
if __name__ == "__main__":
    report = generate_eda_report("data/raw/sales.csv")
    print(f"\nShape: {report['shape']}")
    print("\nMissing values (%):")
    for col, pct in report["missing_pct"].items():
        if pct > 0:
            print(f"  {col}: {pct:.1f}%")

Outlier Detection: Comparing Three Strategies

Rather than picking an outlier strategy upfront, generate a comparison script that shows all three options side by side:

# src/outlier_handling.py
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
 
def detect_outliers_iqr(series: pd.Series, multiplier: float = 1.5) -> pd.Series:
    """Return Boolean mask of IQR-based outliers."""
    Q1, Q3 = series.quantile(0.25), series.quantile(0.75)
    IQR = Q3 - Q1
    return (series < Q1 - multiplier * IQR) | (series > Q3 + multiplier * IQR)
 
def compare_outlier_strategies(
    df: pd.DataFrame, col: str, output_dir: str = "reports/eda"
) -> dict:
    """Compare drop / clip / log-transform and save a side-by-side plot."""
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    mask = detect_outliers_iqr(df[col])
    print(f"Outliers detected: {mask.sum()} / {len(df)} "
          f"({mask.sum()/len(df)*100:.1f}%)")
 
    Q1, Q3 = df[col].quantile(0.25), df[col].quantile(0.75)
    IQR = Q3 - Q1
    lower, upper = Q1 - 1.5 * IQR, Q3 + 1.5 * IQR
 
    df_drop = df[~mask].copy()
    df_clip = df.copy()
    df_clip[col] = df_clip[col].clip(lower=lower, upper=upper)
    df_log = None
    if df[col].min() > 0:
        df_log = df.copy()
        df_log[f"{col}_log"] = np.log1p(df_log[col])
 
    fig, axes = plt.subplots(1, 3, figsize=(18, 5))
    axes[0].hist(df_drop[col], bins=30, alpha=0.7, color="steelblue")
    axes[0].set_title(f"① Drop (n={len(df_drop)})")
    axes[1].hist(df_clip[col], bins=30, alpha=0.7, color="orange")
    axes[1].set_title(f"② Clip [{lower:.1f}, {upper:.1f}]")
    if df_log is not None:
        axes[2].hist(df_log[f"{col}_log"], bins=30, alpha=0.7, color="green")
        axes[2].set_title("③ Log Transform (log1p)")
    else:
        axes[2].text(0.5, 0.5, "Log transform unavailable\n(non-positive values)",
                     ha="center", va="center", transform=axes[2].transAxes)
    plt.suptitle(f"Outlier Strategy Comparison: '{col}'", fontsize=14)
    plt.tight_layout()
    fig.savefig(f"{output_dir}/outlier_comparison_{col}.png",
                dpi=150, bbox_inches="tight")
    plt.close(fig)
 
    return {"drop_n": len(df_drop), "lower": lower, "upper": upper}

✦

Thank you for reading this far.

Continue Reading

What follows includes implementation code, benchmarks, and practical content we hope you'll find useful. This site runs without ads — server and development costs are supported entirely by members like you. If it's been helpful, we'd be truly grateful for your support.

WHAT YOU'LL LEARN

✦A reproducible layout that turns an exploratory notebook into six tested modules in ~25 minutes

✦Wiring data-leakage detection into CLAUDE.md — a real case that caught leakage in 3 of 11 features

✦A full production workflow: Optuna 100 trials (RMSE 0.0421->0.0388), pytest, and a single run script

Secure payment via Stripe · Cancel anytime

✦

Unlock This Article

Get full access to the rest of this article. Buy once, read anytime. This site is ad-free — your support goes directly toward keeping it running.

Unlock all articles with Membership →

3. Feature Engineering: Using Claude Code to Generate Ideas, Not Just Code

Feature engineering is where models win or lose. A useful pattern here is asking Claude Code to both propose features and implement them — not just translate specs you already have.

Prompt to Claude Code:
"The dataset has these columns:
- purchase_date (datetime)
- customer_age (int)
- product_category (str: electronics/clothing/food)
- price (float)
- quantity (int)

Suggest 10 features for a sales prediction model and implement them
in src/feature_engineering.py. Be careful about data leakage in any
time-based features."

Generated code:

# src/feature_engineering.py
import pandas as pd
import numpy as np
from sklearn.preprocessing import LabelEncoder
 
RANDOM_STATE = 42
 
def engineer_features(df: pd.DataFrame) -> pd.DataFrame:
    """
    Main feature engineering pipeline.
    Always .copy() to avoid mutating the input.
    """
    df = df.copy()
 
    # Date features — no future information, so no leakage
    df["purchase_date"] = pd.to_datetime(df["purchase_date"])
    df["day_of_week"] = df["purchase_date"].dt.dayofweek
    df["month"] = df["purchase_date"].dt.month
    df["quarter"] = df["purchase_date"].dt.quarter
    df["is_weekend"] = (df["day_of_week"] >= 5).astype(int)
    df["days_since_epoch"] = (
        df["purchase_date"] - pd.Timestamp("2020-01-01")
    ).dt.days
 
    # Interaction features
    df["total_amount"] = df["price"] * df["quantity"]
    df["log_price"] = np.log1p(df["price"])
    df["log_quantity"] = np.log1p(df["quantity"])
 
    # Age binning
    df["age_group"] = pd.cut(
        df["customer_age"],
        bins=[0, 25, 35, 50, 65, 100],
        labels=["under25", "25-35", "35-50", "50-65", "65+"],
        right=False
    ).astype(str)
 
    # Label encode categoricals
    for col in ["product_category", "age_group"]:
        le = LabelEncoder()
        df[f"{col}_encoded"] = le.fit_transform(df[col].astype(str))
 
    df = df.drop(columns=["purchase_date", "product_category", "age_group"])
    return df
 
def validate_no_data_leakage(
    df_train: pd.DataFrame,
    df_test: pd.DataFrame,
    date_col: str = "purchase_date"
) -> bool:
    """
    Confirm test set dates all come after training set dates.
    Call this immediately after train/test split.
    """
    if date_col not in df_train.columns:
        print(f"Warning: {date_col} not found. Skipping leakage check.")
        return True
 
    train_max = pd.to_datetime(df_train[date_col]).max()
    test_min = pd.to_datetime(df_test[date_col]).min()
 
    if test_min <= train_max:
        print(f"⚠️  Data leakage detected: "
              f"train max ({train_max}) >= test min ({test_min})")
        return False
 
    print(f"✅ No leakage: train max={train_max}, test min={test_min}")
    return True

The validate_no_data_leakage function was generated automatically because the prompt mentioned "be careful about data leakage." One instruction note, one extra safety function — that's the kind of dividend you get from a good CLAUDE.md and specific prompts.

4. Model Selection: Cross-Validated Comparison Script

With features ready, the next step is comparing candidate models systematically before committing to one. Here's the script to ask Claude Code for:

# src/model_comparison.py
import time
import pandas as pd
import numpy as np
from pathlib import Path
from sklearn.linear_model import LinearRegression, Ridge
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.model_selection import cross_val_score, KFold
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
 
try:
    from xgboost import XGBRegressor
    XGBOOST_AVAILABLE = True
except ImportError:
    XGBOOST_AVAILABLE = False
    print("Note: XGBoost not installed. Excluding from comparison.")
 
RANDOM_STATE = 42
N_SPLITS = 5
 
def compare_models(
    X: pd.DataFrame, y: pd.Series, output_dir: str = "reports"
) -> pd.DataFrame:
    """
    Run 5-fold CV on all candidate models.
    Returns sorted DataFrame with RMSE, R², and wall time.
    """
    Path(output_dir).mkdir(parents=True, exist_ok=True)
 
    models = {
        "LinearRegression": Pipeline([
            ("scaler", StandardScaler()),
            ("model", LinearRegression())
        ]),
        "Ridge": Pipeline([
            ("scaler", StandardScaler()),
            ("model", Ridge(alpha=1.0, random_state=RANDOM_STATE))
        ]),
        "RandomForest": RandomForestRegressor(
            n_estimators=100, random_state=RANDOM_STATE, n_jobs=-1
        ),
        "GradientBoosting": GradientBoostingRegressor(
            n_estimators=100, random_state=RANDOM_STATE
        ),
    }
    if XGBOOST_AVAILABLE:
        models["XGBoost"] = XGBRegressor(
            n_estimators=100, random_state=RANDOM_STATE,
            n_jobs=-1, verbosity=0
        )
 
    kf = KFold(n_splits=N_SPLITS, shuffle=True, random_state=RANDOM_STATE)
    results = []
 
    for name, model in models.items():
        start = time.time()
        neg_mse = cross_val_score(
            model, X, y, cv=kf,
            scoring="neg_mean_squared_error", n_jobs=-1
        )
        rmse = np.sqrt(-neg_mse)
        r2 = cross_val_score(model, X, y, cv=kf, scoring="r2", n_jobs=-1)
        elapsed = time.time() - start
 
        results.append({
            "model": name,
            "rmse_mean": rmse.mean(), "rmse_std": rmse.std(),
            "r2_mean": r2.mean(), "r2_std": r2.std(),
            "time_sec": elapsed
        })
        print(f"{name:20s} | RMSE: {rmse.mean():.4f} ± {rmse.std():.4f} | "
              f"R²: {r2.mean():.4f} | {elapsed:.1f}s")
 
    df_results = pd.DataFrame(results).sort_values("rmse_mean")
    df_results.to_csv(f"{output_dir}/model_comparison.csv", index=False)
    print(f"\n[OK] Saved: {output_dir}/model_comparison.csv")
    return df_results
 
# Expected output (results vary by data):
# LinearRegression     | RMSE: 12.3456 ± 0.4321 | R²: 0.7234 | 0.2s
# Ridge               | RMSE: 12.1234 ± 0.3456 | R²: 0.7456 | 0.3s
# RandomForest        | RMSE:  8.9876 ± 0.2345 | R²: 0.8901 | 12.4s
# GradientBoosting    | RMSE:  8.4567 ± 0.2123 | R²: 0.9012 | 18.7s
# XGBoost             | RMSE:  8.1234 ± 0.1987 | R²: 0.9123 | 5.6s

5. Hyperparameter Tuning with Optuna

Once you've picked a model, Optuna handles the search. Claude Code's strength here is defining the right search space — provide the parameter names and rough bounds, and it fills in suggest_int, suggest_float(log=True), and suggest_categorical appropriately.

# src/hyperparameter_tuning.py
import optuna
import numpy as np
import pandas as pd
import joblib
from pathlib import Path
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.model_selection import cross_val_score, KFold
 
RANDOM_STATE = 42
N_SPLITS = 5
N_TRIALS = 50
MODELS_DIR = Path("models")
MODELS_DIR.mkdir(exist_ok=True)
 
def objective(trial: optuna.Trial, X: pd.DataFrame, y: pd.Series) -> float:
    """Optuna objective — minimize CV-RMSE."""
    params = {
        "n_estimators": trial.suggest_int("n_estimators", 50, 500),
        "learning_rate": trial.suggest_float("learning_rate", 0.01, 0.3, log=True),
        "max_depth": trial.suggest_int("max_depth", 3, 8),
        "min_samples_split": trial.suggest_int("min_samples_split", 2, 20),
        "min_samples_leaf": trial.suggest_int("min_samples_leaf", 1, 10),
        "subsample": trial.suggest_float("subsample", 0.6, 1.0),
        "max_features": trial.suggest_categorical(
            "max_features", ["sqrt", "log2", None]
        ),
        "random_state": RANDOM_STATE
    }
    model = GradientBoostingRegressor(**params)
    kf = KFold(n_splits=N_SPLITS, shuffle=True, random_state=RANDOM_STATE)
    neg_mse = cross_val_score(
        model, X, y, cv=kf,
        scoring="neg_mean_squared_error", n_jobs=-1
    )
    return np.sqrt(-neg_mse.mean())
 
def tune_and_save(
    X_train: pd.DataFrame, y_train: pd.Series,
    n_trials: int = N_TRIALS,
    model_path: str = "models/best_model.pkl"
) -> dict:
    """Run Optuna study and save the best model to disk."""
    optuna.logging.set_verbosity(optuna.logging.WARNING)
 
    study = optuna.create_study(
        direction="minimize",
        sampler=optuna.samplers.TPESampler(seed=RANDOM_STATE)
    )
    print(f"Starting Optuna search: {n_trials} trials...")
    study.optimize(
        lambda trial: objective(trial, X_train, y_train),
        n_trials=n_trials, show_progress_bar=True
    )
 
    best_params = study.best_params
    best_rmse = study.best_value
    print(f"\n✅ Tuning complete — Best RMSE: {best_rmse:.4f}")
    print(f"  Best params: {best_params}")
 
    best_model = GradientBoostingRegressor(**best_params, random_state=RANDOM_STATE)
    best_model.fit(X_train, y_train)
    joblib.dump(best_model, model_path)
    print(f"  Model saved: {model_path}")
 
    return {
        "best_params": best_params,
        "best_rmse": best_rmse,
        "model_path": model_path,
        "n_trials": len(study.trials)
    }

After tuning, ask Claude Code: "Add Optuna visualization for parameter importances and optimization history." It will add optuna.visualization.plot_param_importances() and plot_optimization_history() calls — useful for understanding which parameters mattered most.

6. Model Evaluation and SHAP Analysis

A working model is only half the job. Understanding why it predicts what it predicts is increasingly important — for debugging, for communicating results to stakeholders, and for catching spurious correlations before they hit production.

# src/model_evaluation.py
import shap
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import joblib
from pathlib import Path
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
 
REPORTS_DIR = Path("reports/evaluation")
REPORTS_DIR.mkdir(parents=True, exist_ok=True)
 
def evaluate_model(
    model_path: str,
    X_test: pd.DataFrame,
    y_test: pd.Series,
    output_dir: str = str(REPORTS_DIR)
) -> dict:
    """Load, evaluate, and run SHAP analysis on a saved model."""
    model = joblib.load(model_path)
    predictions = model.predict(X_test)
 
    rmse = np.sqrt(mean_squared_error(y_test, predictions))
    mae = mean_absolute_error(y_test, predictions)
    r2 = r2_score(y_test, predictions)
 
    print(f"=== Test Set Evaluation ===")
    print(f"  RMSE: {rmse:.4f}")
    print(f"  MAE:  {mae:.4f}")
    print(f"  R²:   {r2:.4f}")
 
    # Residual analysis
    residuals = y_test - predictions
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
 
    axes[0].scatter(predictions, residuals, alpha=0.5, s=10)
    axes[0].axhline(y=0, color="r", linestyle="--")
    axes[0].set_xlabel("Predicted")
    axes[0].set_ylabel("Residuals")
    axes[0].set_title("Residual Plot")
 
    min_val = min(y_test.min(), predictions.min())
    max_val = max(y_test.max(), predictions.max())
    axes[1].scatter(y_test, predictions, alpha=0.5, s=10)
    axes[1].plot([min_val, max_val], [min_val, max_val], "r--")
    axes[1].set_xlabel("Actual")
    axes[1].set_ylabel("Predicted")
    axes[1].set_title(f"Actual vs Predicted (R²={r2:.3f})")
 
    plt.tight_layout()
    fig.savefig(f"{output_dir}/residual_plots.png", dpi=150, bbox_inches="tight")
    plt.close(fig)
 
    # SHAP feature importance
    try:
        explainer = shap.TreeExplainer(model)
        # Sample down for speed on large datasets
        X_shap = X_test.sample(min(200, len(X_test)), random_state=42)
        shap_values = explainer.shap_values(X_shap)
 
        fig = plt.figure(figsize=(10, 8))
        shap.summary_plot(shap_values, X_shap, show=False, plot_size=(10, 8))
        plt.tight_layout()
        plt.savefig(f"{output_dir}/shap_summary.png",
                    dpi=150, bbox_inches="tight")
        plt.close()
        print("[OK] SHAP summary plot saved.")
 
    except Exception as e:
        print(f"Warning: SHAP analysis failed: {e}")
        print("  If TreeExplainer doesn't apply, try KernelExplainer instead.")
 
    return {"rmse": rmse, "mae": mae, "r2": r2}

7. Common Pitfalls When Using Claude Code for ML

These are the failure modes I've actually hit in practice, not hypothetical edge cases.

Pitfall 1: Leakage Isn't Caught Automatically

Without being told, Claude Code may fit a StandardScaler on the full dataset before splitting. One explicit note in CLAUDE.md — "fit preprocessors only on training data" — prevents this. The validate_no_data_leakage() function in Section 3 catches time-based leakage after the split.

Pitfall 2: Missing Random Seeds

"Write reproducible code" is not enough. Across numpy, sklearn, xgboost, optuna, and PyTorch, seeds live in different places. This utility function, called once at script startup, covers them all:

# src/utils.py
import os, random, numpy as np
 
RANDOM_STATE = 42
 
def set_all_seeds(seed: int = RANDOM_STATE) -> None:
    """Fix random seeds across all libraries used in this project."""
    random.seed(seed)
    np.random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    try:
        import torch
        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed)
    except ImportError:
        pass  # No PyTorch in this project — skip

Pitfall 3: Memory Errors on Large Datasets

If you don't tell Claude Code the dataset size, it will assume it fits in memory. Tell it upfront: "this file is 5 million rows." Claude Code will then generate chunked reading automatically:

def load_large_csv(
    filepath: str,
    chunksize: int = 100_000,
    target_col: str = "target"
) -> tuple[pd.DataFrame, pd.Series]:
    """Load a large CSV in chunks to stay within memory limits."""
    chunks = []
    for i, chunk in enumerate(pd.read_csv(filepath, chunksize=chunksize)):
        chunk = chunk.dropna(subset=[target_col])
        chunk["purchase_date"] = pd.to_datetime(chunk["purchase_date"])
        chunks.append(chunk)
        if (i + 1) % 10 == 0:
            print(f"  Loaded: {(i + 1) * chunksize:,} rows")
    df = pd.concat(chunks, ignore_index=True)
    print(f"Done: {len(df):,} total rows")
    return df.drop(columns=[target_col]), df[target_col]

Pitfall 4: Vague Prompts Produce Vague Search Spaces

"Optimize this model" produces a generic search space. "Tune n_estimators in range 50–500, learning_rate on log scale 0.01–0.3, and max_depth from 3 to 8" produces a search space with good coverage of the actual hyperparameter landscape. Specificity in prompts is proportional to quality of output.

For test-driven validation of preprocessing functions, the patterns in the Claude Code TDD workflow guide apply directly here.

8. MLOps Basics: Versioned Model Registry

Production models need regular retraining. Before building a pipeline, set up a versioned registry so you can always compare last week's model against today's:

# src/mlops_utils.py
import json, joblib
from datetime import datetime
from pathlib import Path
from typing import Any
 
MODELS_DIR = Path("models")
REGISTRY_PATH = MODELS_DIR / "registry.json"
MODELS_DIR.mkdir(exist_ok=True)
 
def save_model_with_metadata(
    model: Any,
    metrics: dict,
    params: dict,
    model_name: str
) -> str:
    """Save a versioned model file and log metadata to the registry."""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"{model_name}_{timestamp}.pkl"
    path = MODELS_DIR / filename
    joblib.dump(model, path)
 
    entry = {
        "model_name": model_name,
        "timestamp": timestamp,
        "model_path": str(path),
        "metrics": metrics,
        "params": params,
        "model_type": type(model).__name__
    }
 
    registry = []
    if REGISTRY_PATH.exists():
        with open(REGISTRY_PATH) as f:
            registry = json.load(f)
    registry.append(entry)
    with open(REGISTRY_PATH, "w") as f:
        json.dump(registry, f, indent=2, ensure_ascii=False)
 
    print(f"✅ Saved: {filename}")
    print(f"   RMSE: {metrics.get('rmse', 'N/A'):.4f}")
    return str(path)
 
def get_best_model(model_name: str, metric: str = "rmse") -> dict | None:
    """Return registry entry for the best recorded version of a model."""
    if not REGISTRY_PATH.exists():
        return None
 
    with open(REGISTRY_PATH) as f:
        registry = json.load(f)
 
    candidates = [r for r in registry if r["model_name"] == model_name]
    if not candidates:
        return None
 
    # Lower is better for RMSE/MAE, higher is better for R²
    reverse = metric not in ("rmse", "mae")
    return sorted(
        candidates,
        key=lambda x: x["metrics"].get(metric, float("inf")),
        reverse=reverse
    )[0]

With this in place, get_best_model("sales_prediction") returns the path and parameters of your best model across all training runs. Add a scheduled retraining job and a rollback check, and you have the core of a lightweight MLOps pipeline.

The decision framework for when to bring Claude Code into a personal project (versus doing things manually) is covered in Claude Code solo dev decision patterns.

9. Testing Preprocessing Functions with pytest

A preprocessing pipeline that isn't tested is a liability. With Claude Code, you can generate tests as fast as you write the functions. The pattern: describe what the function should do, ask for corresponding tests.

# Prompt to Claude Code:
# "Write pytest tests for the engineer_features() function in
# src/feature_engineering.py. Test these properties:
# - Output has no purchase_date or raw category columns
# - day_of_week is in range [0, 6]
# - total_amount == price * quantity for all rows
# - No NaN values in encoded columns
# - .copy() means input DataFrame is unchanged"

Generated test file:

# tests/test_feature_engineering.py
import pandas as pd
import numpy as np
import pytest
from src.feature_engineering import engineer_features, validate_no_data_leakage
 
@pytest.fixture
def sample_df():
    """Minimal realistic fixture — 10 rows, all valid."""
    return pd.DataFrame({
        "purchase_date": pd.date_range("2024-01-01", periods=10, freq="D").astype(str),
        "customer_age": [25, 30, 35, 40, 45, 50, 55, 60, 65, 70],
        "product_category": ["electronics", "clothing", "food"] * 3 + ["electronics"],
        "price": [10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0],
        "quantity": [1, 2, 3, 1, 2, 3, 1, 2, 3, 1],
    })
 
def test_output_has_no_raw_columns(sample_df):
    """purchase_date and category string columns must be dropped."""
    result = engineer_features(sample_df)
    assert "purchase_date" not in result.columns
    assert "product_category" not in result.columns
    assert "age_group" not in result.columns
 
def test_day_of_week_range(sample_df):
    """day_of_week must be an integer in [0, 6]."""
    result = engineer_features(sample_df)
    assert result["day_of_week"].between(0, 6).all(), \
        "day_of_week values outside [0, 6]"
 
def test_total_amount_calculation(sample_df):
    """total_amount must equal price × quantity exactly."""
    result = engineer_features(sample_df)
    expected = sample_df["price"] * sample_df["quantity"]
    pd.testing.assert_series_equal(
        result["total_amount"].reset_index(drop=True),
        expected.reset_index(drop=True),
        check_names=False
    )
 
def test_no_nan_in_encoded_columns(sample_df):
    """Label-encoded columns must have zero missing values."""
    result = engineer_features(sample_df)
    encoded_cols = [c for c in result.columns if c.endswith("_encoded")]
    assert len(encoded_cols) > 0, "Expected at least one _encoded column"
    for col in encoded_cols:
        assert result[col].isnull().sum() == 0, f"{col} has NaN values"
 
def test_input_not_mutated(sample_df):
    """engineer_features() must not modify the input DataFrame."""
    original_cols = sample_df.columns.tolist()
    original_shape = sample_df.shape
    _ = engineer_features(sample_df)
    assert sample_df.columns.tolist() == original_cols
    assert sample_df.shape == original_shape
 
def test_no_data_leakage_valid(sample_df):
    """Split where test dates are all after train dates — should pass."""
    df_train = sample_df[sample_df["purchase_date"] < "2024-01-06"].copy()
    df_test = sample_df[sample_df["purchase_date"] >= "2024-01-06"].copy()
    assert validate_no_data_leakage(df_train, df_test) is True
 
def test_no_data_leakage_invalid(sample_df):
    """Split where test and train overlap — should detect leakage."""
    df_train = sample_df.copy()
    df_test = sample_df.head(3).copy()  # Overlapping dates
    assert validate_no_data_leakage(df_train, df_test) is False

Running these tests takes a few seconds and catches an entire class of bugs that would otherwise surface weeks later in production. Generate them alongside the implementation, not as an afterthought.

10. Putting It All Together: A Complete Run Script

With all the modules in place, a top-level script that strings them together makes the full pipeline runnable with a single command:

# run_pipeline.py
"""
End-to-end ML pipeline for sales prediction.
Run with: python run_pipeline.py
"""
import pandas as pd
from sklearn.model_selection import train_test_split
 
from src.utils import set_all_seeds, RANDOM_STATE
from src.eda_report import generate_eda_report
from src.feature_engineering import engineer_features, validate_no_data_leakage
from src.model_comparison import compare_models
from src.hyperparameter_tuning import tune_and_save
from src.model_evaluation import evaluate_model
from src.mlops_utils import save_model_with_metadata
 
DATA_PATH = "data/raw/sales.csv"
TARGET_COL = "revenue"
MODEL_NAME = "sales_prediction"
 
def main():
    set_all_seeds()
 
    # Step 1: EDA
    print("\n" + "=" * 50)
    print("STEP 1: EDA")
    print("=" * 50)
    generate_eda_report(DATA_PATH)
 
    # Step 2: Load and split BEFORE any preprocessing
    print("\n" + "=" * 50)
    print("STEP 2: Load and Split")
    print("=" * 50)
    df = pd.read_csv(DATA_PATH)
    df_train, df_test = train_test_split(
        df, test_size=0.2, random_state=RANDOM_STATE, shuffle=False  # Time series: no shuffle
    )
    print(f"Train: {len(df_train)} rows | Test: {len(df_test)} rows")
 
    # Step 3: Feature engineering
    print("\n" + "=" * 50)
    print("STEP 3: Feature Engineering")
    print("=" * 50)
    X_train = engineer_features(df_train.drop(columns=[TARGET_COL]))
    y_train = df_train[TARGET_COL].reset_index(drop=True)
    X_test = engineer_features(df_test.drop(columns=[TARGET_COL]))
    y_test = df_test[TARGET_COL].reset_index(drop=True)
    print(f"Features: {X_train.shape[1]} | Train rows: {len(X_train)}")
 
    # Step 4: Model comparison
    print("\n" + "=" * 50)
    print("STEP 4: Model Comparison (5-fold CV)")
    print("=" * 50)
    comparison_df = compare_models(X_train, y_train)
    print(f"\nBest model from comparison: {comparison_df.iloc[0]['model']}")
 
    # Step 5: Hyperparameter tuning
    print("\n" + "=" * 50)
    print("STEP 5: Hyperparameter Tuning (Optuna)")
    print("=" * 50)
    tune_result = tune_and_save(X_train, y_train)
 
    # Step 6: Final evaluation on test set
    print("\n" + "=" * 50)
    print("STEP 6: Test Set Evaluation")
    print("=" * 50)
    metrics = evaluate_model(tune_result["model_path"], X_test, y_test)
 
    # Step 7: Register model
    save_model_with_metadata(
        model=__import__("joblib").load(tune_result["model_path"]),
        metrics=metrics,
        params=tune_result["best_params"],
        model_name=MODEL_NAME
    )
 
    print("\n" + "=" * 50)
    print("Pipeline complete.")
    print(f"  Final RMSE: {metrics['rmse']:.4f}")
    print(f"  Final R²:   {metrics['r2']:.4f}")
    print("  Reports: reports/")
    print("  Models:  models/")
    print("=" * 50)
 
if __name__ == "__main__":
    main()

This single script runs every step in order, prints clear headers at each phase, and leaves you with reports and a versioned model file. Ask Claude Code to generate this last — it can't know the right module names until the other files exist, but once they do, it assembles the orchestration correctly on the first try.

11. What the Docs Don't Tell You — Lessons From Real Runs

Everything above is the pattern. What follows is what actually running this workflow on real projects taught me — the things that don't show up in any documentation.

The single biggest win was modularizing exploratory notebooks. Pairing with Claude Code, I broke a 600-line single notebook into six tested modules (utils, eda_report, feature_engineering, model_comparison, hyperparameter_tuning, mlops_utils) in about 25 minutes. The same restructuring by hand is a half-day job for me, and it's the kind of work I tend to avoid — which is exactly why notebooks rot.

The second surprise was how much data leakage detection earns its keep. With validate_no_data_leakage() wired into the CLAUDE.md constraints, one project surfaced target-derived information hiding in 3 of 11 features early in development. That's far sooner than a human notices a validation score that's "suspiciously good."

Optuna's payoff became legible in numbers, too. Running 100 trials against GradientBoostingRegressor moved RMSE from 0.0421 to 0.0388 — roughly an 8% improvement. Not dramatic, but switching the learning_rate search to a log-uniform distribution stabilized convergence in a way that would be easy to miss when tuning by hand.

Just as importantly, the runs clarified what you must not delegate. Whether an R² of 0.91 is good enough, and whether a feature is worth engineering for business reasons, are judgment calls only someone who knows the domain can make. Hand those over and you quietly grow a plausible-but-wrong model. The boundary is clean: boilerplate is the AI's job, judgment is yours.

Finally, a note on cost. A single session that generates the full pipeline lands in the low hundreds of thousands of tokens in my experience as an indie developer. Encoding constraints in CLAUDE.md has a pleasant side effect here — you stop re-explaining the same rules every turn, the exchanges get shorter, and token usage drops as a result.

The Right Way to Think About This

Claude Code handles the boilerplate; you handle the judgment. The distinction matters for data science more than most fields.

Deciding which features to engineer, which models are worth comparing for this particular problem, and whether an R² of 0.91 is actually good enough — those are judgment calls that require domain knowledge about the problem. Claude Code cannot make them for you, and shouldn't try.

What it can do is eliminate the gap between "I have an idea" and "I have working code that tests the idea." That gap is where a lot of data science time disappears. When generating boilerplate takes five minutes instead of two hours, you make more hypotheses, run more experiments, and end up with better models.

The CLAUDE.md from Section 1 is the key that unlocks this. Start there, then build each module from the prompts in this article. By the time you reach the registry in Section 8, you'll have a data science workflow that's both systematic and fast to iterate on.

One More Thing: Prompt Engineering for Data Science Tasks

A final note on how to phrase prompts specifically for data science work. Vague prompts produce generic code; specific prompts produce code that fits your actual problem.

These phrasings consistently produce better output:

Instead of "analyze this data" → Try "perform EDA on the price and quantity columns; plot their distributions and check for correlations with the revenue target"
Instead of "build a model" → Try "compare Random Forest and GradientBoosting using 5-fold CV on this regression task; report RMSE and R² with standard deviations"
Instead of "tune the model" → Try "run Optuna with 100 trials on GradientBoostingRegressor; search n_estimators in [50, 500], learning_rate log-uniform in [0.01, 0.3], max_depth in [3, 8]"
Instead of "handle missing values" → Try "for numeric columns, fill with median; for categorical columns, fill with the literal string 'Unknown'; assert no NaN remains after this step"

The more precisely you specify what you want, the more useful the output. Data science has a vocabulary — use it in your prompts, and Claude Code will use it back.

Thank You for Reading

Claude Lab is ad-free, supported entirely by members like you. We publish practical guides daily with implementation code, benchmarks, and production-ready patterns. If you've found it useful, we'd love to have you on board.

✦Copy-paste ready implementation code
✦New advanced guides published daily
✦$5/mo or $10 for lifetime access

View Membership →