JSTprove 1.0.0__py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl

python/core/circuit_models/generic_onnx.py

@@ -0,0 +1,231 @@
+from __future__ import annotations
+
+import math
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import numpy as np
+import torch
+
+from python.core import RUST_BINARY_NAME
+from python.core.circuits.errors import (
+    CircuitFileError,
+    CircuitProcessingError,
+    CircuitRunError,
+)
+from python.core.circuits.zk_model_base import ZKModelBase
+
+if TYPE_CHECKING:
+    from python.core.model_processing.converters.onnx_converter import (
+        CircuitParamsDict,
+        ONNXLayerDict,
+    )
+from python.core.model_processing.converters.onnx_converter import (
+    ONNXConverter,
+    ONNXOpQuantizer,
+)
+from python.core.model_processing.onnx_quantizer.layers.base import BaseOpQuantizer
+
+
+class GenericModelONNX(ONNXConverter, ZKModelBase):
+    """
+    A generic ONNX-based Zero-Knowledge (ZK) circuit model wrapper.
+
+    This class provides:
+        - Integration for ONNX model loading, quantization (in `ONNXConverter`)
+          and ZK circuit infrastructure (in `ZKModelBase`).
+        - Support for model quantization via `ONNXOpQuantizer`.
+        - Input/output scaling and formatting utilities for ZK compatibility.
+
+    Attributes
+    ----------
+    name : str
+        Internal identifier for the binary to be run in rust backend.
+    op_quantizer : ONNXOpQuantizer
+        Operator quantizer for applying custom ONNX quantization rules.
+    rescale_config : dict
+        Per-node override for rescaling during quantization.
+        Keys are node names, values are booleans.
+        If not specified, assumption is to rescale each layer
+    model_file_name : str
+        Path to the ONNX model file used for the circuit.
+    scale_base : int
+        Base multiplier for scaling (default: 2).
+    scale_exponent : int
+        Exponent applied to `scale_base` for final scaling factor.
+
+    Parameters
+    ----------
+    model_name : str
+        Name of the model to load (with or without `.onnx` extension).
+
+    Notes
+    -----
+    - The scaling factor (`scale_base ** scale_exponent`) determines how floating point
+      inputs/outputs are represented as integers inside the ZK circuit.
+    - By default, scaling is fixed; dynamic scaling based on model analysis
+      is planned for future implementation.
+    - The quantization logic assumes operators are registered with
+      `ONNXOpQuantizer`.
+    """
+
+    def __init__(
+        self: GenericModelONNX,
+        model_name: str,
+        *,
+        use_find_model: bool = False,
+    ) -> None:
+        try:
+            self.name = RUST_BINARY_NAME
+            self.op_quantizer = ONNXOpQuantizer()
+            self.rescale_config = {}
+            if use_find_model:
+                self.model_file_name = self.find_model(model_name)
+            else:
+                self.model_file_name = model_name
+
+            self.scale_base = 2
+            self.scale_exponent = 18
+            ONNXConverter.__init__(self)
+        except Exception as e:
+
+            msg = f"Failed to initialize GenericModelONNX with model '{model_name}'"
+            raise CircuitFileError(
+                msg,
+                file_path=model_name,
+            ) from e
+
+    def find_model(self: GenericModelONNX, model_name: str) -> str:
+        """Resolve the ONNX model file path.
+
+        Args:
+            model_name (str): Name of the model (with or without `.onnx` extension).
+
+        Returns:
+            str: Full path to the model file.
+        """
+        if ".onnx" not in model_name:
+            model_name = model_name + ".onnx"
+
+        # Check direct path first
+        if Path(model_name).exists():
+            return model_name
+
+        # Check models_onnx directory
+        if "models_onnx" in model_name:
+            if Path(model_name).exists():
+                return model_name
+            models_onnx_path = model_name
+        else:
+            models_onnx_path = f"models_onnx/{model_name}"
+
+        if not Path(models_onnx_path).exists():
+            msg = f"Model file not found: '{model_name}'"
+            raise CircuitFileError(
+                msg,
+                file_path=models_onnx_path,
+            )
+        return models_onnx_path
+
+    def adjust_inputs(self: GenericModelONNX, input_file: str) -> str:
+        """Preprocess and flatten model inputs for the circuit.
+
+        Args:
+            input_file (str): Input data file or array compatible with the model.
+
+        Returns:
+            str: Adjusted input file after reshaping and scaling.
+        """
+        try:
+            input_shape = self.input_shape.copy()
+            shape = self.adjust_shape(input_shape)
+            self.input_shape = [math.prod(shape)]
+            x = super().adjust_inputs(input_file)
+            self.input_shape = input_shape.copy()
+        except Exception as e:
+            msg = f"Failed to adjust inputs for GenericModelONNX: {e}"
+            raise ValueError(msg) from e
+        else:
+            return x
+
+    def get_outputs(
+        self: GenericModelONNX,
+        inputs: np.ndarray | list[int] | torch.Tensor,
+    ) -> torch.Tensor:
+        """Run inference and flatten outputs.
+
+        Args:
+            inputs (List[int]): Preprocessed model inputs.
+
+        Returns:
+            torch.Tensor: Flattened model outputs as a tensor.
+        """
+        try:
+            raw_outputs = super().get_outputs(inputs)
+        except Exception as e:
+            msg = "Failed to get outputs for GenericModelONNX"
+            raise CircuitRunError(
+                msg,
+                operation="get_outputs",
+            ) from e
+        else:
+            return torch.as_tensor(np.array(raw_outputs)).flatten()
+
+    def format_inputs(
+        self: GenericModelONNX,
+        inputs: np.ndarray | list[int] | torch.Tensor,
+    ) -> dict[str, list[int]]:
+        """Format raw inputs into scaled integer tensors for the circuit
+        and transformed into json to be sent to rust backend.
+        Inputs are scaled by `scale_base ** scale_exponent`
+        and converted to long to ensure compatibility with ZK circuits
+
+        Args:
+            inputs (Any): Raw model inputs.
+
+        Returns:
+            Dict[str, List[int]]: Dictionary mapping `input` to scaled integer values.
+        """
+        try:
+            x = {"input": inputs}
+            scaling = BaseOpQuantizer.get_scaling(
+                scale_base=self.scale_base,
+                scale_exponent=self.scale_exponent,
+            )
+            for key in x:  # noqa: PLC0206
+                x[key] = torch.as_tensor(x[key]).flatten().tolist()
+                x[key] = (torch.as_tensor(x[key]) * scaling).long().tolist()
+        except Exception as e:
+            msg = f"Failed to format inputs for GenericModelONNX: {e}"
+            raise CircuitProcessingError(
+                msg,
+                operation="format_inputs",
+                data_type=type(inputs).__name__,
+            ) from e
+        else:
+            return x
+
+    def get_weights(
+        self: GenericModelONNX,
+    ) -> dict[str, list[ONNXLayerDict]]:
+        _, w_and_b, _ = super().get_weights()
+        # Currently want to read these in separately
+        return w_and_b
+
+    def get_architecture(
+        self: GenericModelONNX,
+    ) -> dict[str, list[ONNXLayerDict]]:
+        architecture, _, _ = super().get_weights()
+        # Currently want to read these in separately
+        return architecture
+
+    def get_metadata(
+        self: GenericModelONNX,
+    ) -> CircuitParamsDict:
+        _, _, circuit_params = super().get_weights()
+        # Currently want to read these in separately
+        return circuit_params
+
+
+if __name__ == "__main__":
+    pass

python/core/circuit_models/simple_circuit.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+from random import randint
+
+from python.core.circuits.base import Circuit, RunType
+from python.core.utils.helper_functions import CircuitExecutionConfig, ZKProofSystems
+
+
+class SimpleCircuit(Circuit):
+    def __init__(self: SimpleCircuit) -> None:
+        # Initialize the base class
+        super().__init__()
+
+        # Circuit-specific parameters
+        self.name = "simple_circuit"  # Use exact name that matches the binary
+        self.scale_exponent = 1
+        self.scale_base = 1
+
+        self.input_a = 100
+        self.input_b = 200
+        #############################################################
+        ### NOTE This is not a prg suitable for use in-production ###
+        #############################################################
+        self.nonce = randint(0, 10000000)  # noqa: S311
+
+        self.required_keys = ["value_a", "value_b", "nonce"]
+
+        self.input_shape = [1]
+
+    def get_inputs(self: SimpleCircuit) -> dict[str, int]:
+        """Retrieve the current input values for the circuit.
+
+        Returns:
+            dict[str, int]: A dictionary containing `value_a`, `value_b`, and `nonce`.
+        """
+        return {"value_a": self.input_a, "value_b": self.input_b, "nonce": self.nonce}
+
+    def get_outputs(self: SimpleCircuit, inputs: dict[str, int] | None = None) -> int:
+        """Compute the output of the circuit.
+
+        Args:
+            inputs (dict[str, int], optional):
+                A dictionary containing `value_a`, `value_b`, and `nonce`.
+                If None, uses the instance's default inputs. Defaults to None.
+
+        Returns:
+            int: output of function
+        """
+        if inputs is None:
+            inputs = {
+                "value_a": self.input_a,
+                "value_b": self.input_b,
+                "nonce": self.nonce,
+            }
+        print(  # noqa: T201
+            f"Performing addition operation: {inputs['value_a']} + {inputs['value_b']}",
+        )
+        return inputs["value_a"] + inputs["value_b"]
+
+    def format_inputs(self: SimpleCircuit, inputs: dict[str, int]) -> dict[str, int]:
+        """Format the inputs for the circuit.
+
+        Args:
+            inputs (dict[str, int]): A dictionary containing circuit input values.
+
+        Returns:
+            dict[str, int]: A dictionary containing circuit input values.
+        """
+        return inputs
+
+
+# Example code demonstrating circuit operations
+if __name__ == "__main__":
+    # Create a single circuit instance
+    print("\n--- Creating circuit instance ---")  # noqa: T201
+    circuit = SimpleCircuit()
+
+    print("\n--- Testing different operations ---")  # noqa: T201
+
+    print("\nGetting output again (should use cached value):")  # noqa: T201
+    output_again = circuit.get_outputs()
+    print(f"Circuit output: {output_again}")  # noqa: T201
+
+    # Run another operation
+    print("\nRunning compilation:")  # noqa: T201
+    circuit.base_testing(
+        CircuitExecutionConfig(
+            run_type=RunType.COMPILE_CIRCUIT,
+            dev_mode=True,
+            circuit_path="simple_circuit.txt",
+            input_file="inputs/simple_circuit_input.json",
+            output_file="output/simple_circuit_output.txt",
+            proof_system=ZKProofSystems.Expander,
+        ),
+    )
+
+    # Read the input and output files to verify
+    print("\n--- Verifying input and output files ---")  # noqa: T201
+    print(f"Input file: {circuit._file_info['input_file']}")  # noqa: SLF001, T201
+    print(f"Output file: {circuit._file_info['output_file']}")  # noqa: SLF001, T201
+
+    circuit.base_testing(
+        CircuitExecutionConfig(
+            run_type=RunType.GEN_WITNESS,
+            circuit_path="simple_circuit.txt",
+            input_file="inputs/simple_circuit_input.json",
+            output_file="output/simple_circuit_output.json",
+            write_json=True,
+            proof_system=ZKProofSystems.Expander,
+        ),
+    )
+
+    circuit = SimpleCircuit()
+    circuit.base_testing(
+        CircuitExecutionConfig(
+            run_type=RunType.PROVE_WITNESS,
+            circuit_path="simple_circuit.txt",
+            input_file="inputs/simple_circuit_input.json",
+            output_file="output/simple_circuit_output.json",
+            proof_system=ZKProofSystems.Expander,
+        ),
+    )
+
+    circuit = SimpleCircuit()
+    circuit.base_testing(
+        CircuitExecutionConfig(
+            run_type=RunType.GEN_VERIFY,
+            circuit_path="simple_circuit.txt",
+            input_file="inputs/simple_circuit_input.json",
+            output_file="output/simple_circuit_output.json",
+            proof_system=ZKProofSystems.Expander,
+        ),
+    )