ZKML Prescripciones

1. Genera un modelo dummy

# Verifica si el notebook está en Colab
try:
    # Instala ezkl y onnx si estás en Colab
    import google.colab
    import subprocess
    import sys
    subprocess.check_call([sys.executable, "-m", "pip", "install", "ezkl"])
    subprocess.check_call([sys.executable, "-m", "pip", "install", "onnx"])
except:
    pass # Si no, usa tu setup de ezkl local

# Importa las librería necesarias
from torch import nn
import torch
import ezkl
import os

# Convert a 3-byte hex number to tensor
def hex_to_byte_tensor(hex_value):
    # Ensure the hex value is 6 characters (3 bytes) by padding with zeros
    hex_str = f"{hex_value:06x}"
    return torch.tensor([int(hex_str[i:i+2], 16) for i in range(0, 6, 2)], dtype=torch.uint8)

# Example input value (can be changed as needed)
input_value = 0x047732  # This will be converted to [04, 77, 32]
input_tensor = hex_to_byte_tensor(input_value)
print(f"Input Tensor: {input_tensor}")

# Example input value (can be changed as needed)
input_value = 0x010402  # This will be converted to [04, 77, 32]
input_tensor = hex_to_byte_tensor(input_value)
print(f"Input Tensor: {input_tensor}")

# Define a model that outputs the same 3 bytes that it receives as input
class ThreeByteModel(nn.Module):
    def __init__(self):
        super(ThreeByteModel, self).__init__()

    def forward(self, x):
        # x should be a tensor of shape (batch_size, 3)
        return x

# Instancia el modelo
circuit = ThreeByteModel()

# Example input for the model
dummy_input = input_tensor.unsqueeze(0)  # Add batch dimension

# Generate output
output = circuit(dummy_input)
print(f"Output (3 bytes): {output}")

# Export the model in Onnx format
torch.onnx.export(
    circuit,
    dummy_input,
    "three_byte_model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
)

print("Model saved as 'three_byte_model.onnx'!")

2. Genera el contrato de Solidity

import asyncio
from torch import nn
import ezkl
import os
import json
import torch

# Import the hex_to_byte_tensor function and ThreeByteModel from generate_model.py
def hex_to_byte_tensor(hex_value):
    hex_str = f"{hex_value:06x}"
    return torch.tensor([int(hex_str[i:i+2], 16) for i in range(0, 6, 2)], dtype=torch.uint8)

class ThreeByteModel(nn.Module):
    def __init__(self):
        super(ThreeByteModel, self).__init__()

    def forward(self, x):
        return x

async def generate_contract():
    circuit = ThreeByteModel()
    model_path = 'three_byte_model.onnx'
    compiled_model_path = 'three_byte_model.compiled'
    pk_path = 'test.pk'
    vk_path = 'test.vk'
    settings_path = 'settings.json'
    witness_path = 'witness.json'
    data_path = 'input.json'

    # Create input data
    input_value = 0x040102
    x = hex_to_byte_tensor(input_value)
    x = x.unsqueeze(0)  # Add batch dimension

    # Export the model
    torch.onnx.export(
        circuit,
        x,
        model_path,
        export_params=True,
        opset_version=10,
        do_constant_folding=True,
        input_names=['input'],
        output_names=['output'],
        dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
    )

    # Prepare input data for JSON
    data_array = x.detach().numpy().reshape([-1]).tolist()
    data = dict(input_data=[data_array])
    json.dump(data, open(data_path, 'w'))

    # Setup run arguments
    py_run_args = ezkl.PyRunArgs()
    py_run_args.input_visibility = "public"
    py_run_args.output_visibility = "public"
    py_run_args.param_visibility = "fixed"

    # Generate settings
    res = ezkl.gen_settings(model_path, settings_path, py_run_args=py_run_args)
    assert res == True

    # Calibrate settings
    cal_path = "calibration.json"
    cal_data = dict(input_data=[x.detach().numpy().reshape([-1]).tolist()])
    json.dump(cal_data, open(cal_path, 'w'))
    await ezkl.calibrate_settings(cal_path, model_path, settings_path, "resources")

    # Compile circuit
    res = ezkl.compile_circuit(model_path, compiled_model_path, settings_path)
    assert res == True

    # Get SRS
    await ezkl.get_srs(settings_path)

    # Generate witness
    res = await ezkl.gen_witness(data_path, compiled_model_path, witness_path)
    assert os.path.isfile(witness_path)

    # Setup circuit
    res = ezkl.setup(compiled_model_path, vk_path, pk_path)
    assert res == True

    # Generate proof
    proof_path = 'test.pf'
    res = ezkl.prove(witness_path, compiled_model_path, pk_path, proof_path, "single")
    assert os.path.isfile(proof_path)

    # Verify proof
    res = ezkl.verify(proof_path, settings_path, vk_path)
    assert res == True

    # Generate calldata
    calldata = 'calldata.proof'
    res = ezkl.encode_evm_calldata(proof_path, calldata)
    calldata_hex = "0x" + ''.join(f"{byte:02x}" for byte in res)
    print("Calldata hex (EVM style):")
    print(calldata_hex)

    # Generate Solidity verifier
    abi_path = 'test.abi'
    sol_code_path = 'test.sol'
    output_sol_path = 'verifier.sol'

    res = await ezkl.create_evm_verifier(vk_path, settings_path, sol_code_path, abi_path)
    assert res == True

    with open(sol_code_path, 'r') as source_file:
        with open(output_sol_path, 'w') as dest_file:
            dest_file.write(source_file.read())

    print("Contract generated successfully!")
    return True

async def main():
    await generate_contract()

if __name__ == "__main__":
    asyncio.run(main())

3. Lanza el contrato de recetas médicas

// SPDX-License-Identifier: GPL-3.0
pragma solidity >=0.7.0 <0.9.0;

import "@openzeppelin/contracts/token/ERC1155/ERC1155.sol";

interface Halo2Verifier {
    function verifyProof(bytes calldata proof, uint256[] calldata instances) external view;
}

contract PrescriptionVerifier is ERC1155 {
    address public immutable HALO2_VERIFIER;

    constructor(address halo2Verifier) ERC1155("https://raw.githubusercontent.com/Turupawn/VoidRxMetadata/refs/heads/main/{id}.json") {
        HALO2_VERIFIER = halo2Verifier;
    }

    function extractThreeUintsFromProof(bytes memory proofCalldata) 
        public 
        pure 
        returns (uint256 a, uint256 b, uint256 c) 
    {
        uint256 len = proofCalldata.length;

        // Extract values from proof (0 if position out of bounds)
        a = (len >= 1) ? uint256(uint8(proofCalldata[len - 1])) : 0;
        b = (len >= 33) ? uint256(uint8(proofCalldata[len - 33])) : 0;
        c = (len >= 65) ? uint256(uint8(proofCalldata[len - 65])) : 0;
    }

    function processPrescription(bytes memory proofCalldata) public {
        // Verify proof first
        (bool success, bytes memory data) = HALO2_VERIFIER.call(proofCalldata);
        require(success && data.length == 32 && uint8(data[31]) == 1, "Invalid proof");

        // Extract medication amounts from proof
        (uint256 medA, uint256 medB, uint256 medC) = extractThreeUintsFromProof(proofCalldata);

        // Mint corresponding tokens to sender
        if (medA > 0) _mint(msg.sender, 1, medA, "");
        if (medB > 0) _mint(msg.sender, 2, medB, "");
        if (medC > 0) _mint(msg.sender, 3, medC, "");
    }
}