Metadata-Version: 2.4
Name: koios-model-utils
Version: 1.0.0rc1
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: MIT
License-File: LICENSE
Keywords: industrial,iot,koios,machine-learning,onnx,reinforcement-learning
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT 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: numpy
Requires-Dist: onnx>=1.14
Provides-Extra: dev
Requires-Dist: pytest>=8.0.0; extra == 'dev'
Requires-Dist: ruff>=0.4.0; extra == 'dev'
Provides-Extra: sb3
Requires-Dist: onnxruntime>=1.16; extra == 'sb3'
Requires-Dist: onnxscript; extra == 'sb3'
Requires-Dist: stable-baselines3>=2.0; extra == 'sb3'
Requires-Dist: torch; extra == 'sb3'
Description-Content-Type: text/markdown

# koios-model-utils

Utilities for exporting and annotating ONNX models for the [Koios](https://ai-op.com) IoT platform.

This package helps data scientists prepare trained ML models (especially Stable-Baselines3 reinforcement learning models) for deployment on Koios. It handles ONNX export, metadata embedding, and normalization parameter annotation so the Koios Predict Engine can run inference correctly.

## Installation

```bash
# Metadata embedding only (no torch/SB3 dependency)
pip install koios-model-utils

# Full SB3 export pipeline
pip install koios-model-utils[sb3]
```

## Quick Start

### Embed metadata into an existing ONNX model

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

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

inputs = [
    InputBinding(name="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="symmetric",
        normalization_source="custom",
        custom_minimum=0.0, custom_maximum=100.0,
        clamp_output=True,
    ),
]
training = TrainingMeta(
    scenario_name="tank_temperature",
    algorithm="PPO",
    obs_depth=5,
    sample_rate=1.0,
)

embed_koios_metadata(model, inputs=inputs, outputs=outputs, training=training)
onnx.save(model, "my_model_koios.onnx")
```

### Export an SB3 model to ONNX

```python
from koios_model_utils import export_onnx, build_input_bindings, build_output_bindings, TrainingMeta

inputs = build_input_bindings(
    ["temperature", "level", "flow_rate"],
    descriptions={"temperature": "Tank temp (C)", "level": "Tank level (%)"},
)
outputs = build_output_bindings(
    ["valve_position"],
    action_ranges={"valve_position": (0.0, 100.0)},
)

export_onnx(
    "runs/best_model.zip",
    inputs=inputs,
    outputs=outputs,
    training=TrainingMeta(scenario_name="tank_temperature", algorithm="PPO", obs_depth=5),
)
```

### CLI

```bash
# Export with auto-detected settings
koios-export runs/best_model.zip

# Specify output path and opset version
koios-export runs/best_model.zip --output model.onnx --opset 18
```

## SB3 Training Callbacks

```python
from koios_model_utils import ObsRangeCallback, VecNormalizeSyncCallback

# Track observation ranges during training
obs_range_cb = ObsRangeCallback(run_dir, verbose=1)

# Sync VecNormalize stats with best model checkpoints
vec_sync_cb = VecNormalizeSyncCallback(eval_callback, verbose=1)
```

## API Reference

### Metadata Types

| Class | Description |
|-------|-------------|
| `InputBinding` | Describes one input (observation) binding — name, normalization, failure bounds |
| `OutputBinding` | Describes one output (action) binding — name, range, normalization, clamping |
| `TrainingMeta` | Model-level metadata — algorithm, obs_depth, sample_rate, model_type |

### Functions

| Function | Description |
|----------|-------------|
| `embed_koios_metadata()` | Embed `koios.training` and `koios.bindings` metadata into an ONNX model |
| `build_input_bindings()` | Build input bindings from node names + optional normalization stats |
| `build_output_bindings()` | Build output bindings from action names + ranges |
| `export_onnx()` | Full SB3-to-ONNX export pipeline (requires `[sb3]` extra) |
| `detect_algorithm()` | Detect SB3 algorithm (PPO, SAC, TD3, etc.) from a saved `.zip` model |

### Callbacks (requires `[sb3]` extra)

| Callback | Description |
|----------|-------------|
| `ObsRangeCallback` | Track per-feature min/max of observations during training |
| `VecNormalizeSyncCallback` | Save VecNormalize stats alongside best model checkpoints |

## Development

```bash
# Install with all dependencies
pip install -e ".[sb3,dev]"

# Run tests
make test

# Lint
make lint
```

## License

MIT License - see [LICENSE](LICENSE) for details.
