Metadata-Version: 2.4
Name: koios-model-utils
Version: 1.1.1
Summary: Utilities for exporting and annotating ONNX models for the Koios IoT platform
Project-URL: Homepage, https://www.ai-op.com
Project-URL: Documentation, https://github.com/Ai-Ops-Inc/koios-model-utils
Project-URL: Repository, https://github.com/Ai-Ops-Inc/koios-model-utils
Project-URL: Issues, https://github.com/Ai-Ops-Inc/koios-model-utils/issues
Author-email: "Ai-OPs, Inc." <support@ai-op.com>
License: Apache-2.0
License-File: LICENSE
Keywords: industrial,iot,koios,machine-learning,onnx,reinforcement-learning
Classifier: Development Status :: 5 - Production/Stable
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Typing :: Typed
Requires-Python: >=3.11
Requires-Dist: onnx<1.22,>=1.16
Provides-Extra: dev
Requires-Dist: basedpyright>=1.29.0; extra == 'dev'
Requires-Dist: pytest>=8.0.0; extra == 'dev'
Requires-Dist: ruff>=0.4.0; extra == 'dev'
Description-Content-Type: text/markdown

# koios-model-utils

Annotate ONNX models with the metadata the [Koios](https://ai-op.com) IoT
platform needs to run them.

This library is the canonical writer (and reader) of the `koios.training`
and `koios.bindings` `metadata_props` blocks Koios consumes when a model
is uploaded to its predict engine.  Bring your own ONNX file from any
framework (Stable-Baselines3, PyTorch, TensorFlow/Keras, scikit-learn);
the library only touches metadata.

The wire format is fully documented in [`docs/CONTRACT.md`](docs/CONTRACT.md).

## Installation

```bash
pip install koios-model-utils
```

Only runtime dependency is `onnx>=1.14`.  No torch, no SB3, no CLI —
this is a metadata library.

## Quick start

```python
import onnx
from koios_model_utils import (
    Algorithm,
    InputBinding,
    NormalizationSource,
    NormalizationType,
    OutputBinding,
    TrainingMeta,
    embed_koios_metadata,
)

model = onnx.load("my_model.onnx")

embed_koios_metadata(
    model,
    inputs=[
        InputBinding(name="tank_temperature", description="Tank temperature (C)"),
        InputBinding(name="pressure", description="Vessel pressure (kPa)"),
    ],
    outputs=[
        OutputBinding(
            name="valve_position",
            range_min=0.0,
            range_max=100.0,
            normalization_type=NormalizationType.SYMMETRIC,
            normalization_source=NormalizationSource.CUSTOM,
            custom_minimum=0.0,
            custom_maximum=100.0,
            clamp_output=True,
        ),
    ],
    training=TrainingMeta(
        scenario_name="tank_temperature",
        algorithm=Algorithm.PPO,
        obs_depth=5,
        sample_rate=1.0,
    ),
)

onnx.save(model, "my_model_koios.onnx")
```

The Koios webapp reads both metadata blocks on upload and uses them to
configure `model.sample_rate`, the input/output bindings, normalization
rules, and (for DRL classifiers) the action map.

See [`examples/`](examples/) for runnable scripts.

## Reading metadata back

The mirror of `embed_koios_metadata`:

```python
from koios_model_utils import parse_koios_metadata

parsed = parse_koios_metadata(onnx.load("my_model_koios.onnx"))
print(parsed.training)        # TrainingMeta dataclass (or None)
print(parsed.inputs)          # list[InputBinding]
print(parsed.outputs)         # list[OutputBinding]
print(parsed.has_metadata)    # True if any koios.* prop was present
```

If the JSON has already been pulled out of the ONNX file (e.g. persisted
in a database column), use `parse_koios_metadata_from_dict` and pass
`{"koios.training": {...}, "koios.bindings": {...}}` directly.

## Public API

### Dataclasses

| Class | Purpose |
|---|---|
| `InputBinding` | One observation feature — name, normalization rules, failure bounds |
| `OutputBinding` | One action / output — name, range, normalization, clamping |
| `TrainingMeta` | Model-level metadata — algorithm, sample/scan rate, model type, action map |
| `ActionMapEntry` | One row of a DISCRETE-mode action map (`value` + `label`) |
| `ParsedKoiosMetadata` | Return type of `parse_koios_metadata` |

All dataclasses run their validators in `__post_init__`, so bad
combinations (e.g. `Z_SCORE` + `custom_minimum`) raise at construction.

### Enums

`NormalizationType`, `NormalizationSource`, `ModelType`, `OutputMode`,
`Algorithm`, `FailureRangeMode` — all `enum.StrEnum`, all uppercase wire
values.  See [`enums.py`](src/koios_model_utils/enums.py) for the members.

Pass enum members rather than raw strings — the dataclasses accept
either, but enums catch typos at type-check time.

### Functions

| Function | Purpose |
|---|---|
| `embed_koios_metadata(model, *, inputs, outputs, training=None, output_denormalized=False)` | Write `koios.training` + `koios.bindings` into an ONNX model (in-place) |
| `parse_koios_metadata(model)` | Parse them back into typed dataclasses |
| `parse_koios_metadata_from_dict(raw)` | Same, but from an already-decoded dict |

### Errors

`KoiosMetadataError`, `UnsupportedSchemaVersionError`,
`KoiosMetadataDecodeError` — all `ValueError` subclasses raised by the
parser.

## sample_rate vs scan_rate

`sample_rate` is the interval (seconds) the model was trained at.
The Koios predict engine resamples historical inputs to this rate.

`scan_rate` is how often the predict engine *executes* the model.  When
omitted (the common case), it inherits `sample_rate` — forecasting
models train and run at the same cadence.  Set it explicitly for RL
controllers that need to act faster than the simulator step they were
trained against (e.g. `sample_rate=1.0` history lookback, `scan_rate=0.1`
control loop).

## Discrete (classification / DQN) heads

For models that emit an action index instead of a continuous value, set
`output_mode` and pass an `action_map`:

```python
from koios_model_utils import (
    ActionMapEntry, OutputMode, TrainingMeta,
)

training = TrainingMeta(
    algorithm="DQN",
    output_mode=OutputMode.DISCRETE,
    action_map=[
        ActionMapEntry(value=-1.0, label="Decrease setpoint"),
        ActionMapEntry(value=0.0, label="Hold"),
        ActionMapEntry(value=1.0, label="Increase setpoint"),
    ],
)
```

`action_map` also accepts raw `{"value": ..., "label": ...}` dicts and
coerces them to `ActionMapEntry` at construction — extra keys are
dropped, missing keys raise.

## Contract & versioning

The wire format is at schema_version `1`.  [`docs/CONTRACT.md`](docs/CONTRACT.md)
is the source of truth for every field, every default, and every
server-side policy applied on upload.

The library and the Koios webapp move in lockstep: any wire-format
change here requires a coordinated webapp release.  The CONTRACT lays
out exactly which keys the webapp reads, the forward-compat rules
within v1, and when a v2 bump would be needed.

## Development

```bash
pip install -e ".[dev]"
make test
make lint
```

## License

Apache License 2.0 — see [LICENSE](LICENSE).
