{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Handwritten Formula (HWF)\n",
    "\n",
    "This notebook shows an implementation of [Handwritten Formula](https://arxiv.org/abs/2006.06649). In this task, handwritten images of decimal formulas and their computed results are given, alongwith a domain knowledge base containing information on how to compute the decimal formula. The task is to recognize the symbols (which can be digits or operators '+', '-', '×', '÷') of handwritten images and accurately determine their results.\n",
    "\n",
    "Intuitively, we first use a machine learning model (learning part) to convert the input images to symbols (we call them pseudo-labels), and then use the knowledge base (reasoning part) to calculate the results of these symbols. Since we do not have ground-truth of the symbols, in Abductive Learning, the reasoning part will leverage domain knowledge and revise the initial symbols yielded by the learning part through abductive reasoning. This process enables us to further update the machine learning model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary libraries and modules\n",
    "import os.path as osp\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "from ablkit.bridge import SimpleBridge\n",
    "from ablkit.data.evaluation import ReasoningMetric, SymbolAccuracy\n",
    "from ablkit.learning import ABLModel, BasicNN\n",
    "from ablkit.reasoning import KBBase, Reasoner\n",
    "from ablkit.utils import ABLLogger, print_log\n",
    "\n",
    "from datasets import get_dataset\n",
    "from models.nn import SymbolNet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Working with Data\n",
    "\n",
    "First, we get the training and testing datasets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = get_dataset(train=True, get_pseudo_label=True)\n",
    "test_data = get_dataset(train=False, get_pseudo_label=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Both `train_data` and `test_data` have the same structures: tuples with three components: X (list where each element is a list of images), gt_pseudo_label (list where each element is a list of symbols, i.e., pseudo-labels) and Y (list where each element is the computed result). The length and structures of datasets are illustrated as follows.\n",
    "\n",
    "Note: ``gt_pseudo_label`` is only used to evaluate the performance of the learning part but not to train the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Both train_data and test_data consist of 3 components: X, gt_pseudo_label, Y\")\n",
    "print()\n",
    "train_X, train_gt_pseudo_label, train_Y = train_data\n",
    "print(\n",
    "    f\"Length of X, gt_pseudo_label, Y in train_data: \"\n",
    "    + f\"{len(train_X)}, {len(train_gt_pseudo_label)}, {len(train_Y)}\"\n",
    ")\n",
    "test_X, test_gt_pseudo_label, test_Y = test_data\n",
    "print(\n",
    "    f\"Length of X, gt_pseudo_label, Y in test_data: \"\n",
    "    + f\"{len(test_X)}, {len(test_gt_pseudo_label)}, {len(test_Y)}\"\n",
    ")\n",
    "print()\n",
    "\n",
    "X_0, gt_pseudo_label_0, Y_0 = train_X[0], train_gt_pseudo_label[0], train_Y[0]\n",
    "print(\n",
    "    f\"X is a {type(train_X).__name__}, \"\n",
    "    + f\"with each element being a {type(X_0).__name__} of {type(X_0[0]).__name__}.\"\n",
    ")\n",
    "print(\n",
    "    f\"gt_pseudo_label is a {type(train_gt_pseudo_label).__name__}, \"\n",
    "    + f\"with each element being a {type(gt_pseudo_label_0).__name__} \"\n",
    "    + f\"of {type(gt_pseudo_label_0[0]).__name__}.\"\n",
    ")\n",
    "print(f\"Y is a {type(train_Y).__name__}, \" + f\"with each element being an {type(Y_0).__name__}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ith element of X, gt_pseudo_label, and Y together constitute the ith data example. Here we use two of them (the 1001st and the 3001st) as illstrations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_1000, gt_pseudo_label_1000, Y_1000 = train_X[1000], train_gt_pseudo_label[1000], train_Y[1000]\n",
    "print(f\"X in the 1001st data example (a list of images):\")\n",
    "for i, x in enumerate(X_1000):\n",
    "    plt.subplot(1, len(X_1000), i + 1)\n",
    "    plt.axis(\"off\")\n",
    "    plt.imshow(x.squeeze(), cmap=\"gray\")\n",
    "plt.show()\n",
    "print(\n",
    "    f\"gt_pseudo_label in the 1001st data example (a list of ground truth pseudo-labels): {gt_pseudo_label_1000}\"\n",
    ")\n",
    "print(f\"Y in the 1001st data example (the computed result): {Y_1000}\")\n",
    "print()\n",
    "X_3000, gt_pseudo_label_3000, Y_3000 = train_X[3000], train_gt_pseudo_label[3000], train_Y[3000]\n",
    "print(f\"X in the 3001st data example (a list of images):\")\n",
    "for i, x in enumerate(X_3000):\n",
    "    plt.subplot(1, len(X_3000), i + 1)\n",
    "    plt.axis(\"off\")\n",
    "    plt.imshow(x.squeeze(), cmap=\"gray\")\n",
    "plt.show()\n",
    "print(\n",
    "    f\"gt_pseudo_label in the 3001st data example (a list of ground truth pseudo-labels): {gt_pseudo_label_3000}\"\n",
    ")\n",
    "print(f\"Y in the 3001st data example (the computed result): {Y_3000}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: The symbols in the HWF dataset can be one of digits or operators '+', '-', '×', '÷'. \n",
    "\n",
    "Note: We may see that, in the 1001st data example, the length of the formula is 3, while in the 3001st data example, the length of the formula is 5. In the HWF dataset, the lengths of the formulas are 1, 3, 5, and 7 (Specifically, 10% of the equations have a length of 1, 10% have a length of 3, 20% have a length of 5, and 60% have a length of 7)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the Learning Part"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To build the learning part, we need to first build a machine learning base model. We use SymbolNet, and encapsulate it within a `BasicNN` object to create the base model. `BasicNN` is a class that encapsulates a PyTorch model, transforming it into a base model with an sklearn-style interface. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# class of symbol may be one of ['1', ..., '9', '+', '-', '*', '/'], total of 13 classes\n",
    "cls = SymbolNet(num_classes=13, image_size=(45, 45, 1))\n",
    "loss_fn = nn.CrossEntropyLoss(label_smoothing=0.2)\n",
    "optimizer = torch.optim.Adam(cls.parameters(), lr=0.001)\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "base_model = BasicNN(\n",
    "    model=cls,\n",
    "    loss_fn=loss_fn,\n",
    "    optimizer=optimizer,\n",
    "    device=device,\n",
    "    batch_size=128,\n",
    "    num_epochs=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`BasicNN` offers methods like `predict` and `predict_proba`, which are used to predict the class index and the probabilities of each class for images. As shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_instances = [torch.randn(1, 45, 45) for _ in range(32)]\n",
    "pred_idx = base_model.predict(X=data_instances)\n",
    "print(\n",
    "    f\"Predicted class index for a batch of 32 instances: \"\n",
    "    + f\"{type(pred_idx).__name__} with shape {pred_idx.shape}\"\n",
    ")\n",
    "pred_prob = base_model.predict_proba(X=data_instances)\n",
    "print(\n",
    "    f\"Predicted class probabilities for a batch of 32 instances: \"\n",
    "    + f\"{type(pred_prob).__name__} with shape {pred_prob.shape}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, the base model built above deals with instance-level data (i.e., individual images), and can not directly deal with example-level data (i.e., a list of images comprising the formula). Therefore, we wrap the base model into `ABLModel`, which enables the learning part to train, test, and predict on example-level data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ABLModel(base_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As an illustration, consider this example of training on example-level data using the `predict` method in `ABLModel`. In this process, the method accepts data examples as input and outputs the class labels and the probabilities of each class for all instances within these data examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ablkit.data.structures import ListData\n",
    "\n",
    "# ListData is a data structure provided by ABLkit that can be used to organize data examples\n",
    "data_examples = ListData()\n",
    "# We use the first 1001st and 3001st data examples in the training set as an illustration\n",
    "data_examples.X = [X_1000, X_3000]\n",
    "data_examples.gt_pseudo_label = [gt_pseudo_label_1000, gt_pseudo_label_3000]\n",
    "data_examples.Y = [Y_1000, Y_3000]\n",
    "\n",
    "# Perform prediction on the two data examples\n",
    "# Remind that, in the 1001st data example, the length of the formula is 3,\n",
    "# while in the 3001st data example, the length of the formula is 5.\n",
    "pred_label, pred_prob = model.predict(data_examples)[\"label\"], model.predict(data_examples)[\"prob\"]\n",
    "print(\n",
    "    f\"Predicted class labels for the 100 data examples: a list of length {len(pred_label)}, \\n\"\n",
    "    + f\"the first element is a {type(pred_label[0]).__name__} of shape {pred_label[0].shape}, \"\n",
    "    + f\"and the second element is a {type(pred_label[1]).__name__} of shape {pred_label[1].shape}.\\n\"\n",
    ")\n",
    "print(\n",
    "    f\"Predicted class probabilities for the 100 data examples: a list of length {len(pred_prob)}, \\n\"\n",
    "    f\"the first element is a {type(pred_prob[0]).__name__} of shape {pred_prob[0].shape}, \"\n",
    "    + f\"and the second element is a {type(pred_prob[1]).__name__} of shape {pred_prob[1].shape}.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the Reasoning Part"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the reasoning part, we first build a knowledge base which contains information on how to compute a formula. We build it by creating a subclass of `KBBase`. In the derived subclass, we initialize the `pseudo_label_list` parameter specifying list of possible pseudo-labels, and override the `logic_forward` function defining how to perform (deductive) reasoning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class HwfKB(KBBase):\n",
    "    def __init__(\n",
    "        self, pseudo_label_list=[\"1\", \"2\", \"3\", \"4\", \"5\", \"6\", \"7\", \"8\", \"9\", \"+\", \"-\", \"*\", \"/\"]\n",
    "    ):\n",
    "        super().__init__(pseudo_label_list)\n",
    "\n",
    "    def _valid_candidate(self, formula):\n",
    "        if len(formula) % 2 == 0:\n",
    "            return False\n",
    "        for i in range(len(formula)):\n",
    "            if i % 2 == 0 and formula[i] not in [\"1\", \"2\", \"3\", \"4\", \"5\", \"6\", \"7\", \"8\", \"9\"]:\n",
    "                return False\n",
    "            if i % 2 != 0 and formula[i] not in [\"+\", \"-\", \"*\", \"/\"]:\n",
    "                return False\n",
    "        return True\n",
    "\n",
    "    # Implement the deduction function\n",
    "    def logic_forward(self, formula):\n",
    "        if not self._valid_candidate(formula):\n",
    "            return np.inf\n",
    "        return eval(\"\".join(formula))\n",
    "\n",
    "\n",
    "kb = HwfKB()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The knowledge base can perform logical reasoning (both deductive reasoning and abductive reasoning). Below is an example of performing (deductive) reasoning, and users can refer to [Documentation]() for details of abductive reasoning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pseudo_labels = [\"1\", \"-\", \"2\", \"*\", \"5\"]\n",
    "reasoning_result = kb.logic_forward(pseudo_labels)\n",
    "print(f\"Reasoning result of pseudo-labels {pseudo_labels} is {reasoning_result}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: In addition to building a knowledge base based on `KBBase`, we can also establish a knowledge base with a ground KB using `GroundKB`. The corresponding code can be found in the `main.py` file. Those interested are encouraged to examine it for further insights.\n",
    "\n",
    "Note: Also, when building the knowledge base, we can also set the `max_err` parameter during initialization, which is shown in the `main.py` file. This parameter specifies the upper tolerance limit when comparing the similarity between the reasoning result of pseudo-labels and the ground truth during abductive reasoning, with a default value of 1e-10."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we create a reasoner by instantiating the class ``Reasoner``. Due to the indeterminism of abductive reasoning, there could be multiple candidates compatible with the knowledge base. When this happens, reasoner can minimize inconsistencies between the knowledge base and pseudo-labels predicted by the learning part, and then return only one candidate that has the highest consistency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reasoner = Reasoner(kb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: During creating reasoner, the definition of \"consistency\" can be customized within the `dist_func` parameter. In the code above, we employ a consistency measurement based on confidence, which calculates the consistency between the data example and candidates based on the confidence derived from the predicted probability. In `main.py`, we provide options for utilizing other forms of consistency measurement.\n",
    "\n",
    "Note: Also, during the process of inconsistency minimization, we can leverage [ZOOpt library](https://github.com/polixir/ZOOpt) for acceleration. Options for this are also available in `main.py`. Those interested are encouraged to explore these features."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building Evaluation Metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we set up evaluation metrics. These metrics will be used to evaluate the model performance during training and testing. Specifically, we use `SymbolAccuracy` and `ReasoningMetric`, which are used to evaluate the accuracy of the machine learning model’s predictions and the accuracy of the final reasoning results, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metric_list = [SymbolAccuracy(prefix=\"hwf\"), ReasoningMetric(kb=kb, prefix=\"hwf\")]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bridge Learning and Reasoning\n",
    "\n",
    "Now, the last step is to bridge the learning and reasoning part. We proceed with this step by creating an instance of `SimpleBridge`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bridge = SimpleBridge(model, reasoner, metric_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perform training and testing by invoking the `train` and `test` methods of `SimpleBridge`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build logger\n",
    "print_log(\"Abductive Learning on the HWF example.\", logger=\"current\")\n",
    "log_dir = ABLLogger.get_current_instance().log_dir\n",
    "weights_dir = osp.join(log_dir, \"weights\")\n",
    "\n",
    "bridge.train(train_data, loops=3, segment_size=1000, save_dir=weights_dir)\n",
    "bridge.test(test_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Performance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We present the results of ABL as follows, which include the reasoning accuracy (for different equation lengths in the HWF dataset), training time (to achieve the accuracy using all equation lengths), and average memory usage (using all equation lengths). These results are compared with the following methods:\n",
    "\n",
    "- [**NGS**](https://github.com/liqing-ustc/NGS): A neural-symbolic framework that uses a grammar model and a back-search algorithm to improve its computing process;\n",
    "\n",
    "- [**DeepProbLog**](https://github.com/ML-KULeuven/deepproblog/tree/master): An extension of ProbLog by introducing neural predicates in Probabilistic Logic Programming;\n",
    "\n",
    "- [**DeepStochLog**](https://github.com/ML-KULeuven/deepstochlog/tree/main): A neural-symbolic framework based on stochastic logic program."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<table class=\"tg\" style=\"margin-left: auto; margin-right: auto;\">\n",
    "<thead>\n",
    "  <tr>\n",
    "    <th rowspan=\"2\"></th>\n",
    "    <th colspan=\"5\">Reasoning Accuracy<br><span style=\"font-weight: normal; font-size: smaller;\">(for different equation lengths)</span></th>\n",
    "    <th rowspan=\"2\">Training Time (s)<br><span style=\"font-weight: normal; font-size: smaller;\">(to achieve the Acc. using all lengths)</span></th>\n",
    "    <th rowspan=\"2\">Average Memory Usage (MB)<br><span style=\"font-weight: normal; font-size: smaller;\">(using all lengths)</span></th>\n",
    "  </tr>\n",
    "  <tr>\n",
    "    <th>1</th>\n",
    "    <th>3</th>\n",
    "    <th>5</th>\n",
    "    <th>7</th>\n",
    "    <th>All</th>\n",
    "  </tr>\n",
    "</thead>\n",
    "<tbody>\n",
    "  <tr>\n",
    "    <td>NGS</td>\n",
    "    <td>91.2</td>\n",
    "    <td>89.1</td>\n",
    "    <td>92.7</td>\n",
    "    <td>5.2</td>\n",
    "    <td>98.4</td>\n",
    "    <td>426.2</td>\n",
    "    <td>3705</td>\n",
    "  </tr>\n",
    "  <tr>\n",
    "    <td>DeepProbLog</td>\n",
    "    <td>90.8</td>\n",
    "    <td>85.6</td>\n",
    "    <td>timeout*</td>\n",
    "    <td>timeout</td>\n",
    "    <td>timeout</td>\n",
    "    <td>timeout</td>\n",
    "    <td>4315</td>\n",
    "  </tr>\n",
    "  <tr>\n",
    "    <td>DeepStochLog</td>\n",
    "    <td>92.8</td>\n",
    "    <td>87.5</td>\n",
    "    <td>92.1</td>\n",
    "    <td>timeout</td>\n",
    "    <td>timeout</td>\n",
    "    <td>timeout</td>\n",
    "    <td>4355</td>\n",
    "  </tr>\n",
    "  <tr>\n",
    "    <td>ABL</td>\n",
    "    <td><span style=\"font-weight:bold\">94.0</span></td>\n",
    "    <td><span style=\"font-weight:bold\">89.7</span></td>\n",
    "    <td><span style=\"font-weight:bold\">96.5</span></td>\n",
    "    <td><span style=\"font-weight:bold\">97.2</span></td>\n",
    "    <td><span style=\"font-weight:bold\">99.2</span></td>\n",
    "    <td><span style=\"font-weight:bold\">77.3</span></td>\n",
    "    <td><span style=\"font-weight:bold\">3074</span></td>\n",
    "  </tr>\n",
    "</tbody>\n",
    "</table>\n",
    "<p style=\"font-size: 13px;\">* timeout: need more than 1 hour to execute</p>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "abl",
   "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.8.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "9c8d454494e49869a4ee4046edcac9a39ff683f7d38abf0769f648402670238e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
