fix: fix output node

Author: ChenZiHong-Gavin

graphgen/engine.py

@@ -271,6 +271,8 @@ def execute(self, initial_ds: ray.data.Dataset) -> Dict[str, ray.data.Dataset]:
 
         for node in sorted_nodes:
             self._execute_node(node, initial_ds)
+            if getattr(node, "save_output", False):
+                self.datasets[node.id] = self.datasets[node.id].materialize()
 
         output_nodes = [n for n in sorted_nodes if getattr(n, "save_output", False)]
         return {node.id: self.datasets[node.id] for node in output_nodes}

graphgen/models/evaluator/qa/mtld_evaluator.py

@@ -11,8 +11,8 @@ class MTLDEvaluator(BaseEvaluator):
 
     def __init__(self, threshold: float = 0.72):
         self.nltk_helper = NLTKHelper()
-        self.stopwords_en: Set[str] = set(self.nltk_helper.get_stopwords("english"))
-        self.stopwords_zh: Set[str] = set(self.nltk_helper.get_stopwords("chinese"))
+        self.stopwords_en: Set[str] = set(self.nltk_helper.get_stopwords("en"))
+        self.stopwords_zh: Set[str] = set(self.nltk_helper.get_stopwords("zh"))
         self.threshold = threshold
 
     def evaluate(self, pair: QAPair) -> float:

graphgen/operators/evaluate/evaluate_service.py

@@ -1,4 +1,5 @@
 from typing import Any
+
 import pandas as pd
 
 from graphgen.bases import BaseLLMWrapper, BaseOperator, QAPair
@@ -18,21 +19,32 @@ def __init__(self, working_dir: str = "cache", metrics: list[str] = None, **kwar
         self.metrics = metrics
         self.kwargs = kwargs
         self.evaluators = {}
+        self._init_evaluators()
 
     def _init_evaluators(self):
         for metric in self.metrics:
             if metric == "qa_length":
                 from graphgen.models import LengthEvaluator
+
                 self.evaluators[metric] = LengthEvaluator()
             elif metric == "qa_mtld":
                 from graphgen.models import MTLDEvaluator
-                self.evaluators[metric] = MTLDEvaluator(self.kwargs.get("mtld_params", {}))
+
+                self.evaluators[metric] = MTLDEvaluator(
+                    **self.kwargs.get("mtld_params", {})
+                )
             elif metric == "qa_reward_score":
                 from graphgen.models import RewardEvaluator
-                self.evaluators[metric] = RewardEvaluator(self.kwargs.get("reward_params", {}))
+
+                self.evaluators[metric] = RewardEvaluator(
+                    **self.kwargs.get("reward_params", {})
+                )
             elif metric == "qa_uni_score":
                 from graphgen.models import UniEvaluator
-                self.evaluators[metric] = UniEvaluator(self.kwargs.get("uni_params", {}))
+
+                self.evaluators[metric] = UniEvaluator(
+                    **self.kwargs.get("uni_params", {})
+                )
             else:
                 raise ValueError(f"Unknown metric: {metric}")
 
@@ -44,16 +56,13 @@ async def _process_single(self, item: dict[str, Any]) -> dict[str, Any]:
         try:
             qa_pair = QAPair(
                 question=str(item.get("question", "")),
-                answer=str(item.get("answer", ""))
+                answer=str(item.get("answer", "")),
             )
             if not qa_pair.question or not qa_pair.answer:
                 self.logger.error("Empty question or answer, skipping.")
                 return {}
         except Exception as e:
-            self.logger.error(
-                "Error in QAPair creation: %s",
-                str(e)
-            )
+            self.logger.error("Error in QAPair creation: %s", str(e))
             return {}
 
         for metric, evaluator in self.evaluators.items():
@@ -65,17 +74,33 @@ async def _process_single(self, item: dict[str, Any]) -> dict[str, Any]:
                 else:
                     item[metric] = float(score)
             except Exception as e:
-                self.logger.error(
-                    "Error in %s evaluation: %s",
-                    metric,
-                    str(e)
-                )
+                self.logger.error("Error in %s evaluation: %s", metric, str(e))
                 item[metric] = None
+        return item
+
+    @staticmethod
+    def transform_messages_format(items: list[dict]) -> list[dict]:
+        """
+        Transform from [{'messages': [...]}, ...] to [{'question': '...', 'answer': '...'}, ...]
+        """
+        transformed = []
+        for item in items:
+            messages = item.get("messages", [])
+            question = next(
+                (m["content"] for m in messages if m.get("role") == "user"), ""
+            )
+            answer = next(
+                (m["content"] for m in messages if m.get("role") == "assistant"), ""
+            )
+
+            transformed.append({"question": question, "answer": answer})
+        return transformed
 
     def evaluate(self, items: list[dict[str, Any]]) -> list[dict[str, Any]]:
         if not items:
             return []
 
+        items = self.transform_messages_format(items)
         results = run_concurrent(
             self._process_single,
             items,

graphgen/run.py

@@ -91,18 +91,19 @@ def main():
     results = engine.execute(ds)
 
     for node_id, dataset in results.items():
-        output_path = os.path.join(output_path, f"{node_id}")
-        os.makedirs(output_path, exist_ok=True)
+        logger.info("Saving results for node %s", node_id)
+        node_output_path = os.path.join(output_path, f"{node_id}")
+        os.makedirs(node_output_path, exist_ok=True)
         dataset.write_json(
-            output_path,
+            node_output_path,
             filename_provider=NodeFilenameProvider(node_id),
             pandas_json_args_fn=lambda: {
                 "force_ascii": False,
                 "orient": "records",
                 "lines": True,
             },
         )
-        logger.info("Node %s results saved to %s", node_id, output_path)
+        logger.info("Node %s results saved to %s", node_id, node_output_path)
 
     save_config(os.path.join(output_path, "config.yaml"), config)
     logger.info("GraphGen completed successfully. Data saved to %s", output_path)