Concepts 🚀 Project 🧮 Labs

Full benchmark engine — putting it all together

The complete benchmarking function: load .pkl → detect model type → load dataset → predict → time → metrics → return JSON. This is what OpenBenchML does.

script.py v3.1

import joblib
import pickle
import time
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split, cross_validate, KFold
from sklearn.metrics import (
    mean_absolute_error, mean_squared_error, r2_score,
    accuracy_score, precision_score, recall_score, f1_score
)

# ════════════════════════════════════════════════════════════════════
#  THE COMPLETE BENCHMARK ENGINE FUNCTION
#  Works with ANY .pkl model — sklearn, xgboost, custom.
# ════════════════════════════════════════════════════════════════════

def benchmark_model(model_bytes: bytes, X, y, run_cv: bool = True) -> dict:
    """Benchmark any uploaded .pkl model.

    Args:
        model_bytes: raw bytes from UploadFile.read()
        X: features (numpy array or DataFrame)
        y: target (numpy array or Series)
        run_cv: whether to run 5-fold cross-validation

    Returns:
        dict with metrics, timing, model info
    """
    result = {}

    # ── 1. Load the model ──
    try:
        model = pickle.loads(model_bytes)
    except Exception as e:
        return {"status": "error", "message": f"Failed to load model: {e}"}

    result["model_type"] = type(model).__name__
    result["model_module"] = type(model).__module__

    # ── 2. Duck-type check ──
    if not hasattr(model, "predict"):
        return {"status": "error", "message": "Not a valid ML model (no predict method)"}

    # ── 3. Split data ──
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )
    result["train_samples"] = len(X_train)
    result["test_samples"] = len(X_test)

    # ── 4. Time the prediction ──
    start = time.perf_counter()
    y_pred = model.predict(X_test)
    inference_time = time.perf_counter() - start

    result["inference_time_sec"] = inference_time
    result["latency_per_sample_sec"] = inference_time / len(X_test)
    result["throughput_samples_per_sec"] = len(X_test) / inference_time

    # ── 5. Detect task type (regression vs classification) ──
    is_regression = y_pred.dtype.kind == "f"
    result["task_type"] = "regression" if is_regression else "classification"

    # ── 6. Compute metrics ──
    if is_regression:
        result["MAE"] = float(mean_absolute_error(y_test, y_pred))
        result["MSE"] = float(mean_squared_error(y_test, y_pred))
        result["RMSE"] = float(np.sqrt(result["MSE"]))
        result["R2"] = float(r2_score(y_test, y_pred))
    else:
        result["accuracy"] = float(accuracy_score(y_test, y_pred))
        result["precision"] = float(precision_score(y_test, y_pred, average="weighted", zero_division=0))
        result["recall"] = float(recall_score(y_test, y_pred, average="weighted", zero_division=0))
        result["f1"] = float(f1_score(y_test, y_pred, average="weighted", zero_division=0))

    # ── 7. Cross-validation (optional) ──
    if run_cv:
        cv = KFold(n_splits=5, shuffle=True, random_state=42)
        if is_regression:
            scoring = {"MAE": "neg_mean_absolute_error", "R2": "r2"}
        else:
            scoring = {"accuracy": "accuracy", "f1": "f1_weighted"}

        try:
            cv_results = cross_validate(model, X, y, cv=cv, scoring=scoring)
            for metric in scoring:
                key = f"CV_{metric}"
                result[key] = float(cv_results[f"test_{metric}"].mean())
        except Exception as e:
            result["cv_error"] = str(e)

    # ── 8. Feature importances (if available) ──
    if hasattr(model, "feature_importances_"):
        result["feature_importances"] = model.feature_importances_.tolist()
    elif hasattr(model, "coef_"):
        result["coefficients"] = np.array(model.coef_).ravel().tolist()

    result["status"] = "success"
    return result


# ════════════════════════════════════════════════════════════════════
#  TEST THE ENGINE WITH 2 DIFFERENT MODEL TYPES
# ════════════════════════════════════════════════════════════════════
from sklearn.linear_model import LinearRegression, LogisticRegression
from sklearn.ensemble import RandomForestRegressor
from sklearn.datasets import make_regression, make_classification

# Test 1: Regression model
print('=' * 60)
print('TEST 1: Regression (LinearRegression)')
print('=' * 60)
Xr, yr = make_regression(n_samples=500, n_features=5, noise=15, random_state=42)
lr = LinearRegression().fit(Xr, yr)
model_bytes_r = pickle.dumps(lr)
result_r = benchmark_model(model_bytes_r, Xr, yr)
for k, v in result_r.items():
    if isinstance(v, list):
        print(f'  {k}: [{len(v)} values]')
    else:
        print(f'  {k}: {v}')

# Test 2: Classification model
print('\n' + '=' * 60)
print('TEST 2: Classification (LogisticRegression)')
print('=' * 60)
Xc, yc = make_classification(n_samples=500, n_features=5, n_informative=3, random_state=42)
clf = LogisticRegression().fit(Xc, yc)
model_bytes_c = pickle.dumps(clf)
result_c = benchmark_model(model_bytes_c, Xc, yc)
for k, v in result_c.items():
    if isinstance(v, list):
        print(f'  {k}: [{len(v)} values]')
    else:
        print(f'  {k}: {v}')

Python output click Run to execute

Click "Run" to execute the code on the left.
The output will appear here.

🧠 How it works

This is the COMPLETE benchmark engine — exactly what OpenBenchML does. 8 steps: load → validate → split → time → detect type → metrics → CV → importances. KEY INSIGHT: the function takes model BYTES (not a file path) — that's what you get from UploadFile.read(). Duck typing throughout: hasattr(model, 'predict'), hasattr(model, 'feature_importances_'). Works with LinearRegression, RandomForest, XGBoost, LightGBM, ANY model with .predict(). Auto-detects regression vs classification by checking y_pred.dtype. Returns a dict that FastAPI can serialize to JSON directly. The 2 tests at the bottom prove it works with BOTH regression AND classification models — same function, no changes.

✏️ Try these changes

Change CV folds from 5 to 10
→ more folds, slower, slightly more accurate
Add a 'training_time' field by re-fitting the model + timing it
→ measures how long training takes
Replace LogisticRegression with RandomForestClassifier
→ same function works — that's the point
Add error handling for when X has wrong number of features
→ model.predict crashes if features don't match

← Previous lab

HTML form — upload .pkl + select dataset

Next lab →

Save benchmark results + compare models