Metadata-Version: 2.4
Name: wp21_train
Version: 0.12.0
Summary: Framework that provides tools allowing for monitoring, training and evaluating ML models
Home-page: https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework.git
Author: ATLAS NextGen WP2.1
Author-email: ngtwp21@cern.ch
Project-URL: Source, https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework
Project-URL: Documentation, https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework/-/blob/train_orchestrator_release/docs/README.md
Project-URL: Examples, https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework/-/tree/train_orchestrator_release/examples
Project-URL: Issues, https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework/-/issues
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: numpy<2.0,>=1.21
Requires-Dist: PyYAML<7.0,>=6.0
Requires-Dist: uproot<6.0,>=5.0
Requires-Dist: awkward<3.0,>=2.5
Requires-Dist: lz4<5.0,>=4.3
Requires-Dist: xxhash<4.0,>=3.4
Requires-Dist: tqdm>=4.67.1
Requires-Dist: pybind11>=3.0.1
Requires-Dist: vector>=1.7.0
Requires-Dist: pyarrow>=18.0.0
Requires-Dist: scipy<2.0,>=1.10
Provides-Extra: ml-utils
Requires-Dist: wp21-ml-utils; extra == "ml-utils"
Requires-Dist: tensorflow>=2.13; extra == "ml-utils"
Requires-Dist: scikit-learn>=1.2; extra == "ml-utils"
Provides-Extra: torch
Requires-Dist: torch>=2.1; extra == "torch"
Provides-Extra: brevitas
Requires-Dist: torch<2.14,>=2.1; extra == "brevitas"
Requires-Dist: brevitas<0.14,>=0.13.2; extra == "brevitas"
Provides-Extra: pytorch
Requires-Dist: torch>=2.1; extra == "pytorch"
Provides-Extra: optuna
Requires-Dist: optuna>=3.0; extra == "optuna"
Provides-Extra: hgq
Requires-Dist: tensorflow>=2.13; extra == "hgq"
Requires-Dist: hgq==0.2.6; extra == "hgq"
Provides-Extra: hgq2
Requires-Dist: HGQ2>=0.1.5; extra == "hgq2"
Requires-Dist: keras>=3.0.0; extra == "hgq2"
Requires-Dist: tensorflow>=2.16; extra == "hgq2"
Provides-Extra: tensorflow
Requires-Dist: tensorflow>=2.13; extra == "tensorflow"
Requires-Dist: tf2onnx>=1.14; extra == "tensorflow"
Provides-Extra: xgboost
Requires-Dist: xgboost>=3.1.2; extra == "xgboost"
Provides-Extra: tf-tools
Requires-Dist: tf2onnx>=1.14; extra == "tf-tools"
Provides-Extra: plotting
Requires-Dist: matplotlib>=3.7; extra == "plotting"
Provides-Extra: physics
Requires-Dist: pandas>=2.0; extra == "physics"
Requires-Dist: matplotlib>=3.7; extra == "physics"
Provides-Extra: legacy
Requires-Dist: pandas>=2.0; extra == "legacy"
Requires-Dist: matplotlib>=3.7; extra == "legacy"
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == "dev"
Requires-Dist: pytest-cov>=4.0; extra == "dev"
Requires-Dist: ruff>=0.4; extra == "dev"
Requires-Dist: mypy>=1.5; extra == "dev"
Dynamic: author
Dynamic: author-email
Dynamic: classifier
Dynamic: description
Dynamic: description-content-type
Dynamic: home-page
Dynamic: project-url
Dynamic: provides-extra
Dynamic: requires-dist
Dynamic: requires-python
Dynamic: summary

# WP21 Train

Current development version: **0.12.0**. Legacy APIs remain available.

[GitLab source repository](https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework)
· [PyPI package](https://pypi.org/project/wp21-train/)
· [Documentation](https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework/-/blob/train_orchestrator_release/docs/README.md)
· [Report an issue](https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework/-/issues)

WP21 Train orchestrates the training and evaluation operations used by the WP21
group. A trainer adapts a model backend; `TrainingRunner` runs it once or searches
its parameters with Optuna. The package also provides backend-independent
classification and trigger evaluation tools.

## Changes in this version

The orchestration interface introduced in 0.10.0 keeps older project/physics utilities in
[legacy_support](https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework/-/blob/train_orchestrator_release/docs/legacy-support.md), while preserving their old imports.
Model list scans use `length_range` (legacy length inputs remain readable).
ML Utils models can load [separate HGQ2 configuration](https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework/-/blob/train_orchestrator_release/docs/hgq2-configuration.md),
including scannable quantizer/beta settings. Keras objective adapters preserve
optional sample weights from `(inputs, targets, sample_weight)` batches.

Optuna model-build failures are recorded and skipped up to a configurable limit.
The pipeline ObjectiveTuner also counts trials rejected by checkpoint cuts;
it retains an eligible winner if the failure limit is reached.
See [Optuna settings and defaults](https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework/-/blob/train_orchestrator_release/docs/optuna.md) before starting a study.
The [custom layer/model hook](https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework/-/blob/train_orchestrator_release/docs/custom-objects.md) registers project-owned Python classes before building models.
Configure it with `meta.custom_objects_factory: custom_layers.py:get_custom_objects`.

## Installation

Python 3.10 or newer is required. Start with an isolated environment:

```bash
python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
```

Install the latest published package from PyPI:

```bash
python -m pip install wp21_train
```

Choose the extras for your workflow. These are alternatives; you do not need to
install every backend. Omit `optuna` if you only need standalone training.

| Workflow | Installation |
| --- | --- |
| WP21 ML Utils models and Optuna | `python -m pip install 'wp21_train[ml-utils,optuna]'` |
| TensorFlow/Keras and Optuna | `python -m pip install 'wp21_train[tensorflow,optuna]'` |
| HGQ2 and Optuna | `python -m pip install 'wp21_train[HGQ2,optuna]'` |
| PyTorch and Optuna | `python -m pip install 'wp21_train[torch,optuna]'` |
| Brevitas quantization and Optuna | `python -m pip install 'wp21_train[brevitas,optuna]'` |

For a CPU-only Torch installation, install Torch from its CPU wheel index first.
The `tensorflow` extra includes the historical tf2onnx dependency for legacy
export; native `.keras` saves do not need ONNX. Use separate environments for
legacy HGQ and modern HGQ2 when their dependencies conflict.

Check the version installed in your active environment:

```bash
python -c 'from importlib.metadata import version; print(version("wp21_train"))'
```

Core imports, configuration, event contracts and array evaluation do not import
TensorFlow, Keras, Torch, HGQ, ML Utils or Optuna. Importing a concrete trainer
without its backend raises an installation hint. The base installation includes
numerical/data dependencies. Fixed runs do not require Optuna.

### Install the development version or contribute

The interfaces documented here target development version **0.12.0** on
`ixiotidi/release_v0.12.0`. Installing from PyPI selects the published release,
which may differ. To use this development branch, clone the GitLab repository
and install it in editable mode with your chosen extras:

```bash
git clone --branch ixiotidi/release_v0.12.0 https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework.git
cd wp21_training_framework
python -m pip install -e '.[tensorflow,optuna,dev]'
```

For example, replace `tensorflow` with `brevitas` for quantized Torch training.
GitLab access may require your CERN account. The documentation links below point
to this development branch. Cloning also provides the examples, which
are not installed as part of the Python package. Record the commit and dependency
versions when reproducing an experiment.

## Modular use and compatibility

The components can be used separately. Call an objective on arrays, call a
trainer directly, or use TrainingRunner when search orchestration is useful.
There is no mandatory centralized Project class or runner around standalone
training/evaluation. Each adapter still requires its documented model/data input.

Version 0.10.0 introduced `wp21_train.objectives.Objective`, a backend-independent base
for user-defined metrics, plus generic `MultiClassAUC` and `TriggerEfficiency`.
Prediction bindings and optional Keras/ML Utils epoch adapters reuse the same
calculations. See [generic objectives and complete examples](https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework/-/blob/train_orchestrator_release/docs/objectives.md).

The current dummy jet-tagger's imports, objective factory, params.yml and model.yml
remain supported without migration. Its legacy MultiClassAUC and ObjectiveTuner
path is retained. Updating the package does not switch the framework to the new
runner. The training image must install the new package version to use new APIs.

## Supported interfaces

| Trainer | Model input | Training data | Native output |
| --- | --- | --- | --- |
| `MLUtilsTrainer` | WP21 ML Utils model.yml referenced by params.yml | Callable returning unbatched train/validation tf.data datasets | model.keras, resolved model.yml, result.json |
| `KerasTrainer` | Callable factory or .keras path | Object with x_train, y_train, x_val, y_val | model.keras, result.json |
| `HGQ2Trainer` | HGQ2 factory or registered HGQ2 .keras model | Same array interface as Keras | model.keras, result.json |
| `TorchTrainer` | Module factory, optionally a state_dict checkpoint | Same attributes containing Torch tensors | model.pt, result.json |
| `BrevitasTrainer` | Quantized Torch factory receiving trial parameters | Same Torch tensor interface | model.pt, result.json |

All five implement `Trainer`: `prepare`, `fit`, `train`, `predict`, and `save`.
Factories must construct fresh models. A saved Keras model is recompiled with a
fresh optimizer; Torch checkpoint loading requires the matching factory.
ML Utils handles its own supported quantization configuration, including use_hgq.
There is no automatic dispatch between these adapters yet.

Legacy `TrainerBase`, `LossTuner`, `ObjectiveTuner`, HGQ 0.2, XGBoost, export and
physics utilities remain available. They are not all migrated to the new
interface. HGQ2 and legacy HGQ use different APIs/environments.

## Input and output reference

See [input and output types](https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework/-/blob/train_orchestrator_release/docs/api-types.md)
for array shapes, data types, units, model factories and returned objects.
Class docstrings describe constructor inputs, defaults, shapes and lifecycle
behavior, including legacy utilities. Inspect them directly from Python:

```python
from wp21_train.training import TrainingRunner, TrainingResult
from wp21_train.evaluation import turn_on_curve

help(TrainingRunner)
help(TrainingResult)
help(turn_on_curve)
```

## Configuration and training

The existing YAML/JSON `data`/`meta` structure is preserved. For native trainers,
`data` contains parameter choices. Save this example as `params.yml`:

```yaml
data:
  lr: [0.001, 0.0001]
  batch_size: [64]
  epochs: [10]
  optimiser: [adam]
  loss: [sparse_categorical_crossentropy]
  metrics: [accuracy]
meta:
  model_config: model.yml
  search: optuna
  sampler: tpe
  trials: 4
  objective: val_loss
  direction: minimize
  seed: 42
  callbacks:
    early_stopping:
      monitor: val_loss
      mode: minimize
      patience: 2
      min_delta: 0.0
      restore_best_weights: true
```

MLUtilsTrainer instead reads architecture and dynamic scan specifications from
meta.model_config. Other trainers currently receive their model source explicitly
in Python; model_config does not route them automatically. Paths supplied to a
trainer resolve relative to its configuration file. No new Project subclass is
required.

```python
from wp21_train.training import TrainingRunner
from wp21_train.training.trainers import KerasTrainer

# Supply dataset with x_train, y_train, x_val and y_val attributes.
# Supply model_factory returning a fresh compiled-or-uncompiled Keras model.
# A runnable example is provided in examples/mnist/run.py (see below).
trainer = KerasTrainer('params.yml', dataset, model_factory)
runner  = TrainingRunner(trainer)
results = runner.run()
winner  = runner.best_result
if winner is not None:
    trainer.save('output/winner', winner)
```

`sampler` accepts tpe/bayesian, random, or grid. Python callers can inject an
Optuna BaseSampler instance. `search: grid` and `search: random` retain legacy
flat-scan ordering and no-replacement semantics; explicit Optuna random sampling
can repeat choices. Callable model-defined grid searches need finite
meta.grid_space choices. See [search orchestration](https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework/-/blob/train_orchestrator_release/docs/search-orchestration.md).

To train once, pass concrete values to `trainer.train(parameters)` or
`runner.run(parameters)`. Missing/fixed/none search settings require singleton
choices. A fixed ML Utils model automatically bypasses search. No optimization
objective is required for fixed runs.

Results retain the model, parameters, configuration, history, scalar metrics and
trial context. Searches use the metric named by meta.objective or an explicit
`objective(result)` callable. The winner is an already trained result; saving it
does not retrain it. Retaining every trial model can consume substantial memory.

## Brevitas and Optuna

BrevitasTrainer reuses the Torch lifecycle with a factory receiving concrete
Optuna parameters and fixed quantizer settings. Weight/activation bit widths
can be scanned alongside ordinary training parameters. It supports standalone
QAT, generic objectives, early stopping, pruning and native checkpoint reloads.
See [Brevitas factory/configuration examples](https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework/-/blob/train_orchestrator_release/docs/brevitas.md). PTQ, quantized
ONNX/QONNX export and framework synthesis integration are separate capabilities.

## Callbacks and evaluation

Common training supports fresh per-trial event callbacks and explicit early
stopping across all five adapters. Native adapters apply these policies through
the common train/fit(context) interface. Legacy calls keep their historical
callback behavior. The ML Utils legacy callback factory remains supported.

```python
from wp21_train.callbacks import TrainingCallback

class Progress(TrainingCallback):
    def on_epoch_end(self, event):
        print(event.context.trial_number, event.epoch, dict(event.metrics))

trainer = KerasTrainer('params.yml', dataset, model_factory,
                      event_callbacks_factory=lambda context: [Progress()])
```

Restored weights and result metrics refer to the same epoch. Only model weights
are restored, not optimizer state. Optional Optuna pruning hooks report epoch
metrics to a trial. See [callbacks and evaluation](https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework/-/blob/train_orchestrator_release/docs/callbacks-and-evaluation.md)
for pruning, event ordering and supported policy options.

Evaluation consumes explicit score arrays; it does not assume a jet tagger,
a class index or a probability conversion:

```python
from wp21_train.evaluation import binary_auc, rate_curve, turn_on_curve

auc   = binary_auc(binary_labels, positive_class_scores)
rates = rate_curve(background_scores, thresholds, input_rate_hz=40_000_000)
turnon = turn_on_curve(reference_pt, positive_class_scores, threshold=0.5,
                      bins=[0, 20, 40, 60, 100])
rates.save('output/rate.json')
turnon.save('output/turnon.json')
```

Rate normalization is explicit in Hz. Selections use score >= threshold.
Nonnegative event weights are supported; empty turn-on bins serialize as null.
`BinaryAUCObjective` connects prediction to scalar search scoring, with explicit
output selection for named/multiclass outputs. `trigger_rate` returns a scalar
rate at a chosen threshold. Curves are separate artifacts, not scalar objectives.
There are no inferred physics selections or weighted-error prescriptions.

## Examples

[MNIST examples](https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework/-/blob/train_orchestrator_release/examples/mnist/README.md) cover all five adapters. From the
repository root after cloning and installing the corresponding extras above:

```bash
python examples/mnist/run.py --backend keras --download --orchestrated
python examples/mnist/run.py --backend torch --data mnist.npz --standalone
```

Downloads happen only with --download. The default
example path exercises legacy APIs; --standalone uses the common trainer and
--orchestrated uses TrainingRunner.

## Scope and documentation

Automatic model routing, framework/MLflow pipeline migration, distributed study
resume and multi-objective search are not implemented by this interface yet.
Native checkpoints support model reuse/inference, not exact optimizer resume.
Save project code, input data and dependency versions alongside model artifacts
when reproducibility is required.

See the [documentation index](https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework/-/blob/train_orchestrator_release/docs/README.md) for API contracts, input/output types and
compatibility details. The [historical API reference](https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework/-/blob/train_orchestrator_release/docs/legacy-api.md) retains
the previous README's legacy examples and utilities.

## Developers and contributors

The developers credited in the previous README are:

- **Ioannis Xiotidis** — [ioannis.xiotidis@cern.ch](mailto:ioannis.xiotidis@cern.ch)
- **Pawel Mucha** — [pawel.mucha@cern.ch](mailto:pawel.mucha@cern.ch)
- **Vila Andela Petrovic** — [vila.andela.petrovic@cern.ch](mailto:vila.andela.petrovic@cern.ch)
- **David Reikher** — [david.reikher@cern.ch](mailto:david.reikher@cern.ch)
