stabilization-uplift / stabilization-uplift.py
IlVarsh's picture
Add Stabilization Uplift calculator and evaluate module
883ccfd verified
Raw History Blame Contribute Delete
3.98 kB
"""Stabilization Score (SS) and Stabilization Uplift (SU) as a Hugging Face `evaluate` module.
The metric logic lives in the `stabilization-uplift` package on PyPI; this module only wraps it.
"""
import datasets
import evaluate
from stabilization_uplift import stabilization_score, stabilization_uplift
_CITATION = """\
@inproceedings{varshavskiy2025mitigating,
title = {Mitigating Model Drift in Developing Economies Using Synthetic Data and Outliers},
author = {Varshavskiy, Ilyas and Boboeva, Bonu and Khalilbekov, Shuhrat and Azimi, Azizjon and Shulgin, Sergey and Nizamitdinov, Akhlitdin and S{\\'a}ez de Oc{\\'a}riz Borde, Haitz},
booktitle = {NeurIPS 2025 Workshop on Generative AI in Finance},
year = {2025},
url = {https://openreview.net/forum?id=zfTaFD0B5Z}
}
"""
_DESCRIPTION = """\
Stabilization Score (SS) and Stabilization Uplift (SU) evaluate how stable a classifier stays under a sudden
distribution shift, such as a macroeconomic shock. SS measures how much one model's ROC AUC changed between the
base (pre-shock) and shock periods, normalized by the severity of the shift. SU measures whether model B
(e.g. trained with synthetic outliers) is more stable and not worse than model A (a baseline) under the shock.
"""
_KWARGS_DESCRIPTION = """
Each row is one comparison of model B with model A.
Args:
auc_base_A (`list` of `float`): ROC AUC of model A on the base (pre-shock) test set, in (0, 1].
auc_shock_A (`list` of `float`): ROC AUC of model A on the shock test set, in (0, 1].
auc_base_B (`list` of `float`): ROC AUC of model B on the base test set, in (0, 1].
auc_shock_B (`list` of `float`): ROC AUC of model B on the shock test set, in (0, 1].
dist_shift (`list` of `float`): distribution shift between the base and shock data, normally in [0, 1].
Compute it with `stabilization_uplift.distribution_shift` (pip install "stabilization-uplift[shift]").
Returns:
stabilization_score_A (`list` of `float`): SS of model A for each row, in [0.5, 1]. 1 means no change in AUC.
stabilization_score_B (`list` of `float`): SS of model B for each row, in [0.5, 1].
stabilization_uplift (`list` of `float`): SU of B over A for each row, in [0, 1). 0 means no uplift.
Examples:
>>> metric = evaluate.load("zyplai/stabilization-uplift")
>>> results = metric.compute(auc_base_A=[0.80], auc_shock_A=[0.70],
... auc_base_B=[0.80], auc_shock_B=[0.81], dist_shift=[0.2])
>>> {k: [round(v, 3) for v in vals] for k, vals in results.items()}
{'stabilization_score_A': [0.915], 'stabilization_score_B': [0.992], 'stabilization_uplift': [0.725]}
"""
_FEATURE_NAMES = ["auc_base_A", "auc_shock_A", "auc_base_B", "auc_shock_B", "dist_shift"]
@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
class StabilizationUplift(evaluate.Metric):
def _info(self):
return evaluate.MetricInfo(
description=_DESCRIPTION,
citation=_CITATION,
inputs_description=_KWARGS_DESCRIPTION,
features=datasets.Features({name: datasets.Value("float64") for name in _FEATURE_NAMES}),
homepage="https://pypi.org/project/stabilization-uplift/",
codebase_urls=["https://github.com/zypl-ai/stabilization_uplift"],
reference_urls=["https://openreview.net/forum?id=zfTaFD0B5Z", "https://arxiv.org/abs/2510.09294"],
)
def _compute(self, auc_base_A, auc_shock_A, auc_base_B, auc_shock_B, dist_shift):
rows = list(zip(auc_base_A, auc_shock_A, auc_base_B, auc_shock_B, dist_shift))
return {
"stabilization_score_A": [stabilization_score(a_base, a_shock, shift) for a_base, a_shock, _, _, shift in rows],
"stabilization_score_B": [stabilization_score(b_base, b_shock, shift) for _, _, b_base, b_shock, shift in rows],
"stabilization_uplift": [stabilization_uplift(*row) for row in rows],
}