Metadata-Version: 2.4
Name: wp21_train
Version: 0.9.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
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: torch
Requires-Dist: torch>=2.1; extra == "torch"
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"
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: 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: provides-extra
Dynamic: requires-dist
Dynamic: requires-python
Dynamic: summary

# wp21_train

A modular, extensible Python framework designed for managing data parsing, serialization, and metadata tracking in machine learning workflows — especially for hardware-aware applications such as HLS and AIE profiling. Built with physicists and hardware engineers in mind, `wp21_train` integrates support for CERN ROOT I/O, common formats like JSON and Pickle, and parsing of Xilinx toolchain outputs.

---

## 🚀 Features

- ✅ **Unified interface** for reading/writing training data and metadata  
- ✅ Supports **JSON**, **Pickle**, **YAML** and **ROOT** formats  
- ✅ Parsers for:
  - **HLS reports** (Vivado HLS)
  - **AIE profiling reports** (Vitis AI Engine)
  - **ATHENA configuration**
- ✅ A **training callback interface** to log key events and outputs  
- ✅ Lightweight **type-to-symbol conversion utility**
- ✅ **Optuna-based hyperparameter optimization**
- ✅ **ML Utils model training integration**
- ✅ **Quantization Aware Training with HGQ2 and model checkpoint selection**  
- ✅ Versioned with easy integration (`__version__`)

---

## 📦 Included Modules

```
+-----------------------------------------------+--------------------------------------------------------------+-----------------------------------------+
| Module                                        | Description                                                  | Notes / Deps                            |
+-----------------------------------------------+--------------------------------------------------------------+-----------------------------------------+
| wp21_train/savers/json_adapter.py             | JSON-based serialization                                     | stdlib json                             |
| wp21_train/savers/pickle_adapter.py           | Pickle-based serialization                                   | stdlib pickle                           |
| wp21_train/savers/root_adapter.py             | ROOT I/O serialization                                       | uproot, awkward (no ROOT needed)        |
| wp21_train/savers/yml_adapter.py              | YAML I/O serialization                                       | PyYAML (yaml)                           |
| wp21_train/parser/hls_parser.py               | XML parsing of HLS synthesis reports                         |                                         |
| wp21_train/parser/aie_parser.py               | XML parsing of AIE runtime profiling                         |                                         |
| wp21_train/parser/athena_parser.py            | Parsing of ATHENA configuration                              |                                         |
| wp21_train/callbacks/base_callback.py         | Base callback for training pipelines                         |                                         |
| wp21_train/training/searchers/base_search.py  | Base class for search strategies                             |                                         |
| wp21_train/training/searchers/grid_search.py  | Grid search implementation                                   |                                         |
| wp21_train/training/searchers/random_search.py| Random search implementation                                 |                                         |
| wp21_train/training/trainers/base.py          | Common trainer utilities                                     |                                         |
| wp21_train/training/trainers/keras_base.py    | Keras trainer base                                           | tensorflow / keras                      |
| wp21_train/training/trainers/keras_trainer.py | Concrete Keras trainer                                       | tensorflow / keras                      |
| wp21_train/training/trainers/torch_base.py    | PyTorch trainer base                                         | torch                                   |
| wp21_train/training/trainers/torch_trainer.py | Concrete PyTorch trainer                                     | torch                                   |
| wp21_train/training/trainers/hgq_trainer.py   | Concrete HGQ trainer                                         | HGQ                                     |
| wp21_train/training/trainers/hgq2_trainer.py  | Concrete hgq2 trainer                                        | hgq2                                    |
| wp21_train/training/trainers/xgb_trainer.py   | Concrete XGBoost trainer                                     | xgboost                                 |
| wp21_train/training/evaluator.py              | Generic evaluation helpers                                   |                                         |
| wp21_train/training/tuner.py                  | High-level tuning orchestration                              |                                         |
| wp21_train/utils/utility.py                   | Type-shortening / helpers for metadata tagging               |                                         |
| wp21_train/utils/logger.py                    | Simple logging of info/warn/error                            | (renamed from logging.py)               |
| wp21_train/utils/version.py                   | Package versioning (__version__)                             |                                         |
| wp21_train/__init__.py                        | Package exports                                              |                                         |
| wp21_train/training/searchers/optuna_search.py   | Optuna-based hyperparameter search                        | optuna                                  |
| wp21_train/training/trainers/ml_utils_trainer.py | ML Utils trainer backend                                  | wp21_ml_utils                           |
| wp21_train/selection/selector.py                 | Final model/checkpoint selection                          |                                         |
| wp21_train/selection/finalizer.py                | Finalizes and saves the selected ML Utils model           | wp21_ml_utils                           | 
| wp21_train/callbacks/training_callbacks.py       | Callbacks factory for QAT (EBOPs, Parto, BetaScheduler)   | HGQ, TensorFlow                         |
+-----------------------------------------------+--------------------------------------------------------------+-----------------------------------------+
```

## 🔧 Installation

### From PyPI

```bash
pip install wp21_train
```

---

## 📁 Example Usage

### 🔄 JSON / Pickle / YAML / ROOT Adapters

```python
from wp21_train.savers import json_adapter, pickle_adapter, yml_adapter, root_adapter

adapter = json_adapter("results", dump_data=my_data, dump_meta=my_metadata)
adapter.write_data()

meta, data = adapter.read_data()
```

### 🧠 HLS Parser

```python
from wp21_train.parsers import hls_parser

parser = hls_parser("hls_report.xml")
print(parser._data)        # extracted info
print(parser._meta_data)   # associated metadata
```

### ⚙️ AIE Parser

```python
from wp21_train.parsers import aie_parser

parser = aie_parser("aie_profile.xml")
print(parser._data)
```

### ⚙️ ATHENA Parser

```python
from wp21_train.parsers import athena_parser

parser = athena_parser(data=data_from_adapter, metadata=meta_from_adapter, nevents=10000)
print(parser.config)
print(parser.environment)
```

### 📋 Training Callback

```python
from wp21_train.callbacks import base_callback

cb = base_callback(project_name="FastML4Jets")
cb.on_train_begin()
# training loop here
cb.on_train_end()
```

### 🌲 XGBoost Trainer

```python
from xgboost import XGBClassifier
import wp21_train as train
from wp21_train.training.trainers import XGBTrainer

def xgb_model(**params):
    params.setdefault("objective", "binary:logistic")
    params.setdefault("eval_metric", ["logloss", "auc"])
    return XGBClassifier(**params)

trainer = XGBTrainer("training_config", dataset, xgb_model)
tuner = train.training.LossTuner("hat", trainer)
tuner.run()

evaluator = train.training.Evaluator(
    tuner.trials,
    trainer,
    meta_path="meta",
    out_path="out",
)
evaluator.save_info(verbose=3)
```

The XGBoost trainer assumes a scikit-learn-like estimator, such as `XGBClassifier` or `XGBRegressor`, with `fit(...)`, `evals_result()`, and `save_model(...)`/`get_booster().save_model(...)` methods. It stores all metrics returned by `evals_result()` using their native names, such as `logloss`, `auc`, `rmse`, and their `val_` counterparts. For `LossTuner`, only XGBoost `logloss` is exposed as `loss`/`val_loss`. To control which additional metrics are written to `history.json`, set `history_metrics` in the config `meta` section:

```json
{
  "meta": {
    "search": "random",
    "trials": 6,
    "history_metrics": ["auc"]
  }
}
```


### 🔍 ML Utils and Optuna Training

`wp21_train` supports training and hyperparameter optimization of models defined with `wp21_ml_utils` through `MLUtilsTrainer`.

The ML Utils workflow uses two YAML configuration files:

- `params.yml` contains pipeline-level settings such as dataset paths, search strategy, number of trials, optimization direction, and optional selection cuts.
- `model.yml` contains the ML Utils model architecture, optimiser, training settings, HGQ configuration, and optional hyperparameter search definitions.

For Optuna searches, `direction` must be explicitly defined as either `maximize` or `minimize`.

Example `params.yml`:

```yaml
data: {}

meta:
  model_config: path/to/model.yml

  data_paths:
    train_path: /path/to/train.npz
    val_path: /path/to/val.npz
    test_path: /path/to/test.npz

  env_name: tf_v3
  save_path: out/

  search: optuna
  trials: 10
  sampler: bayesian
  objective: mean_auc
  direction: maximize

  cuts:
    min_score: 0.93
    max_ebops: 20000
```

Example `model.yml`:

```yaml
inputs:
  jets:
    shape: [16]

layers:
  backbone:
    class: DenseLayers
    inputs: jets
    params:
      hidden_layer_sizes:
        type: list
        length: [1, 4]
        values: [16, 32, 64, 128]

      activation: relu
      use_hgq: true

  output:
    class: Dense
    inputs: backbone
    params:
      units: 5
      activation: softmax

outputs:
  output:
    loss: categorical_crossentropy

optimiser:
  class: Adam
  params:
    learning_rate:
      type: float
      values:
        - 0.01
        - 0.005
        - 0.001
        - 0.0005
        - 0.0001

training:
  batch_size:
    type: int
    values: [64, 128, 256, 512]

  epochs: 30

hgq_config:
  quantizer_scopes:
    - place: all
      default_q_type: kbi
      overflow_mode: SAT_SYM
      heterogeneous_axis: []
```

`MLUtilsTrainer` interprets the nested hyperparameter search definitions in `model.yml` and provides the sampled parameters to `OptunaSearch`.

`OptunaSearch` remains generic and also supports traditional flat search-space dictionaries used by other trainers.

If `model.yml` contains no hyperparameter search definitions, it is treated as a fixed configuration and the model is trained once without running Optuna.

The dataset paths are provided through `meta.data_paths`.

For quantization-aware training, Pareto checkpoints are produced during training and evaluated by `Selector`.

If selection cuts are provided under `meta.cuts`, the best checkpoint satisfying the requested constraints is selected.

If no cuts are provided, the best-scoring checkpoint from each trial is considered and the checkpoint with the lowest EBOP cost among those candidates is selected as the final model.

The selected model is reconstructed and finalized by `MLUtilsFinalizer`.

The final artifacts are stored as:

- `winner.yml`
- `winner.keras`

#### QAT callbacks and EBOP scheduling

For HGQ-enabled models, `MLUtilsCallbackFactory` builds the callbacks required for quantization-aware training.

HGQ is enabled at the model layer level with `use_hgq: true`. EBOP accounting and the initial penalty coefficient are configured through `hgq_config`:

```yaml
layers:
  backbone:
    class: DenseLayers
    inputs: jets
    params:
      use_hgq: true

hgq_config:
  layer:
    enable_ebops: true
    beta0: 0.0
```

The evolution of the EBOP penalty during training is configured separately through `qat.beta_schedule`:

```yaml
qat:
  beta_schedule:
    - [0, 0.0, linear]
    - [3, 1.0e-7, log]
    - [18, 1.0e-4, constant]
```

Each schedule entry specifies the epoch, target beta value, and interpolation mode. `MLUtilsCallbackFactory` converts this configuration into an HGQ `PieceWiseSchedule` and attaches it to training through `BetaScheduler`.

The callback factory is created and passed to `MLUtilsTrainer`:

```python
from callbacks.multi_class_auc import MultiClassAUC
from wp21_train.callbacks.training_callbacks import MLUtilsCallbackFactory
from wp21_train.training.trainers.ml_utils_trainer import MLUtilsTrainer


def make_objective(valid_ds, resolved_config):
    output_node = next(iter(resolved_config["outputs"]))

    n_classes = resolved_config[
        "layers"
    ][output_node]["params"]["units"]

    return MultiClassAUC(
        valid_ds=valid_ds,
        output_node=output_node,
        n_classes=n_classes,
    )


callbacks_factory = MLUtilsCallbackFactory(
    objective_factory=make_objective,
    pareto_root="out/qat_pareto",
)

trainer = MLUtilsTrainer(
    params_config="params.yml",
    dataset=dataset,
    callbacks_factory=callbacks_factory,
)
```

For an HGQ-enabled model, the factory attaches the following callbacks:

- `BetaScheduler`, which applies the configured `qat.beta_schedule`;
- `FreeEBOPs`, which makes the model EBOP cost available during training;
- the user-provided objective callback;
- `ParetoFront`, which tracks the trade-off between predictive performance and EBOP cost and saves Pareto-optimal checkpoints.

Pareto checkpoints are stored separately for each trial:

```text
out/qat_pareto/
  trial_0/
  trial_1/
  ...
```

Checkpoint filenames contain the epoch, objective score, and EBOP value:

```text
epoch=<epoch>_score=<score>_ebops=<ebops>.weights.h5
```

The objective side of the Pareto front follows the configured optimization `direction`, while EBOPs are always minimized.

The saved Pareto checkpoints are later inspected by `Selector`, which applies the optional constraints from `meta.cuts` and chooses the final deployment candidate.

The MultiClassAUC is an example of a custom callable class for scoring a user should provide.


### 🧬 Type Utility

```python
from wp21_train.utils.utility import get_short_type

print(get_short_type(42))     # 'd'
print(get_short_type("abc"))  # 's'
```

### 🧬 Logging Utility

```python
from wp21_train.utils.logging import log_message

log_message("error", f"Provided number of events ({nevents}) is not an integer.")
```

---

## 🧪 Testing

```bash
pytest tests/
```

---

## 📜 Requirements

- Python ≥ 3.7  
- `uproot` (for reading ROOT files)  
- `xml.etree.ElementTree` (standard lib, for HLS/AIE parsing)  
- **CERN ROOT** (installed and configured) if you use `.root` I/O
- PyYAML ≥ 6.0.0, < 7.0.0 (for reading yaml files)
- awkward (data is represented as awkward arrays)
- pybind11 and openmp - used to implement slow bits in C++

---

## ⚠️ Note About ROOT

This package supports `.root` file serialization and reading **via CERN ROOT**. If you intend to use this feature, ensure that ROOT is installed and properly sourced in your environment. You can install ROOT via Conda:

```bash
conda install -c conda-forge root
```

Or follow the official installation guide:  
https://root.cern/install/

---

## 🧠 Versioning

The current package version is defined in:

```python
from wp21_train.utils.version import __version__
```

---

## 🔖 License

This project is licensed under the **MIT License**. See the [LICENSE](LICENSE) file for details.

---

## 👤 Author

**Ioannis Xiotidis**  
Email: [ioannis.xiotidis@cern.ch](mailto:ioannis.xiotidis@cern.ch)

**Pawel Mucha**  
Email: [pawel.mucha@cern.ch](mailto:pawel.mucha@cern.ch)

**Vila Andela Petrovic**  
Email: [vila.andela.petrovic@cern.ch](mailto:vila.andela.petrovic@cern.ch)

**David Reikher**  
Email: [david.reikher@cern.ch](mailto:david.reikher@cern.ch)

---
