{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use Case: Encrypted Inference with ONNX Model\n",
    "\n",
    "In this tutorial, we will demonstrate how we can run encrypted machine learning inference from a model saved as an [ONNX](https://onnx.ai/) file.\n",
    "\n",
    "For simplicity, we will start by prototyping this computation between the different parties locally using the `pm.LocalMooseRuntime`. Then we will execute the same Moose computation over the network with `pm.GrpcMooseRuntime`. You can also check additional gRPC example [here](https://github.com/tf-encrypted/moose/tree/main/pymoose/examples/grpc)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pathlib\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from sklearn.datasets import make_classification\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "from onnxmltools.convert import convert_sklearn\n",
    "from skl2onnx.common import data_types as onnx_dtypes\n",
    "\n",
    "import pymoose as pm\n",
    "\n",
    "random_state = 5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use Case\n",
    "\n",
    "You are a healthcare AI startup who has trained a model to diagnose heart disease. You would like to serve this model to a hospital to help doctors diagnose potential heart disease for their patients. However, the patients' data is too sensitive to be shared with the AI startup. For this reason, you would like to encrypt the patient's data and run the model on it.\n",
    "\n",
    "In this tutorial, we will perform the following steps:\n",
    "- Train a model with [Scikit-Learn](https://scikit-learn.org/stable/)\n",
    "- Convert the trained model to ONNX.\n",
    "- Convert the model from ONNX to a Moose computation.\n",
    "- Run encrypted inference by evaluating the Moose computation.\n",
    "\n",
    "### Training\n",
    "\n",
    "For this tutorial, we use a synthetic dataset. The dataset contains 10 features (`X`). Each record is labeled (`y`) by 0 or 1 (heart disease or not). For the model, we train a logistic regression with [Scikit-Learn](https://scikit-learn.org/stable/). But you could experiment with other models such as `XGBClassifier` from [XGBoost](https://xgboost.readthedocs.io/en/stable/index.html#) or even [Multi-layer Perceptron](https://scikit-learn.org/stable/modules/neural_networks_supervised.html#multi-layer-perceptron).\n",
    "\n",
    "Once the model is trained, you can convert it to ONNX which is a format to represent machine learning models. Since it's a Scikit-Learn model, you can convert it to ONNX with `convert_sklearn` from [ONNXMLTools](https://github.com/onnx/onnxmltools). Use `convert_xgboost` for XGBoost models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_samples = 1000\n",
    "n_features = 10\n",
    "n_classes = 2\n",
    "\n",
    "# Generate synthetic dataset\n",
    "X, y = make_classification(\n",
    "    n_samples=n_samples,\n",
    "    n_features=n_features,\n",
    "    n_classes=n_classes,\n",
    "    random_state=random_state,\n",
    ")\n",
    "\n",
    "# Split dataset between training and testing datasets\n",
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    X, y, test_size=0.2, random_state=random_state\n",
    ")\n",
    "\n",
    "# Train logistic regression\n",
    "lg = LogisticRegression()\n",
    "lg.fit(X_train, y_train)\n",
    "\n",
    "# Convert scikit-learn model to ONNX\n",
    "initial_type = (\"float_input\", onnx_dtypes.FloatTensorType([None, n_features]))\n",
    "onnx_proto = convert_sklearn(lg, initial_types=[initial_type])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert ONNX to Moose Predictor\n",
    "\n",
    "PyMoose provides several predictor classes to translate an ONNX model into a PyMoose DSL program.\n",
    "\n",
    "It currently supports 8 types of model predictors:\n",
    "- `linear_regressor.LinearRegressor`: for models such as linear regression, ridge regression, etc.\n",
    "- `linear_regressor.LinearClassifier`: for models such as logistic regression, classifier using ridge regression, etc.\n",
    "- `tree_ensemble.TreeEnsembleRegressor` for models such as XGBoost regressor, random forest regressor, etc.\n",
    "- `tree_ensemble.TreeEnsembleClassifier`: for models such as XGBoost classifier, random forest classifier, etc.\n",
    "- `multilayer_perceptron_predictor.MLPRegressor`: for multi-layer perceptron regressor models. \n",
    "- `multilayer_perceptron_predictor.MLPClassifier`: for multi-layer perceptron classifier models.\n",
    "- `neural_network_predictor.NeuralNetwork`: for feed forward neural network from PyTorch and TensorFlow.\n",
    "\n",
    "Because the trained model is a logistic regression, we should use the class `linear_regressor.LinearClassifier`. The class has a method `from_onnx` which will parse the ONNX file. More specifically, it will extract the model weights, intercepts and the post transform (e.g sigmoid or softmax etc.). The returned object is callable. When called, it will compute the forward pass of the logistic regression."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor = pm.predictors.LinearClassifier.from_onnx(onnx_proto)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On the `LinearClassifier` object, there is a `host_placements` property. As you can see, when instantiating the object, there are three host placements which have been created automatically: `alice`, `bob` and `carole`. These three players are grouped under the replicated placement to perform the encrypted inference. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "List of host placements: (HostPlacementExpression(name='alice'), HostPlacementExpression(name='bob'), HostPlacementExpression(name='carole'))\n"
     ]
    }
   ],
   "source": [
    "print(\"List of host placements:\", predictor.host_placements)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Moose Computation\n",
    "\n",
    "For this example, `alice` will play the role  of the hospital. \n",
    "\n",
    "The Moose computation below performs the following steps:\n",
    "- Loads patient's data in plaintext from alice's (hospital) storage.\n",
    "- Secret share (encrypts) the patient's data.\n",
    "- Computes logistic regression on secret shared data.\n",
    "- Reveals the prediction only to alice (hospital) and saves it into its storage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "@pm.computation\n",
    "def moose_predictor_computation():\n",
    "    # Alice (the hospital in our use case) load the patients' data in plaintext\n",
    "    # Then the data gets converted from float to fixed-point\n",
    "    with predictor.alice:\n",
    "        x = pm.load(\"x\", dtype=pm.float64)\n",
    "        x_fixed = pm.cast(x, dtype=pm.predictors.predictor_utils.DEFAULT_FIXED_DTYPE)\n",
    "    # The patients' data gets secret shared when moving from host placement\n",
    "    # to replicated placement.\n",
    "    # Then compute the logistic regression on secret shared data\n",
    "    with predictor.replicated:\n",
    "        y_pred = predictor(x_fixed, pm.predictors.predictor_utils.DEFAULT_FIXED_DTYPE)\n",
    "\n",
    "    # The predictions gets revealed only to Alice (the hospital)\n",
    "    # Convert the data from fixed-point to floats and save the data in the storage\n",
    "    with predictor.alice:\n",
    "        y_pred = pm.cast(y_pred, dtype=pm.float64)\n",
    "        y_pred = pm.save(\"y_pred\", y_pred)\n",
    "    return y_pred"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate the computation\n",
    "\n",
    "For simplicity, we will use `LocalMooseRuntime` to locally simulate this computation running across hosts. To do so, we need to provide: the Moose computation, the list of host identities to simulate, and a mapping of the data stored by each simulated host. \n",
    "\n",
    "Since the hospital is represented by `alice`, we will place the patients' data in `alice` storage.\n",
    "\n",
    "Once you have instantiated the `LocalMooseRuntime` with the identities and additional storage mapping and the runtime set as default, you can simply call the Moose computatio to evaluate it. If you prefer, you can also evaluate the computation with `runtime.evaluate_computation(computation=multiparty_correlation, arguments={})`. We can also provide arguments to the computation if needed, but we don't have any in this example. Note that the output of `evaluate_computation` is an empty dictionary, since this function's output operation `pm.save` returns the Unit type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "executive_storage = {\"alice\": {\"x\": X_test}, \"bob\": {}, \"carole\": {}}\n",
    "identities = [plc.name for plc in predictor.host_placements]\n",
    "\n",
    "runtime = pm.LocalMooseRuntime(identities, storage_mapping=executive_storage)\n",
    "runtime.set_default()\n",
    "\n",
    "_ = moose_predictor_computation()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Results\n",
    "\n",
    "Once the computation is done, we can extract the results. The predictions have been stored in alice's storage. We can extract the value from the storage with `read_value_from_storage`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred = runtime.read_value_from_storage(\"alice\", \"y_pred\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this simulated setting, we can validate that the results on encrypted data match the computation on plaintext data. To do so, we compute the logistic regression prediction with [Scikit-Learn](https://scikit-learn.org/stable/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "expected_result = lg.predict_proba(X_test)\n",
    "np.testing.assert_almost_equal(y_pred, expected_result, decimal=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Nice! You were able to compute the inference on patients' data while keeping the data encrypted during the entire process."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run Computation over the Network with gRPC\n",
    "\n",
    "To run the same computation over the network, you need to launch a gRPC worker at the right endpoints for each party (alice, bob and carole)). You can launch the three workers as follow:\n",
    "\n",
    "```\n",
    "cargo run --bin comet -- --identity localhost:50000 --port 50000\n",
    "cargo run --bin comet -- --identity localhost:50001 --port 50001\n",
    "cargo run --bin comet -- --identity localhost:50002 --port 50002\n",
    "```\n",
    "\n",
    "For the data, we will use the numpy files saved in the `data` folder of this tutorial containing the patients data (`x_test.npy`).\n",
    "\n",
    "For the Moose computation, we will use is exact the same computation as the one used with `LocalMooseRuntime` except for the key of the load operations we will provide the actual file path. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "_DATA_DIR = pathlib.Path().parent / \"data\"\n",
    "x_path = str((_DATA_DIR / \"x_test.npy\").resolve())\n",
    "\n",
    "\n",
    "@pm.computation\n",
    "def moose_predictor_computation():\n",
    "    # Alice (the hospital in our use case) load the patients' data in plaintext\n",
    "    # Then the data gets converted from float to fixed-point\n",
    "    with predictor.alice:\n",
    "        x = pm.load(x_path, dtype=pm.float64)\n",
    "        x_fixed = pm.cast(x, dtype=pm.predictors.predictor_utils.DEFAULT_FIXED_DTYPE)\n",
    "    # The patients' data gets secret shared when moving from host placement\n",
    "    # to replicated placement.\n",
    "    # Then compute the logistic regression on secret shared data\n",
    "    with predictor.replicated:\n",
    "        y_pred = predictor(x_fixed, pm.predictors.predictor_utils.DEFAULT_FIXED_DTYPE)\n",
    "\n",
    "    # The predictions gets revealed only to Alice (the hospital)\n",
    "    # Convert the data from fixed-point to floats and return the prediction\n",
    "    # to Alice who's launching the computation. You could also save the result\n",
    "    # to her storage with `pm.save`.\n",
    "    with predictor.alice:\n",
    "        y_pred = pm.cast(y_pred, dtype=pm.float64)\n",
    "\n",
    "    return y_pred"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the runtime, we will use ` pm.GrpcMooseRuntime` this time. As an argument, we need to provide a mapping between the players identity and the gRPC host address. Once you set the runtime as default, you can call the moose computation to compute the encrypted prediction. You could also evaluate the computation with `runtime.evaluate_computation(computation=multiparty_correlation, arguments={})` if you prefer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "skip-execution"
    ]
   },
   "outputs": [],
   "source": [
    "role_map = {\n",
    "    predictor.alice: \"localhost:50000\",\n",
    "    predictor.bob: \"localhost:50001\",\n",
    "    predictor.carole: \"localhost:50002\",\n",
    "}\n",
    "\n",
    "runtime = pm.GrpcMooseRuntime(role_map)\n",
    "runtime.set_default()\n",
    "\n",
    "grpc_y_pred = moose_predictor_computation()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can finally comfirm that we get same predictions when running this computation with gRPC 🎉!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "skip-execution"
    ]
   },
   "outputs": [],
   "source": [
    "np.testing.assert_almost_equal(grpc_y_pred[0][\"output_0\"], y_pred, decimal=2)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "moose",
   "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.10.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "377d2daf6546a2dd25b0fbd54faa498522b36d8a6837c782d0bfb30722f17a46"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}