{
    "cells": [
        {
            "attachments": {},
            "cell_type": "markdown",
            "id": "cf69bb3f-94e6-4dba-92cd-ce08df117d67",
            "metadata": {},
            "source": [
                "## EZKL Jupyter Notebook Demo with Keras\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "a27b0cd9",
            "metadata": {},
            "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\", \"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",
                "from keras.models import Sequential\n",
                "from keras.layers import Dense, Dropout, Activation, Flatten\n",
                "from keras.layers import Convolution2D, MaxPooling2D\n",
                "import logging\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.DEBUG)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "95613ee9",
            "metadata": {},
            "outputs": [],
            "source": [
                "\n",
                "\n",
                "# Defines the model\n",
                "# we got convs, we got relu, we got linear layers, max pooling layers etc... \n",
                "# What else could one want ????\n",
                "\n",
                "model = Sequential()\n",
                "model.add(Convolution2D(2, (3,3), activation='relu', input_shape=(28,28,1)))\n",
                "model.add(Convolution2D(2, (3,3), activation='relu'))\n",
                "model.add(MaxPooling2D(pool_size=(2,2)))\n",
                "model.add(Dropout(0.25))\n",
                "model.add(Flatten())\n",
                "model.add(Dense(128, activation='relu'))\n",
                "model.add(Dropout(0.5))\n",
                "model.add(Dense(10, activation='softmax'))\n",
                "\n",
                "\n",
                "# Train the model as you like here (skipped for brevity)\n",
                "\n",
                "\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "b37637c4",
            "metadata": {},
            "outputs": [],
            "source": [
                "model_path = os.path.join('network.onnx')\n",
                "compiled_model_path = os.path.join('network.compiled')\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')"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "82db373a",
            "metadata": {},
            "outputs": [],
            "source": [
                "\n",
                "import numpy as np\n",
                "import tf2onnx\n",
                "import tensorflow as tf\n",
                "\n",
                "\n",
                "shape = [1, 28, 28]\n",
                "# After training, export to onnx (network.onnx) and create a data file (input.json)\n",
                "x = 0.1*np.random.rand(1,*shape)\n",
                "\n",
                "spec = tf.TensorSpec([1, 28, 28, 1], tf.float32, name='input_0')\n",
                "\n",
                "\n",
                "tf2onnx.convert.from_keras(model, input_signature=[spec], inputs_as_nchw=['input_0'], opset=12, output_path=model_path)\n",
                "\n",
                "data_array = x.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"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "d5e374a2",
            "metadata": {},
            "outputs": [],
            "source": [
                "!RUST_LOG=trace\n",
                "\n",
                "# TODO: Dictionary outputs\n",
                "res = ezkl.gen_settings(model_path, settings_path)\n",
                "assert res == True\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "cal_path = os.path.join(\"calibration.json\")\n",
                "\n",
                "data_array = (0.1*np.random.rand(20,*shape)).reshape([-1]).tolist()\n",
                "\n",
                "data = dict(input_data = [data_array])\n",
                "\n",
                "# Serialize data into file:\n",
                "json.dump(data, open(cal_path, 'w'))\n",
                "\n",
                "\n",
                "res = ezkl.calibrate_settings(data_path, model_path, settings_path, \"resources\")\n",
                "assert res == True\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "b6e051d5",
            "metadata": {},
            "outputs": [],
            "source": [
                "res = ezkl.compile_circuit(model_path, compiled_model_path, settings_path)\n",
                "assert res == True"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "8b74dcee",
            "metadata": {},
            "outputs": [],
            "source": [
                "# srs path\n",
                "res = ezkl.get_srs(settings_path = settings_path)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "18c8b7c7",
            "metadata": {},
            "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,
            "id": "b1c561a8",
            "metadata": {},
            "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",
                "\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,
            "id": "c384cbc8",
            "metadata": {},
            "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",
                "        \n",
                "        \"single\",\n",
                "    )\n",
                "\n",
                "print(res)\n",
                "assert os.path.isfile(proof_path)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "76f00d41",
            "metadata": {},
            "outputs": [],
            "source": [
                "# VERIFY IT\n",
                "res = ezkl.verify(\n",
                "        proof_path,\n",
                "        settings_path,\n",
                "        vk_path,\n",
                "        \n",
                "    )\n",
                "\n",
                "assert res == True\n",
                "print(\"verified\")"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3 (ipykernel)",
            "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.13"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}
