From 52f8451d56475b5ac0a30786b4d56c788679a561 Mon Sep 17 00:00:00 2001
From: troyyyyy <49091847+troyyyyy@users.noreply.github.com>
Date: Mon, 6 Mar 2023 08:02:17 +0000
Subject: [PATCH] Update prolog_KB

---
 abl/abducer/abducer_base.py |  8 ++++----
 abl/abducer/kb.py           | 15 ++-------------
 examples/example.py         |  4 ++--
 3 files changed, 8 insertions(+), 19 deletions(-)

diff --git a/abl/abducer/abducer_base.py b/abl/abducer/abducer_base.py
index 777533a..cb50ee3 100644
--- a/abl/abducer/abducer_base.py
+++ b/abl/abducer/abducer_base.py
@@ -158,7 +158,7 @@ class AbducerBase(abc.ABC):
 
 
 if __name__ == '__main__':
-    from kb import add_KB, add_prolog_KB, HWF_KB, HED_prolog_KB
+    from kb import add_KB, prolog_KB, HWF_KB
     
     prob1 = [[0, 0.99, 0.01, 0, 0, 0, 0, 0, 0, 0], [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]]
     prob2 = [[0, 0, 0.01, 0, 0, 0, 0, 0.99, 0, 0], [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]]
@@ -191,7 +191,7 @@ if __name__ == '__main__':
     print()
     
 
-    kb = add_prolog_KB()
+    kb = prolog_KB(pseudo_label_list=list(range(10)), pl_file='../examples/datasets/mnist_add/add.pl')
     abd = AbducerBase(kb, 'confidence')
     res = abd.abduce(([1, 1], prob1, 8), max_address_num=2, require_more_address=0)
     print(res)
@@ -205,7 +205,7 @@ if __name__ == '__main__':
     print(res)
     print()
 
-    kb = add_prolog_KB()
+    kb = prolog_KB(pseudo_label_list=list(range(10)), pl_file='../examples/datasets/mnist_add/add.pl')
     abd = AbducerBase(kb, 'confidence', zoopt=True)
     res = abd.abduce(([1, 1], prob1, 8), max_address_num=2, require_more_address=0)
     print(res)
@@ -260,7 +260,7 @@ if __name__ == '__main__':
     print(res)
     print()
 
-    kb = HED_prolog_KB()
+    kb = prolog_KB(pseudo_label_list=[1, 0, '+', '='], pl_file='../examples/datasets/hed/hed.pl')
     abd = AbducerBase(kb, zoopt=True, multiple_predictions=True)
     consist_exs = [[1, 1, '+', 0, '=', 1, 1], [1, '+', 1, '=', 1, 0], [0, '+', 0, '=', 0]]
     inconsist_exs = [[1, '+', 0, '=', 0], [1, '=', 1, '=', 0], [0, '=', 0, '=', 1, 1]]
diff --git a/abl/abducer/kb.py b/abl/abducer/kb.py
index 71bdcbe..c5fc675 100644
--- a/abl/abducer/kb.py
+++ b/abl/abducer/kb.py
@@ -225,9 +225,10 @@ class add_KB(ClsKB):
 
 
 class prolog_KB(KBBase):
-    def __init__(self, pseudo_label_list):
+    def __init__(self, pseudo_label_list, pl_file):
         super().__init__(pseudo_label_list)
         self.prolog = pyswip.Prolog()
+        self.prolog.consult(pl_file)
 
     def logic_forward(self, pseudo_labels):
         result = list(self.prolog.query("logic_forward(%s, Res)." % pseudo_labels))[0]['Res']
@@ -278,18 +279,6 @@ class prolog_KB(KBBase):
             candidates.append(candidate)
         return candidates
 
-
-class add_prolog_KB(prolog_KB):
-    def __init__(self, pseudo_label_list=list(range(10))):
-        super().__init__(pseudo_label_list)
-        self.prolog.consult('../datasets/mnist_add/add.pl')
-
-
-class HED_prolog_KB(prolog_KB):
-    def __init__(self, pseudo_label_list=[0, 1, '+', '=']):
-        super().__init__(pseudo_label_list)
-        self.prolog.consult('./datasets/hed/learn_add.pl')
-
     def consist_rule(self, exs, rules):
         rules = str(rules).replace("\'","")
         return len(list(self.prolog.query("eval_inst_feature(%s, %s)." % (exs, rules)))) != 0
diff --git a/examples/example.py b/examples/example.py
index 30c50fc..ee0b898 100644
--- a/examples/example.py
+++ b/examples/example.py
@@ -20,7 +20,7 @@ from models.wabl_models import DecisionTree, WABLBasicModel
 
 from multiprocessing import Pool
 from abducer.abducer_base import AbducerBase
-from abducer.kb import add_KB, HWF_KB, HED_prolog_KB
+from abducer.kb import add_KB, HWF_KB, prolog_KB
 from datasets.mnist_add.get_mnist_add import get_mnist_add
 from datasets.hwf.get_hwf import get_hwf
 from datasets.hed.get_hed import get_hed, split_equation
@@ -33,7 +33,7 @@ def run_test():
     # kb = HWF_KB(True)
     # abducer = AbducerBase(kb)
 
-    kb = HED_prolog_KB()
+    kb = prolog_KB(pseudo_label_list=[1, 0, '+', '='], pl_file='../examples/datasets/hed/hed.pl')
     abducer = AbducerBase(kb, zoopt=True, multiple_predictions=True)
 
     recorder = logger()