feat(graphgen): use average loss when choosing nodes & edges

Author: ChenZiHong-Gavin

graphgen/operators/split_graph.py

@@ -26,9 +26,10 @@ async def _get_node_info(
 
 def _get_level_n_edges_by_max_width(
     edge_adj_list: dict,
+    node_dict: dict,
     edges: list,
-    src_id: str,
-    tgt_id: str,
+    nodes,
+    src_edge: tuple,
     max_depth: int,
     bidirectional: bool,
     max_extra_edges: int,
@@ -39,15 +40,18 @@ def _get_level_n_edges_by_max_width(
     n is decided by max_depth in traverse_strategy
 
     :param edge_adj_list
+    :param node_dict
     :param edges
-    :param src_id
-    :param tgt_id
+    :param nodes
+    :param src_edge
     :param max_depth
     :param bidirectional
     :param max_extra_edges
     :param edge_sampling
     :return: level n edges
     """
+    src_id, tgt_id, _ = src_edge
+
     level_n_edges = []
 
     start_nodes = {tgt_id} if not bidirectional else {src_id, tgt_id}
@@ -66,7 +70,8 @@ def _get_level_n_edges_by_max_width(
             break
 
         if len(candidate_edges) >= max_extra_edges:
-            candidate_edges = _sort_edges(candidate_edges, edge_sampling)[:max_extra_edges]
+            er_tuples = [([nodes[node_dict[edge[0]]], nodes[node_dict[edge[1]]]], edge) for edge in candidate_edges]
+            candidate_edges = _sort_edges(er_tuples, edge_sampling)[:max_extra_edges]
             for edge in candidate_edges:
                 level_n_edges.append(edge)
                 edge[2]["visited"] = True
@@ -138,7 +143,8 @@ def _get_level_n_edges_by_max_tokens(
         if not candidate_edges:
             break
 
-        candidate_edges = _sort_edges(candidate_edges, edge_sampling)
+        er_tuples = [([nodes[node_dict[edge[0]]], nodes[node_dict[edge[1]]]], edge) for edge in candidate_edges]
+        candidate_edges = _sort_edges(er_tuples, edge_sampling)
         for edge in candidate_edges:
             max_tokens -= edge[2]["length"]
             if not edge[0] in temp_nodes:
@@ -166,22 +172,24 @@ def _get_level_n_edges_by_max_tokens(
     return level_n_edges
 
 
-def _sort_edges(edges: list, edge_sampling: str) -> list:
+def _sort_edges(er_tuples: list, edge_sampling: str) -> list:
     """
     Sort edges with edge sampling strategy
 
-    :param edges: total edges
+    :param er_tuples: [(nodes:list, edge:tuple)]
     :param edge_sampling: edge sampling strategy (random, min_loss, max_loss)
     :return: sorted edges
     """
     if edge_sampling == "random":
-        random.shuffle(edges)
+        er_tuples = random.sample(er_tuples, len(er_tuples))
     elif edge_sampling == "min_loss":
-        edges = sorted(edges, key=lambda x: x[2]["loss"])
+        er_tuples = sorted(er_tuples, key=lambda x: sum(node[1]["loss"] for node in x[0]) + x[1][2]["loss"])
     elif edge_sampling == "max_loss":
-        edges = sorted(edges, key=lambda x: x[2]["loss"], reverse=True)
+        er_tuples = sorted(er_tuples, key=lambda x: sum(node[1]["loss"] for node in x[0]) + x[1][2]["loss"],
+                           reverse=True)
     else:
         raise ValueError(f"Invalid edge sampling: {edge_sampling}")
+    edges = [edge for _, edge in er_tuples]
     return edges
 
 async def get_batches_with_strategy(
@@ -199,8 +207,6 @@ async def get_batches_with_strategy(
     max_depth = traverse_strategy.max_depth
     edge_sampling = traverse_strategy.edge_sampling
 
-    edges = _sort_edges(edges, edge_sampling)
-
     # 构建临接矩阵
     edge_adj_list = defaultdict(list)
     node_dict = {}
@@ -220,6 +226,9 @@ async def get_cached_node_info(node_id: str) -> dict:
     for i, (node_name, _) in enumerate(nodes):
         node_dict[node_name] = i
 
+    er_tuples = [([nodes[node_dict[edge[0]]], nodes[node_dict[edge[1]]]], edge) for edge in edges]
+    edges = _sort_edges(er_tuples, edge_sampling)
+
     for edge in tqdm_async(edges, desc="Preparing batches"):
         if "visited" in edge[2] and edge[2]["visited"]:
             continue
@@ -238,7 +247,7 @@ async def get_cached_node_info(node_id: str) -> dict:
 
         if expand_method == "max_width":
             level_n_edges = _get_level_n_edges_by_max_width(
-                edge_adj_list, edges, src_id, tgt_id, max_depth,
+                edge_adj_list, node_dict, edges, nodes, edge, max_depth,
                 traverse_strategy.bidirectional, traverse_strategy.max_extra_edges,
                 edge_sampling
             )
@@ -260,7 +269,6 @@ async def get_cached_node_info(node_id: str) -> dict:
 
         processing_batches.append((_process_nodes, _process_edges))
 
-    l
     # isolate nodes
     isolated_node_strategy = traverse_strategy.isolated_node_strategy
     if isolated_node_strategy == "add":