fix: update anchor_bfs_partitioner & fix vqa_generator (#166)

Author: ChenZiHong-Gavin

graphgen/models/generator/vqa_generator.py

@@ -106,7 +106,7 @@ def format_generation_results(
                                 {"text": v["question"], "image": v.get("img_path", "")}
                             ],
                         },
-                        {"from": "gpt", "value": v["answer"]},
+                        {"from": "gpt", "value": [{"text": v["answer"]}]},
                     ]
                 }
                 for item in results
@@ -122,7 +122,10 @@ def format_generation_results(
                                 {"text": v["question"], "image": v.get("img_path", "")}
                             ],
                         },
-                        {"role": "assistant", "content": v["answer"]},
+                        {
+                            "role": "assistant",
+                            "content": [{"type": "text", "text": v["answer"]}],
+                        },
                     ]
                 }
                 for item in results

graphgen/models/partitioner/anchor_bfs_partitioner.py

@@ -37,9 +37,6 @@ def partition(
         **kwargs: Any,
     ) -> Iterable[Community]:
         nodes = g.get_all_nodes()  # List[tuple[id, meta]]
-        edges = g.get_all_edges()  # List[tuple[u, v, meta]]
-
-        adj, _ = self._build_adjacency_list(nodes, edges)
 
         anchors: Set[str] = self._pick_anchor_ids(nodes)
         if not anchors:
@@ -55,7 +52,7 @@ def partition(
             if seed_node in used_n:
                 continue
             comm_n, comm_e = self._grow_community(
-                seed_node, adj, max_units_per_community, used_n, used_e
+                seed_node, g, max_units_per_community, used_n, used_e
             )
             if comm_n or comm_e:
                 yield Community(id=seed_node, nodes=comm_n, edges=comm_e)
@@ -77,15 +74,15 @@ def _pick_anchor_ids(
     @staticmethod
     def _grow_community(
         seed: str,
-        adj: dict[str, List[str]],
+        g: BaseGraphStorage,
         max_units: int,
         used_n: set[str],
         used_e: set[frozenset[str]],
     ) -> Tuple[List[str], List[Tuple[str, str]]]:
         """
         Grow a community from the seed node using BFS.
         :param seed: seed node id
-        :param adj: adjacency list
+        :param g: graph storage
         :param max_units: maximum number of units (nodes + edges) in the community
         :param used_n: set of used node ids
         :param used_e: set of used edge keys
@@ -105,7 +102,7 @@ def _grow_community(
                 used_n.add(it)
                 comm_n.append(it)
                 cnt += 1
-                for nei in adj[it]:
+                for nei in g.get_neighbors(it):
                     e_key = frozenset((it, nei))
                     if e_key not in used_e:
                         queue.append((EDGE_UNIT, e_key))