feat: support partitioning the graph into quintuples

Author: beanbun

examples/generate/generate_masked_fill_in_blank_qa/masked_fill_in_blank_config.yaml

@@ -38,7 +38,7 @@ nodes:
     dependencies:
       - build_kg
     params:
-      method: triple 
+      method: quintuple
 
   - id: generate
     op_name: generate

graphgen/models/__init__.py

@@ -15,13 +15,13 @@
         AtomicGenerator,
         CoTGenerator,
         FillInBlankGenerator,
+        MaskedFillInBlankGenerator,
         MultiAnswerGenerator,
         MultiChoiceGenerator,
         MultiHopGenerator,
         QuizGenerator,
         TrueFalseGenerator,
         VQAGenerator,
-        MaskedFillInBlankGenerator,
     )
     from .kg_builder import LightRAGKGBuilder, MMKGBuilder
     from .llm import HTTPClient, OllamaClient, OpenAIClient
@@ -31,6 +31,7 @@
         DFSPartitioner,
         ECEPartitioner,
         LeidenPartitioner,
+        QuintuplePartitioner,
         TriplePartitioner,
     )
     from .reader import (
@@ -90,6 +91,7 @@
     "ECEPartitioner": ".partitioner",
     "LeidenPartitioner": ".partitioner",
     "TriplePartitioner": ".partitioner",
+    "QuintuplePartitioner": ".partitioner",
     # Reader
     "CSVReader": ".reader",
     "JSONReader": ".reader",

graphgen/models/generator/masked_fill_in_blank_generator.py

@@ -96,22 +96,35 @@ async def generate(
         if not context:
             return []
 
-        nodes, edge = batch
-        assert (
-            len(nodes) == 2
-        ), "MaskedFillInBlankGenerator currently only supports triples, which should has 2 nodes."
-        assert (
-            len(edge) == 1
-        ), "MaskedFillInBlankGenerator currently only supports triples, which should has 1 edge."
+        nodes, edges = batch
 
-        node1, node2 = nodes
-        mask_node = random.choice([node1, node2])
-        mask_node_name = mask_node[1]["entity_name"].strip("'\" \n\r\t")
+        assert len(nodes) == 3, (
+            "MaskedFillInBlankGenerator currently only supports quintuples that has 3 nodes, "
+            f"but got {len(nodes)} nodes."
+        )
+        assert len(edges) == 2, (
+            "MaskedFillInBlankGenerator currently only supports quintuples that has 2 edges, "
+            f"but got {len(edges)} edges."
+        )
 
+        node1, node2, node3 = nodes
+        mask_node = random.choice([node1, node2, node3])
+        mask_node_name = mask_node[1]["entity_name"].strip("'\" \n\r\t")
         mask_pattern = re.compile(re.escape(mask_node_name), re.IGNORECASE)
-        masked_context = mask_pattern.sub("___", context)
-        # For accuracy, extract the actual replaced text from the context as the ground truth
-        gth = re.search(mask_pattern, context).group(0)
+
+        match = re.search(mask_pattern, context)
+        if match:
+            gth = match.group(0)
+            masked_context = mask_pattern.sub("___", context)
+        else:
+            logger.debug(
+                "Regex Match Failed!\n"
+                "Expected name of node: %s\n"
+                "Actual context: %s\n",
+                mask_node_name,
+                context,
+            )
+            return []
 
         logger.debug("masked_context: %s", masked_context)
         qa_pairs = {

graphgen/models/partitioner/__init__.py

@@ -3,4 +3,5 @@
 from .dfs_partitioner import DFSPartitioner
 from .ece_partitioner import ECEPartitioner
 from .leiden_partitioner import LeidenPartitioner
+from .quintuple_partitioner import QuintuplePartitioner
 from .triple_partitioner import TriplePartitioner

graphgen/models/partitioner/quintuple_partitioner.py

@@ -0,0 +1,74 @@
+import random
+from collections import deque
+from typing import Any, Iterable, Set
+
+from graphgen.bases import BaseGraphStorage, BasePartitioner
+from graphgen.bases.datatypes import Community
+
+random.seed(42)
+
+
+class QuintuplePartitioner(BasePartitioner):
+    """
+    quintuple Partitioner that partitions the graph into multiple distinct quintuple (node, edge, node, edge, node).
+    1. Automatically ignore isolated points.
+    2. In each connected component, yield quintuples in the order of BFS.
+    """
+
+    def partition(
+        self,
+        g: BaseGraphStorage,
+        **kwargs: Any,
+    ) -> Iterable[Community]:
+        nodes = [n[0] for n in g.get_all_nodes()]
+        random.shuffle(nodes)
+
+        visited_nodes: Set[str] = set()
+        used_edges: Set[frozenset[str]] = set()
+
+        for seed in nodes:
+            if seed in visited_nodes:
+                continue
+
+            # start BFS in a connected component
+            queue = deque([seed])
+            visited_nodes.add(seed)
+
+            while queue:
+                u = queue.popleft()
+
+                # collect all neighbors connected to node u via unused edges
+                available_neighbors = []
+                for v in g.get_neighbors(u):
+                    edge_key = frozenset((u, v))
+                    if edge_key not in used_edges:
+                        available_neighbors.append(v)
+
+                    # standard BFS queue maintenance
+                    if v not in visited_nodes:
+                        visited_nodes.add(v)
+                        queue.append(v)
+
+                random.shuffle(available_neighbors)
+
+                # every two neighbors paired with the center node u creates one quintuple
+                # Note: If available_neighbors has an odd length, the remaining edge
+                # stays unused for now. It may be matched into a quintuple later
+                # when its other endpoint is processed as a center node.
+                for i in range(0, len(available_neighbors) // 2 * 2, 2):
+                    v1 = available_neighbors[i]
+                    v2 = available_neighbors[i + 1]
+
+                    edge1 = frozenset((u, v1))
+                    edge2 = frozenset((u, v2))
+
+                    used_edges.add(edge1)
+                    used_edges.add(edge2)
+
+                    v1_s, v2_s = sorted((v1, v2))
+
+                    yield Community(
+                        id=f"{v1_s}-{u}-{v2_s}",
+                        nodes=[v1_s, u, v2_s],
+                        edges=[tuple(sorted((v1_s, u))), tuple(sorted((u, v2_s)))],
+                    )

graphgen/operators/partition/partition_service.py

@@ -61,6 +61,10 @@ def __init__(
             from graphgen.models import TriplePartitioner
 
             self.partitioner = TriplePartitioner()
+        elif method == "quintuple":
+            from graphgen.models import QuintuplePartitioner
+
+            self.partitioner = QuintuplePartitioner()
         else:
             raise ValueError(f"Unsupported partition method: {method}")