Spaces:
Running
Running
| import json | |
| import os | |
| from typing import Iterable | |
| import pandas as pd | |
| from graphgen.bases import BaseGraphStorage, BaseKVStorage, BaseOperator, BaseTokenizer | |
| from graphgen.common import init_storage | |
| from graphgen.models import ( | |
| AnchorBFSPartitioner, | |
| BFSPartitioner, | |
| DFSPartitioner, | |
| ECEPartitioner, | |
| LeidenPartitioner, | |
| Tokenizer, | |
| ) | |
| from graphgen.utils import logger | |
| class PartitionService(BaseOperator): | |
| def __init__( | |
| self, | |
| working_dir: str = "cache", | |
| graph_backend: str = "kuzu", | |
| kv_backend: str = "rocksdb", | |
| **partition_kwargs, | |
| ): | |
| super().__init__(working_dir=working_dir, op_name="partition_service") | |
| self.kg_instance: 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", | |
| ) | |
| tokenizer_model = os.getenv("TOKENIZER_MODEL", "cl100k_base") | |
| self.tokenizer_instance: BaseTokenizer = Tokenizer(model_name=tokenizer_model) | |
| self.partition_kwargs = partition_kwargs | |
| def process(self, batch: pd.DataFrame) -> Iterable[pd.DataFrame]: | |
| # this operator does not consume any batch data | |
| # but for compatibility we keep the interface | |
| _ = batch.to_dict(orient="records") | |
| self.kg_instance.reload() | |
| self.chunk_storage.reload() | |
| yield from self.partition() | |
| def partition(self) -> Iterable[pd.DataFrame]: | |
| method = self.partition_kwargs["method"] | |
| method_params = self.partition_kwargs["method_params"] | |
| if method == "bfs": | |
| logger.info("Partitioning knowledge graph using BFS method.") | |
| partitioner = BFSPartitioner() | |
| elif method == "dfs": | |
| logger.info("Partitioning knowledge graph using DFS method.") | |
| partitioner = DFSPartitioner() | |
| elif method == "ece": | |
| logger.info("Partitioning knowledge graph using ECE method.") | |
| # TODO： before ECE partitioning, we need to: | |
| # 1. 'quiz' and 'judge' to get the comprehension loss if unit_sampling is not random | |
| # 2. pre-tokenize nodes and edges to get the token length | |
| self._pre_tokenize() | |
| partitioner = ECEPartitioner() | |
| elif method == "leiden": | |
| logger.info("Partitioning knowledge graph using Leiden method.") | |
| partitioner = LeidenPartitioner() | |
| elif method == "anchor_bfs": | |
| logger.info("Partitioning knowledge graph using Anchor BFS method.") | |
| partitioner = AnchorBFSPartitioner( | |
| anchor_type=method_params.get("anchor_type"), | |
| anchor_ids=set(method_params.get("anchor_ids", [])) | |
| if method_params.get("anchor_ids") | |
| else None, | |
| ) | |
| else: | |
| raise ValueError(f"Unsupported partition method: {method}") | |
| communities = partitioner.partition(g=self.kg_instance, **method_params) | |
| for community in communities: | |
| batch = partitioner.community2batch(community, g=self.kg_instance) | |
| batch = self._attach_additional_data_to_node(batch) | |
| yield pd.DataFrame( | |
| { | |
| "nodes": [batch[0]], | |
| "edges": [batch[1]], | |
| } | |
| ) | |
| def _pre_tokenize(self) -> None: | |
| """Pre-tokenize all nodes and edges to add token length information.""" | |
| logger.info("Starting pre-tokenization of nodes and edges...") | |
| nodes = self.kg_instance.get_all_nodes() | |
| edges = self.kg_instance.get_all_edges() | |
| # Process nodes | |
| for node_id, node_data in nodes: | |
| if "length" not in node_data: | |
| try: | |
| description = node_data.get("description", "") | |
| tokens = self.tokenizer_instance.encode(description) | |
| node_data["length"] = len(tokens) | |
| self.kg_instance.update_node(node_id, node_data) | |
| except Exception as e: | |
| logger.warning("Failed to tokenize node %s: %s", node_id, e) | |
| node_data["length"] = 0 | |
| # Process edges | |
| for u, v, edge_data in edges: | |
| if "length" not in edge_data: | |
| try: | |
| description = edge_data.get("description", "") | |
| tokens = self.tokenizer_instance.encode(description) | |
| edge_data["length"] = len(tokens) | |
| self.kg_instance.update_edge(u, v, edge_data) | |
| except Exception as e: | |
| logger.warning("Failed to tokenize edge %s-%s: %s", u, v, e) | |
| edge_data["length"] = 0 | |
| # Persist changes | |
| self.kg_instance.index_done_callback() | |
| logger.info("Pre-tokenization completed.") | |
| def _attach_additional_data_to_node(self, batch: tuple) -> tuple: | |
| """ | |
| Attach additional data from chunk_storage to nodes in the batch. | |
| :param batch: tuple of (nodes_data, edges_data) | |
| :return: updated batch with additional data attached to nodes | |
| """ | |
| nodes_data, edges_data = batch | |
| for node_id, node_data in nodes_data: | |
| entity_type = (node_data.get("entity_type") or "").lower() | |
| if not entity_type: | |
| continue | |
| source_ids = [ | |
| sid.strip() | |
| for sid in node_data.get("source_id", "").split("<SEP>") | |
| if sid.strip() | |
| ] | |
| # Handle images | |
| if "image" in entity_type: | |
| image_chunks = [ | |
| data | |
| for sid in source_ids | |
| if "image" in sid.lower() | |
| and (data := self.chunk_storage.get_by_id(sid)) | |
| ] | |
| if image_chunks: | |
| # The generator expects a dictionary with an 'img_path' key, not a list of captions. | |
| # We'll use the first image chunk found for this node. | |
| node_data["image_data"] = json.loads(image_chunks[0]["content"]) | |
| logger.debug("Attached image data to node %s", node_id) | |
| return nodes_data, edges_data | |