{
    "cells": [
        {
            "attachments": {
                "image-3.png": {
                    "image/png": "[encoded data, 144816 characters]"
                },
                "image.png": {
                    "image/png": "[encoded data, 13700 characters]"
                }
            },
            "cell_type": "markdown",
            "id": "cf69bb3f-94e6-4dba-92cd-ce08df117d67",
            "metadata": {},
            "source": [
                "## LightBGM models\n",
                "\n",
                "\n",
                "LightBGM based models are slightly finicky to get into a suitable onnx format. By default most tree based models will export into something that looks like this: \n",
                "\n",
                "\n",
                "![image.png](attachment:image.png)\n",
                "\n",
                "\n",
                "Processing such nodes can be difficult and error prone. It would be much better if the operations of the tree were represented as a proper graph, possibly ... like this: \n",
                "\n",
                "\n",
                "![image-3.png](attachment:image-3.png)\n",
                "\n",
                "\n",
                "This notebook showcases how to do that using the `hummingbird` python package ! "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "a60b90d6",
            "metadata": {},
            "outputs": [],
            "source": [
                "!python -m pip install hummingbird_ml"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "95613ee9",
            "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\", \"onnx\"])\n",
                "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"hummingbird-ml\"])\n",
                "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"lightgbm\"])\n",
                "\n",
                "# rely on local installation of ezkl if the notebook is not in colab\n",
                "except:\n",
                "    pass\n",
                "\n",
                "\n",
                "# here we create and (potentially train a model)\n",
                "\n",
                "# make sure you have the dependencies required here already installed\n",
                "import json\n",
                "import numpy as np\n",
                "from sklearn.datasets import load_iris\n",
                "from sklearn.model_selection import train_test_split\n",
                "from lightgbm import LGBMClassifier as Gbc\n",
                "import torch\n",
                "import ezkl\n",
                "import os\n",
                "from torch import nn\n",
                "from hummingbird.ml import convert\n",
                "\n",
                "NUM_CLASSES = 3\n",
                "\n",
                "iris = load_iris()\n",
                "X, y = iris.data, iris.target\n",
                "X = X.astype(np.float32)\n",
                "X_train, X_test, y_train, y_test = train_test_split(X, y)\n",
                "clr = Gbc(n_estimators=12)\n",
                "clr.fit(X_train, y_train)\n",
                "\n",
                "# convert to torch\n",
                "\n",
                "\n",
                "torch_gbt = convert(clr, 'torch', X_test[:1])\n",
                "\n",
                "print(torch_gbt)\n",
                "# assert predictions from torch are = to sklearn\n",
                "diffs = []\n",
                "\n",
                "for i in range(len(X_test)):\n",
                "    torch_pred = torch_gbt.predict(torch.tensor(X_test[i].reshape(1, -1)))\n",
                "    sk_pred = clr.predict(X_test[i].reshape(1, -1))\n",
                "    diffs.append(torch_pred != sk_pred[0])\n",
                "\n",
                "print(\"num diff: \", sum(diffs))\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": [
                "# !!!!!!!!!!!!!!!!! This cell will flash a warning about onnx runtime compat but it is fine !!!!!!!!!!!!!!!!!!!!!\n",
                "\n",
                "\n",
                "# export to onnx format\n",
                "\n",
                "\n",
                "# Input to the model\n",
                "shape = X_train.shape[1:]\n",
                "x = torch.rand(1, *shape, requires_grad=False)\n",
                "torch_out = torch_gbt.predict(x)\n",
                "# Export the model\n",
                "torch.onnx.export(torch_gbt.model,               # model being run\n",
                "                  # model input (or a tuple for multiple inputs)\n",
                "                  x,\n",
                "                  # where to save the model (can be a file or file-like object)\n",
                "                  \"network.onnx\",\n",
                "                  export_params=True,        # store the trained parameter weights inside the model file\n",
                "                  opset_version=18,          # the ONNX version to export the model to\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",
                "d = ((x).detach().numpy()).reshape([-1]).tolist()\n",
                "\n",
                "data = dict(input_shapes=[shape],\n",
                "            input_data=[d],\n",
                "            output_data=[(o).reshape([-1]).tolist() for o in torch_out])\n",
                "\n",
                "# Serialize data into file:\n",
                "json.dump(data, open(\"input.json\", 'w'))\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "d5e374a2",
            "metadata": {},
            "outputs": [],
            "source": [
                "run_args = ezkl.PyRunArgs()\n",
                "run_args.variables = [(\"batch_size\", 1)]\n",
                "\n",
                "# TODO: Dictionary outputs\n",
                "res = ezkl.gen_settings(model_path, settings_path, py_run_args=run_args)\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 = (torch.randn(20, *shape).detach().numpy()).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": "3aa4f090",
            "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)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "18c8b7c7",
            "metadata": {},
            "outputs": [],
            "source": [
                "# now generate the witness file \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",
                "\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",
                "\n",
                "res = ezkl.verify(\n",
                "        proof_path,\n",
                "        settings_path,\n",
                "        vk_path,\n",
                "        \n",
                "    )\n",
                "\n",
                "assert res == True\n",
                "print(\"verified\")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "1f50a8d0",
            "metadata": {},
            "outputs": [],
            "source": [
                "# VERIFY IT\n",
                "\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.15"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}
