feat: add kg_structure evaluation

Author: ChenZiHong-Gavin

examples/evaluate/evaluate_kg/kg_evaluation_config.yaml

@@ -27,7 +27,7 @@ nodes:
     op_name: build_kg
     type: map_batch
     dependencies:
-      - chunk_documents
+      - chunk
     execution_params:
       replicas: 1
       batch_size: 128
@@ -40,6 +40,6 @@ nodes:
       - build_kg
     params:
       metrics:
-        - kg_accuracy
-        - kg_consistency
         - kg_structure
+#        - kg_accuracy
+#        - kg_consistency

examples/evaluate/evaluate_qa/qa_evaluation_config.yaml

@@ -92,7 +92,7 @@ nodes:
       metrics:
         - qa_length
         - qa_mtld
-        # - qa_reward_score
-        # - qa_uni_score
+        - qa_reward_score
+        - qa_uni_score
       mtld_params:
         threshold: 0.7

graphgen/bases/base_storage.py

@@ -1,5 +1,6 @@
+from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Generic, TypeVar, Union
+from typing import Dict, Generic, List, Set, TypeVar, Union
 
 T = TypeVar("T")
 
@@ -45,52 +46,90 @@ def reload(self):
         raise NotImplementedError
 
 
-class BaseGraphStorage(StorageNameSpace):
+class BaseGraphStorage(StorageNameSpace, ABC):
+    @abstractmethod
+    def is_directed(self) -> bool:
+        pass
+
+    @abstractmethod
     def has_node(self, node_id: str) -> bool:
         raise NotImplementedError
 
+    @abstractmethod
     def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
         raise NotImplementedError
 
+    @abstractmethod
     def node_degree(self, node_id: str) -> int:
         raise NotImplementedError
 
-    def edge_degree(self, src_id: str, tgt_id: str) -> int:
-        raise NotImplementedError
+    @abstractmethod
+    def get_all_node_degrees(self) -> Dict[str, int]:
+        pass
 
+    def get_isolated_nodes(self) -> List[str]:
+        return [
+            node_id
+            for node_id, degree in self.get_all_node_degrees().items()
+            if degree == 0
+        ]
+
+    @abstractmethod
     def get_node(self, node_id: str) -> Union[dict, None]:
         raise NotImplementedError
 
+    @abstractmethod
     def update_node(self, node_id: str, node_data: dict[str, str]):
         raise NotImplementedError
 
+    @abstractmethod
     def get_all_nodes(self) -> Union[list[tuple[str, dict]], None]:
         raise NotImplementedError
 
+    @abstractmethod
+    def get_node_count(self) -> int:
+        pass
+
+    @abstractmethod
     def get_edge(self, source_node_id: str, target_node_id: str) -> Union[dict, None]:
         raise NotImplementedError
 
+    @abstractmethod
     def update_edge(
         self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
     ):
         raise NotImplementedError
 
+    @abstractmethod
     def get_all_edges(self) -> Union[list[tuple[str, str, dict]], None]:
         raise NotImplementedError
 
+    @abstractmethod
+    def get_edge_count(self) -> int:
+        pass
+
+    @abstractmethod
     def get_node_edges(self, source_node_id: str) -> Union[list[tuple[str, str]], None]:
         raise NotImplementedError
 
+    @abstractmethod
     def upsert_node(self, node_id: str, node_data: dict[str, str]):
         raise NotImplementedError
 
+    @abstractmethod
     def upsert_edge(
         self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
     ):
         raise NotImplementedError
 
+    @abstractmethod
     def delete_node(self, node_id: str):
         raise NotImplementedError
 
+    @abstractmethod
     def reload(self):
         raise NotImplementedError
+
+    @abstractmethod
+    def get_connected_components(self, undirected: bool = True) -> List[Set[str]]:
+        raise NotImplementedError

graphgen/common/init_storage.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Union
+from typing import Any, Dict, List, Set, Union
 
 import ray
 
@@ -68,6 +68,21 @@ def __init__(self, backend: str, working_dir: str, namespace: str):
     def index_done_callback(self):
         return self.graph.index_done_callback()
 
+    def is_directed(self) -> bool:
+        return self.graph.is_directed()
+
+    def get_all_node_degrees(self) -> Dict[str, int]:
+        return self.graph.get_all_node_degrees()
+
+    def get_node_count(self) -> int:
+        return self.graph.get_node_count()
+
+    def get_edge_count(self) -> int:
+        return self.graph.get_edge_count()
+
+    def get_connected_components(self, undirected: bool = True) -> List[Set[str]]:
+        return self.graph.get_connected_components(undirected)
+
     def has_node(self, node_id: str) -> bool:
         return self.graph.has_node(node_id)
 
@@ -165,6 +180,21 @@ def __init__(self, actor_handle: ray.actor.ActorHandle):
     def index_done_callback(self):
         return ray.get(self.actor.index_done_callback.remote())
 
+    def is_directed(self) -> bool:
+        return ray.get(self.actor.is_directed.remote())
+
+    def get_all_node_degrees(self) -> Dict[str, int]:
+        return ray.get(self.actor.get_all_node_degrees.remote())
+
+    def get_node_count(self) -> int:
+        return ray.get(self.actor.get_node_count.remote())
+
+    def get_edge_count(self) -> int:
+        return ray.get(self.actor.get_edge_count.remote())
+
+    def get_connected_components(self, undirected: bool = True) -> List[Set[str]]:
+        return ray.get(self.actor.get_connected_components.remote(undirected))
+
     def has_node(self, node_id: str) -> bool:
         return ray.get(self.actor.has_node.remote(node_id))
 
@@ -239,10 +269,14 @@ def create_storage(backend: str, working_dir: str, namespace: str):
         try:
             actor_handle = ray.get_actor(actor_name)
         except ValueError:
-            actor_handle = ray.remote(actor_class).options(
-                name=actor_name,
-                get_if_exists=True,
-            ).remote(backend, working_dir, namespace)
+            actor_handle = (
+                ray.remote(actor_class)
+                .options(
+                    name=actor_name,
+                    get_if_exists=True,
+                )
+                .remote(backend, working_dir, namespace)
+            )
             ray.get(actor_handle.ready.remote())
         return proxy_class(actor_handle)
 

graphgen/models/__init__.py

@@ -1,8 +1,10 @@
 from .evaluator import (
-    KGQualityEvaluator,
+    AccuracyEvaluator,
+    ConsistencyEvaluator,
     LengthEvaluator,
     MTLDEvaluator,
     RewardEvaluator,
+    StructureEvaluator,
     UniEvaluator,
 )
 from .generator import (

graphgen/models/evaluator/__init__.py

@@ -1,7 +1,2 @@
+from .kg import AccuracyEvaluator, ConsistencyEvaluator, StructureEvaluator
 from .qa import LengthEvaluator, MTLDEvaluator, RewardEvaluator, UniEvaluator
-from .kg import (
-    AccuracyEvaluator,
-    ConsistencyEvaluator,
-    KGQualityEvaluator,
-    StructureEvaluator,
-)

graphgen/models/evaluator/kg/__init__.py

@@ -9,12 +9,10 @@
 
 from .accuracy_evaluator import AccuracyEvaluator
 from .consistency_evaluator import ConsistencyEvaluator
-from .kg_quality_evaluator import KGQualityEvaluator
 from .structure_evaluator import StructureEvaluator
 
 __all__ = [
     "AccuracyEvaluator",
     "ConsistencyEvaluator",
-    "KGQualityEvaluator",
     "StructureEvaluator",
 ]