zk-stats-demo/public/assets/template.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install torch zkstats==0.1.10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "\n",
    "from zkstats.core import (\n",
    "    prover_gen_settings,\n",
    "    prover_gen_proof,\n",
    "    setup,\n",
    "    verifier_verify,\n",
    ")\n",
    "\n",
    "cwd = os.getcwd()\n",
    "\n",
    "# FIXME: fill this in with the path to your data\n",
    "data_path = f\"{cwd}/data.json\"\n",
    "\n",
    "# Paths to the output files\n",
    "output_dir = f\"{cwd}/out\"\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "model_onnx_path = f\"{output_dir}/model.onnx\"\n",
    "compiled_model_path = f\"{output_dir}/model.compiled\"\n",
    "\n",
    "pk_path = f\"{output_dir}/model.pk\"\n",
    "vk_path = f\"{output_dir}/model.vk\"\n",
    "proof_path = f\"{output_dir}/model.pf\"\n",
    "settings_path = f\"{output_dir}/settings.json\"\n",
    "witness_path = f\"{output_dir}/witness.json\"\n",
    "comb_data_path = f\"{output_dir}/comb_data.json\"\n",
    "precal_witness_path = f\"{output_dir}/precal_witness.json\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data provided shared by the data provider beforehand\n",
    "\n",
    "def get_data_shape(data_path: str) -> dict[str, int]:\n",
    "    \"\"\"\n",
    "    Get the shape of the data from the data file.\n",
    "\n",
    "    Parameters:\n",
    "    - data_path (str): The path to the data file.\n",
    "\n",
    "    Returns:\n",
    "    - shape_info (dict): A dictionary where keys are column names and values are the number of elements (shape).\n",
    "    \"\"\"\n",
    "    with open(data_path, 'r') as f:\n",
    "        data = json.load(f)\n",
    "    shape_info = {col: len(data[col]) for col in data.keys()}\n",
    "    return shape_info\n",
    "\n",
    "data_shape = get_data_shape(data_path)\n",
    "print(data_shape)\n",
    "data_commitment_path = f\"{output_dir}/data_commitment.json\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## User-defined Computation\n",
    "\n",
    "A computation should be of type `TComputation`. For example, the following code snippet defines a computation that computes the sum of the private data.\n",
    "\n",
    "```python\n",
    "def computation(state: State, args: Args):\n",
    "    x = args[\"x\"]\n",
    "    y = args[\"y\"]\n",
    "    return state.mean(x), state.mean(y)\n",
    "```\n",
    "\n",
    "FIXME: The following code snippet is entirely from the user. You MUST check\n",
    "1. the code only performs zkstats-related operations.\n",
    "2. the computation must not leak any information about the private data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is just a dummy computation. Replace it with user's computation\n",
    "import torch\n",
    "from zkstats.computation import State, Args\n",
    "\n",
    "def computation(state: State, args: Args):\n",
    "    x = args[\"x\"]\n",
    "    y = args[\"y\"]\n",
    "    return state.mean(x), state.mean(y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate settings and setup with user's computation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from zkstats.computation import computation_to_model\n",
    "\n",
    "used_columns, state, model = computation_to_model(computation, precal_witness_path, data_shape, isProver=True)\n",
    "prover_gen_settings(\n",
    "    data_path,\n",
    "    used_columns,\n",
    "    comb_data_path,\n",
    "    model,\n",
    "    model_onnx_path,\n",
    "    [7],\n",
    "    \"resources\",\n",
    "    settings_path,\n",
    ")\n",
    "\n",
    "# Determine which srs to use with the logrows in the settings.json\n",
    "with open(settings_path, \"r\") as f:\n",
    "    settings = json.load(f)\n",
    "logrows = int(settings[\"run_args\"][\"logrows\"])\n",
    "srs_path = f'~/.ezkl/srs/kzg{logrows}.srs'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "setup(model_onnx_path, compiled_model_path, settings_path, vk_path, pk_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate proof with your data and user's computation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prover_gen_proof(\n",
    "    model_onnx_path,\n",
    "    comb_data_path,\n",
    "    witness_path,\n",
    "    compiled_model_path,\n",
    "    settings_path,\n",
    "    proof_path,\n",
    "    pk_path,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Verify the proof to ensure it is correct\n",
    "NOTE: The following section is to illustrate what should be done on the user (data consumer) side. This step is not required by the data provider."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# verify the proof\n",
    "res = verifier_verify(proof_path, settings_path, vk_path, used_columns, data_commitment_path)\n",
    "print(f\"!@# res: {res}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Print the file paths. You should share the following files back to the user for them to verify the proof. You **SHOULD NOT** share more files otherwise data might be leaked."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Model onnx:\\t\\t\", model_onnx_path)\n",
    "print(\"Settings:\\t\\t\", settings_path)\n",
    "print(\"Proof:\\t\\t\\t\", proof_path)\n",
    "print(\"Verification key:\\t\", vk_path)\n",
    "print(\"Srs path:\\t\\t\", srs_path)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}