ABLkit/examples/hed/hed_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os.path as osp\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from zoopt import Dimension, Objective, Opt, Parameter\n",
    "\n",
    "from abl.evaluation import ReasoningMetric, SymbolMetric\n",
    "from abl.learning import ABLModel, BasicNN\n",
    "from abl.reasoning import PrologKB, Reasoner\n",
    "from abl.utils import ABLLogger, print_log, reform_list\n",
    "from examples.hed.datasets.get_hed import get_hed, split_equation\n",
    "from examples.hed.hed_bridge import HEDBridge\n",
    "from examples.models.nn import SymbolNet"
   ]
  },
  {
   "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\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Logic Part"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize knowledge base and abducer\n",
    "class HedKB(PrologKB):\n",
    "    def __init__(self, pseudo_label_list, pl_file):\n",
    "        super().__init__(pseudo_label_list, pl_file)\n",
    "\n",
    "    def consist_rule(self, exs, rules):\n",
    "        rules = str(rules).replace(\"'\", \"\")\n",
    "        return len(list(self.prolog.query(\"eval_inst_feature(%s, %s).\" % (exs, rules)))) != 0\n",
    "\n",
    "    def abduce_rules(self, pred_res):\n",
    "        prolog_result = list(self.prolog.query(\"consistent_inst_feature(%s, X).\" % pred_res))\n",
    "        if len(prolog_result) == 0:\n",
    "            return None\n",
    "        prolog_rules = prolog_result[0][\"X\"]\n",
    "        rules = [rule.value for rule in prolog_rules]\n",
    "        return rules\n",
    "\n",
    "\n",
    "class HedReasoner(Reasoner):\n",
    "    def revise_at_idx(self, data_example):\n",
    "        revision_idx = np.where(np.array(data_example.flatten(\"revision_flag\")) != 0)[0]\n",
    "        candidate = self.kb.revise_at_idx(\n",
    "            data_example.pred_pseudo_label, data_example.Y, data_example.X, revision_idx\n",
    "        )\n",
    "        return candidate\n",
    "\n",
    "    def zoopt_revision_score(self, symbol_num, data_example, sol):\n",
    "        revision_flag = reform_list(\n",
    "            list(sol.get_x().astype(np.int32)), data_example.pred_pseudo_label\n",
    "        )\n",
    "        data_example.revision_flag = revision_flag\n",
    "\n",
    "        lefted_idxs = [i for i in range(len(data_example.pred_idx))]\n",
    "        candidate_size = []\n",
    "        max_consistent_idxs = []\n",
    "        while lefted_idxs:\n",
    "            idxs = []\n",
    "            idxs.append(lefted_idxs.pop(0))\n",
    "            max_candidate_idxs = []\n",
    "            found = False\n",
    "            for idx in range(-1, len(data_example.pred_idx)):\n",
    "                if (not idx in idxs) and (idx >= 0):\n",
    "                    idxs.append(idx)\n",
    "                candidates, _ = self.revise_at_idx(data_example[idxs])\n",
    "                if len(candidates) == 0:\n",
    "                    if len(idxs) > 1:\n",
    "                        idxs.pop()\n",
    "                else:\n",
    "                    if len(idxs) > len(max_candidate_idxs):\n",
    "                        found = True\n",
    "                        max_candidate_idxs = idxs.copy()\n",
    "            removed = [i for i in lefted_idxs if i in max_candidate_idxs]\n",
    "            if found:\n",
    "                removed.insert(0, idxs[0])\n",
    "                candidate_size.append(len(removed))\n",
    "                max_consistent_idxs = max_candidate_idxs.copy()\n",
    "                lefted_idxs = [i for i in lefted_idxs if i not in max_candidate_idxs]\n",
    "        candidate_size.sort()\n",
    "        score = 0\n",
    "        import math\n",
    "\n",
    "        for i in range(0, len(candidate_size)):\n",
    "            score -= math.exp(-i) * candidate_size[i]\n",
    "        return score, max_consistent_idxs\n",
    "    \n",
    "    def _zoopt_get_solution(self, symbol_num, data_example, max_revision_num):\n",
    "        dimension = Dimension(size=symbol_num, regs=[[0, 1]] * symbol_num, tys=[False] * symbol_num)\n",
    "        objective = Objective(\n",
    "            lambda sol: self.zoopt_revision_score(symbol_num, data_example, sol)[0],\n",
    "            dim=dimension,\n",
    "            constraint=lambda sol: self._constrain_revision_num(sol, max_revision_num),\n",
    "        )\n",
    "        parameter = Parameter(budget=200, intermediate_result=False, autoset=True)\n",
    "        solution = Opt.min(objective, parameter)\n",
    "        return solution\n",
    "\n",
    "    def abduce(self, data_example):\n",
    "        symbol_num = data_example.elements_num(\"pred_pseudo_label\")\n",
    "        max_revision_num = self._get_max_revision_num(self.max_revision, symbol_num)\n",
    "\n",
    "        solution = self._zoopt_get_solution(symbol_num, data_example, max_revision_num)\n",
    "        _, max_candidate_idxs = self.zoopt_revision_score(symbol_num, data_example, solution)\n",
    "\n",
    "        abduced_pseudo_label = [[] for _ in range(len(data_example))]\n",
    "\n",
    "        if len(max_candidate_idxs) > 0:\n",
    "            candidates, _ = self.revise_at_idx(data_example[max_candidate_idxs])\n",
    "            for i, idx in enumerate(max_candidate_idxs):\n",
    "                abduced_pseudo_label[idx] = candidates[0][i]\n",
    "        data_example.abduced_pseudo_label = abduced_pseudo_label\n",
    "        return abduced_pseudo_label\n",
    "\n",
    "    def abduce_rules(self, pred_res):\n",
    "        return self.kb.abduce_rules(pred_res)\n",
    "\n",
    "\n",
    "kb = HedKB(pseudo_label_list=[1, 0, \"+\", \"=\"], pl_file=\"./datasets/learn_add.pl\")\n",
    "reasoner = HedReasoner(kb, dist_func=\"hamming\", use_zoopt=True, max_revision=10)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Machine Learning Part"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build necessary components for BasicNN\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\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build BasicNN\n",
    "# The function of BasicNN is to wrap NN models into the form of an sklearn estimator\n",
    "base_model = BasicNN(\n",
    "    cls,\n",
    "    loss_fn,\n",
    "    optimizer,\n",
    "    device,\n",
    "    batch_size=32,\n",
    "    num_epochs=1,\n",
    "    save_interval=1,\n",
    "    stop_loss=None,\n",
    "    save_dir=weights_dir,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build ABLModel\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": [
    "### Metric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up metrics\n",
    "metric_list = [SymbolMetric(prefix=\"hed\"), ReasoningMetric(kb=kb, 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, reasoner, metric_list)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_train_data = get_hed(train=True)\n",
    "train_data, val_data = split_equation(total_train_data, 3, 1)\n",
    "test_data = get_hed(train=False)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train and Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bridge.pretrain(\"./weights\")\n",
    "bridge.train(train_data, val_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.18"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "fb6f4ceeabb9a733f366948eb80109f83aedf798cc984df1e68fb411adb27d58"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}