refactor: remove KGQualityEvaluator and restructure KG evaluation integration

Author: CHERRY-ui8

graphgen/models/__init__.py

@@ -1,5 +1,4 @@
 from .evaluator import (
-    KGQualityEvaluator,
     LengthEvaluator,
     MTLDEvaluator,
     RewardEvaluator,

graphgen/models/evaluator/__init__.py

@@ -2,6 +2,5 @@
 from .kg import (
     AccuracyEvaluator,
     ConsistencyEvaluator,
-    KGQualityEvaluator,
     StructureEvaluator,
 )

graphgen/models/evaluator/kg/README.md

@@ -4,12 +4,13 @@ This module provides comprehensive quality evaluation for knowledge graphs built
 
 ## Module Structure
 
-The evaluation functionality has been split into modular components:
+The evaluation functionality is organized into modular components:
 
 - **`accuracy_evaluator.py`**: Entity/relation extraction quality evaluation using LLM-as-a-Judge
 - **`consistency_evaluator.py`**: Attribute value conflict detection
 - **`structure_evaluator.py`**: Graph structural robustness metrics
-- **`kg_quality_evaluator.py`**: Main evaluator class that integrates all modules
+
+The evaluation components are integrated in `graphgen/operators/evaluate/evaluate_kg.py`, which provides functions to create and use these evaluators.
 
 ## Features
 

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",
 ]

graphgen/models/evaluator/kg/kg_quality_evaluator.py

@@ -2,70 +2,86 @@
 
 from dotenv import load_dotenv
 
-from graphgen.models import KGQualityEvaluator
+from graphgen.bases import BaseGraphStorage, BaseKVStorage, BaseLLMWrapper
+from graphgen.common import init_llm, init_storage
+from graphgen.models.evaluator.kg.accuracy_evaluator import AccuracyEvaluator
+from graphgen.models.evaluator.kg.consistency_evaluator import ConsistencyEvaluator
+from graphgen.models.evaluator.kg.structure_evaluator import StructureEvaluator
 from graphgen.utils import logger
 
 # Load environment variables
 load_dotenv()
 
 
-def evaluate_accuracy(evaluator: KGQualityEvaluator) -> Dict[str, Any]:
-    """Evaluate accuracy of entity and relation extraction.
-    
-    Args:
-        evaluator: KGQualityEvaluator instance
+class KGEvaluators:
+    def __init__(
+        self,
+        working_dir: str = "cache",
+        graph_backend: str = "kuzu",
+        kv_backend: str = "rocksdb",
+        **kwargs
+    ):
+        # Initialize storage
+        self.graph_storage: BaseGraphStorage = init_storage(
+            backend=graph_backend, working_dir=working_dir, namespace="graph"
+        )
+        self.chunk_storage: BaseKVStorage = init_storage(
+            backend=kv_backend, working_dir=working_dir, namespace="chunk"
+        )
 
-    Returns:
-        Dictionary containing entity_accuracy and relation_accuracy metrics.
-    """
+        # Initialize LLM client
+        self.llm_client: BaseLLMWrapper = init_llm("synthesizer")
+        
+        # Initialize individual evaluators
+        self.accuracy_evaluator = AccuracyEvaluator(
+            graph_storage=self.graph_storage,
+            chunk_storage=self.chunk_storage,
+            llm_client=self.llm_client,
+        )
+        
+        self.consistency_evaluator = ConsistencyEvaluator(
+            graph_storage=self.graph_storage,
+            chunk_storage=self.chunk_storage,
+            llm_client=self.llm_client,
+        )
+        
+        # Structure evaluator doesn't need chunk_storage or llm_client
+        structure_params = kwargs.get("structure_params", {})
+        self.structure_evaluator = StructureEvaluator(
+            graph_storage=self.graph_storage,
+            **structure_params
+        )
+        
+        logger.info("KG evaluators initialized")
+
+
+def evaluate_accuracy(evaluators: KGEvaluators) -> Dict[str, Any]:
     logger.info("Running accuracy evaluation...")
-    results = evaluator.evaluate_accuracy()
+    results = evaluators.accuracy_evaluator.evaluate()
     logger.info("Accuracy evaluation completed")
     return results
 
 
-def evaluate_consistency(evaluator: KGQualityEvaluator) -> Dict[str, Any]:
-    """Evaluate consistency by detecting semantic conflicts.
-    
-    Args:
-        evaluator: KGQualityEvaluator instance
-        
-    Returns:
-        Dictionary containing consistency metrics including conflict_rate and conflicts.
-    """
+def evaluate_consistency(evaluators: KGEvaluators) -> Dict[str, Any]:
     logger.info("Running consistency evaluation...")
-    results = evaluator.evaluate_consistency()
+    results = evaluators.consistency_evaluator.evaluate()
     logger.info("Consistency evaluation completed")
     return results
 
 
-def evaluate_structure(evaluator: KGQualityEvaluator) -> Dict[str, Any]:
-    """Evaluate structural robustness of the graph.
-    
-    Args:
-        evaluator: KGQualityEvaluator instance
-        
-    Returns:
-        Dictionary containing structural metrics including noise_ratio, largest_cc_ratio, etc.
-    """
+def evaluate_structure(evaluators: KGEvaluators) -> Dict[str, Any]:
     logger.info("Running structural robustness evaluation...")
-    results = evaluator.evaluate_structure()
+    results = evaluators.structure_evaluator.evaluate()
     logger.info("Structural robustness evaluation completed")
     return results
 
 
-def evaluate_all(evaluator: KGQualityEvaluator) -> Dict[str, Any]:
-    """Run all evaluations (accuracy, consistency, structure).
-    
-    Args:
-        evaluator: KGQualityEvaluator instance
-        
-    Returns:
-        Dictionary containing all evaluation results with keys: accuracy, consistency, structure.
-    """
+def evaluate_all(evaluators: KGEvaluators) -> Dict[str, Any]:
     logger.info("Running all evaluations...")
-    results = evaluator.evaluate_all()
+    results = {
+        "accuracy": evaluate_accuracy(evaluators),
+        "consistency": evaluate_consistency(evaluators),
+        "structure": evaluate_structure(evaluators),
+    }
     logger.info("All evaluations completed")
     return results
-
-

graphgen/operators/evaluate/evaluate_kg.py

@@ -1,10 +1,15 @@
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Optional, Union
 
 import pandas as pd
 
-from graphgen.bases import BaseLLMWrapper, BaseOperator, QAPair
-from graphgen.common import init_llm
-from graphgen.models import KGQualityEvaluator
+from graphgen.bases import BaseOperator, QAPair
+from graphgen.operators.evaluate.evaluate_kg import (
+    KGEvaluators,
+    evaluate_accuracy,
+    evaluate_all,
+    evaluate_consistency,
+    evaluate_structure,
+)
 from graphgen.utils import logger, run_concurrent
 
 
@@ -23,7 +28,6 @@ def __init__(
         **kwargs
     ):
         super().__init__(working_dir=working_dir, op_name="evaluate_service")
-        self.llm_client: BaseLLMWrapper = init_llm("synthesizer")
         self.metrics = metrics or []
         self.kwargs = kwargs
         self.graph_backend = graph_backend
@@ -35,7 +39,7 @@ def __init__(
 
         # Initialize evaluators
         self.qa_evaluators = {}
-        self.kg_evaluator = None
+        self.kg_evaluators: Optional[KGEvaluators] = None
 
         self._init_evaluators()
 
@@ -65,16 +69,15 @@ def _init_evaluators(self):
             else:
                 raise ValueError(f"Unknown QA metric: {metric}")
 
-        # Initialize KG evaluator if KG metrics are specified
+        # Initialize KG evaluators if KG metrics are specified
         if self.kg_metrics:
             kg_params = self.kwargs.get("kg_params", {})
-            self.kg_evaluator = KGQualityEvaluator(
+            self.kg_evaluators = KGEvaluators(
                 working_dir=self.working_dir,
                 graph_backend=self.graph_backend,
                 kv_backend=self.kv_backend,
                 **kg_params
             )
-            logger.info("KG evaluator initialized")
 
     async def _process_single(self, item: dict[str, Any]) -> dict[str, Any]:
         try:
@@ -140,17 +143,17 @@ def _evaluate_qa(self, items: list[dict[str, Any]]) -> list[dict[str, Any]]:
         return results
 
     def _evaluate_kg(self) -> Dict[str, Any]:
-        if not self.kg_evaluator:
-            logger.warning("No KG evaluator initialized, skipping KG evaluation")
+        if not self.kg_evaluators:
+            logger.warning("No KG evaluators initialized, skipping KG evaluation")
             return {}
 
         results = {}
 
         # Map metric names to evaluation functions
         kg_metric_map = {
-            "kg_accuracy": self.kg_evaluator.evaluate_accuracy,
-            "kg_consistency": self.kg_evaluator.evaluate_consistency,
-            "kg_structure": self.kg_evaluator.evaluate_structure,
+            "kg_accuracy": evaluate_accuracy,
+            "kg_consistency": evaluate_consistency,
+            "kg_structure": evaluate_structure,
         }
 
         # Run KG evaluations based on metrics
@@ -159,7 +162,7 @@ def _evaluate_kg(self) -> Dict[str, Any]:
                 logger.info("Running %s evaluation...", metric)
                 metric_key = metric.replace("kg_", "")  # Remove "kg_" prefix
                 try:
-                    results[metric_key] = kg_metric_map[metric]()
+                    results[metric_key] = kg_metric_map[metric](self.kg_evaluators)
                 except Exception as e:
                     logger.error("Error in %s evaluation: %s", metric, str(e))
                     results[metric_key] = {"error": str(e)}
@@ -169,7 +172,7 @@ def _evaluate_kg(self) -> Dict[str, Any]:
         # If no valid metrics were found, run all evaluations
         if not results:
             logger.info("No valid KG metrics found, running all evaluations")
-            results = self.kg_evaluator.evaluate_all()
+            results = evaluate_all(self.kg_evaluators)
 
         return results