Metadata-Version: 2.4
Name: scikit-rank
Version: 0.1.0
Summary: Scikit-learn-Compatible Neural Ranking Models for Tabular Recommender Systems
License-File: LICENSE
Requires-Python: <3.14,>=3.12
Requires-Dist: accelerate>=1.14.0
Requires-Dist: einops>=0.8.2
Requires-Dist: muon-optimizer>=0.1.0
Requires-Dist: numpy>=2.4.6
Requires-Dist: pandas>=3.0.3
Requires-Dist: polars>=1.41.2
Requires-Dist: pyarrow>=19.0.0
Requires-Dist: pytorch-ignite>=0.5.4
Requires-Dist: schedulefree>=1.4.1
Requires-Dist: scikit-learn>=1.8.0
Requires-Dist: scipy>=1.17.1
Requires-Dist: torch>=2.12.0
Requires-Dist: tqdm>=4.68.2
Description-Content-Type: text/markdown

# 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](https://arxiv.org/abs/2008.13535)
(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`.

```python
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`](https://docs.astral.sh/uv/) and targets Python ≥ 3.12.

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

Or add it to your own project:

```bash
uv add scikit-rank
```

## Quick start

### Classification

```python
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

```python
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.

```python
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

```python
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

```python
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`](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:

```bash
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/`](exps/) directory contains reproducible benchmarks and the scripts,
configs, and hyperparameters behind the published results:

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

See [`exps/README.md`](exps/README.md) for dataset sources and how to recreate the
temporal-split datasets.

## Development

If you have [devenv](https://devenv.sh/) support, run `devenv shell` from the
repository root to instantiate the development environment with the required
modules and tools.

```bash
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
```
