fix: add param min_units_per_community

Author: ChenZiHong-Gavin

.pylintrc

@@ -308,7 +308,7 @@ max-public-methods=20
 max-returns=6
 
 # Maximum number of statements in function / method body.
-max-statements=50
+max-statements=60
 
 # Minimum number of public methods for a class (see R0903).
 min-public-methods=2

graphgen/configs/aggregated_config.yaml

@@ -14,6 +14,7 @@ partition: # graph partition configuration
   method: ece # ece is a custom partition method based on comprehension loss
   method_params:
     max_units_per_community: 20 # max nodes and edges per community
+    min_units_per_community: 5 # min nodes and edges per community
     max_tokens_per_community: 10240 # max tokens per community
     unit_sampling: max_loss # edge sampling strategy, support: random, max_loss, min_loss
 generate:

graphgen/models/partitioner/ece_partitioner.py

@@ -1,7 +1,7 @@
 import asyncio
 import random
 from dataclasses import dataclass
-from typing import Any, Dict, List, Set, Tuple
+from typing import Any, Dict, List, Optional, Set, Tuple
 
 from tqdm.asyncio import tqdm as tqdm_async
 
@@ -54,6 +54,7 @@ async def partition(
         self,
         g: BaseGraphStorage,
         max_units_per_community: int = 10,
+        min_units_per_community: int = 1,
         max_tokens_per_community: int = 10240,
         unit_sampling: str = "random",
         **kwargs: Any,
@@ -75,7 +76,9 @@ async def partition(
 
         all_units = self._sort_units(all_units, unit_sampling)
 
-        async def _grow_community(seed_unit: Tuple[str, Any, dict]) -> Community:
+        async def _grow_community(
+            seed_unit: Tuple[str, Any, dict]
+        ) -> Optional[Community]:
             nonlocal used_n, used_e
 
             community_nodes: Dict[str, dict] = {}
@@ -135,6 +138,9 @@ async def _add_unit(u):
                     if await _add_unit(nb):
                         await queue.put(nb)
 
+            if len(community_nodes) + len(community_edges) < min_units_per_community:
+                return None
+
             return Community(
                 id=len(communities),
                 nodes=list(community_nodes.keys()),
@@ -145,6 +151,8 @@ async def _add_unit(u):
             utype, uid, _ = unit
             if (utype == "n" and uid in used_n) or (utype == "e" and uid in used_e):
                 continue
-            communities.append(await _grow_community(unit))
+            comm = await _grow_community(unit)
+            if comm is not None:
+                communities.append(comm)
 
         return communities

graphgen/operators/__init__.py

@@ -1,8 +1,3 @@
-from graphgen.operators.partition.traverse_graph import (
-    traverse_graph_for_aggregated,
-    traverse_graph_for_multi_hop,
-)
-
 from .build_kg import build_kg
 from .generate import generate_qas
 from .judge import judge_statement