From 12952e0d4d9c5fafd277d1e0cf3a0cf5247861ea Mon Sep 17 00:00:00 2001
From: Gao Enhao <gaoeh@lamda.nju.edu.cn>
Date: Sun, 11 Jun 2023 22:43:32 +0800
Subject: [PATCH] [ENH] run hed example successfully

---
 examples/hed/hed_example.ipynb | 94 +++++++++++++++++++++++-----------
 1 file changed, 63 insertions(+), 31 deletions(-)

diff --git a/examples/hed/hed_example.ipynb b/examples/hed/hed_example.ipynb
index d562852..2e5add3 100644
--- a/examples/hed/hed_example.ipynb
+++ b/examples/hed/hed_example.ipynb
@@ -10,17 +10,14 @@
     "import torch.nn as nn\n",
     "import torch\n",
     "\n",
-    "from abl.reasoning.reasoner import ReasonerBase\n",
-    "from abl.reasoning.kb import prolog_KB\n",
-    "\n",
-    "from abl.utils.plog import logger\n",
-    "from abl.learning.basic_nn import BasicNN\n",
-    "from abl.learning.abl_model import ABLModel\n",
-    "from abl.utils.utils import reform_idx\n",
+    "from abl.reasoning import ReasonerBase, prolog_KB\n",
+    "from abl.learning import BasicNN, ABLModel\n",
+    "from abl.evaluation import SymbolMetric, ABLMetric\n",
+    "from abl.utils import ABLLogger, reform_idx\n",
     "\n",
+    "from examples.hed.hed_bridge import HEDBridge\n",
     "from models.nn import SymbolNet\n",
-    "from datasets.get_hed import get_hed, split_equation\n",
-    "import framework_hed"
+    "from datasets.get_hed import get_hed, split_equation"
    ]
   },
   {
@@ -30,7 +27,7 @@
    "outputs": [],
    "source": [
     "# Initialize logger\n",
-    "recorder = logger()"
+    "logger = ABLLogger.get_instance(\"abl\")"
    ]
   },
   {
@@ -64,7 +61,7 @@
     "        rules = [rule.value for rule in prolog_rules]\n",
     "        return rules\n",
     "        \n",
-    "    \n",
+    "\n",
     "kb = HED_prolog_KB(pseudo_label_list=[1, 0, '+', '='], pl_file='./datasets/learn_add.pl')\n",
     "\n",
     "class HED_Abducer(ReasonerBase):\n",
@@ -113,6 +110,24 @@
     "        for i in range(0, len(candidate_size)):\n",
     "            score -= math.exp(-i) * candidate_size[i]\n",
     "        return score\n",
+    "    \n",
+    "    def abduce(self, data, max_revision=-1, require_more_revision=0):\n",
+    "        batch_pred_label, batch_pred_prob, batch_pred_pseudo_label, batch_y = data\n",
+    "\n",
+    "        solution = self.zoopt_get_solution(\n",
+    "            batch_pred_label, batch_pred_pseudo_label, batch_pred_prob, batch_y, max_revision\n",
+    "        )\n",
+    "        batch_revision_idx = reform_idx(solution.astype(np.int32), batch_pred_label)\n",
+    "        \n",
+    "        batch_abduced_pseudo_label = []\n",
+    "        for pred_pseudo_label, pred_prob, revision_idx in zip(batch_pred_pseudo_label, batch_pred_prob, batch_revision_idx):\n",
+    "            candidates = self.revise_by_idx([pred_pseudo_label], None, list(np.nonzero(np.array(revision_idx))[0]))\n",
+    "            if len(candidates) == 0:\n",
+    "                batch_abduced_pseudo_label.append([])\n",
+    "            else:\n",
+    "                batch_abduced_pseudo_label.append(candidates[0][0])\n",
+    "                # batch_abduced_pseudo_label.append(self._get_one_candidate(pred_pseudo_label, pred_prob, candidates)[0])\n",
+    "        return batch_abduced_pseudo_label\n",
     "\n",
     "    def abduce_rules(self, pred_res):\n",
     "        return self.kb.abduce_rules(pred_res)\n",
@@ -144,16 +159,6 @@
     "optimizer = torch.optim.RMSprop(cls.parameters(), lr=0.001, weight_decay=1e-6)"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Pretrain NN classifier\n",
-    "framework_hed.hed_pretrain(kb, cls, recorder)"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -168,19 +173,30 @@
     "    optimizer,\n",
     "    device,\n",
     "    save_interval=1,\n",
-    "    save_dir=recorder.save_dir,\n",
+    "    save_dir=logger.save_dir,\n",
     "    batch_size=32,\n",
     "    num_epochs=1,\n",
-    "    recorder=recorder,\n",
     ")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize ABL model\n",
+    "# The main function of the ABL model is to serialize data and \n",
+    "# provide a unified interface for different machine learning models\n",
+    "model = ABLModel(base_model)"
+   ]
+  },
   {
    "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Use ABL model to join two parts"
+    "### Metric"
    ]
   },
   {
@@ -189,7 +205,25 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model = ABLModel(base_model)"
+    "# Add metric\n",
+    "metric = [SymbolMetric(prefix=\"hed\"), ABLMetric(prefix=\"hed\")]"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Bridge Machine Learning and Logic Reasoning"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "bridge = HEDBridge(model, abducer, metric)"
    ]
   },
   {
@@ -216,19 +250,17 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Train and save"
+    "### Train and Test"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "model, mapping = framework_hed.train_with_rule(model, abducer, train_data, val_data, select_num=10, min_len=5, max_len=8)\n",
-    "framework_hed.hed_test(model, abducer, mapping, train_data, test_data, min_len=5, max_len=8)\n",
-    "\n",
-    "recorder.dump()"
+    "bridge.pretrain(\"./weights/\")\n",
+    "bridge.train(train_data, val_data)"
    ]
   }
  ],