{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Handwritten Equation Decipherment (HED)\n",
    "\n",
    "This notebook shows an implementation of [Handwritten Equation Decipherment](https://proceedings.neurips.cc/paper_files/paper/2019/file/9c19a2aa1d84e04b0bd4bc888792bd1e-Paper.pdf). In this task, the handwritten equations are given, which consist of sequential pictures of characters. The equations are generated with unknown operation rules from images of symbols ('0', '1', '+' and '='), and each equation is associated with a label indicating whether the equation is correct (i.e., positive) or not (i.e., negative). Also, we are given a knowledge base which involves the structure of the equations and a recursive definition of bit-wise operations. The task is to learn from a training set of above mentioned equations and then to predict labels of unseen equations. \n",
    "\n",
    "Intuitively, we first use a machine learning model (learning part) to obtain the pseudo-labels ('0', '1', '+' and '=') for the observed pictures. We then use the knowledge base (reasoning part) to perform abductive reasoning so as to yield ground hypotheses as possible explanations to the observed facts, suggesting some pseudo-labels to be revised. This process enables us to further update the machine learning model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary libraries and modules\n",
    "import os.path as osp\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "from ablkit.learning import ABLModel, BasicNN\n",
    "from ablkit.utils import ABLLogger, print_log\n",
    "\n",
    "from bridge import HedBridge\n",
    "from consistency_metric import ConsistencyMetric\n",
    "from datasets import get_dataset, split_equation\n",
    "from models.nn import SymbolNet\n",
    "from reasoning import HedKB, HedReasoner"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Working with Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we get the datasets of handwritten equations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_train_data = get_dataset(train=True)\n",
    "train_data, val_data = split_equation(total_train_data, 3, 1)\n",
    "test_data = get_dataset(train=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The datasets are shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Equations in the dataset are organized by equation length, from 5 to 26.\n",
      "\n",
      "For each equation length, there are 225 true equations and 225 false equations in the training set.\n",
      "For each equation length, there are 75 true equations and 75 false equations in the validation set.\n",
      "For each equation length, there are 300 true equations and 300 false equations in the test set.\n"
     ]
    }
   ],
   "source": [
    "true_train_equation = train_data[1]\n",
    "false_train_equation = train_data[0]\n",
    "print(\n",
    "    f\"Equations in the dataset are organized by equation length, \"\n",
    "    + f\"from {min(train_data[0].keys())} to {max(train_data[0].keys())}.\"\n",
    ")\n",
    "print()\n",
    "\n",
    "true_train_equation_with_length_5 = true_train_equation[5]\n",
    "false_train_equation_with_length_5 = false_train_equation[5]\n",
    "print(\n",
    "    f\"For each equation length, there are {len(true_train_equation_with_length_5)} \"\n",
    "    + f\"true equations and {len(false_train_equation_with_length_5)} false equations \"\n",
    "    + f\"in the training set.\"\n",
    ")\n",
    "\n",
    "true_val_equation = val_data[1]\n",
    "false_val_equation = val_data[0]\n",
    "true_val_equation_with_length_5 = true_val_equation[5]\n",
    "false_val_equation_with_length_5 = false_val_equation[5]\n",
    "print(\n",
    "    f\"For each equation length, there are {len(true_val_equation_with_length_5)} \"\n",
    "    + f\"true equations and {len(false_val_equation_with_length_5)} false equations \"\n",
    "    + f\"in the validation set.\"\n",
    ")\n",
    "\n",
    "true_test_equation = test_data[1]\n",
    "false_test_equation = test_data[0]\n",
    "true_test_equation_with_length_5 = true_test_equation[5]\n",
    "false_test_equation_with_length_5 = false_test_equation[5]\n",
    "print(\n",
    "    f\"For each equation length, there are {len(true_test_equation_with_length_5)} \"\n",
    "    + f\"true equations and {len(false_test_equation_with_length_5)} false equations \"\n",
    "    + f\"in the test set.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As illustrations, we show four equations in the training dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First true equation with length 5 in the training dataset:\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAgQAAABpCAYAAABF9zs7AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8pXeV/AAAACXBIWXMAAA9hAAAPYQGoP6dpAAAIG0lEQVR4nO3dS6hN7xsH8OVwXI9CLgllwoDCKaEMXKYumYjcyiUiGTF1KUNRKANKSjERuU0xUCYuE4aS0nEr90PH7T97f4/zP/tc995r77U/n9E31naefp21e3/Ps953Dfr79+/fDABoaE15FwAA5M+CAACwIAAALAgAgMyCAADILAgAgMyCAADILAgAgCzLhvT2wkGDBlWyjoZVjnOh2traylBJeTU3N6ccf3c6OjryKKdfJk+ePKDPu2cqoxz3zLt378pQCZ1NmDBhQJ93z1RGb+8ZHQIAwIIAAOjDyAB60tT03/ry2bNnKbe3t6c8Z86clLUHAWqHDgEAYEEAABgZUEYjR45M+cCBAyk/fPgw5efPn6ccdyIAkC8dAgDAggAAMDKgjOLhFyNGjEi5paUlj3IA6AMdAgDAggAAMDIAgAG5du1ayj9+/Eh5/fr1OVTTfzoEAIAFAQBQ0JHBuXPnUt6+fXvK+/btS/nUqVNVrQmoHfE9Gg7I6rufP3+mXI7XUdejefPmpbxq1aqU47tbtm7dmvL58+erUtdA6BAAABYEAECBRgax7bdy5cqU//z5k/LVq1erWhMU1eHDh7v887t373aZa83bt29TXrFiRcpDhw7No5y68O3bt5QvXryYcmtrax7l5GLYsGEpx/e1RPFQtm3btqVsZAAA1AULAgCgOCOD+NTwhAkTurzmzZs31SoHal6ptv+hQ4dSXrZsWcpxBBCv6c1nO38e6tHEiRNTXrduXY/Xf/z4sYLVlJ8OAQBgQQAAWBAAAFmBniEAuhafFSg1+y/lzp07/f65nT8bn/PJW5wFP3nyJL9C6lQ8qbCjoyPHSmrb0aNH8y6hT3QIAAALAgDAyAAKY+nSpSkPpNVfKbGmzlsSqy2+kEfLm/6II7Cmpv/+3zqejltLY7Le0CEAACwIAAAjA6hr1RwTHDlyJOW4c6G3NcTrYsu+3tqqNK65c+emHH+H45jgxo0bKT969Kg6hZWJDgEAYEEAABRoZLBx48a8S4CKi233LKvMmKDUaKCU+NKi+Nm+HoIEtW7z5s09XvP9+/eU620Hiw4BAGBBAAAUaGQwevToLv/89evXKcenQqFexDHBQNvwsb2f9+FAtW7w4MFd5qKLT8z/+vUrx0pqw8yZM1NeuHBhjpVUng4BAGBBAAAUaGRQyu3bt1MuavsrnqOdp1hHqcNm4jW1UnctirsHOu8s6Ku+7hpoZC0tLSmfOnUq5ZMnT6Y8atSoqtZUDfHJ+NWrV6d87NixlL98+VLVmmrFuHHjUp46dWqP11+/fr2S5VSUb2QAwIIAACjQyGD8+PF5l5Cbtra2vEvIsizLRo4cmXKpAzniro8hQ2r712/SpEl5l9BvcQdB3FkA9M3x48dTLjUKffz4ccpxTF1vdAgAAAsCAKBAI4NG1tramncJ/yceFBUPdVmwYEHKtX5QVLXri7sJ+rOzoJpjgrhbYcmSJSkPdEdE3r5+/Zryjh07Ut61a1ce5eQiHkzUqDsLpk2blvLYsWNTLvWdEO+3T58+VayuStMhAAAsCACAAo0M3r9/n3cJudm0aVPeJWRZ9u+ugdhC+/DhQ8obNmxIudQTu42qP+32eOjQQMYEpcYV5RwHlHpNcq36/ft3l5niW7RoUcozZszo8fqXL19Wspyq0SEAACwIAIACjQwa2ZkzZ/IuIcuyf8+BX758ecpxZHD69OmUa/1gomrrz6uNY0u/1HsKBvrK5HLxuuX6E5+qjwePDR06tMtrGlV810U90yEAACwIAAAjg0KolcNDYuvw27dvKceDTj5//pxyc3NzdQrrpzgCqYa466K3bdiBHmZUbnEnwb179/IrhG51foVzqR0/I0aMSDmO++7fv59yHB9cuHChXCWSAx0CAMCCAAAwMqCM2tvbUz5x4kTKW7duTbmpyRq0N0qNDzo/qR93EFRzZFBqNFBqpwPlF98REnOpcVM8XGnNmjX//N2vX796/HlPnz5N+e3btynHe7ooI4N4/zXSAWq+nQEACwIAoEAjg+HDh+ddQsOLuwlmz57d5TVr165N+caNGyn//PmzcoXVue5aluU6dKg3uwPiNZV+vXKl9aZFnqfY9o/3VXzqP94/Bw8eTHnMmDFd/pvx33n16lXJnxd3CO3Zsyfls2fPpvz9+/du66938b9HIx28pEMAAFgQAAAFGhnMnz8/7xII4gggtirj08kMXHyNcCO0+stl1qxZeZfQrTgmmjJlSsrxif54yNeuXbtS3r17d8qx/R913u1TaldLfN9IPEiskdrovTF9+vSUX7x4kVsdA6VDAABYEAAAFgQAQFagZwhKbSOKM7D+vDyG8nJSYXl5PqB/Wltb8y6hW/F769KlSynHEwnj91k8hTB+Fw4bNqxsNdX6Vs1yitsqOzo6Uo4vcor27t2b8v79+ytXWIX5dgYALAgAgAKNDB48eJDyunXrUt6yZUvKO3fuTNnJeNC4bt26lXcJ3Sp1cmDcwkvl3Lx5M+XHjx+nvHDhwi6vj6dG1jMdAgDAggAAKNDI4PLlyynHU7uuXLmSciM9JQuU9vXr17xLoE4sXrw47xKqRocAALAgAAAKNDKIL82ZPXt2jpXQWXt7e94lANADHQIAwIIAACjQyIDaEnd0XLhwocdrAMiXDgEAYEEAABgZUCHxLPZ58+Z1eY33SQDUDh0CAMCCAADIskF/Y28XAGhIOgQAgAUBAGBBAABkFgQAQGZBAABkFgQAQGZBAABkFgQAQGZBAABkWfY/nLpHqMndLi0AAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 640x480 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First true equation with length 8 in the training dataset:\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 14800 characters]",
      "text/plain": [
       "<Figure size 640x480 with 8 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First false equation with length 5 in the training dataset:\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAgQAAABpCAYAAABF9zs7AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8pXeV/AAAACXBIWXMAAA9hAAAPYQGoP6dpAAAK+klEQVR4nO3dWUiU7RvH8afSStMyC0KKNqJ90UKrg7CkrGgTNLAVIiiio1YokizwpKBogRYjolJC6iQL8iAqCj3JSlHrSIooW2hBzT37n/y5urR5Xsfx2Wbm+zn6vfLMzN3ozHtzXc993/3+/PnzxwAAAGGtv9sDAAAA7mNCAAAAmBAAAAAmBAAAwGBCAAAADCYEAADAYEIAAAAMJgQAAMAwjAh/L+zXr5+d4whbVuwLVVdXZ8FIrBUZGSlZ/+20tbW5MZyAJCQk9OnxfGbsYcVn5vv37xaMBN3Fx8f36fF8Zuzh72eGCgEAAGBCAAAAetEyAHrSv//f+WVNTY3kpqYmybNnz5ZMeRAAvIMKAQAAYEIAAABoGcBC0dHRkg8cOCC5vLxccm1trWS9EgFwUkxMTI/X6JZWc3Oz5Pb2dp/XAMGOCgEAAGBCAAAAgrxlkJubK/no0aM+r6Gk5xy9+UVUVJRkf8qzgJNOnTrV4zUtLS2SMzMzJc+aNUtyR0eH5M7Ozi6Pb2xs7MsQAcdRIQAAAEwIAABAkLcMzNoEjx8/dnYgAILKsWPHerxGrybQ3ynTp0+XrNsEcXFxXR6fl5cnWbcWQo1egRFO9FknO3bskLx//37Jb9++lbx582bJFRUV9g4uQFQIAAAAEwIAABCELYNHjx65PQQgJCxYsEByWVmZZF3aLCgocHRMTnnw4EGP1+gVSjrv3btX8uvXryV3X01TWVkpOSIi6L5q/5NugZw9e1Zy97ZJsBo9erTkq1evStb/bp2XLl0q+f3795JHjhwpubS0VPL8+fMlV1VVWTBia1AhAAAATAgAAECQtAz0BkSLFy/u8Xp9jd4sR98pvGTJEgtGBgSvVatWSdZ3yyclJUkO1ZbBvHnzAn5sUVGRZL150cCBA7tcp8/w2L17t2R95kew0iswdA4VFy9elJyWlib51q1bkhcuXCg5OztbcnFxsWT9/5mSkhLJukWXnJws+c2bN30Zdp9RIQAAAEwIAABAkLQMUlNTLXkes1YC5x0gHE2ZMsXtIbimra0t4McOHz7cr+uWLVsmWd9JHgrfN/r7M1RWUMTGxkqeO3eu5JMnT0o+fPhwr55Tr4rbs2eP5HPnzknWLXHdenADFQIAAMCEAAAAMCEAAABGkNxD4M9SQ728w+x6s8OQuJ8AgL/0Ek1/RUZG2jASWEnvuKgPLnr37p0lz5+fny/5+PHjktPT0yUPHTq0y2Pq6+steW1/USEAAABMCAAAgIdbBv4cYqTPNNe7EOqs6Z+bPb9eAqIzEMrMDvIBwkViYqLk1tZWyXrnwb7QS131MtRFixb5zIZhGPfv37fktf1FhQAAADAhAAAAHm4Z+LOyoLclfd0y0GVRvcpAr0TQOyRyGBJCmf4M6AyEsokTJ0qeMGGC5Lt370r++PGj5a974sQJyd3bBG6iQgAAAJgQAAAAD7cMnKTbAXr1gW5bdG9hmK1kAOB9Q4YMkaw3GmpqapLMaovQZ7YpUEFBga2v++HDB1ufP1BUCAAAABMCAADgsZaBPysL9GZEVjHb1IiWAcJReXm520OwXUZGhuRp06ZJ1nd/t7S02D6Ojo4OyYGckYC+0ZsR/f79W/KvX79cGE3Xv0vDYGMiAADgAiYEAAAg+FoGdnvy5IlkPZ7uRyd76ZyD/v29Ma/T4zC7Q1tf45Vxo+vv6+nTpy6OxBkvXryQXF1dLVn/2+3aoEmXo/X3ysqVKyXrvfTN6PFFRET4zGwy9d+ysrIk//z5U/LDhw9dGM2/xx87jW9kAADAhAAAAHisZdC9LO+L3Xf36+f3ZzxeUFdX5/YQDMMwjOjoaMn6qE/t06dPknVp04tGjRrl9hAcE26l5bFjx/r8uRPvg/6cnDlzRnJOTo5k3TIwa7/pVRC7du2SvH37dslO3y0fFxfn6OtZSX832W3r1q2S9e/3+vXrjo3BFyoEAACACQEAAHC5ZRDIqgK7WwZeWOnQW0lJSW4P4R+xsbGSBwwYIDklJUWy18vUXh9fIHQbJJyP9NZnlrhp8ODBkg8ePCi5pqZG8sCBAyXrlTmfP3+WrFsPOjvtx48frr22v6ZOnSo5KipKcklJiWNjyMzMlOyl7xkqBAAAgAkBAAAIkpaBHecXhJLNmze7PQTDMLquGtCtHV1G3Lhxo2SOl3WeLj/Hx8e7OBJ3ma2CcZpeTZCXlyfZbNOuQYMGSS4rK5N86dIlyboMjn9NmjRJsn6v9FkGdtCtU3389rdv3ySXlpbaOoaeUCEAAABMCAAAgMc2JkJgLly44PYQDMMwjJiYGMlpaWmSdcvg/Pnzkr2+MVGo02XppqYmyXaXTvGXbps1Nzf3eH1jY6PkmTNnSr5y5YpkL9217kVjxozx+fPCwkLLX0u3Ca5duyZZt+v27dsn2e1VGlQIAAAAEwIAAEDLwDCMrpuUmK188PImLg0NDW4PwTCMrqVKvYd6Z2en5Pr6esmRkZHODCxAugUSKlavXi1Z/16Ki4sle+VsDPxLtxja29t9ZjcFw8oV3c6025o1ayRv2rRJclVVleSbN286Np6eUCEAAABMCAAAgMstg+7nEpgdN6x/rh9jdq5Bbm6uz5+npqZK9mdTJH9eC3/pO9VPnz4tedu2bZLNNlyBM/SmLEA4snsVxuTJkyXfuHFDst4MKyMjQ/LXr19tHU9v8O0MAACYEAAAAI+1DPR/m5X07T62VI/ByysLvEjftT5jxgyf16xfv16yvrPdK3dJhyu9sU240Xfuc76GOf35DmZmv2O9+qCioqJXz6mPFS8qKpKsz57Q57jU1tb26vmdQoUAAAAwIQAAAB7bmEgfc+zv0ciB0q0B/bqsJrCGbgHoUuOXL1/cGA7+T5dL9YoPvUd+ONDvg777+/v37z6vCSf6Lvxhw4ZJjouLk6zfs2BrJehjo7OysiRnZ2dL1ucO6PMFEhMTJa9bt07yoUOHJOsjxp8/fy759u3bgQ/aIVQIAAAAEwIAAMCEAAAAGB67h0D373X/Tu886M9ug9wf4F3sVOgu3R8Ott6vlfTBWi9fvpS8fPlyyaF4uJWZAQMGSNZ/F1u2bJG8YcMGyWPGjJEcbO/T5cuXJefk5EhOTk6WXFlZKVnvwDpu3DjJ+l6B1tZWyfqwop07d1owYufw7QwAAJgQAAAAj7UMzJgdVgQgcPpM9urqahdH4jxdFh8xYoRkXSLXu8yFGt0iMIyuB+yUlJRILiwslJyfny/53r17klNSUuwYom1+/folee3atZKPHDkiOT093edjdSvhzp07kvWuq69evbJimK6gQgAAAJgQAACAIGkZALCGvqN+wYIFkhsaGtwYjmv0Tprjx4+XrO9A1ysyQk1ERNev/vfv30tOSkqSXFpaKlmXxUPFs2fPJK9YscLFkXgDFQIAAMCEAAAAGEa/P37WxcL1oA+7WVGWrKurs2Ak1tIbv8yZM8fnNfrMcV3C9YqEhIQ+PZ7PjD2s+MzoQ4zQdcMwfYiRvmNeH9STlpYmOT4+3udjA8Fnxh7+fmaoEAAAACYEAACAloHrQrVloP9e9AY42syZMyV78Y5uWgbeRMvAXvr91fv16zZgS0uLZL3Jk24fBILPjD1oGQAAAL8xIQAAAGxMBHvoElViYqLPa7y4sgAId7psrz+jfF5DHxUCAADAhAAAAPRilQEAAAhdVAgAAAATAgAAwIQAAAAYTAgAAIDBhAAAABhMCAAAgMGEAAAAGEwIAACAwYQAAAAYhvE/mo9c6U3R9SgAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 640x480 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First false equation with length 8 in the training dataset:\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 12072 characters]",
      "text/plain": [
       "<Figure size 640x480 with 8 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "true_train_equation_with_length_5 = true_train_equation[5]\n",
    "true_train_equation_with_length_8 = true_train_equation[8]\n",
    "print(f\"First true equation with length 5 in the training dataset:\")\n",
    "for i, x in enumerate(true_train_equation_with_length_5[0]):\n",
    "    plt.subplot(1, 5, i + 1)\n",
    "    plt.axis(\"off\")\n",
    "    plt.imshow(x.squeeze(), cmap=\"gray\")\n",
    "plt.show()\n",
    "print(f\"First true equation with length 8 in the training dataset:\")\n",
    "for i, x in enumerate(true_train_equation_with_length_8[0]):\n",
    "    plt.subplot(1, 8, i + 1)\n",
    "    plt.axis(\"off\")\n",
    "    plt.imshow(x.squeeze(), cmap=\"gray\")\n",
    "plt.show()\n",
    "\n",
    "false_train_equation_with_length_5 = false_train_equation[5]\n",
    "false_train_equation_with_length_8 = false_train_equation[8]\n",
    "print(f\"First false equation with length 5 in the training dataset:\")\n",
    "for i, x in enumerate(false_train_equation_with_length_5[0]):\n",
    "    plt.subplot(1, 5, i + 1)\n",
    "    plt.axis(\"off\")\n",
    "    plt.imshow(x.squeeze(), cmap=\"gray\")\n",
    "plt.show()\n",
    "print(f\"First false equation with length 8 in the training dataset:\")\n",
    "for i, x in enumerate(false_train_equation_with_length_8[0]):\n",
    "    plt.subplot(1, 8, i + 1)\n",
    "    plt.axis(\"off\")\n",
    "    plt.imshow(x.squeeze(), cmap=\"gray\")\n",
    "plt.show()"
   ]
  },
  {
   "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": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# class of symbol may be one of ['0', '1', '+', '='], total of 4 classes\n",
    "cls = SymbolNet(num_classes=4)\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.RMSprop(cls.parameters(), lr=0.001, weight_decay=1e-4)\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "base_model = BasicNN(\n",
    "    cls,\n",
    "    loss_fn,\n",
    "    optimizer,\n",
    "    device=device,\n",
    "    batch_size=32,\n",
    "    num_epochs=1,\n",
    "    stop_loss=None,\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 equation). 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": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ABLModel(base_model)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the Reasoning Part"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the reasoning part, we first build a knowledge base. As mentioned before, the knowledge base in this task involves the structure of the equations and a recursive definition of bit-wise operations, which are defined in Prolog file ``reasoning/BK.pl``  and ``reasoning/learn_add.pl``, respectively.  Specifically, the knowledge about the structure of equations is a set of DCG rules recursively define that a digit is a sequence of ‘0’ and ‘1’, and equations share the structure of X+Y=Z, though the length of X, Y and Z can be varied. The knowledge about bit-wise operations is a recursive  logic program, which reversely calculates X+Y, i.e., it operates on  X and Y digit-by-digit and from the last digit to the first.\n",
    "\n",
    "The knowledge base is already built in ``HedKB``.  ``HedKB`` is derived from class ``PrologKB``, and is built upon the aformentioned Prolog files. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "kb = HedKB()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: Please notice that, the specific rules for calculating the operations are undefined in the knowledge base, i.e., results of '0+0', '0+1' and '1+1' could be '0', '1', '00', '01' or even '10'. The missing calculation rules are required to be learned from the data. Therefore, `HedKB` incorporates methods for abducing rules from data. Users interested can refer to the specific implementation of `HedKB`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we create a 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. \n",
    "\n",
    "In this task, we create the reasoner by instantiating the class `HedReasoner`, which is a reasoner derived from `Reasoner` and tailored specifically for this task. `HedReasoner` leverages [ZOOpt library](https://github.com/polixir/ZOOpt) for acceleration, and has designed a specific strategy to better harness ZOOpt’s capabilities. Additionally, methods for abducing rules from data have been incorporated. Users interested can refer to the specific implementation of `HedReasoner` in `reasoning/reasoning.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "reasoner = HedReasoner(kb, dist_func=\"hamming\", use_zoopt=True, max_revision=10)"
   ]
  },
  {
   "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 = [ConsistencyMetric(kb=kb)]"
   ]
  },
  {
   "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 `HedBridge`, which is derived from `SimpleBridge` and tailored specific for this task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bridge = HedBridge(model, reasoner, metric_list)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perform pretraining, training and testing by invoking the `pretrain`, `train` and `test` methods of `HedBridge`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build logger\n",
    "print_log(\"Abductive Learning on the HED example.\", logger=\"current\")\n",
    "\n",
    "# Retrieve the directory of the Log file and define the directory for saving the model weights.\n",
    "log_dir = ABLLogger.get_current_instance().log_dir\n",
    "weights_dir = osp.join(log_dir, \"weights\")\n",
    "\n",
    "bridge.pretrain(weights_dir)\n",
    "bridge.train(train_data, val_data, save_dir=weights_dir)\n",
    "bridge.test(test_data)"
   ]
  }
 ],
 "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": "fb6f4ceeabb9a733f366948eb80109f83aedf798cc984df1e68fb411adb27d58"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
