{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n8QlFzjPRIGN"
      },
      "source": [
        "# EZKL DEMO\n",
        "\n",
        "**Learning Objectives**\n",
        "1. Learn some basic AI/ML techniques by training a toy model in pytorch to perform classification\n",
        "2. Convert the toy model into zk circuit with ezkl to do provable inference\n",
        "3. Create a solidity verifier and deploy it on Remix (you can deploy it however you like but we will use Remix as it's quite easy to setup)\n",
        "\n",
        "\n",
        "**Important Note**: You might want to avoid calling \"Run All\". There's some file locking issue with Colab which can cause weird bugs. To mitigate this issue you should run cell by cell on Colab."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dx81GOIySIpa"
      },
      "source": [
        "# Step 1: Training a toy model\n",
        "\n",
        "For this demo we will use a toy data set called the Iris dataset to demonstrate how training can be performed. The Iris dataset is a collection of Iris flowers and is one of the earliest dataset used to validate classification methodologies.\n",
        "\n",
        "[More info in the dataset](https://archive.ics.uci.edu/dataset/53/iris)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JhHE2WMvS9NP"
      },
      "source": [
        "First, we will need to import all the various dependencies required to train the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gvQ5HL1bTDWF"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "from sklearn.datasets import load_iris\n",
        "from sklearn.model_selection import train_test_split\n",
        "from sklearn.metrics import accuracy_score, precision_score, recall_score\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.autograd import Variable\n",
        "import tqdm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Op9SHfZHUkaR"
      },
      "source": [
        "Inspect the dataset. Note that for the Iris dataset we have 3 targets.\n",
        "\n",
        "0 = Iris-setosa\n",
        "\n",
        "1 = Iris-versicolor\n",
        "\n",
        "2 = Iris-virginica"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 424
        },
        "id": "C4XXA1hoU30c",
        "outputId": "4fbd47ec-88d1-4ef7-baee-3e3894cc29db"
      },
      "outputs": [],
      "source": [
        "iris = load_iris()\n",
        "dataset = pd.DataFrame(\n",
        "    data= np.c_[iris['data'], iris['target']],\n",
        "    columns= iris['feature_names'] + ['target'])\n",
        "dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I8RargmGTWN2"
      },
      "source": [
        "Next, we can begin defining the neural net model. For this dataset we will use a small fully connected neural net.\n",
        "\n",
        "<br />\n",
        "\n",
        "**Note:**\n",
        "For the 1st layer we use 4x20, because there are 4 features we want as inputs. After which we add a ReLU.\n",
        "\n",
        "For the 2nd layer we use 20x20, then add a ReLU.\n",
        "\n",
        "And for the last layer we use 20x3, because there are 3 classes we want to classify, then add a ReLU.\n",
        "\n",
        "The last ReLU function gives us an array of 3 elements where the position of the largest value gives us the target that we want to classify.\n",
        "\n",
        "For example, if we get [0, 0.001, 0.002] as the output of the last ReLU. As, 0.002 is the largest value, the inferred value is 2.\n",
        "\n",
        "\n",
        "![image.png](data:image/png;base64,[encoded data, 88192 characters])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dIdQ9U3yTKtP"
      },
      "outputs": [],
      "source": [
        "class Model(nn.Module):\n",
        "    # define nn\n",
        "    def __init__(self):\n",
        "        super(Model, self).__init__()\n",
        "        self.fc1 = nn.Linear(4, 20)\n",
        "        self.fc2 = nn.Linear(20, 20)\n",
        "        self.fc3 = nn.Linear(20, 3)\n",
        "        self.relu = nn.ReLU()\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.fc1(x)\n",
        "        x = self.relu(x)\n",
        "        x = self.fc2(x)\n",
        "        x = self.relu(x)\n",
        "        x = self.fc3(x)\n",
        "        x = self.relu(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "# Initialize Model\n",
        "model = Model()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SfC03XLNXDPZ"
      },
      "source": [
        "We will now need to split the dataset into a training set and testing set for ML. This is done fairly easily with the `train_test_split` helper function from sklearn."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "agmbEdmfUO1-",
        "outputId": "87766edd-50db-48af-aa5d-3f4fc164f8b7"
      },
      "outputs": [],
      "source": [
        "train_X, test_X, train_y, test_y = train_test_split(\n",
        "    dataset[dataset.columns[0:4]].values, # use columns 0-4 as X\n",
        "    dataset.target, # use target as y\n",
        "    test_size=0.2 # use 20% of data for testing\n",
        ")\n",
        "\n",
        "# Uncomment for sanity checks\n",
        "# print(\"train_X: \", train_X)\n",
        "# print(\"test_X: \", test_X)\n",
        "print(\"train_y: \", train_y)\n",
        "print(\"test_y: \", test_y)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_FrQXhAGZGS3"
      },
      "source": [
        "We can now define the parameters for training, we will use the [Cross Entropy Loss](https://machinelearningmastery.com/cross-entropy-for-machine-learning/) and [Stochastic Gradient Descent Optimizer](https://en.wikipedia.org/wiki/Stochastic_gradient_descent)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9PjADXnuXbdk",
        "outputId": "81602926-c386-4f68-a9ee-ae2d5837fe47"
      },
      "outputs": [],
      "source": [
        "# our loss function\n",
        "loss_fn = nn.CrossEntropyLoss()\n",
        "\n",
        "# our optimizer\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr=0.01)\n",
        "\n",
        "\n",
        "# use 800 EPOCHS\n",
        "EPOCHS = 800\n",
        "\n",
        "# Convert training data to pytorch variables\n",
        "train_X = Variable(torch.Tensor(train_X).float())\n",
        "test_X = Variable(torch.Tensor(test_X).float())\n",
        "train_y = Variable(torch.Tensor(train_y.values).long())\n",
        "test_y = Variable(torch.Tensor(test_y.values).long())\n",
        "\n",
        "\n",
        "loss_list     = np.zeros((EPOCHS,))\n",
        "accuracy_list = np.zeros((EPOCHS,))\n",
        "\n",
        "\n",
        "# we use tqdm for nice loading bars\n",
        "for epoch in tqdm.trange(EPOCHS):\n",
        "\n",
        "    # To train, we get a prediction from the current network\n",
        "    predicted_y = model(train_X)\n",
        "\n",
        "    # Compute the loss to see how bad or good we are doing\n",
        "    loss = loss_fn(predicted_y, train_y)\n",
        "\n",
        "    # Append the loss to keep track of our performance\n",
        "    loss_list[epoch] = loss.item()\n",
        "\n",
        "    # Afterwards, we will need to zero the gradients to reset\n",
        "    optimizer.zero_grad()\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "    # Calculate the accuracy, call torch.no_grad() to prevent updating gradients\n",
        "    # while calculating accuracy\n",
        "    with torch.no_grad():\n",
        "        y_pred = model(test_X)\n",
        "        correct = (torch.argmax(y_pred, dim=1) == test_y).type(torch.FloatTensor)\n",
        "        accuracy_list[epoch] = correct.mean()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 546
        },
        "id": "2fHJAgvwboCe",
        "outputId": "513c73b7-2663-4bb3-f7b4-cae208940070"
      },
      "outputs": [],
      "source": [
        "# Plot the Accuracy and Loss\n",
        "\n",
        "# import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "plt.style.use('ggplot')\n",
        "\n",
        "\n",
        "fig, (ax1, ax2) = plt.subplots(2, figsize=(12, 6), sharex=True)\n",
        "\n",
        "ax1.plot(accuracy_list)\n",
        "ax1.set_ylabel(\"Accuracy\")\n",
        "ax2.plot(loss_list)\n",
        "ax2.set_ylabel(\"Loss\")\n",
        "ax2.set_xlabel(\"epochs\");"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "djB-UtvgYbF2"
      },
      "source": [
        "## Congratulations! You've just trained a neural network\n",
        "\n",
        "**Exercise:** The model provided is very simplistic, what are other ways the model can be improved upon?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JgtwrbMZcgla"
      },
      "source": [
        "# Step 2: ZK the Neural Network\n",
        "\n",
        "Now that we have the Neural Network trained, we can use ezkl to easily ZK our model.\n",
        "\n",
        "To proceed we will now need to install `ezkl`\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C_YiqknhdDwN"
      },
      "outputs": [],
      "source": [
        "# check if notebook is in colab\n",
        "try:\n",
        "    import google.colab\n",
        "    import subprocess\n",
        "    import sys\n",
        "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"ezkl\"])\n",
        "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"onnx\"])\n",
        "\n",
        "# rely on local installation of ezkl if the notebook is not in colab\n",
        "except:\n",
        "    pass\n",
        "\n",
        "import os\n",
        "import json\n",
        "import ezkl"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-b_z_d2FdVTB"
      },
      "source": [
        "Next, we will need to export the neural network to a `.onnx` file. ezkl reads this `.onnx` file and converts it into a circuit which then allows you to generate proofs as well as verify proofs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YeKWP0tFeCpq"
      },
      "outputs": [],
      "source": [
        "# Specify all the files we need\n",
        "\n",
        "model_path = os.path.join('network.onnx')\n",
        "compiled_model_path = os.path.join('network.ezkl')\n",
        "pk_path = os.path.join('test.pk')\n",
        "vk_path = os.path.join('test.vk')\n",
        "settings_path = os.path.join('settings.json')\n",
        "\n",
        "witness_path = os.path.join('witness.json')\n",
        "data_path = os.path.join('input.json')\n",
        "cal_data_path = os.path.join('cal_data.json')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cQeNw_qndQ8g",
        "outputId": "2d40f14e-7fbb-4377-e9ee-0e7678edb2ce"
      },
      "outputs": [],
      "source": [
        "# After training, export to onnx (network.onnx) and create a data file (input.json)\n",
        "\n",
        "# create a random input\n",
        "x = test_X[0].reshape(1, 4)\n",
        "\n",
        "# Flips the neural net into inference mode\n",
        "model.eval()\n",
        "\n",
        "# Export the model\n",
        "torch.onnx.export(model,                     # model being run\n",
        "                  x,                         # model input (or a tuple for multiple inputs)\n",
        "                  model_path,                # where to save the model (can be a file or file-like object)\n",
        "                  export_params=True,        # store the trained parameter weights inside the model file\n",
        "                  opset_version=10,          # the ONNX version to export the model to\n",
        "                  do_constant_folding=True,  # whether to execute constant folding for optimization\n",
        "                  input_names = ['input'],   # the model's input names\n",
        "                  output_names = ['output'], # the model's output names\n",
        "                  dynamic_axes={'input' : {0 : 'batch_size'},    # variable length axes\n",
        "                                'output' : {0 : 'batch_size'}})\n",
        "\n",
        "data_array = ((x).detach().numpy()).reshape([-1]).tolist()\n",
        "\n",
        "data = dict(input_data = [data_array])\n",
        "\n",
        "    # Serialize data into file:\n",
        "json.dump(data, open(data_path, 'w'))\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9P4x79hIeiLO"
      },
      "source": [
        "After which we can proceed to generate the settings file for `ezkl` and run calibrate settings to find the optimal settings for `ezkl`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cY25BIyreIX8"
      },
      "outputs": [],
      "source": [
        "!RUST_LOG=trace\n",
        "# TODO: Dictionary outputs\n",
        "res = ezkl.gen_settings(model_path, settings_path)\n",
        "assert res == True\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# use the test set to calibrate the circuit\n",
        "cal_data = dict(input_data = test_X.flatten().tolist())\n",
        "\n",
        "# Serialize calibration data into file:\n",
        "json.dump(data, open(cal_data_path, 'w'))\n",
        "\n",
        "# Optimize for resources, we cap logrows at 12 to reduce setup and proving time, at the expense of accuracy\n",
        "# You may want to increase the max logrows if accuracy is a concern\n",
        "res = ezkl.calibrate_settings(cal_data_path, model_path, settings_path, \"resources\", max_logrows = 12, scales = [2])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MFmPMBQ1jYao"
      },
      "source": [
        "Next, we will compile the model. The compilation step allow us to generate proofs faster."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "De5XtpGUerkZ"
      },
      "outputs": [],
      "source": [
        "res = ezkl.compile_circuit(model_path, compiled_model_path, settings_path)\n",
        "assert res == True"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UbkuSVKljmhA"
      },
      "source": [
        "Before we can setup the circuit params, we need a SRS (Structured Reference String). The SRS is used to generate the proofs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "amaTcWG6f2GI"
      },
      "outputs": [],
      "source": [
        "res = ezkl.get_srs( settings_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y92p3GhVj1Jd"
      },
      "source": [
        "Now run setup, this will generate a proving key (pk) and verification key (vk). The proving key is used for proving while the verification key is used for verificaton."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fdsteit9jzfK"
      },
      "outputs": [],
      "source": [
        "res = ezkl.setup(\n",
        "        compiled_model_path,\n",
        "        vk_path,\n",
        "        pk_path,\n",
        "    )\n",
        "\n",
        "\n",
        "assert res == True\n",
        "assert os.path.isfile(vk_path)\n",
        "assert os.path.isfile(pk_path)\n",
        "assert os.path.isfile(settings_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QYlqpP3jkExm"
      },
      "source": [
        "Now, we can generate a proof and verify the proof as a sanity check. We will use the \"evm\" transcript. This will allow us to provide proofs to the EVM."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yoz5Vks5kaHI"
      },
      "outputs": [],
      "source": [
        "# Generate the Witness for the proof\n",
        "\n",
        "# now generate the witness file\n",
        "witness_path = os.path.join('witness.json')\n",
        "\n",
        "res = ezkl.gen_witness(data_path, compiled_model_path, witness_path)\n",
        "assert os.path.isfile(witness_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eKkFBZX1kBdE",
        "outputId": "48c67e19-a491-4515-f09c-a560df8c3834"
      },
      "outputs": [],
      "source": [
        "# Generate the proof\n",
        "\n",
        "proof_path = os.path.join('proof.json')\n",
        "\n",
        "proof = ezkl.prove(\n",
        "        witness_path,\n",
        "        compiled_model_path,\n",
        "        pk_path,\n",
        "        proof_path,\n",
        "        \"single\",\n",
        "    )\n",
        "\n",
        "print(proof)\n",
        "assert os.path.isfile(proof_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DuuH-qcOkQf1",
        "outputId": "375fdd63-1c0b-4c3c-eddd-f890a752923c"
      },
      "outputs": [],
      "source": [
        "# verify our proof\n",
        "\n",
        "res = ezkl.verify(\n",
        "        proof_path,\n",
        "        settings_path,\n",
        "        vk_path,\n",
        "    )\n",
        "\n",
        "assert res == True\n",
        "print(\"verified\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TOSRigalkwH-"
      },
      "source": [
        "## Congratulations! You have just turned your Neural Network into a Halo2 Circuit!\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "flrg3NOGwsJh"
      },
      "source": [
        "\n",
        "# Part 3: Deploying the Verifier\n",
        "Now that we have the circuit setup, we can proceed to deploy the verifier onchain.\n",
        "\n",
        "We will need to setup `solc=0.8.20` for this."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CVqMeMYqktvl",
        "outputId": "60ef81a5-867e-4a27-a0a1-0a492244e7f7"
      },
      "outputs": [],
      "source": [
        "# check if notebook is in colab\n",
        "try:\n",
        "    import google.colab\n",
        "    import subprocess\n",
        "    import sys\n",
        "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"solc-select\"])\n",
        "    !solc-select install 0.8.20\n",
        "    !solc-select use 0.8.20\n",
        "    !solc --version\n",
        "\n",
        "# rely on local installation if the notebook is not in colab\n",
        "except:\n",
        "    pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HRHvkMjVlfWU"
      },
      "source": [
        "With solc in our environment we can now create the evm verifier."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gYlw20VZkva7"
      },
      "outputs": [],
      "source": [
        "sol_code_path = os.path.join('Verifier.sol')\n",
        "abi_path = os.path.join('Verifier.abi')\n",
        "\n",
        "res = ezkl.create_evm_verifier(\n",
        "        vk_path,\n",
        "        \n",
        "        settings_path,\n",
        "        sol_code_path,\n",
        "        abi_path\n",
        "    )\n",
        "\n",
        "assert res == True\n",
        "assert os.path.isfile(sol_code_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jQSAVMvxrBQD",
        "outputId": "691484fa-ef21-4b40-e179-9d2d90abd3d0"
      },
      "outputs": [],
      "source": [
        "onchain_input_array = []\n",
        "\n",
        "# using a loop\n",
        "# avoiding printing last comma\n",
        "formatted_output = \"[\"\n",
        "for i, value in enumerate(proof[\"instances\"]):\n",
        "    for j, field_element in enumerate(value):\n",
        "        onchain_input_array.append(ezkl.string_to_felt(field_element))\n",
        "        formatted_output += str(onchain_input_array[-1])\n",
        "        if j != len(value) - 1:\n",
        "            formatted_output += \", \"\n",
        "    formatted_output += \"]\"\n",
        "\n",
        "# This will be the values you use onchain\n",
        "# copy them over to remix and see if they verify\n",
        "# What happens when you change a value?\n",
        "print(\"pubInputs: \", formatted_output)\n",
        "print(\"proof: \", \"0x\" + proof[\"proof\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zrzPxPvZmX9b"
      },
      "source": [
        "We will exit colab for the next steps. At the left of colab you can see a folder icon. Click on that.\n",
        "\n",
        "\n",
        "You should see a `Verifier.sol`. Right-click and save it locally.\n",
        "\n",
        "Now go to [https://remix.ethereum.org](https://remix.ethereum.org).\n",
        "\n",
        "Create a new file within remix and copy the verifier code over.\n",
        "\n",
        "Finally, compile the code and deploy. For the demo you can deploy to the test environment within remix.\n",
        "\n",
        "If everything works, you would have deployed your verifer onchain! Copy the values in the cell above to the respective fields to test if the verifier is working.\n",
        "\n",
        "**Note that right now this setup accepts random values!**\n",
        "\n",
        "This might not be great for some applications. For that we will want to use a data attested verifier instead. [See this tutorial.](https://github.com/zkonduit/ezkl/blob/main/examples/notebooks/data_attest.ipynb)\n",
        "\n",
        "## Congratulations for making it this far!\n",
        "\n",
        "If you have followed the whole tutorial, you would have deployed a neural network inference model onchain! That's no mean feat!"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.9.15"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
