From d6d29d632b894c9856bc724f42aa502cee4f36e0 Mon Sep 17 00:00:00 2001
From: Gao Enhao <gaoeh@lamda.nju.edu.cn>
Date: Sat, 11 Nov 2023 22:32:04 +0800
Subject: [PATCH] [MNT] resolve some comments in kbs

---
 abl/reasoning/base_kb.py         |  9 ++----
 abl/reasoning/ground_kb.py       |  5 ++--
 abl/reasoning/prolog_based_kb.py | 49 ++++++++++++++++++++++++++++++++
 abl/reasoning/search_based_kb.py | 19 ++++++-------
 examples/hwf/hwf_kb.py           |  2 +-
 5 files changed, 62 insertions(+), 22 deletions(-)
 create mode 100644 abl/reasoning/prolog_based_kb.py

diff --git a/abl/reasoning/base_kb.py b/abl/reasoning/base_kb.py
index 560f3d4..4822e45 100644
--- a/abl/reasoning/base_kb.py
+++ b/abl/reasoning/base_kb.py
@@ -1,16 +1,11 @@
 from abc import ABC, abstractmethod
-from typing import List, Tuple, Union
-
-import numpy as np
 
 from ..structures import ListData
 
 
 class BaseKB(ABC):
-    @abstractmethod
-    def logic_forward(self, data_sample: ListData):
-        """Placeholder for the forward reasoning of the knowledge base."""
-        pass
+    def __init__(self, pseudo_label_list) -> None:
+        self.pseudo_label_list = pseudo_label_list
 
     @abstractmethod
     def abduce_candidates(self, data_sample: ListData):
diff --git a/abl/reasoning/ground_kb.py b/abl/reasoning/ground_kb.py
index 95ee093..ed32efd 100644
--- a/abl/reasoning/ground_kb.py
+++ b/abl/reasoning/ground_kb.py
@@ -7,8 +7,8 @@ from .base_kb import BaseKB
 
 
 class GroundKB(BaseKB, ABC):
-    def __init__(self, pseudo_label_list):
-        self.pseudo_label_list = pseudo_label_list
+    def __init__(self, pseudo_label_list: List) -> None:
+        super().__init__(pseudo_label_list)
         self.base = self.construct_base()
 
     @abstractmethod
@@ -19,7 +19,6 @@ class GroundKB(BaseKB, ABC):
     def get_key(self, data_sample: ListData) -> Hashable:
         pass
 
-    @abstractmethod
     def key2candidates(self, key: Hashable) -> List[List[Any]]:
         return self.base[key]
 
diff --git a/abl/reasoning/prolog_based_kb.py b/abl/reasoning/prolog_based_kb.py
new file mode 100644
index 0000000..9bce03b
--- /dev/null
+++ b/abl/reasoning/prolog_based_kb.py
@@ -0,0 +1,49 @@
+from abc import ABC, abstractmethod
+from typing import Any, Generator, List, Tuple, Union
+
+import numpy as np
+import pyswip
+
+from abl.structures import ListData
+
+from ..structures import ListData
+from .base_kb import BaseKB
+
+
+class PrologBasedKB(BaseKB, ABC):
+    def __init__(self, pseudo_label_list, pl_file):
+        self.pseudo_label_list = pseudo_label_list
+        self.prolog = pyswip.Prolog()
+        self.prolog.consult(pl_file)
+
+    def logic_forward(
+        self, data_sample: ListData, revision_idx: Union[List, Tuple, np.ndarray] = None
+    ) -> Generator[Union[Any, pyswip.Variable, list, dict, None], Any, None]:
+        return self.prolog.query(self.to_query(data_sample, revision_idx))
+
+    @abstractmethod
+    def to_query(self, data_sample: ListData, revision_idx: Union[List, Tuple, np.ndarray] = None):
+        pass
+
+    @abstractmethod
+    def postprocess(
+        self, query_res, data_sample: ListData, revision_idx: Union[List, Tuple, np.ndarray]
+    ):
+        return list(query_res)
+
+    @abstractmethod
+    def filter_candidates(
+        self,
+        data_sample: ListData,
+        candidates: List[List[Any]],
+        max_revision_num: int,
+        require_more_revision: int = 0,
+    ) -> List[List[Any]]:
+        return candidates
+
+    def revise_at_idx(self, data_sample: ListData, revision_idx: Union[List, Tuple, np.ndarray]):
+        query_res = self.logic_forward(data_sample, revision_idx)
+        return self.postprocess(query_res, data_sample, revision_idx)
+
+    def abduce_candidates(self, data_sample: ListData):
+        return super().abduce_candidates(data_sample)
diff --git a/abl/reasoning/search_based_kb.py b/abl/reasoning/search_based_kb.py
index 0a37beb..4598edd 100644
--- a/abl/reasoning/search_based_kb.py
+++ b/abl/reasoning/search_based_kb.py
@@ -36,25 +36,22 @@ class SearchBasedKB(BaseKB, ABC):
         use_cache: bool = True,
         cache_root: Optional[str] = None,
     ) -> None:
-        self.pseudo_label_list = pseudo_label_list
+        super().__init__(pseudo_label_list)
         self.search_strategy = search_strategy
         self.use_cache = use_cache
-        if self.use_cache and getattr(self, "get_key") is getattr(SearchBasedKB, "get_key", None):
-            raise NotImplementedError("If use_cache is True, get_key should be implemented.")
+        if self.use_cache:
+            if not hasattr(self, "get_key"):
+                raise NotImplementedError("If use_cache is True, get_key should be implemented.")
+            key_func = self.get_key
+        else:
+            key_func = lambda x: x
         self.cache = Cache[ListData, List[List[Any]]](
             func=self._abduce_by_search,
             cache=use_cache,
             cache_root=cache_root,
-            key_func=lambda x: self.get_key(x),
+            key_func=key_func,
         )
 
-    @abstractmethod
-    def get_key(self, data_sample: ListData):
-        """
-        If 'use_cache' is set to 'True', this method should be implemented.
-        """
-        pass
-
     @abstractmethod
     def entail(self, data_sample: ListData, y: Any):
         """Placeholder for entail."""
diff --git a/examples/hwf/hwf_kb.py b/examples/hwf/hwf_kb.py
index 5da7b36..821a79f 100644
--- a/examples/hwf/hwf_kb.py
+++ b/examples/hwf/hwf_kb.py
@@ -16,7 +16,7 @@ class HWF_KB(GroundKB):
         self,
         pseudo_label_list=["1", "2", "3", "4", "5", "6", "7", "8", "9", "+", "-", "times", "div"],
         GKB_len_list=[1, 3, 5, 7],
-        max_err=1e-3,
+        max_err=1e-10,
     ):
         self.GKB_len_list = GKB_len_list
         self.max_err = max_err