wip: move storage logic to baseOperator

Author: ChenZiHong-Gavin

graphgen/bases/base_operator.py

@@ -1,7 +1,7 @@
 import inspect
 import os
 from abc import ABC, abstractmethod
-from typing import Iterable, Union
+from typing import Iterable, Union, Tuple
 
 import pandas as pd
 import ray
@@ -64,12 +64,17 @@ def __call__(
             if to_process.empty:
                 return
 
-            docs = to_process.to_dict(orient="records")
-            result = self.process(docs)
+            data = to_process.to_dict(orient="records")
+            result, meta_update = self.process(data)
             if inspect.isgenerator(result):
-                yield from result
+                is_first = True
+                for res in result:
+                    yield pd.DataFrame([res])
+                    self.store([res], meta_update if is_first else {})
+                    is_first = False
             else:
-                yield result
+                yield pd.DataFrame(result)
+                self.store(result, meta_update)
         finally:
             CURRENT_LOGGER_VAR.reset(logger_token)
 
@@ -130,5 +135,11 @@ def store(self, results: list, meta_update: dict):
         self.kv_storage.index_done_callback()
 
     @abstractmethod
-    def process(self, batch: list) -> Union[pd.DataFrame, Iterable[pd.DataFrame]]:
-        pass
+    def process(self, batch: list) -> Tuple[Union[list, Iterable[list]], dict]:
+        """
+        Process the input batch and return the result.
+        :param batch
+        :return:
+            result: DataFrame of processed documents
+            meta_update: dict of meta data to be updated
+        """

graphgen/operators/build_kg/build_kg_service.py

@@ -1,4 +1,4 @@
-import pandas as pd
+from typing import Tuple
 
 from graphgen.bases import BaseGraphStorage, BaseLLMWrapper, BaseOperator
 from graphgen.bases.datatypes import Chunk
@@ -27,7 +27,7 @@ def __init__(
         self.build_kwargs = build_kwargs
         self.max_loop: int = int(self.build_kwargs.get("max_loop", 3))
 
-    def process(self, batch: list) -> pd.DataFrame:
+    def process(self, batch: list) -> Tuple[list, dict]:
         """
         Build knowledge graph (KG) and merge into kg_instance
         """
@@ -99,8 +99,4 @@ def process(self, batch: list) -> pd.DataFrame:
             source_ids = edge.get("source_id", "").split("<SEP>")
             for source_id in source_ids:
                 meta_updates.setdefault(source_id, []).append(str(trace_id))
-        self.store(
-            results,
-            meta_updates,
-        )
-        return pd.DataFrame(results)
+        return results, meta_updates

graphgen/operators/chunk/chunk_service.py

@@ -1,8 +1,6 @@
 import os
 from functools import lru_cache
-from typing import Union
-
-import pandas as pd
+from typing import Union, Tuple
 
 from graphgen.bases import BaseOperator
 from graphgen.models import (
@@ -51,7 +49,7 @@ def __init__(
         self.tokenizer_instance: Tokenizer = Tokenizer(model_name=tokenizer_model)
         self.chunk_kwargs = chunk_kwargs
 
-    def process(self, batch: list) -> pd.DataFrame:
+    def process(self, batch: list) -> Tuple[list, dict]:
         results = []
         meta_updates = {}
         for doc in batch:
@@ -72,17 +70,15 @@ def process(self, batch: list) -> pd.DataFrame:
                         else len(text_chunk),
                         "language": doc_language,
                     }
-                    chunk["_trace_id"] = self.generate_trace_id(chunk)
+                    chunk["_trace_id"] = self.get_trace_id(chunk)
                     results.append(chunk)
                     meta_updates.setdefault(doc["_trace_id"], []).append(
                         chunk["_trace_id"]
                     )
             else:
                 # other types of documents(images, sequences) are not chunked
-                doc["_trace_id"] = self.generate_trace_id(doc)
-                results.append(doc)
-        self.store(
-            results,
-            meta_updates,
-        )
-        return pd.DataFrame(results)
+                data = doc.copy()
+                data["_trace_id"] = self.get_trace_id(data)
+                results.append(data)
+                meta_updates.setdefault(doc["_trace_id"], []).append(data["_trace_id"])
+        return results, meta_updates

graphgen/operators/extract/extract_service.py

@@ -1,6 +1,5 @@
 import json
-
-import pandas as pd
+from typing import Tuple
 
 from graphgen.bases import BaseLLMWrapper, BaseOperator, Chunk
 from graphgen.common import init_llm
@@ -26,7 +25,7 @@ def __init__(
         else:
             raise ValueError(f"Unsupported extraction method: {self.method}")
 
-    def process(self, batch: list) -> pd.DataFrame:
+    def process(self, batch: list) -> Tuple[list, dict]:
         logger.info("Start extracting information from %d items", len(batch))
         chunks = [Chunk.from_dict(item["_trace_id"], item) for item in batch]
         results = run_concurrent(

graphgen/operators/generate/generate_service.py

@@ -1,5 +1,4 @@
-import pandas as pd
-
+from typing import Tuple
 from graphgen.bases import BaseKVStorage, BaseLLMWrapper, BaseOperator
 from graphgen.common import init_llm, init_storage
 from graphgen.utils import logger, run_concurrent
@@ -80,11 +79,9 @@ def __init__(
         else:
             raise ValueError(f"Unsupported generation mode: {method}")
 
-    def process(self, batch: list[dict]) -> pd.DataFrame:
+    def process(self, batch: list) -> Tuple[list, dict]:
         """
         Generate question-answer pairs based on nodes and edges.
-        :param batch
-        :return: QA pairs
         """
         logger.info("[Generation] mode: %s, batches: %d", self.method, len(batch))
         triples = [(item["nodes"], item["edges"]) for item in batch]
@@ -106,11 +103,7 @@ def process(self, batch: list[dict]) -> pd.DataFrame:
                 res = self.generator.format_generation_results(
                     qa_pair, output_data_format=self.data_format
                 )
-                res["_trace_id"] = self.generate_trace_id(res)
+                res["_trace_id"] = self.get_trace_id(res)
                 final_results.append(res)
                 meta_updates.setdefault(input_trace_id, []).append(res["_trace_id"])
-        self.store(
-            final_results,
-            meta_updates,
-        )
-        return pd.DataFrame(final_results)
+        return final_results, meta_updates

graphgen/operators/judge/judge_service.py

@@ -1,7 +1,6 @@
+from typing import Tuple
 import math
 
-import pandas as pd
-
 from graphgen.bases import BaseGraphStorage, BaseLLMWrapper, BaseOperator
 from graphgen.common import init_llm, init_storage
 from graphgen.templates import STATEMENT_JUDGEMENT_PROMPT
@@ -44,7 +43,7 @@ async def _process_single_judge(self, item: dict) -> dict:
             item["loss"] = -math.log(0.1)
         return item
 
-    def process(self, batch: list) -> pd.DataFrame:
+    def process(self, batch: list) -> Tuple[list, dict]:
         """
         Judge the description in the item and compute the loss.
         """
@@ -78,10 +77,9 @@ def process(self, batch: list) -> pd.DataFrame:
                 edge_data["loss"] = loss
                 self.graph_storage.update_edge(edge_source, edge_target, edge_data)
 
-            result["_trace_id"] = self.generate_trace_id(result)
+            result["_trace_id"] = self.get_trace_id(result)
             to_store.append(result)
             meta_update.setdefault(input_trace_id, []).append(result["_trace_id"])
         self.graph_storage.index_done_callback()
-        self.store(to_store, meta_update)
 
-        return pd.DataFrame(results)
+        return results, meta_update

graphgen/operators/partition/partition_service.py

@@ -1,7 +1,5 @@
 import os
-from typing import Iterable
-
-import pandas as pd
+from typing import Iterable, Tuple
 
 from graphgen.bases import BaseGraphStorage, BaseOperator, BaseTokenizer
 from graphgen.common import init_storage
@@ -24,7 +22,9 @@ def __init__(
         graph_backend: str = "kuzu",
         **partition_kwargs,
     ):
-        super().__init__(working_dir=working_dir, op_name="partition")
+        super().__init__(
+            working_dir=working_dir, kv_backend=kv_backend, op_name="partition"
+        )
         self.kg_instance: BaseGraphStorage = init_storage(
             backend=graph_backend,
             working_dir=working_dir,
@@ -55,7 +55,7 @@ def __init__(
         else:
             raise ValueError(f"Unsupported partition method: {method}")
 
-    def process(self, batch: pd.DataFrame) -> Iterable[pd.DataFrame]:
+    def process(self, batch: list) -> Tuple[Iterable[list], dict]:
         # this operator does not consume any batch data
         # but for compatibility we keep the interface
         self.kg_instance.reload()
@@ -64,19 +64,22 @@ def process(self, batch: pd.DataFrame) -> Iterable[pd.DataFrame]:
             g=self.kg_instance, **self.method_params
         )
 
-        count = 0
-        for community in communities:
-            count += 1
-            batch = self.partitioner.community2batch(community, g=self.kg_instance)
-            # batch = self._attach_additional_data_to_node(batch)
+        def generator():
+            count = 0
+            for community in communities:
+                count += 1
+                batch = self.partitioner.community2batch(community, g=self.kg_instance)
+                # batch = self._attach_additional_data_to_node(batch)
+
+                result = {
+                    "nodes": batch[0],
+                    "edges": batch[1],
+                }
+                result["_trace_id"] = self.get_trace_id(result)
+                yield result
+            logger.info("Total communities partitioned: %d", count)
 
-            result = {
-                "nodes": batch[0],
-                "edges": batch[1],
-            }
-            result["_trace_id"] = self.generate_trace_id(result)
-            yield pd.DataFrame([result])
-        logger.info("Total communities partitioned: %d", count)
+        return generator(), {}
 
     # def _attach_additional_data_to_node(self, batch: tuple) -> tuple:
     #     """

graphgen/operators/quiz/quiz_service.py

@@ -1,4 +1,4 @@
-import pandas as pd
+from typing import Tuple
 
 from graphgen.bases import BaseGraphStorage, BaseLLMWrapper, BaseOperator
 from graphgen.common import init_llm, init_storage
@@ -50,7 +50,7 @@ async def _process_single_quiz(self, item: tuple) -> dict | None:
             logger.error("Error when quizzing description %s: %s", item, e)
             return None
 
-    def process(self, batch: list) -> pd.DataFrame:
+    def process(self, batch: list) -> Tuple[list, dict]:
         """
         Get all nodes and edges and quiz their descriptions using QuizGenerator.
         """
@@ -67,7 +67,7 @@ def process(self, batch: list) -> pd.DataFrame:
                 edge_key = (edge["src_id"], edge["tgt_id"])
                 items.append((input_id, edge["description"], edge_key))
         if not items:
-            return pd.DataFrame()
+            return [], {}
 
         logger.info("Total descriptions to quiz: %d", len(items))
         results = run_concurrent(
@@ -83,11 +83,8 @@ def process(self, batch: list) -> pd.DataFrame:
         for (input_id, _, _), quiz_data in zip(items, results):
             if quiz_data is None:
                 continue
-            quiz_data["_trace_id"] = self.generate_trace_id(quiz_data)
+            quiz_data["_trace_id"] = self.get_trace_id(quiz_data)
             final_results.append(quiz_data)
             meta_update[input_id] = [quiz_data["_trace_id"]]
 
-        if final_results:
-            self.store(final_results, meta_update)
-
-        return pd.DataFrame(final_results)
+        return final_results, meta_update