mlbundle 0.1.0a1__py3-none-any.whl

mlbundle/_onnx/description.py

@@ -0,0 +1,252 @@
+"""Native graph signatures, output contracts and original label claims."""
+
+from dataclasses import dataclass
+from typing import Any, cast
+
+from pydantic import JsonValue
+
+from mlbundle.canonical import canonical_dumps
+from mlbundle.claims import ClassLabels, ModelClaims
+from mlbundle.contract import ModelKind
+from mlbundle.environment import OnnxInfo, Opset
+from mlbundle.errors import InvalidModelError, LimitExceededError
+from mlbundle.limits import HYPERPARAMETER_BYTES, MAX_CLASSES
+
+from .signatures import _batch, _inputs, signature
+
+LABEL_OPS = frozenset(
+    {"ZipMap", "LinearClassifier", "TreeEnsembleClassifier", "SVMClassifier"}
+)
+
+
+@dataclass(frozen=True)
+class Description:
+    claims: ModelClaims
+    inputs: tuple[dict[str, Any], ...]
+    output: str
+    maps: bool
+
+
+def _ancestors(graph: Any, output: str) -> list[Any]:
+    producers = {name: node for node in graph.node for name in node.output if name}
+    pending, seen, result = [output], set(), []
+    while pending:
+        name = pending.pop()
+        if name in seen:
+            continue
+        seen.add(name)
+        node = producers.get(name)
+        if node is not None:
+            result.append(node)
+            pending.extend(node.input)
+    return result
+
+
+def _labels(nodes: list[Any]) -> ClassLabels | None:
+    found = []
+    for node in nodes:
+        if node.domain != "ai.onnx.ml" or node.op_type not in LABEL_OPS:
+            continue
+        for attr in node.attribute:
+            if attr.name == "classlabels_strings":
+                found.append([value.decode("utf-8") for value in attr.strings])
+            elif attr.name in ("classlabels_ints", "classlabels_int64s"):
+                found.append(list(attr.ints))
+    if not found:
+        return None
+    if any(
+        canonical_dumps(labels) != canonical_dumps(found[0]) for labels in found[1:]
+    ):
+        raise InvalidModelError("ONNX native label records disagree.")
+    return cast(ClassLabels, [found[0]])
+
+
+def _native_labels(graph: Any, output: str) -> ClassLabels | None:
+    producers = {name: node for node in graph.node for name in node.output if name}
+    original = output
+    for _ in range(len(graph.node)):
+        node = producers.get(output)
+        if node is None:
+            return None
+        if node.domain == "ai.onnx.ml" and node.op_type in LABEL_OPS:
+            if node.op_type == "ZipMap" or (
+                len(node.output) > 1 and output == node.output[1]
+            ):
+                return _labels([node])
+            break
+        if (
+            node.domain in ("", "ai.onnx")
+            and node.op_type in ("Identity", "Cast")
+            and len(node.input) == 1
+        ):
+            output = node.input[0]
+        else:
+            break
+    if any(
+        node.domain == "ai.onnx.ml" and node.op_type in LABEL_OPS
+        for node in _ancestors(graph, original)
+    ):
+        raise InvalidModelError(
+            "Unsupported transformation of ONNX labeled probabilities."
+        )
+    return None
+
+
+def _softmax(graph: Any, output: str) -> bool:
+    producers = {name: node for node in graph.node for name in node.output if name}
+    for _ in range(len(graph.node)):
+        node = producers.get(output)
+        if node is None or node.domain not in ("", "ai.onnx"):
+            return False
+        if node.op_type == "Softmax":
+            axis = next((attr.i for attr in node.attribute if attr.name == "axis"), -1)
+            return axis in (1, -1)
+        if node.op_type not in ("Identity", "Cast") or len(node.input) != 1:
+            return False
+        output = node.input[0]
+    return False
+
+
+def _probability(value: dict[str, Any]) -> bool:
+    if value["dtype"].startswith("sequence<map<"):
+        return True
+    shape = value["shape"]
+    return (
+        value["dtype"] in ("float32", "float64")
+        and shape is not None
+        and len(shape) == 2
+        and _batch(shape)
+        and (not isinstance(shape[1], int) or 1 <= shape[1] <= MAX_CLASSES)
+    )
+
+
+def _label_output(value: dict[str, Any]) -> bool:
+    shape = value["shape"]
+    return (
+        value["dtype"] in ("int64", "string")
+        and shape is not None
+        and _batch(shape)
+        and (len(shape) == 1 or (len(shape) == 2 and shape[1] == 1))
+    )
+
+
+def _regression_width(graph: Any, output: str) -> None:
+    producers = {name: node for node in graph.node for name in node.output if name}
+    for _ in range(len(graph.node)):
+        node = producers.get(output)
+        if node is None:
+            return
+        if node.domain == "ai.onnx.ml" and node.op_type in (
+            "LinearRegressor",
+            "TreeEnsembleRegressor",
+        ):
+            attr_name = "targets" if node.op_type == "LinearRegressor" else "n_targets"
+            width = next(
+                (attr.i for attr in node.attribute if attr.name == attr_name), 1
+            )
+            if width != 1:
+                raise InvalidModelError(
+                    "ONNX native regression operator must expose one target."
+                )
+            return
+        if (
+            node.domain not in ("", "ai.onnx")
+            or node.op_type not in ("Identity", "Cast")
+            or len(node.input) != 1
+        ):
+            return
+        output = node.input[0]
+
+
+def _outputs(
+    graph: Any, outputs: list[dict[str, Any]]
+) -> tuple[ModelKind, dict[str, Any], ClassLabels | None, int | None]:
+    if len(outputs) == 2:
+        candidates = [
+            value
+            for index, value in enumerate(outputs)
+            if _probability(value) and _label_output(outputs[1 - index])
+        ]
+    elif len(outputs) == 1:
+        value = outputs[0]
+        evidence = _native_labels(graph, value["name"]) is not None or _softmax(
+            graph, value["name"]
+        )
+        candidates = [value] if _probability(value) and evidence else []
+    else:
+        candidates = []
+    if len(candidates) == 1:
+        value = candidates[0]
+        labels = _native_labels(graph, value["name"])
+        count = value["shape"][1] if value["shape"] is not None else None
+        if not isinstance(count, int):
+            count = len(labels[0]) if labels else None
+        if count is None or (labels is not None and len(labels[0]) != count):
+            raise InvalidModelError(
+                "ONNX probability width and native labels disagree."
+            )
+        return ModelKind.CLASSIFIER, value, labels, count
+    if len(outputs) == 1:
+        value = outputs[0]
+        shape = value["shape"]
+        if (
+            value["dtype"] in ("float32", "float64")
+            and shape is not None
+            and _batch(shape)
+            and (len(shape) == 1 or (len(shape) == 2 and shape[1] == 1))
+        ):
+            _regression_width(graph, value["name"])
+            return ModelKind.REGRESSOR, value, None, None
+    raise InvalidModelError(
+        "ONNX kind or probability/regression output layout is indeterminate."
+    )
+
+
+def describe(model: Any) -> Description:
+    graph = model.graph
+    if not 1 <= len(graph.output) <= 2:
+        raise InvalidModelError(
+            "ONNX requires one regression or probability output, optionally with labels."
+        )
+    inputs = _inputs(graph)
+    outputs = [signature(value) for value in graph.output]
+    kind, output, labels, count = _outputs(graph, outputs)
+    width = sum(value["shape"][1] for value in inputs)
+    metadata = {item.key: item.value for item in model.metadata_props}
+    if len(metadata) != len(model.metadata_props):
+        raise InvalidModelError("Duplicate ONNX metadata keys.")
+    projection = {
+        "inputs": list(inputs),
+        "outputs": outputs,
+        "metadata_props": metadata,
+    }
+    try:
+        canonical_dumps(projection, max_bytes=HYPERPARAMETER_BYTES, max_depth=16)
+    except LimitExceededError:
+        raise LimitExceededError(
+            "ONNX hyperparameters (inputs/outputs/metadata_props) exceed the "
+            f"{HYPERPARAMETER_BYTES}-byte or depth-16 limit."
+        ) from None
+    claims = ModelClaims(
+        estimator_class="ONNX",
+        kind=kind,
+        hyperparameters=cast(dict[str, JsonValue], projection),
+        n_features=width,
+        n_classes=None if count is None else [count],
+        n_outputs=1,
+        feature_names=[value["name"] for value in inputs] if len(inputs) > 1 else None,
+        feature_dtypes=[
+            value["dtype"] for value in inputs for _ in range(value["shape"][1])
+        ],
+        class_labels=labels,
+        onnx=OnnxInfo(
+            opsets=[
+                Opset(domain=item.domain, version=item.version)
+                for item in sorted(model.opset_import, key=lambda item: item.domain)
+            ],
+            ir_version=model.ir_version,
+            producer_name=model.producer_name,
+            producer_version=model.producer_version,
+        ),
+    )
+    return Description(claims, inputs, output["name"], output["shape"] is None)

mlbundle/_onnx/prediction.py

@@ -0,0 +1,136 @@
+"""Numeric tabular feeds and validated tensor/ZipMap outputs, entirely in memory."""
+
+import math
+import warnings
+from typing import Any, cast
+
+import numpy as np
+
+from mlbundle.adapters import ProbabilityOutput, RegressionOutput
+from mlbundle.contract import ModelKind
+from mlbundle.errors import InvalidModelError, PredictionInputError
+
+from .description import Description
+
+_InputShapeError = PredictionInputError
+
+
+def _array(data: Any, dtype: str) -> np.ndarray:
+    raw = np.asarray(data, dtype=object)
+    if np.dtype(dtype).kind in "iu":
+        limits = np.iinfo(dtype)
+        for value in raw.flat:
+            if (
+                isinstance(value, (bool, np.bool_))
+                or not isinstance(value, (int, float, np.integer, np.floating))
+                or not math.isfinite(value)
+                or int(value) != value
+                or not limits.min <= int(value) <= limits.max
+            ):
+                raise InvalidModelError(
+                    "ONNX integer inputs require whole numbers within the dtype range."
+                )
+    elif any(
+        isinstance(value, (bool, np.bool_))
+        or not isinstance(value, (int, float, np.integer, np.floating))
+        for value in raw.flat
+    ):
+        raise InvalidModelError("ONNX inputs require finite numeric values.")
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", RuntimeWarning)
+        values = np.asarray(raw, dtype=dtype)
+    if not np.isfinite(values).all():
+        raise InvalidModelError("ONNX inputs require finite numeric values.")
+    return values
+
+
+class Predictor:
+    def __init__(self, session: Any, description: Description) -> None:
+        self.session = session
+        self.description = description
+        self.claims = description.claims
+
+    def _feeds(self, data: object) -> tuple[dict[str, np.ndarray], int]:
+        frame = None
+        if any(
+            base.__module__.split(".")[0] == "pandas" for base in type(data).__mro__
+        ):
+            import pandas as pd  # noqa: PLC0415
+
+            if isinstance(data, pd.DataFrame):
+                frame = data
+        values = np.asarray(data, dtype=object)
+        if values.ndim != 2 or values.shape[1] != self.claims.n_features:
+            raise _InputShapeError(
+                "Prediction input must have the model's feature width."
+            )
+        rows = values.shape[0]
+        feeds = {}
+        for index, value in enumerate(self.description.inputs):
+            shape = value["shape"]
+            if isinstance(shape[0], int) and rows != shape[0]:
+                raise _InputShapeError(
+                    "Prediction row count must match the ONNX static batch."
+                )
+            if len(self.description.inputs) == 1:
+                column = values
+            elif frame is not None:
+                if not frame.columns.is_unique or value["name"] not in frame.columns:
+                    raise InvalidModelError(
+                        "ONNX DataFrame columns must match model input names."
+                    )
+                column = frame[[value["name"]]].to_numpy(dtype=object)
+            else:
+                column = values[:, index : index + 1]
+            feeds[value["name"]] = _array(column, value["dtype"])
+        return feeds, rows
+
+    def _run(self, data: object) -> tuple[Any, int]:
+        try:
+            feeds, rows = self._feeds(data)
+            return self.session.run([self.description.output], feeds)[0], rows
+        except (InvalidModelError, _InputShapeError):
+            raise
+        except Exception:
+            raise InvalidModelError(
+                "ONNX prediction failed or input values are invalid."
+            ) from None
+
+    def predict(self, data: object) -> RegressionOutput:
+        if self.claims.kind != ModelKind.REGRESSOR:
+            raise ValueError("predict requires a regressor.")
+        raw, rows = self._run(data)
+        try:
+            values = np.asarray(raw, dtype=float)
+            if values.shape == (rows, 1):
+                values = values[:, 0]
+            if values.shape != (rows,) or not np.isfinite(values).all():
+                raise ValueError
+            return cast(RegressionOutput, values.tolist())
+        except Exception:
+            raise InvalidModelError("Invalid ONNX regression output.") from None
+
+    def predict_proba(self, data: object) -> ProbabilityOutput:
+        if self.claims.kind != ModelKind.CLASSIFIER:
+            raise ValueError("predict_proba requires a classifier.")
+        raw, rows = self._run(data)
+        try:
+            if self.description.maps:
+                labels = (self.claims.class_labels or [[]])[0]
+                if not isinstance(raw, list) or any(
+                    not isinstance(row, dict) or set(row) != set(labels) for row in raw
+                ):
+                    raise ValueError
+                raw = [[row[label] for label in labels] for row in raw]
+            values = np.asarray(raw, dtype=float)
+            if (
+                values.shape != (rows, (self.claims.n_classes or [0])[0])
+                or not np.isfinite(values).all()
+                or np.any(values < 0)
+                or np.any(values > 1)
+                or not np.allclose(values.sum(axis=1), 1.0, rtol=1e-6, atol=1e-5)
+            ):
+                raise ValueError
+            return cast(ProbabilityOutput, values.tolist())
+        except Exception:
+            raise InvalidModelError("Invalid ONNX probability output.") from None

mlbundle/_onnx/protobuf.py

@@ -0,0 +1,111 @@
+"""Wire bounds before protobuf constructs objects; no tensor decoding or file I/O."""
+
+from typing import Any
+
+from mlbundle.errors import InvalidModelError, LimitExceededError
+from mlbundle.limits import ArchiveLimits
+
+from ._wire_runs import packed_varints, repeated
+
+# Each field/header or packed scalar consumes at least one distinct wire byte.
+MAX_FIELDS = ArchiveLimits().max_member_bytes
+MAX_DEPTH = 32
+
+
+def _varint(data: bytes, position: int, end: int) -> tuple[int, int]:
+    value = 0
+    for shift in range(0, 70, 7):
+        if position >= end:
+            raise InvalidModelError("Truncated ONNX protobuf.")
+        byte = data[position]
+        position += 1
+        value |= (byte & 127) << shift
+        if byte < 128:
+            if shift == 63 and byte > 1:
+                break
+            return value, position
+    raise InvalidModelError("Invalid ONNX protobuf integer.")
+
+
+def validate_wire(data: bytes, descriptor: Any) -> None:
+    if not data:
+        raise InvalidModelError("Empty ONNX model.")
+    if len(data) > ArchiveLimits().max_member_bytes:
+        raise LimitExceededError("ONNX model exceeds the byte limit.")
+    budget = [0]
+    _message(data, 0, len(data), descriptor, depth=0, budget=budget)
+
+
+def _message(data, position, end, descriptor, *, depth, budget):
+    if depth > MAX_DEPTH:
+        raise LimitExceededError("ONNX protobuf nesting exceeds the limit.")
+    seen = set()
+    while position < end:
+        budget[0] += 1
+        if budget[0] > MAX_FIELDS:
+            raise LimitExceededError("ONNX protobuf field count exceeds the limit.")
+        start = position
+        tag, position = _varint(data, position, end)
+        number, wire = tag >> 3, tag & 7
+        field = descriptor.fields_by_number.get(number)
+        if field is None or (not field.is_repeated and number in seen):
+            raise InvalidModelError("Unknown or duplicate ONNX protobuf field.")
+        seen.add(number)
+        value_start = position
+        position = _field(data, position, end, field, wire, depth=depth, budget=budget)
+        if field.is_repeated and (wire != 2 or field.type in (9, 12)):
+            position, count = repeated(
+                data, start, end, value_start=value_start, stop=position, wire=wire
+            )
+            budget[0] += count - 1
+            if budget[0] > MAX_FIELDS:
+                raise LimitExceededError("ONNX protobuf field count exceeds the limit.")
+
+
+def _field(data, position, end, field, wire, *, depth, budget):
+    expected = (
+        2
+        if field.type in (9, 11, 12)
+        else {1: 1, 2: 5, 6: 1, 7: 5, 15: 5, 16: 1}.get(field.type, 0)
+    )
+    packed = field.is_repeated and expected != 2 and wire == 2
+    if wire != expected and not packed:
+        raise InvalidModelError("Invalid ONNX protobuf wire type.")
+    if wire == 0:
+        return _varint(data, position, end)[1]
+    if wire in (1, 5):
+        stop = position + (8 if wire == 1 else 4)
+    else:
+        size, position = _varint(data, position, end)
+        stop = position + size
+    if stop > end:
+        raise InvalidModelError("Truncated ONNX protobuf field.")
+    if field.type == 11:
+        _message(
+            data, position, stop, field.message_type, depth=depth + 1, budget=budget
+        )
+    elif packed:
+        _packed(data, position, stop, expected, budget)
+    return stop
+
+
+def _packed(data, position, end, wire, budget):
+    if wire in (1, 5):
+        size = 8 if wire == 1 else 4
+        if (end - position) % size:
+            raise InvalidModelError("Invalid packed ONNX protobuf field.")
+        budget[0] += (end - position) // size
+    else:
+        while position < end:
+            stop, count = packed_varints(data, position, end)
+            if not count:
+                stop = _varint(data, position, end)[1]
+                count = 1
+            position = stop
+            budget[0] += count
+            if budget[0] > MAX_FIELDS:
+                raise LimitExceededError(
+                    "ONNX protobuf scalar count exceeds the limit."
+                )
+    if budget[0] > MAX_FIELDS:
+        raise LimitExceededError("ONNX protobuf scalar count exceeds the limit.")

mlbundle/_onnx/session.py

@@ -0,0 +1,87 @@
+"""Fixed runtime settings and agreement with its inferred input/output metadata."""
+
+from typing import Any, cast
+
+import onnxruntime as ort
+
+from mlbundle.errors import InvalidModelError
+
+from .description import Description
+
+ORT_TYPES = {
+    "float32": "tensor(float)",
+    "float64": "tensor(double)",
+    "int32": "tensor(int32)",
+    "int64": "tensor(int64)",
+    "string": "tensor(string)",
+}
+
+
+def _agree(expected: dict[str, Any], actual: Any) -> None:
+    dtype = expected["dtype"]
+    if dtype.startswith("sequence<map<"):
+        dtype = dtype.removeprefix("sequence<map<").removesuffix(">>")
+        key, value = dtype.split(",")
+        native_type = f"seq(map({key},{ORT_TYPES[value]}))"
+    else:
+        native_type = ORT_TYPES[dtype]
+    if expected["name"] != actual.name or native_type != actual.type:
+        raise InvalidModelError("ONNX graph and runtime input/output types disagree.")
+    shape = expected["shape"]
+    if shape is not None and (
+        actual.shape is None
+        or len(shape) != len(actual.shape)
+        or any(
+            isinstance(a, int) and a != b
+            for a, b in zip(shape, actual.shape, strict=True)
+        )
+    ):
+        raise InvalidModelError(
+            "ONNX graph and runtime input/output dimensions disagree."
+        )
+
+
+def create_session(data: bytes, description: Description) -> Any:
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    options.inter_op_num_threads = 1
+    options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
+    options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
+    options.enable_cpu_mem_arena = False
+    options.enable_mem_pattern = False
+    options.enable_profiling = False
+    options.optimized_model_filepath = ""
+    options.log_severity_level = 4
+    options.add_session_config_entry("session.load_model_format", "ONNX")
+    options.add_session_config_entry("session.disable_prepacking", "1")
+    session = ort.InferenceSession(
+        data,
+        sess_options=options,
+        providers=["CPUExecutionProvider"],
+        read_config_from_model=False,
+    )
+    inputs, outputs = session.get_inputs(), session.get_outputs()
+    expected_outputs = cast(
+        list[dict[str, Any]], description.claims.hyperparameters["outputs"]
+    )
+    if len(inputs) != len(description.inputs) or len(outputs) != len(expected_outputs):
+        raise InvalidModelError("ONNX graph and runtime input/output counts disagree.")
+    for expected, actual in zip(description.inputs, inputs, strict=True):
+        _agree(expected, actual)
+    for expected, actual in zip(expected_outputs, outputs, strict=True):
+        _agree(expected, actual)
+        if actual.name == description.output and expected["shape"] is not None:
+            width = (
+                1
+                if description.claims.n_classes is None
+                else description.claims.n_classes[0]
+            )
+            if (
+                len(actual.shape) == 2
+                and isinstance(actual.shape[1], int)
+                and actual.shape[1] != width
+            ):
+                raise InvalidModelError(
+                    "ONNX runtime prediction width disagrees with native claims."
+                )
+    return session

mlbundle/_onnx/signatures.py

@@ -0,0 +1,63 @@
+"""Supported tabular graph signatures and feature input layouts."""
+
+from typing import Any
+
+from mlbundle.errors import InvalidModelError
+from mlbundle.limits import MAX_FEATURES
+
+DTYPES = {1: "float32", 11: "float64", 6: "int32", 7: "int64"}
+
+
+def signature(value: Any) -> dict[str, Any]:
+    typ = value.type
+    if typ.HasField("tensor_type"):
+        tensor = typ.tensor_type
+        if not tensor.HasField("shape"):
+            raise InvalidModelError("ONNX input/output ranks must be declared.")
+        dtype = {**DTYPES, 8: "string"}.get(tensor.elem_type)
+        if dtype is None:
+            raise InvalidModelError("Unsupported ONNX input/output dtype.")
+        shape = [
+            dim.dim_value if dim.HasField("dim_value") else dim.dim_param or None
+            for dim in tensor.shape.dim
+        ]
+        return {"name": value.name, "dtype": dtype, "shape": shape}
+    if typ.HasField("sequence_type"):
+        element = typ.sequence_type.elem_type
+        if element.HasField("map_type"):
+            mapping = element.map_type
+            dtype = DTYPES.get(mapping.value_type.tensor_type.elem_type)
+            if mapping.key_type in (7, 8) and dtype in ("float32", "float64"):
+                key = "int64" if mapping.key_type == 7 else "string"
+                return {
+                    "name": value.name,
+                    "dtype": f"sequence<map<{key},{dtype}>>",
+                    "shape": None,
+                }
+    raise InvalidModelError("Unsupported ONNX input/output layout.")
+
+
+def _batch(shape: list[Any]) -> bool:
+    return bool(shape) and (not isinstance(shape[0], int) or shape[0] == 1)
+
+
+def _inputs(graph: Any) -> tuple[dict[str, Any], ...]:
+    inputs = tuple(signature(value) for value in graph.input)
+    for value in inputs:
+        shape = value["shape"]
+        if (
+            value["dtype"] not in DTYPES.values()
+            or shape is None
+            or len(shape) != 2
+            or not _batch(shape)
+            or not isinstance(shape[1], int)
+            or not 1 <= shape[1] <= MAX_FEATURES
+        ):
+            raise InvalidModelError(
+                "ONNX inputs require numeric matrices with fixed feature width and dynamic or one-row batches."
+            )
+        if len(inputs) > 1 and shape[1] != 1:
+            raise InvalidModelError(
+                "Multiple ONNX inputs require one column per input."
+            )
+    return inputs