wp21-train 0.3.0

Framework that provides tools allowing for monitoring training of ML models with the aim to be executed in firmware developed under the ATLAS NextGen WP21 project

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
• ✅ Versioned with easy integration (__version__)

---

📦 Included Modules

Module	Description
savers/json_adapter.py	JSON-based serialization
savers/pickle_adapter.py	Pickle-based serialization
savers/root_adapter.py	ROOT I/O serialization (requires ROOT installed)
savers/yml_adapter.py	YML I/O serialization (requires PyYAML installed)
parsers/hls_parser.py	XML parsing of HLS synthesis reports
parsers/aie_parser.py	XML parsing of AIE runtime profiling
parsers/athena_parser.py	Parsing of ATHENA configuration
callbacks/base_callback.py	Base callback for training pipelines
utils/utility.py	Type-shortening utility for metadata tagging
utils/logging.py	Simple logging of info warnings and errors
utils/version.py	Package versioning (__version__)

---

🔧 Installation

From PyPI

pip install wp21_train

---

📁 Example Usage

🔄 JSON / Pickle / YAML / ROOT Adapters

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

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

from wp21_train.parsers import aie_parser

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

⚙️ ATHENA Parser

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

from wp21_train.callbacks import base_callback

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

🧬 Type Utility

from wp21_train.utils.utility import get_short_type

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

🧬 Logging Utility

from wp21_train.utils.logging import log_message

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

---

🧪 Testing

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)

---

⚠️ 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:

conda install -c conda-forge root

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

---

🧠 Versioning

The current package version is defined in:

from wp21_train.utils.version import __version__

---

🔖 License

This project is licensed under the MIT License. See the LICENSE file for details.

---

👤 Author

Ioannis Xiotidis
Email: [e-mail]

Pawel Mucha Email: [e-mail]

Vila Andela Petrovic Email: [e-mail]

---