scikit-rank 0.1.0

Scikit-learn-Compatible Neural Ranking Models for Tabular Recommender Systems

scikit-rank

sklearn-compatible DCNv2 estimators for tabular ranking, classification, and regression.

scikit-rank provides neural, drop-in alternatives to LightGBM / CatBoost / XGBoost for tabular problems. The core model is a DCNv2 (Deep & Cross Network v2) wrapped in three scikit-learn estimators with a familiar fit / predict / predict_proba API, so they slot straight into existing sklearn pipelines, GridSearchCV, clone, and get_params / set_params.

from scikit_rank import DCNClassifier, DCNRegressor, DCNRanker

Features

• Three estimators, one backbone. DCNClassifier (binary / multiclass), DCNRegressor, and DCNRanker (pairwise / listwise learning-to-rank) share the same DCNv2 core and training loop.
• sklearn-native ergonomics. Flat __init__ hyperparameters, fit(X, y, group=..., eval_set=...), early stopping, custom eval metrics, and full get_params / set_params / clone support — mirroring the LightGBM / XGBoost sklearn wrappers.
• Flexible inputs. Accepts numpy.ndarray, pandas.DataFrame, polars.DataFrame, and polars.LazyFrame. String categoricals and NaNs are handled natively (no manual encoding required). When X is a polars (Lazy)Frame, y and group may be passed as column names.
• Out-of-core lazy training. A polars.LazyFrame is preprocessed and streamed through a temporary Arrow IPC file, so datasets larger than memory can be trained without materializing the full frame. Lazy training is numerically equivalent to the in-memory path (see tests/test_lazy_eager_parity.py).
• Rich feature encoders. Per-feature and hashed (Unified Embedding / multihash) categorical encoders, piecewise-linear (PLE) and quantile numeric encoders, and named external embedding streams (e.g. precomputed entity vectors).
• Ranking losses. BCE, BPR, LambdaRank, listwise/softmax, ordinal (CORAL), and more — selectable by string spec, e.g. loss="bpr:sampling=all_pairs".
• Scales with 🤗 Accelerate. Mixed precision, gradient accumulation/clipping, EMA weights, LR schedulers, and multi-GPU / DDP training (including sharded, group-aware ranking evaluation).

Installation

The project uses uv and targets Python ≥ 3.12.

uv sync            # full environment (torch, accelerate, scikit-rank)

Or add it to your own project:

uv add scikit-rank

Quick start

Classification

import polars as pl
from scikit_rank import DCNClassifier

X = pl.DataFrame({"num": [0.1, 1.2, -0.3], "cat": ["a", "b", "a"]})
y = [0, 1, 0]

clf = DCNClassifier(
    hidden_units=[256, 128],
    cross_layers=3,
    epochs=10,
    batch_size=1024,
    num_features=["num"],
    cat_features=["cat"],
    random_state=0,
)
clf.fit(X, y)

proba = clf.predict_proba(X)   # (n_samples, n_classes)
labels = clf.predict(X)

Regression

from scikit_rank import DCNRegressor

reg = DCNRegressor(epochs=10, num_features=["num"], cat_features=["cat"]).fit(X, y_reg)
pred = reg.predict(X)

Ranking

Pass a group (query / impression id) so the loss and metrics are computed per group. group can be an array or, for (Lazy)Frames, a column name.

from scikit_rank import DCNRanker

ranker = DCNRanker(
    loss="listwise",
    hidden_units=[256, 128],
    cross_layers=3,
    epochs=10,
    num_features=["num"],
    cat_features=["cat"],
    random_state=0,
)
ranker.fit(train_df, y="click", group="impression_id")
scores = ranker.predict(test_df)   # higher score = ranked higher

Lazy / out-of-core training

lf = pl.scan_parquet("large_dataset.parquet")   # never fully materialized
clf = DCNClassifier(
    num_features=["x1", "x2"],
    cat_features=["cat"],
    chunk_rows=100_000,          # Arrow streaming chunk size
    epochs=5,
).fit(lf, y="target")

Custom evaluation metric + early stopping

from sklearn.metrics import roc_auc_score

clf = DCNClassifier(
    epochs=100,
    early_stopping_rounds=5,
    eval_metric=lambda y_true, y_pred: roc_auc_score(y_true, y_pred),
    eval_metric_name="auc",
    eval_metric_direction="max",
)
clf.fit(X_train, y_train, eval_set=(X_val, y_val))
print(clf.history_)   # per-epoch train/val metrics

Demo notebook

demo.ipynb is an interactive Jupyter notebook that showcases the library's capabilities end to end. It demonstrates preparing tabular data, configuring and training the sklearn-compatible DCNv2 estimators, and inspecting their predictions and fitted state. Run it after installing the project environment:

uv sync
uv run jupyter lab demo.ipynb

Key hyperparameters

Group	Parameters
Architecture	hidden_units, cross_layers, cross_rank, structure (stacked/parallel), cross_type, use_moe, use_inner_cross_layers
Embeddings / encoders	embedding_dim, cat_encoder, num_encoder, multihash_features, embedding_features, normalize_numeric, ple_n_bins
Optimization	loss, lr, weight_decay, optimizer, epochs, batch_size, lr_scheduler, grad_clip_norm, embedding_regularizer, ema_decay
Evaluation	eval_metric, eval_metric_name, eval_metric_direction, eval_metric_group_aware, early_stopping_rounds
Data / runtime	num_features, cat_features, chunk_rows, random_state, accelerator_config, verbose

Multi-GPU / device settings are passed through accelerator_config (e.g. accelerator_config={"cpu": True} or mixed-precision / DDP options from 🤗 Accelerate).

Experiments & benchmarks

The exps/ directory contains reproducible benchmarks and the scripts, configs, and hyperparameters behind the published results:

• exps/bars/ — BARS-CTR benchmark: DCNv2 vs. LightGBM / CatBoost / XGBoost / FuxiCTR on Criteo and Avazu.
• exps/temporal_article/ — temporal-split experiments (Avazu-Time, MIND-small) with the best models and their exact hyperparameters from the paper.

See exps/README.md for dataset sources and how to recreate the temporal-split datasets.

Development

If you have devenv support, run devenv shell from the repository root to instantiate the development environment with the required modules and tools.

devenv shell                 # optional: enter the configured development shell
uv sync                      # install project and development dependencies
uv run python -m pytest -q   # run the test suite
uv run ruff check            # lint