examples/notebooks/mean_postgres.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Mean of ERC20 transfer amounts\n",
        "\n",
        "This notebook shows how to calculate the mean of ERC20 transfer amounts, pulling data in from a Postgres database. First we install and get the necessary libraries running. \n",
        "The first of which is [shovel](https://indexsupply.com/shovel/docs/#getting-started), which is a library that allows us to pull data from the Ethereum blockchain into a Postgres database.\n",
        "\n",
        "Make sure you install postgres if needed https://indexsupply.com/shovel/docs/#getting-started. \n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "import getpass\n",
        "import json\n",
        "import time\n",
        "import subprocess\n",
        "\n",
        "# swap out for the relevant linux/amd64, darwin/arm64, darwin/amd64, windows/amd64\n",
        "os.system(\"curl -LO https://indexsupply.net/bin/1.0/linux/amd64/shovel\")\n",
        "os.system(\"chmod +x shovel\")\n",
        "\n",
        "\n",
        "os.environ[\"PG_URL\"] = \"postgres://\" + getpass.getuser() + \":@localhost:5432/shovel\"\n",
        "\n",
        "# create a config.json file with the following contents\n",
        "config = {\n",
        "  \"pg_url\": \"$PG_URL\",\n",
        "  \"eth_sources\": [\n",
        "    {\"name\": \"mainnet\", \"chain_id\": 1, \"url\": \"https://ethereum-rpc.publicnode.com\"},\n",
        "    {\"name\": \"base\", \"chain_id\": 8453, \"url\": \"https://base-rpc.publicnode.com\"}\n",
        "  ],\n",
        "  \"integrations\": [{\n",
        "    \"name\": \"usdc_transfer\",\n",
        "    \"enabled\": True,\n",
        "    \"sources\": [{\"name\": \"mainnet\"}, {\"name\": \"base\"}],\n",
        "    \"table\": {\n",
        "      \"name\": \"usdc\",\n",
        "      \"columns\": [\n",
        "        {\"name\": \"log_addr\",  \"type\": \"bytea\"},\n",
        "        {\"name\": \"block_num\", \"type\": \"numeric\"},\n",
        "        {\"name\": \"f\",         \"type\": \"bytea\"},\n",
        "        {\"name\": \"t\",         \"type\": \"bytea\"},\n",
        "        {\"name\": \"v\",         \"type\": \"numeric\"}\n",
        "      ]\n",
        "    },\n",
        "    \"block\": [\n",
        "      {\"name\": \"block_num\", \"column\": \"block_num\"},\n",
        "      {\n",
        "        \"name\": \"log_addr\",\n",
        "        \"column\": \"log_addr\",\n",
        "        \"filter_op\": \"contains\",\n",
        "        \"filter_arg\": [\n",
        "          \"a0b86991c6218b36c1d19d4a2e9eb0ce3606eb48\",\n",
        "          \"833589fCD6eDb6E08f4c7C32D4f71b54bdA02913\"\n",
        "        ]\n",
        "      }\n",
        "    ],\n",
        "    \"event\": {\n",
        "      \"name\": \"Transfer\",\n",
        "      \"type\": \"event\",\n",
        "      \"anonymous\": False,\n",
        "      \"inputs\": [\n",
        "        {\"indexed\": True,  \"name\": \"from\",  \"type\": \"address\", \"column\": \"f\"},\n",
        "        {\"indexed\": True,  \"name\": \"to\",    \"type\": \"address\", \"column\": \"t\"},\n",
        "        {\"indexed\": False, \"name\": \"value\", \"type\": \"uint256\", \"column\": \"v\"}\n",
        "      ]\n",
        "    }\n",
        "  }]\n",
        "}\n",
        "\n",
        "# write the config to a file\n",
        "with open(\"config.json\", \"w\") as f:\n",
        "  f.write(json.dumps(config))\n",
        "\n",
        "\n",
        "# print the two env variables\n",
        "os.system(\"echo $PG_URL\")\n",
        "\n",
        "os.system(\"createdb -h localhost -p 5432 shovel\")\n",
        "\n",
        "os.system(\"echo shovel is now installed. starting:\")\n",
        "\n",
        "command = [\"./shovel\", \"-config\", \"config.json\"]\n",
        "proc = subprocess.Popen(command)\n",
        "\n",
        "os.system(\"echo shovel started.\")\n",
        "\n",
        "time.sleep(10)\n",
        "\n",
        "# after we've fetched some data -- kill the process\n",
        "proc.terminate()\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2wIAHwqH2_mo"
      },
      "source": [
        "**Import Dependencies**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9Byiv2Nc2MsK"
      },
      "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",
        "\n",
        "# rely on local installation of ezkl if the notebook is not in colab\n",
        "except:\n",
        "    pass\n",
        "\n",
        "import ezkl\n",
        "import torch\n",
        "import datetime\n",
        "import pandas as pd\n",
        "import requests\n",
        "import json\n",
        "import os\n",
        "\n",
        "import logging\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)\n",
        "\n",
        "print(\"ezkl version: \", ezkl.__version__)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "osjj-0Ta3E8O"
      },
      "source": [
        "**Create Computational Graph**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "x1vl9ZXF3EEW",
        "outputId": "bda21d02-fe5f-4fb2-8106-f51a8e2e67aa"
      },
      "outputs": [],
      "source": [
        "from torch import nn\n",
        "import torch\n",
        "\n",
        "\n",
        "class Model(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Model, self).__init__()\n",
        "\n",
        "    # x is a time series \n",
        "    def forward(self, x):\n",
        "        return [torch.mean(x)]\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "circuit = Model()\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "x = 0.1*torch.rand(1,*[1,5], requires_grad=True)\n",
        "\n",
        "# # print(torch.__version__)\n",
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "print(device)\n",
        "\n",
        "circuit.to(device)\n",
        "\n",
        "# Flips the neural net into inference mode\n",
        "circuit.eval()\n",
        "\n",
        "# Export the model\n",
        "torch.onnx.export(circuit,               # model being run\n",
        "                      x,                   # model input (or a tuple for multiple inputs)\n",
        "                      \"lol.onnx\",            # 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=11,          # 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",
        "# export(circuit, input_shape=[1, 20])\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E3qCeX-X5xqd"
      },
      "source": [
        "**Set Data Source and Get Data**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6RAMplxk5xPk",
        "outputId": "bd2158fe-0c00-44fd-e632-6a3f70cdb7c9"
      },
      "outputs": [],
      "source": [
        "import getpass\n",
        "# make an input.json file from the df above\n",
        "input_filename = os.path.join('input.json')\n",
        "\n",
        "pg_input_file = dict(input_data = {\n",
        "  \"host\": \"localhost\",\n",
        "  # make sure you replace this with your own username\n",
        "  \"user\": getpass.getuser(),\n",
        "  \"dbname\": \"shovel\",\n",
        "  \"password\": \"\",\n",
        "  \"query\": \"SELECT v FROM usdc ORDER BY block_num DESC LIMIT 5\",\n",
        "  \"port\": \"5432\",\n",
        "})\n",
        "\n",
        "json_formatted_str = json.dumps(pg_input_file, indent=2)\n",
        "print(json_formatted_str)\n",
        "\n",
        "\n",
        "    # Serialize data into file:\n",
        "json.dump(pg_input_file, open(input_filename, 'w' ))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# this corresponds to 4 batches\n",
        "calibration_filename = os.path.join('calibration.json')\n",
        "\n",
        "pg_cal_file = dict(input_data = {\n",
        "  \"host\": \"localhost\",\n",
        "  # make sure you replace this with your own username\n",
        "  \"user\": getpass.getuser(),\n",
        "  \"dbname\": \"shovel\",\n",
        "  \"password\": \"\",\n",
        "  \"query\": \"SELECT v FROM usdc ORDER BY block_num DESC LIMIT 20\",\n",
        "  \"port\": \"5432\",\n",
        "})\n",
        "\n",
        "    # Serialize data into file:\n",
        "json.dump( pg_cal_file, open(calibration_filename, 'w' ))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eLJ7oirQ_HQR"
      },
      "source": [
        "**EZKL Workflow**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rNw0C9QL6W88"
      },
      "outputs": [],
      "source": [
        "import subprocess\n",
        "import os\n",
        "\n",
        "onnx_filename = os.path.join('lol.onnx')\n",
        "compiled_filename = os.path.join('lol.compiled')\n",
        "settings_filename = os.path.join('settings.json')\n",
        "\n",
        "# Generate settings using ezkl\n",
        "res = ezkl.gen_settings(onnx_filename, settings_filename)\n",
        "\n",
        "assert res == True\n",
        "\n",
        "res = await ezkl.calibrate_settings(input_filename, onnx_filename, settings_filename, \"resources\")\n",
        "\n",
        "assert res == True"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "\n",
        "ezkl.compile_circuit(onnx_filename, compiled_filename, settings_filename)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4MmE9SX66_Il",
        "outputId": "16403639-66a4-4280-ac7f-6966b75de5a3"
      },
      "outputs": [],
      "source": [
        "# generate settings\n",
        "\n",
        "\n",
        "# show the settings.json\n",
        "with open(\"settings.json\") as f:\n",
        "  data = json.load(f)\n",
        "  json_formatted_str = json.dumps(data, indent=2)\n",
        "\n",
        "  print(json_formatted_str)\n",
        "\n",
        "assert os.path.exists(\"settings.json\")\n",
        "assert os.path.exists(\"input.json\")\n",
        "assert os.path.exists(\"lol.onnx\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fULvvnK7_CMb"
      },
      "outputs": [],
      "source": [
        "pk_path = os.path.join('test.pk')\n",
        "vk_path = os.path.join('test.vk')\n",
        "\n",
        "\n",
        "# setup the proof\n",
        "res = ezkl.setup(\n",
        "        compiled_filename,\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_filename)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "witness_path = \"witness.json\"\n",
        "\n",
        "# generate the witness\n",
        "res = await ezkl.gen_witness(\n",
        "        input_filename,\n",
        "        compiled_filename,\n",
        "        witness_path\n",
        "    )\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Oog3j6Kd-Wed",
        "outputId": "5839d0c1-5b43-476e-c2f8-6707de562260"
      },
      "outputs": [],
      "source": [
        "# prove the zk circuit\n",
        "# GENERATE A PROOF\n",
        "proof_path = os.path.join('test.pf')\n",
        "\n",
        "\n",
        "proof = ezkl.prove(\n",
        "        witness_path,\n",
        "        compiled_filename,\n",
        "        pk_path,\n",
        "        proof_path,\n",
        "        \"single\"\n",
        "    )\n",
        "\n",
        "\n",
        "print(\"proved\")\n",
        "\n",
        "assert os.path.isfile(proof_path)\n",
        "\n"
      ]
    }
  ],
  "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.12.2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}