koios-model-utils 1.0.0rc1__tar.gz

koios_model_utils-1.0.0rc1/.github/workflows/release.yml

@@ -0,0 +1,64 @@
+name: Release
+
+on:
+  push:
+    tags: ["v*"]
+
+env:
+  FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - run: pip install -e ".[dev]"
+      - run: ruff check src/ tests/
+      - run: ruff format --check src/ tests/
+      - run: pytest
+
+  build:
+    needs: test
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - run: pip install build
+      - run: python -m build
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+  publish-pypi:
+    needs: build
+    runs-on: ubuntu-latest
+    environment: release
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@release/v1
+
+  github-release:
+    needs: build
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+      - uses: softprops/action-gh-release@v2
+        with:
+          files: dist/*
+          generate_release_notes: true

koios_model_utils-1.0.0rc1/.gitignore

@@ -0,0 +1,54 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual environments
+.venv/
+venv/
+env/
+ENV/
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# macOS
+.DS_Store
+**/.DS_Store
+
+# Testing
+.pytest_cache/
+.coverage
+htmlcov/
+.tox/
+
+# Dev tooling
+uv.lock
+
+# Environment
+.env
+.env.local
+
+# Logs
+*.log

koios_model_utils-1.0.0rc1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024-2026 Ai-OPs, Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

koios_model_utils-1.0.0rc1/Makefile

@@ -0,0 +1,22 @@
+.PHONY: setup test lint format clean
+
+setup:
+	pip install -e ".[sb3]"
+
+setup-minimal:
+	pip install -e .
+
+test:
+	python3 -m pytest tests/ -v $(ARGS)
+
+lint:
+	ruff check src/ tests/
+
+format:
+	ruff format src/ tests/
+	ruff check --fix src/ tests/
+
+clean:
+	rm -rf build/ dist/ *.egg-info src/*.egg-info
+	find . -type d -name __pycache__ -exec rm -rf {} +
+	find . -type f -name "*.pyc" -delete

koios_model_utils-1.0.0rc1/PKG-INFO

@@ -0,0 +1,176 @@
+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.

koios_model_utils-1.0.0rc1/README.md

@@ -0,0 +1,141 @@
+# 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.

koios_model_utils-1.0.0rc1/pyproject.toml

@@ -0,0 +1,69 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "koios-model-utils"
+version = "1.0.0rc1"
+description = "Utilities for exporting and annotating ONNX models for the Koios IoT platform"
+readme = "README.md"
+requires-python = ">=3.11"
+license = {text = "MIT"}
+authors = [
+    {name = "Ai-OPs, Inc.", email = "support@ai-op.com"}
+]
+keywords = ["koios", "onnx", "iot", "industrial", "machine-learning", "reinforcement-learning"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Typing :: Typed",
+]
+
+dependencies = [
+    "numpy",
+    "onnx>=1.14",
+]
+
+[project.optional-dependencies]
+sb3 = [
+    "stable-baselines3>=2.0",
+    "torch",
+    "onnxruntime>=1.16",
+    "onnxscript",
+]
+dev = [
+    "pytest>=8.0.0",
+    "ruff>=0.4.0",
+]
+
+[project.scripts]
+koios-export = "koios_model_utils.export:main"
+
+[project.urls]
+Homepage = "https://www.ai-op.com"
+Documentation = "https://github.com/Ai-Ops-Inc/koios-model-utils"
+Repository = "https://github.com/Ai-Ops-Inc/koios-model-utils"
+Issues = "https://github.com/Ai-Ops-Inc/koios-model-utils/issues"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/koios_model_utils"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py311"
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "N", "W", "UP"]
+ignore = ["E501"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+pythonpath = ["src"]
+addopts = ["-v", "--tb=short"]

koios_model_utils-1.0.0rc1/src/koios_model_utils/__init__.py

@@ -0,0 +1,47 @@
+"""koios-model-utils — utilities for exporting and annotating ONNX models for the Koios IoT platform."""
+
+__version__ = "1.0.0rc1"
+
+from koios_model_utils.metadata import (
+    InputBinding,
+    OutputBinding,
+    TrainingMeta,
+    build_input_bindings,
+    build_output_bindings,
+    embed_koios_metadata,
+)
+
+# SB3-dependent exports — only available with koios-model-utils[sb3]
+try:
+    from koios_model_utils.callbacks import (
+        ObsRangeCallback,
+        VecNormalizeSyncCallback,
+    )
+    from koios_model_utils.export import (
+        OnnxActorWrapper,
+        OnnxDQNWrapper,
+        detect_algorithm,
+        export_onnx,
+        load_vec_normalize_stats,
+    )
+except ImportError:
+    pass
+
+__all__ = [
+    # metadata (always available)
+    "InputBinding",
+    "OutputBinding",
+    "TrainingMeta",
+    "embed_koios_metadata",
+    "build_input_bindings",
+    "build_output_bindings",
+    # export (requires sb3 extra)
+    "export_onnx",
+    "detect_algorithm",
+    "load_vec_normalize_stats",
+    "OnnxActorWrapper",
+    "OnnxDQNWrapper",
+    # callbacks (requires sb3 extra)
+    "ObsRangeCallback",
+    "VecNormalizeSyncCallback",
+]