style: fix formatting issues

Author: beanbun

graphgen/models/generator/masked_fill_in_blank_generator.py

@@ -1,5 +1,5 @@
-import re
 import random
+import re
 from typing import Any, Optional
 
 from graphgen.bases import BaseGenerator
@@ -8,6 +8,7 @@
 
 random.seed(42)
 
+
 class MaskedFillInBlankGenerator(BaseGenerator):
     """
     Masked Fill-in-blank Generator follows a TWO-STEP process:
@@ -94,18 +95,22 @@ async def generate(
         context = self.parse_rephrased_text(response)
         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."
+        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."
 
         node1, node2 = nodes
         mask_node = random.choice([node1, node2])
-        mask_node_name = mask_node[1]["entity_name"].strip('\'" \n\r\t')
+        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 (keeping the original case)
+        # For accuracy, extract the actual replaced text from the context as the ground truth
         gth = re.search(mask_pattern, context).group(0)
 
         logger.debug("masked_context: %s", masked_context)
@@ -114,4 +119,3 @@ async def generate(
             "answer": gth,
         }
         return [qa_pairs]
-

graphgen/models/partitioner/triple_partitioner.py

@@ -7,6 +7,7 @@
 
 random.seed(42)
 
+
 class TriplePartitioner(BasePartitioner):
     """
     Triple Partitioner that partitions the graph into multiple distinct triples (node, edge, node).
@@ -28,30 +29,30 @@ def partition(
         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()
-                
+
                 for v in g.get_neighbors(u):
                     edge_key = frozenset((u, v))
-                    
+
                     # if this edge has not been used, a new triple has been found
                     if edge_key not in used_edges:
                         used_edges.add(edge_key)
-                        
+
                         # use the edge name to ensure the uniqueness of the ID
                         u_sorted, v_sorted = sorted((u, v))
                         yield Community(
                             id=f"{u_sorted}-{v_sorted}",
                             nodes=[u_sorted, v_sorted],
-                            edges=[(u_sorted, v_sorted)]
+                            edges=[(u_sorted, v_sorted)],
                         )
-                    
+
                     # continue to BFS
                     if v not in visited_nodes:
                         visited_nodes.add(v)
-                        queue.append(v)
+                        queue.append(v)