feat: add KGQualityEvaluator and integrate into EvaluateService for KG evaluations

Author: CHERRY-ui8

graphgen/models/__init__.py

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

graphgen/models/evaluator/__init__.py

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

graphgen/models/evaluator/kg/__init__.py

@@ -9,10 +9,12 @@
 
 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

@@ -0,0 +1,79 @@
+from typing import Any, Dict
+
+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
+
+
+class KGQualityEvaluator:
+    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"
+        )
+        
+        # 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("KGQualityEvaluator initialized")
+
+    def evaluate_accuracy(self) -> Dict[str, Any]:
+        logger.info("Running accuracy evaluation...")
+        results = self.accuracy_evaluator.evaluate()
+        logger.info("Accuracy evaluation completed")
+        return results
+
+    def evaluate_consistency(self) -> Dict[str, Any]:
+        logger.info("Running consistency evaluation...")
+        results = self.consistency_evaluator.evaluate()
+        logger.info("Consistency evaluation completed")
+        return results
+
+    def evaluate_structure(self) -> Dict[str, Any]:
+        logger.info("Running structural robustness evaluation...")
+        results = self.structure_evaluator.evaluate()
+        logger.info("Structural robustness evaluation completed")
+        return results
+
+    def evaluate_all(self) -> Dict[str, Any]:
+        logger.info("Running all KG evaluations...")
+        results = {
+            "accuracy": self.evaluate_accuracy(),
+            "consistency": self.evaluate_consistency(),
+            "structure": self.evaluate_structure(),
+        }
+        logger.info("All KG evaluations completed")
+        return results
+

graphgen/operators/evaluate/evaluate_service.py

@@ -1,10 +1,11 @@
-from typing import Any
+from typing import Any, Dict, List, Union
 
 import pandas as pd
 
 from graphgen.bases import BaseLLMWrapper, BaseOperator, QAPair
 from graphgen.common import init_llm
-from graphgen.utils import run_concurrent
+from graphgen.models import KGQualityEvaluator
+from graphgen.utils import logger, run_concurrent
 
 
 class EvaluateService(BaseOperator):
@@ -13,40 +14,67 @@ class EvaluateService(BaseOperator):
     2. QA Quality Evaluation
     """
 
-    def __init__(self, working_dir: str = "cache", metrics: list[str] = None, **kwargs):
+    def __init__(
+        self,
+        working_dir: str = "cache",
+        metrics: list[str] = None,
+        graph_backend: str = "kuzu",
+        kv_backend: str = "rocksdb",
+        **kwargs
+    ):
         super().__init__(working_dir=working_dir, op_name="evaluate_service")
         self.llm_client: BaseLLMWrapper = init_llm("synthesizer")
-        self.metrics = metrics
+        self.metrics = metrics or []
         self.kwargs = kwargs
-        self.evaluators = {}
+        self.graph_backend = graph_backend
+        self.kv_backend = kv_backend
+        
+        # Separate QA and KG metrics
+        self.qa_metrics = [m for m in self.metrics if m.startswith("qa_")]
+        self.kg_metrics = [m for m in self.metrics if m.startswith("kg_")]
+        
+        # Initialize evaluators
+        self.qa_evaluators = {}
+        self.kg_evaluator = None
+        
         self._init_evaluators()
 
     def _init_evaluators(self):
-        for metric in self.metrics:
+        """Initialize QA and KG evaluators based on metrics."""
+        # Initialize QA evaluators
+        for metric in self.qa_metrics:
             if metric == "qa_length":
                 from graphgen.models import LengthEvaluator
 
-                self.evaluators[metric] = LengthEvaluator()
+                self.qa_evaluators[metric] = LengthEvaluator()
             elif metric == "qa_mtld":
                 from graphgen.models import MTLDEvaluator
-
-                self.evaluators[metric] = MTLDEvaluator(
+                self.qa_evaluators[metric] = MTLDEvaluator(
                     **self.kwargs.get("mtld_params", {})
                 )
             elif metric == "qa_reward_score":
                 from graphgen.models import RewardEvaluator
-
-                self.evaluators[metric] = RewardEvaluator(
+                self.qa_evaluators[metric] = RewardEvaluator(
                     **self.kwargs.get("reward_params", {})
                 )
             elif metric == "qa_uni_score":
                 from graphgen.models import UniEvaluator
-
-                self.evaluators[metric] = UniEvaluator(
+                self.qa_evaluators[metric] = UniEvaluator(
                     **self.kwargs.get("uni_params", {})
                 )
             else:
-                raise ValueError(f"Unknown metric: {metric}")
+                raise ValueError(f"Unknown QA metric: {metric}")
+        
+        # Initialize KG evaluator if KG metrics are specified
+        if self.kg_metrics:
+            kg_params = self.kwargs.get("kg_params", {})
+            self.kg_evaluator = KGQualityEvaluator(
+                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:
@@ -61,7 +89,7 @@ async def _process_single(self, item: dict[str, Any]) -> dict[str, Any]:
             self.logger.error("Error in QAPair creation: %s", str(e))
             return {}
 
-        for metric, evaluator in self.evaluators.items():
+        for metric, evaluator in self.qa_evaluators.items():
             try:
                 score = evaluator.evaluate(qa_pair)
                 if isinstance(score, dict):
@@ -92,18 +120,98 @@ def transform_messages_format(items: list[dict]) -> list[dict]:
             transformed.append({"question": question, "answer": answer})
         return transformed
 
-    def evaluate(self, items: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    def _evaluate_qa(self, items: list[dict[str, Any]]) -> list[dict[str, Any]]:
         if not items:
             return []
 
+        if not self.qa_evaluators:
+            logger.warning("No QA evaluators initialized, skipping QA evaluation")
+            return []
+
         items = self.transform_messages_format(items)
         results = run_concurrent(
             self._process_single,
             items,
-            desc="Evaluating items",
+            desc="Evaluating QA items",
             unit="item",
         )
 
         results = [item for item in results if item]
+        return results
 
+    def _evaluate_kg(self) -> Dict[str, Any]:
+        if not self.kg_evaluator:
+            logger.warning("No KG evaluator 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,
+        }
+        
+        # Run KG evaluations based on metrics
+        for metric in self.kg_metrics:
+            if metric in kg_metric_map:
+                logger.info("Running %s evaluation...", metric)
+                metric_key = metric.replace("kg_", "")  # Remove "kg_" prefix
+                try:
+                    results[metric_key] = kg_metric_map[metric]()
+                except Exception as e:
+                    logger.error("Error in %s evaluation: %s", metric, str(e))
+                    results[metric_key] = {"error": str(e)}
+            else:
+                logger.warning("Unknown KG metric: %s, skipping", metric)
+        
+        # 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()
+        
         return results
+
+    def evaluate(
+        self, items: list[dict[str, Any]] = None
+    ) -> Union[List[Dict[str, Any]], Dict[str, Any]]:
+        # Determine evaluation type
+        has_qa_metrics = len(self.qa_metrics) > 0
+        has_kg_metrics = len(self.kg_metrics) > 0
+        
+        # If items provided and QA metrics exist, do QA evaluation
+        if items is not None and has_qa_metrics:
+            return self._evaluate_qa(items)
+        
+        # If KG metrics exist, do KG evaluation
+        if has_kg_metrics:
+            return self._evaluate_kg()
+        
+        # If no metrics specified, try to infer from context
+        if items is not None:
+            logger.warning("No QA metrics specified but items provided, skipping evaluation")
+            return []
+        else:
+            logger.warning("No metrics specified, skipping evaluation")
+            return {}
+
+    def process(self, batch: pd.DataFrame) -> pd.DataFrame:
+        has_qa_metrics = len(self.qa_metrics) > 0
+        has_kg_metrics = len(self.kg_metrics) > 0
+        
+        # QA evaluation: process batch items
+        if has_qa_metrics:
+            items = batch.to_dict(orient="records")
+            results = self._evaluate_qa(items)
+            return pd.DataFrame(results)
+        
+        # KG evaluation: evaluate from storage
+        if has_kg_metrics:
+            results = self._evaluate_kg()
+            # Convert dict to DataFrame (single row)
+            return pd.DataFrame([results])
+        
+        # No metrics specified
+        logger.warning("No metrics specified, returning empty DataFrame")
+        return pd.DataFrame()