Module: ezkl_utils

Library containing utility functions related to EZKL and the EZKL Proof Service.

EZKLProof

Bases: BaseModel

All of the artifacts needed to verify a proof.

Attributes:

Name	Type	Description
ezkl_proof	Dict[str, Any]	Proof dictionary as generated by ezkl
witness	Dict[str, Any]	Witness used to generate the proof
verify_key	Optional[bytes]	Verification key used to verify the proof. Available if a new setup has been done, otherwise None
input	Optional[List[float]]	Input data
output	Optional[List[float]]	Output data
flops	Optional[float]	Number of FLOPS in the model
ezkl_proof_path	Optional[Path]	Path to the proof file
witness_path	Optional[Path]	Path to the witness file
settings_path	Optional[Path]	Path to the settings file
vk_path	Optional[Path]	Path to the verification

Source code in src/infernet_ml/zk/ezkl/ezkl_utils.py

class EZKLProof(BaseModel):
    """
    All of the artifacts needed to verify a proof.

    Attributes:
        ezkl_proof (Dict[str, Any]): Proof dictionary as generated by ezkl
        witness (Dict[str, Any]): Witness used to generate the proof
        verify_key (Optional[bytes]): Verification key used to verify the proof.
            Available if a new setup has been done, otherwise None

        input (Optional[List[float]]): Input data
        output (Optional[List[float]]): Output data
        flops (Optional[float]): Number of FLOPS in the model

        ezkl_proof_path (Optional[Path]): Path to the proof file
        witness_path (Optional[Path]): Path to the witness file
        settings_path (Optional[Path]): Path to the settings file
        vk_path (Optional[Path]): Path to the verification
    """

    ezkl_proof: Dict[str, Any]
    witness: Dict[str, Any]

    verify_key: Optional[bytes] = None

    input: Optional[RitualVector] = None

    output: Optional[RitualVector] = None

    flops: Optional[float] = None

    ezkl_proof_path: Optional[Path]
    witness_path: Optional[Path]
    settings_path: Optional[Path]
    vk_path: Optional[Path] = None

    @classmethod
    def from_proof_and_verify_key(
        cls,
        proof: Dict[str, Any],
        verify_key: bytes | str,
        proof_directory: Path = DEFAULT_PROOFS_DIRECTORY,
    ) -> "EZKLProof":
        """
        Create an EZKLProof object from the proof and verification key.

        Args:
            proof (Dict[str, Any]): Proof dictionary
            verify_key (bytes | str): Verification key
            proof_directory (Path, optional): Directory to store the proof.
                Defaults to ~/.cache/ritual/proofs.

        Returns:
            EZKLProof: EZKLProof object
        """
        verify_key = (
            bytes.fromhex(verify_key.removeprefix("0x"))
            if isinstance(verify_key, str)
            else verify_key
        )
        proof_dir, proof_path = get_proof_dir_and_path(proof, proof_directory)
        return cls(
            ezkl_proof=proof,
            verify_key=verify_key,
            witness={},
            ezkl_proof_path=proof_path,
            witness_path=None,
            settings_path=None,
        )

    @classmethod
    def from_files(
        cls,
        proof_path: Path,
        witness_path: Path,
        settings_path: Path,
        flops: Optional[float] = None,
        vk_path: Optional[Path] = None,
    ) -> "EZKLProof":
        """
        Create an EZKLProof object from the files generated by EZKL.

        Args:
            proof_path (Path): Path to the proof file
            witness_path (Path): Path to the witness file
            settings_path (Path): Path to the settings file
            flops (Optional[float], optional): Number of FLOPS in the model.
            vk_path (Optional[Path], optional): Path to the verification key file.
                Defaults to None.

        Returns:
            EZKLProof: EZKLProof object
        """
        with open(proof_path, "r") as f:
            proof = json.loads(f.read())
            public_inputs = proof.get("pretty_public_inputs", {})
            inputs = public_inputs.get("rescaled_inputs")
            inputs = numpy.array(inputs, dtype=float) if inputs else None
            outputs = public_inputs.get("rescaled_outputs")
            outputs = numpy.array(outputs, dtype=float) if outputs else None

        with open(witness_path, "r") as f:
            witness = json.load(f)
            _input, _output = extract_io_from_proof_execution(settings_path, witness)
        if vk_path is not None:
            with open(vk_path, "rb") as f:
                vk = f.read()
        else:
            vk = None

        return cls(
            ezkl_proof=proof,
            verify_key=vk,
            witness=witness,
            input=RitualVector.from_numpy(inputs),
            output=RitualVector.from_numpy(outputs),
            flops=flops,
            ezkl_proof_path=proof_path,
            witness_path=witness_path,
            vk_path=vk_path,
            settings_path=settings_path,
        )

from_files(proof_path, witness_path, settings_path, flops=None, vk_path=None) classmethod

Create an EZKLProof object from the files generated by EZKL.

Parameters:

Name	Type	Description	Default
proof_path	Path	Path to the proof file	required
witness_path	Path	Path to the witness file	required
settings_path	Path	Path to the settings file	required
flops	Optional[float]	Number of FLOPS in the model.	None
vk_path	Optional[Path]	Path to the verification key file. Defaults to None.	None

Returns:

Name	Type	Description
EZKLProof	EZKLProof	EZKLProof object

Source code in src/infernet_ml/zk/ezkl/ezkl_utils.py

@classmethod
def from_files(
    cls,
    proof_path: Path,
    witness_path: Path,
    settings_path: Path,
    flops: Optional[float] = None,
    vk_path: Optional[Path] = None,
) -> "EZKLProof":
    """
    Create an EZKLProof object from the files generated by EZKL.

    Args:
        proof_path (Path): Path to the proof file
        witness_path (Path): Path to the witness file
        settings_path (Path): Path to the settings file
        flops (Optional[float], optional): Number of FLOPS in the model.
        vk_path (Optional[Path], optional): Path to the verification key file.
            Defaults to None.

    Returns:
        EZKLProof: EZKLProof object
    """
    with open(proof_path, "r") as f:
        proof = json.loads(f.read())
        public_inputs = proof.get("pretty_public_inputs", {})
        inputs = public_inputs.get("rescaled_inputs")
        inputs = numpy.array(inputs, dtype=float) if inputs else None
        outputs = public_inputs.get("rescaled_outputs")
        outputs = numpy.array(outputs, dtype=float) if outputs else None

    with open(witness_path, "r") as f:
        witness = json.load(f)
        _input, _output = extract_io_from_proof_execution(settings_path, witness)
    if vk_path is not None:
        with open(vk_path, "rb") as f:
            vk = f.read()
    else:
        vk = None

    return cls(
        ezkl_proof=proof,
        verify_key=vk,
        witness=witness,
        input=RitualVector.from_numpy(inputs),
        output=RitualVector.from_numpy(outputs),
        flops=flops,
        ezkl_proof_path=proof_path,
        witness_path=witness_path,
        vk_path=vk_path,
        settings_path=settings_path,
    )

from_proof_and_verify_key(proof, verify_key, proof_directory=DEFAULT_PROOFS_DIRECTORY) classmethod

Create an EZKLProof object from the proof and verification key.

Parameters:

Name	Type	Description	Default
proof	Dict[str, Any]	Proof dictionary	required
verify_key	bytes | str	Verification key	required
proof_directory	Path	Directory to store the proof. Defaults to ~/.cache/ritual/proofs.	DEFAULT_PROOFS_DIRECTORY

Returns:

Name	Type	Description
EZKLProof	EZKLProof	EZKLProof object

Source code in src/infernet_ml/zk/ezkl/ezkl_utils.py

@classmethod
def from_proof_and_verify_key(
    cls,
    proof: Dict[str, Any],
    verify_key: bytes | str,
    proof_directory: Path = DEFAULT_PROOFS_DIRECTORY,
) -> "EZKLProof":
    """
    Create an EZKLProof object from the proof and verification key.

    Args:
        proof (Dict[str, Any]): Proof dictionary
        verify_key (bytes | str): Verification key
        proof_directory (Path, optional): Directory to store the proof.
            Defaults to ~/.cache/ritual/proofs.

    Returns:
        EZKLProof: EZKLProof object
    """
    verify_key = (
        bytes.fromhex(verify_key.removeprefix("0x"))
        if isinstance(verify_key, str)
        else verify_key
    )
    proof_dir, proof_path = get_proof_dir_and_path(proof, proof_directory)
    return cls(
        ezkl_proof=proof,
        verify_key=verify_key,
        witness={},
        ezkl_proof_path=proof_path,
        witness_path=None,
        settings_path=None,
    )

ProofGenerationArgs

Bases: BaseModel

Arguments for generating the proof.

Attributes:

Name	Type	Description
compiled_path	Path	Path to the compiled model
settings_path	Path	Path to the settings

Source code in src/infernet_ml/zk/ezkl/ezkl_utils.py

class ProofGenerationArgs(BaseModel):
    """
    Arguments for generating the proof.

    Attributes:
        compiled_path (Path): Path to the compiled model
        settings_path (Path): Path to the settings
    """

    compiled_path: Path
    settings_path: Path

generate_proof(artifact_files, input_data, output_data=None, prover_key=None, proof_directory=DEFAULT_PROOFS_DIRECTORY) async

Generate a proof for the model. If a verifier

Parameters:

Name	Type	Description	Default
artifact_files	EZKLArtifact	Paths to the various artifacts generated by EZKL	required
input_data	ONNXInput	Input data for the model	required
output_data	Optional[ndarray]	Output data for the model. Defaults to None.	None
prover_key	Optional[bytes]	Provide a prover key to generate the proof. Defaults to None. If it is not provided, a new circuit setup will happen & a new prover key will be generated.	None
proof_directory	Optional[Path]	Directory to store the proof. Defaults to ~/.cache/ritual/proofs.	DEFAULT_PROOFS_DIRECTORY

Returns:

Name	Type	Description
EZKLProof	EZKLProof	Proof and verification key

Source code in src/infernet_ml/zk/ezkl/ezkl_utils.py

async def generate_proof(
    artifact_files: EZKLArtifact,
    input_data: ONNXInput,
    output_data: Optional[np.ndarray] = None,  # type: ignore
    prover_key: Optional[bytes] = None,
    proof_directory: Path = DEFAULT_PROOFS_DIRECTORY,
) -> EZKLProof:
    """
    Generate a proof for the model. If a verifier

    Args:
        artifact_files (EZKLArtifact): Paths to the various artifacts generated by
            EZKL
        input_data (ONNXInput): Input data for the model
        output_data (Optional[np.ndarray], optional): Output data for the model.
            Defaults to None.
        prover_key (Optional[bytes], optional): Provide a prover key to generate the
            proof. Defaults to None. If it is not provided, a new circuit setup will
            happen & a new prover key will be generated.
        proof_directory (Optional[Path], optional): Directory to store the proof.
            Defaults to ~/.cache/ritual/proofs.

    Returns:
        EZKLProof: Proof and verification key
    """
    settings_path = artifact_files.settings_path
    compiled_model_path = artifact_files.compiled_model_path
    srs_path = artifact_files.srs_path
    vk_path = artifact_files.verifier_key_path
    pk_path = settings_path.parent / "test.pk"
    prover_key = prover_key or artifact_files.prover_key_path.read_bytes()

    if srs_path is None:
        srs_path = settings_path.parent / "srs_file.srs"
        res = await ezkl.get_srs(settings_path, srs_path=srs_path)
        assert res is True

    if prover_key is None:
        # if pk is not provided, then we'll generate both pk & vk
        log.info("proving key not provided, performing a new ZK setup")
        vk_path = settings_path.parent / "test.vk"
        res = ezkl.setup(
            compiled_model_path, vk_path=vk_path, pk_path=pk_path, srs_path=srs_path
        )
        assert res is True
        assert os.path.isfile(vk_path)
        assert os.path.isfile(pk_path)
    else:
        log.info("proving key provided using the existing setup")
        pk_path.write_bytes(prover_key)

    proof_dir, proof_path = get_proof_dir_and_path({}, proof_directory)

    os.makedirs(proof_dir, exist_ok=True)

    witness_path = proof_dir / "witness.json"
    data_path = proof_dir / "data.json"

    input_data_ = list(input_data.values())[0].reshape([-1]).tolist()

    data = dict(input_data=[input_data_])

    if output_data is not None:
        data["output_data"] = [output_data.reshape([-1]).tolist()]

    json.dump(data, open(data_path, "w"))

    await ezkl.gen_witness(
        data_path, model=compiled_model_path, output=witness_path, srs_path=srs_path
    )

    assert os.path.isfile(witness_path)

    ONNXModelAnalyzer(artifact_files.onnx_path).calculate_flops()

    res = ezkl.prove(
        witness=witness_path,
        model=compiled_model_path,
        pk_path=pk_path,
        proof_path=proof_path,
        proof_type="single",
        srs_path=srs_path,
    )

    assert res, "unable to generate proof"

    assert os.path.isfile(proof_path)

    if vk_path:
        log.debug("vk available, verifying the proof")
        verify_success = ezkl.verify(
            proof_path=proof_path,
            settings_path=settings_path,
            vk_path=vk_path,
            srs_path=srs_path,
        )
        assert verify_success, "unable to verify generated proof"

    if artifact_files.onnx_path:
        flops = get_onnx_flops(artifact_files.onnx_path)
    else:
        flops = None

    return EZKLProof.from_files(
        proof_path=proof_path,
        witness_path=witness_path,
        settings_path=settings_path,
        flops=flops,
        vk_path=vk_path,
    )

generate_proof_from_repo_id(repo_id, input_vector, hf_token=None) async

Generate a proof for the model using the provided repo_id.

Parameters:

Name	Type	Description	Default
repo_id	str | RitualRepoId	Repo ID for the model	required
input_vector	FloatNumpy	Input vector for the model	required
hf_token	Optional[str]	Hugging Face API token, optional	None

Returns:

Name	Type	Description
EZKLProof	EZKLProof	Proof and verification

Source code in src/infernet_ml/zk/ezkl/ezkl_utils.py

async def generate_proof_from_repo_id(
    repo_id: str | RitualRepoId,
    input_vector: np.ndarray[Any, Any],
    hf_token: Optional[str] = None,
) -> EZKLProof:
    """
    Generate a proof for the model using the provided repo_id.

    Args:
        repo_id (str | RitualRepoId): Repo ID for the model
        input_vector (FloatNumpy): Input vector for the model
        hf_token (Optional[str]): Hugging Face API token, optional

    Returns:
        EZKLProof: Proof and verification
    """
    artifacts: EZKLArtifact = RitualArtifactManager.from_repo(
        EZKLArtifact, repo_id, hf_token=hf_token
    ).artifact

    names = get_onnx_input_names(artifacts.onnx_path)
    first = names[0]

    return await generate_proof(
        artifact_files=artifacts,
        prover_key=artifacts.prover_key_path.read_bytes(),
        input_data={first: input_vector},
    )

get_proof_dir_and_path(proof, proofs_directory=DEFAULT_PROOFS_DIRECTORY)

Get the proof directory and path for the proof. Prune old items in the proof directory to avoid filling up the disk.

Parameters:

Name	Type	Description	Default
proof	Dict[str, Any]	Proof dictionary	required
proofs_directory	Path	Directory to store the proof. Defaults to ~/.cache/ritual/proofs.	DEFAULT_PROOFS_DIRECTORY

Returns:

Type	Description
tuple[Path, Path]	tuple[Path, Path]: Proof directory and path

Source code in src/infernet_ml/zk/ezkl/ezkl_utils.py

def get_proof_dir_and_path(
    proof: Dict[str, Any],
    proofs_directory: Path = DEFAULT_PROOFS_DIRECTORY,
) -> tuple[Path, Path]:
    """
    Get the proof directory and path for the proof. Prune old items in the proof
    directory to avoid filling up the disk.

    Args:
        proof (Dict[str, Any]): Proof dictionary
        proofs_directory (Path, optional): Directory to store the proof.
            Defaults to ~/.cache/ritual/proofs.

    Returns:
        tuple[Path, Path]: Proof directory and path
    """
    proof_dir = get_proof_dir(proofs_directory)
    proof_path = proof_dir / "proof.json"
    proof_path.write_text(json.dumps(proof))
    return proof_dir, proof_path

verify_proof(proof, artifact_files) async

Verify a proof using the provided artifacts.

Parameters:

Name	Type	Description	Default
proof	EZKLProof	Proof to verify	required
artifact_files	EZKLArtifact	Paths to the various artifacts generated	required

Returns:

Name	Type	Description
bool	bool	True if the proof is valid, False otherwise

Source code in src/infernet_ml/zk/ezkl/ezkl_utils.py

async def verify_proof(proof: EZKLProof, artifact_files: EZKLArtifact) -> bool:
    """
    Verify a proof using the provided artifacts.

    Args:
        proof (EZKLProof): Proof to verify
        artifact_files (EZKLArtifact): Paths to the various artifacts generated
        by EZKL

    Returns:
        bool: True if the proof is valid, False otherwise
    """

    with tempfile.TemporaryDirectory() as _dir:
        temp_dir = Path(_dir)
        vk_path = temp_dir / "test.vk"
        assert proof.verify_key
        vk_path.write_bytes(proof.verify_key)

        settings_path = artifact_files.settings_path
        if not artifact_files.srs_path or not artifact_files.srs_path.exists():
            res = await ezkl.get_srs(settings_path, srs_path=artifact_files.srs_path)
            assert res is True

        verify_success = ezkl.verify(
            proof_path=proof.ezkl_proof_path,
            settings_path=settings_path,
            vk_path=vk_path,
            srs_path=artifact_files.srs_path,
        )

        return bool(verify_success)