wip: add checkpoint

Author: ChenZiHong-Gavin

graphgen/bases/base_generator.py

@@ -42,46 +42,36 @@ async def generate(
 
     @staticmethod
     def format_generation_results(
-        results: list[dict], output_data_format: str
-    ) -> list[dict[str, Any]]:
+        result: dict, output_data_format: str
+    ) -> dict[str, Any]:
+        question = result.get("question", "")
+        answer = result.get("answer", "")
+        if "options" in result and result["options"]:
+            options = result["options"]
+            options_str = "\n".join(
+                [f"{key}. {options[key]}" for key in sorted(options.keys())]
+            )
+            question += f"\nOptions:\n{options_str}"
 
-        flat_results = []
-        for qa_data in results:
-            question = qa_data.get("question", "")
-            answer = qa_data.get("answer", "")
-            if "options" in qa_data and qa_data["options"]:
-                options = qa_data["options"]
-                options_str = "\n".join(
-                    [f"{key}. {options[key]}" for key in sorted(options.keys())]
-                )
-                question += f"\nOptions:\n{options_str}"
+        if output_data_format == "Alpaca":
+            return {
+                "instruction": question,
+                "input": "",
+                "output": answer,
+            }
 
-            if output_data_format == "Alpaca":
-                flat_results.append(
-                    {
-                        "instruction": question,
-                        "input": "",
-                        "output": answer,
-                    }
-                )
-            elif output_data_format == "Sharegpt":
-                flat_results.append(
-                    {
-                        "conversations": [
-                            {"from": "human", "value": question},
-                            {"from": "gpt", "value": answer},
-                        ]
-                    }
-                )
-            elif output_data_format == "ChatML":
-                flat_results.append(
-                    {
-                        "messages": [
-                            {"role": "user", "content": question},
-                            {"role": "assistant", "content": answer},
-                        ]
-                    }
-                )
-            else:
-                raise ValueError(f"Unknown output data format: {output_data_format}")
-        return flat_results
+        if output_data_format == "Sharegpt":
+            return {
+                "conversations": [
+                    {"from": "human", "value": question},
+                    {"from": "gpt", "value": answer},
+                ]
+            }
+        if output_data_format == "ChatML":
+            return {
+                "messages": [
+                    {"role": "user", "content": question},
+                    {"role": "assistant", "content": answer},
+                ]
+            }
+        raise ValueError(f"Unknown output data format: {output_data_format}")

graphgen/bases/base_operator.py

@@ -8,12 +8,22 @@
 
 
 class BaseOperator(ABC):
-    def __init__(self, working_dir: str = "cache", op_name: str = None):
+    def __init__(
+        self,
+        working_dir: str = "cache",
+        kv_backend: str = "rocksdb",
+        op_name: str = None,
+    ):
         # lazy import to avoid circular import
+        from graphgen.common import init_storage
         from graphgen.utils import set_logger
 
         log_dir = os.path.join(working_dir, "logs")
         self.op_name = op_name or self.__class__.__name__
+        self.working_dir = working_dir
+        self.kv_storage = init_storage(
+            backend=kv_backend, working_dir=working_dir, namespace=self.op_name
+        )
 
         try:
             ctx = ray.get_runtime_context()
@@ -45,17 +55,80 @@ def __call__(
 
         logger_token = CURRENT_LOGGER_VAR.set(self.logger)
         try:
-            result = self.process(batch)
+            self.kv_storage.reload()
+            to_process, recovered = self.split(batch)
+            # yield recovered chunks first
+            if not recovered.empty:
+                yield recovered
+
+            if to_process.empty:
+                return
+
+            docs = to_process.to_dict(orient="records")
+            result = self.process(docs)
             if inspect.isgenerator(result):
                 yield from result
             else:
                 yield result
         finally:
             CURRENT_LOGGER_VAR.reset(logger_token)
 
-    @abstractmethod
-    def process(self, batch):
-        raise NotImplementedError("Subclasses must implement the process method.")
-
     def get_logger(self):
         return self.logger
+
+    def get_meta_forward(self):
+        return self.kv_storage.get_by_id("_meta_forward") or {}
+
+    def get_meta_inverse(self):
+        return self.kv_storage.get_by_id("_meta_inverse") or {}
+
+    def get_trace_id(self, content: dict) -> str:
+        from graphgen.utils import compute_dict_hash
+
+        return compute_dict_hash(content, prefix=f"{self.op_name}-")
+
+    def split(self, batch: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
+        """
+        Split the input batch into to_process & processed based on _meta data in KV_storage
+        :param batch
+        :return:
+            to_process: DataFrame of documents to be chunked
+            recovered: Result DataFrame of already chunked documents
+        """
+        meta_forward = self.get_meta_forward()
+        meta_ids = set(meta_forward.keys())
+        mask = batch["_trace_id"].isin(meta_ids)
+        to_process = batch[~mask]
+        processed = batch[mask]
+
+        if processed.empty:
+            return to_process, pd.DataFrame()
+
+        all_ids = [
+            pid for tid in processed["_trace_id"] for pid in meta_forward.get(tid, [])
+        ]
+
+        recovered_chunks = self.kv_storage.get_by_ids(all_ids)
+        recovered_chunks = [c for c in recovered_chunks if c is not None]
+        return to_process, pd.DataFrame(recovered_chunks)
+
+    def store(self, results: list, meta_update: dict):
+        batch = {res["_trace_id"]: res for res in results}
+        self.kv_storage.upsert(batch)
+
+        # update forward meta
+        forward_meta = self.get_meta_forward()
+        forward_meta.update(meta_update)
+        self.kv_storage.update({"_meta_forward": forward_meta})
+
+        # update inverse meta
+        inverse_meta = self.get_meta_inverse()
+        for k, v_list in meta_update.items():
+            for v in v_list:
+                inverse_meta[v] = k
+        self.kv_storage.update({"_meta_inverse": inverse_meta})
+        self.kv_storage.index_done_callback()
+
+    @abstractmethod
+    def process(self, batch: list) -> Union[pd.DataFrame, Iterable[pd.DataFrame]]:
+        pass

graphgen/models/extractor/schema_guided_extractor.py

@@ -1,7 +1,7 @@
 import json
 from typing import Dict, List
 
-from graphgen.bases import BaseExtractor, BaseLLMWrapper
+from graphgen.bases import BaseExtractor, BaseLLMWrapper, Chunk
 from graphgen.templates import SCHEMA_GUIDED_EXTRACTION_PROMPT
 from graphgen.utils import compute_dict_hash, detect_main_language, logger
 
@@ -59,9 +59,9 @@ def build_prompt(self, text: str) -> str:
         )
         return prompt
 
-    async def extract(self, chunk: dict) -> dict:
+    async def extract(self, chunk: Chunk) -> dict:
         _chunk_id = chunk.get("_chunk_id", "")
-        text = chunk.get("content", "")
+        text = chunk.content
 
         prompt = self.build_prompt(text)
         response = await self.llm_client.generate_answer(prompt)

graphgen/operators/build_kg/build_kg_service.py

@@ -1,5 +1,3 @@
-from typing import List
-
 import pandas as pd
 
 from graphgen.bases import BaseGraphStorage, BaseLLMWrapper, BaseOperator
@@ -13,31 +11,27 @@
 
 class BuildKGService(BaseOperator):
     def __init__(
-        self, working_dir: str = "cache", graph_backend: str = "kuzu", **build_kwargs
+        self,
+        working_dir: str = "cache",
+        kv_backend: str = "rocksdb",
+        graph_backend: str = "kuzu",
+        **build_kwargs
     ):
-        super().__init__(working_dir=working_dir, op_name="build_kg_service")
+        super().__init__(
+            working_dir=working_dir, kv_backend=kv_backend, op_name="build_kg"
+        )
         self.llm_client: BaseLLMWrapper = init_llm("synthesizer")
         self.graph_storage: BaseGraphStorage = init_storage(
             backend=graph_backend, working_dir=working_dir, namespace="graph"
         )
         self.build_kwargs = build_kwargs
         self.max_loop: int = int(self.build_kwargs.get("max_loop", 3))
 
-    def process(self, batch: pd.DataFrame) -> pd.DataFrame:
-        docs = batch.to_dict(orient="records")
-        docs = [Chunk.from_dict(doc["_chunk_id"], doc) for doc in docs]
-
-        # consume the chunks and build kg
-        nodes, edges = self.build_kg(docs)
-        return pd.DataFrame(
-            [{"node": node, "edge": []} for node in nodes]
-            + [{"node": [], "edge": edge} for edge in edges]
-        )
-
-    def build_kg(self, chunks: List[Chunk]) -> tuple:
+    def process(self, batch: list) -> pd.DataFrame:
         """
         Build knowledge graph (KG) and merge into kg_instance
         """
+        chunks = [Chunk.from_dict(doc["_trace_id"], doc) for doc in batch]
         text_chunks = [chunk for chunk in chunks if chunk.type == "text"]
         mm_chunks = [
             chunk
@@ -75,4 +69,38 @@ def build_kg(self, chunks: List[Chunk]) -> tuple:
         self.graph_storage.index_done_callback()
         logger.info("Knowledge graph building completed.")
 
-        return nodes, edges
+        meta_updates = {}
+        results = []
+        for node in nodes:
+            if not node:
+                continue
+            trace_id = node["entity_name"]
+            results.append(
+                {
+                    "_trace_id": trace_id,
+                    "node": node,
+                    "edge": {},
+                }
+            )
+            source_ids = node.get("source_id", "").split("<SEP>")
+            for source_id in source_ids:
+                meta_updates.setdefault(source_id, []).append(trace_id)
+        for edge in edges:
+            if not edge:
+                continue
+            trace_id = frozenset((edge["src_id"], edge["tgt_id"]))
+            results.append(
+                {
+                    "_trace_id": str(trace_id),
+                    "node": {},
+                    "edge": edge,
+                }
+            )
+            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)

graphgen/operators/build_kg/build_text_kg.py

@@ -30,6 +30,7 @@ def build_text_kg(
         desc="[2/4]Extracting entities and relationships from chunks",
         unit="chunk",
     )
+    results = [res for res in results if res]
 
     nodes = defaultdict(list)
     edges = defaultdict(list)