{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6ypZa6jg0rZy"
      },
      "source": [
        "## Mnist Clan-ssifier ;)\n",
        "\n",
        "Here we demonstrate how to use the EZKL package to build an MNIST classifier for on-chain handrawn digit recognition.\n",
        "The proofs get submitted to a contract that assigns the users account to a digit clan (0-9). The contract keeps track of the member count of each clan. The clan with the most members is the winner!\n",
        "\n",
        "![zk-gaming-diagram-transformed](https://file.notion.so/f/f/f9535faf-4480-4499-9059-a48ba240eaa9/cd13414a-ecd8-4b8f-90a1-8a2311baa278/Untitled.png?id=365d66ee-e653-4ec3-8eb6-6d2b6306455a&table=block&spaceId=f9535faf-4480-4499-9059-a48ba240eaa9&expirationTimestamp=1701568800000&signature=VJ9p3YsOjYjeLxmkVEWOJw_3VmM6IBkTYxMwQUFKeus&downloadName=Untitled.png)\n",
        "> **A typical ZK application flow**. For all the image classifictiton hackers out there — this is an fairly straight forward example. A user computes a ZKML-proof that they have calculated a valid classification of a hand drawn digit from a MNIST trained lenet model. They submit this proof to a verifier contract which governs a set of clans, along with the output values of the model (length 10 tensor whereby the index with the max value represented the prediction), and the clan count updates according the lenets model's prediction. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Fm2DMmXW0rZ0"
      },
      "outputs": [],
      "source": [
        "# check if notebook is in colab\n",
        "try:\n",
        "    # install ezkl\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\", \"torch\"])\n",
        "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"torchvision\"])\n",
        "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"tf2onnx\"])\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",
        "# make sure you have the dependencies required here already installed\n",
        "import ezkl\n",
        "import os\n",
        "import json\n",
        "import time\n",
        "import random\n",
        "import logging\n",
        "\n",
        "\n",
        "# uncomment for more descriptive logging\n",
        "FORMAT = '%(levelname)s %(name)s %(asctime)-15s %(filename)s:%(lineno)d %(message)s'\n",
        "logging.basicConfig(format=FORMAT)\n",
        "logging.getLogger().setLevel(logging.INFO)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jbRkAJLz0rZ1"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "class LeNet(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(LeNet, self).__init__()\n",
        "        # Convolutional encoder\n",
        "        self.conv1 = nn.Conv2d(1, 6, 5)  # 1 input channel, 6 output channels, 5x5 kernel\n",
        "        self.conv2 = nn.Conv2d(6, 16, 5) # 6 input channels, 16 output channels, 5x5 kernel\n",
        "\n",
        "        # Fully connected layers / Dense block\n",
        "        self.fc1 = nn.Linear(16 * 4 * 4, 120) \n",
        "        self.fc2 = nn.Linear(120, 84)         # 120 inputs, 84 outputs\n",
        "        self.fc3 = nn.Linear(84, 10)          # 84 inputs, 10 outputs (number of classes)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Convolutional block\n",
        "        x = F.avg_pool2d(F.sigmoid(self.conv1(x)), (2, 2)) # Convolution -> Sigmoid -> Avg Pool\n",
        "        x = F.avg_pool2d(F.sigmoid(self.conv2(x)), (2, 2)) # Convolution -> Sigmoid -> Avg Pool\n",
        "\n",
        "        # Flattening\n",
        "        x = x.view(x.size(0), -1)\n",
        "\n",
        "        # Fully connected layers\n",
        "        x = F.sigmoid(self.fc1(x))\n",
        "        x = F.sigmoid(self.fc2(x))\n",
        "        x = self.fc3(x)  # No activation function here, will use CrossEntropyLoss later\n",
        "        return x\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8yYNrvus0rZ2",
        "outputId": "d3ae2380-fa03-4368-9bea-46c38a65ddca"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import os\n",
        "import torch\n",
        "from torchvision.datasets import mnist\n",
        "from torch.nn import CrossEntropyLoss\n",
        "from torch.optim import Adam  # Import Adam\n",
        "from torch.utils.data import DataLoader\n",
        "from torchvision.transforms import ToTensor\n",
        "\n",
        "def normalize_img(image, label):\n",
        "  return torch.round(image), label\n",
        "\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "batch_size = 256\n",
        "train_dataset = mnist.MNIST(root='./train', train=True, transform=ToTensor(), download=True)\n",
        "test_dataset = mnist.MNIST(root='./test', train=False, transform=ToTensor(), download=True)\n",
        "train_loader = DataLoader(train_dataset, batch_size=batch_size)\n",
        "test_loader = DataLoader(test_dataset, batch_size=batch_size)\n",
        "model = LeNet().to(device)\n",
        "adam = Adam(model.parameters())  # Using Adam with a learning rate of 1e-3\n",
        "loss_fn = CrossEntropyLoss()\n",
        "all_epoch = 25\n",
        "prev_acc = 0\n",
        "for current_epoch in range(all_epoch):\n",
        "    model.train()\n",
        "    for idx, (train_x, train_label) in enumerate(train_loader):\n",
        "        train_x = train_x.to(device)\n",
        "        # normalize the image to 0 or 1 to reflect the inputs from the drawing board\n",
        "        train_x = train_x.round()\n",
        "        train_label = train_label.to(device)\n",
        "        adam.zero_grad()  # Use adam optimizer\n",
        "        predict_y = model(train_x.float())\n",
        "        loss = loss_fn(predict_y, train_label.long())\n",
        "        loss.backward()\n",
        "        adam.step()  # Use adam optimizer\n",
        "    all_correct_num = 0\n",
        "    all_sample_num = 0\n",
        "    model.eval()\n",
        "\n",
        "    for idx, (test_x, test_label) in enumerate(test_loader):\n",
        "        test_x = test_x.to(device)\n",
        "         # normalize the image to 0 or 1 to reflect the inputs from the drawing board\n",
        "        test_x = test_x.round()\n",
        "        test_label = test_label.to(device)\n",
        "        predict_y = model(test_x.float()).detach()\n",
        "        predict_y = torch.argmax(predict_y, dim=-1)\n",
        "        current_correct_num = predict_y == test_label\n",
        "        all_correct_num += np.sum(current_correct_num.to('cpu').numpy(), axis=-1)\n",
        "        all_sample_num += current_correct_num.shape[0]\n",
        "    acc = all_correct_num / all_sample_num\n",
        "    print('test accuracy: {:.3f}'.format(acc), flush=True)\n",
        "    if not os.path.isdir(\"models\"):\n",
        "        os.mkdir(\"models\")\n",
        "    torch.save(model, 'models/mnist_{:.3f}.pkl'.format(acc))\n",
        "    prev_acc = acc\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "from matplotlib import colors\n",
        "from scipy import stats\n",
        "n_bins = 100\n",
        "# whats the smallest parameter in this model ?\n",
        "minimum_abs_val = 100000\n",
        "params_abs = []\n",
        "for param in model.parameters():\n",
        "  if param.abs().min() < minimum_abs_val:\n",
        "    minimum_abs_val = param.abs().min()\n",
        "  params_abs.extend(param.abs().detach().numpy().flatten().tolist())\n",
        "\n",
        "print(minimum_abs_val)\n",
        "\n",
        "xx = np.linspace(0, 0.2, 1000)\n",
        "\n",
        "kde = stats.gaussian_kde(params_abs)\n",
        "fig, ax = plt.subplots(figsize = (6,4))\n",
        "\n",
        "ax.set_xlim(0, 0.2)\n",
        "\n",
        "# N is the count in each bin, bins is the lower-limit of the bin\n",
        "N, bins, patches = ax.hist(params_abs, bins=n_bins, density = True, alpha =0.65)\n",
        "ax.plot(xx, kde(xx))\n",
        "\n",
        "\n",
        "ax.set_ylim(0, 14)\n",
        "ax.set_yticklabels([])\n",
        "ax.set_ylabel(\"\")\n",
        "plt.style.use(\"bmh\")\n",
        "ax.set_xticks([1.0/128.0, 1.0/32.0, 1.0/16.0])\n",
        "ax.set_xticklabels([\"1/2^7\", \"1/2^6\", \"1/2^5\"])\n",
        "ax.grid(False)\n",
        "# Calculate percentiles\n",
        "quant_5, quant_25, quant_50, quant_75, quant_95 = np.quantile(params_abs, 0.05), np.quantile(params_abs, 0.25), np.quantile(params_abs, 0.5), np.quantile(params_abs, 0.75), np.quantile(params_abs, 0.95)\n",
        "\n",
        "# [quantile, opacity, length]\n",
        "quants = [[quant_5, 0.6, 0.16], [quant_25, 0.8, 0.26], [quant_50, 1, 0.36],  [quant_75, 0.8, 0.46], [quant_95, 0.6, 0.56]]\n",
        "\n",
        "# Plot the lines with a loop\n",
        "for i in quants:\n",
        "    ax.axvline(i[0], alpha = i[1], ymax = i[2], linestyle = \":\")\n",
        "\n",
        "# Annotations\n",
        "ax.text(quant_5, 14 * 0.17, \"5th\", size = 10, alpha = 0.8)\n",
        "ax.text(quant_25, 14 * 0.27, \"25th\", size = 10, alpha = 0.85)\n",
        "ax.text(quant_50, 14 * 0.37, \"50th\", size = 10, alpha = 1)\n",
        "ax.text(quant_75, 14 * 0.47, \"75th\", size = 10, alpha = 0.85)\n",
        "ax.text(quant_95, 14 * 0.57, \"95th Percentile\", size = 10, alpha =.8)\n",
        "\n",
        "ax.set_title(\"Absolute value of parameters\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tmGo25eb0rZ3"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "model_path = os.path.join('network_lenet.onnx')\n",
        "compiled_model_path = os.path.join('network.compiled')\n",
        "pk_path = os.path.join('key.pk')\n",
        "vk_path = os.path.join('key.vk')\n",
        "settings_path = os.path.join('settings.json')\n",
        "witness_path = os.path.join('witness.json')\n",
        "data_path = os.path.join('input.json')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RfFjaXrM0rZ3"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import json\n",
        "\n",
        "model.eval()  # Set the model to evaluation mode\n",
        "\n",
        "# # Fetch a single data point from the train_dataset\n",
        "# # Ensure train_dataset is already loaded and accessible\n",
        "train_data_point, _ = next(iter(train_dataset))\n",
        "train_data_point = train_data_point.unsqueeze(0)  # Add a batch dimension\n",
        "\n",
        "# Verify the device (CPU or CUDA) and transfer the data point to the same device as the model\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "train_data_point = train_data_point.to(device)\n",
        "\n",
        "# # Export the model to ONNX format\n",
        "torch.onnx.export(model, train_data_point, model_path, export_params=True, opset_version=12, do_constant_folding=True, input_names=['input_0'], output_names=['output'])\n",
        "\n",
        "# Convert the tensor to numpy array and reshape it for JSON serialization\n",
        "x = train_data_point.cpu().detach().numpy().reshape([-1]).tolist()\n",
        "data = {'input_data': [x]}\n",
        "with open('input.json', 'w') as f:\n",
        "    json.dump(data, f)\n",
        "\n",
        "print(f\"Model exported to {model_path} and input data saved to input.json\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dS-yXte30rZ3"
      },
      "outputs": [],
      "source": [
        "import ezkl\n",
        "\n",
        "run_args = ezkl.PyRunArgs()\n",
        "run_args.input_visibility = \"private\"\n",
        "run_args.param_visibility = \"fixed\"\n",
        "run_args.output_visibility = \"public\"\n",
        "run_args.num_inner_cols = 2\n",
        "run_args.variables = [(\"batch_size\", 1)]\n",
        "\n",
        "# Capture set of data points\n",
        "num_data_points = 8\n",
        "\n",
        "# Fetch 30 data points from the train_dataset\n",
        "data_points = []\n",
        "for i, (data_point, _) in enumerate(train_dataset):\n",
        "    if i >= num_data_points:\n",
        "        break\n",
        "    data_points.append(data_point)\n",
        "\n",
        "# Stack the data points to create a batch\n",
        "train_data_batch = torch.stack(data_points)\n",
        "\n",
        "# Add a batch dimension if not already present\n",
        "if train_data_batch.dim() == 3:\n",
        "    train_data_batch = train_data_batch.unsqueeze(0)\n",
        "\n",
        "x = train_data_batch.cpu().detach().numpy().reshape([-1]).tolist()\n",
        "\n",
        "data = dict(input_data = [x])\n",
        "\n",
        "cal_path = os.path.join('cal_data.json')\n",
        "\n",
        "# Serialize data into file:\n",
        "json.dump( data, open(cal_path, 'w' ))\n",
        "\n",
        "!RUST_LOG=trace\n",
        "# TODO: Dictionary outputs\n",
        "res = ezkl.gen_settings(model_path, settings_path, py_run_args=run_args)\n",
        "assert res == True\n",
        "\n",
        "res = ezkl.calibrate_settings(cal_path, model_path, settings_path, \"resources\", scales=[2,7])\n",
        "assert res == True"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pGb1nj720rZ3"
      },
      "outputs": [],
      "source": [
        "res = ezkl.compile_circuit(model_path, compiled_model_path, settings_path)\n",
        "assert res == True"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "umpLxVAI0rZ3"
      },
      "outputs": [],
      "source": [
        "# srs path\n",
        "res = ezkl.get_srs(settings_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "syLy2Kt90rZ3"
      },
      "outputs": [],
      "source": [
        "# now generate the witness file\n",
        "witness_path = \"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": {
        "id": "65u_ObBQ0rZ4"
      },
      "outputs": [],
      "source": [
        "res = ezkl.mock(witness_path, compiled_model_path)\n",
        "assert res == True"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c8GIoMD40rZ4"
      },
      "outputs": [],
      "source": [
        "\n",
        "# HERE WE SETUP THE CIRCUIT PARAMS\n",
        "# WE GOT KEYS\n",
        "# WE GOT CIRCUIT PARAMETERS\n",
        "# EVERYTHING ANYONE HAS EVER NEEDED FOR ZK\n",
        "\n",
        "res = ezkl.setup(\n",
        "        compiled_model_path,\n",
        "        vk_path,\n",
        "        pk_path,\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": "code",
      "execution_count": null,
      "metadata": {
        "id": "vkIutAhR0rZ4"
      },
      "outputs": [],
      "source": [
        "# GENERATE A PROOF\n",
        "\n",
        "\n",
        "proof_path = os.path.join('test.pf')\n",
        "\n",
        "res = ezkl.prove(\n",
        "        witness_path,\n",
        "        compiled_model_path,\n",
        "        pk_path,\n",
        "        proof_path,\n",
        "        \"single\",\n",
        "    )\n",
        "\n",
        "print(res)\n",
        "assert os.path.isfile(proof_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KmpQBZUT0rZ4"
      },
      "outputs": [],
      "source": [
        "# VERIFY IT\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": "CZ2CYQdm0rZ4"
      },
      "source": [
        "We can now create an EVM / `.sol` verifier that can be deployed on chain to verify submitted proofs using a view function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jASymUWQ0rZ4"
      },
      "outputs": [],
      "source": [
        "\n",
        "abi_path = 'test.abi'\n",
        "sol_code_path = 'test_1.sol'\n",
        "\n",
        "res = ezkl.create_evm_verifier(\n",
        "        vk_path,\n",
        "        settings_path,\n",
        "        sol_code_path,\n",
        "        abi_path,\n",
        "    )\n",
        "assert res == True"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-icN2yw60rZ5"
      },
      "source": [
        "## Verify on the evm"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sTLsLwaC0rZ5"
      },
      "outputs": [],
      "source": [
        "# Make sure anvil is running locally first\n",
        "# run with $ anvil -p 3030\n",
        "# we use the default anvil node here\n",
        "import json\n",
        "\n",
        "address_path = os.path.join(\"address.json\")\n",
        "\n",
        "res = ezkl.deploy_evm(\n",
        "    address_path,\n",
        "    sol_code_path,\n",
        "    'http://127.0.0.1:3030'\n",
        ")\n",
        "\n",
        "assert res == True\n",
        "\n",
        "with open(address_path, 'r') as file:\n",
        "    addr = file.read().rstrip()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PrSyQK4B0rZ5"
      },
      "outputs": [],
      "source": [
        "# make sure anvil is running locally\n",
        "# $ anvil -p 3030\n",
        "\n",
        "res = ezkl.verify_evm(\n",
        "    addr,\n",
        "    proof_path,\n",
        "    \"http://127.0.0.1:3030\"\n",
        ")\n",
        "assert res == True"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "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
}
