fix: fix fetching img_path in vqa_generator

Author: ChenZiHong-Gavin

graphgen/models/generator/vqa_generator.py

@@ -80,7 +80,7 @@ async def generate(
         for node in nodes:
             node_data = node[1]
             if "images" in node_data and node_data["images"]:
-                img_path = node_data["images"]
+                img_path = node_data["images"]["img_path"]
                 for qa in qa_pairs.values():
                     qa["img_path"] = img_path
         result.update(qa_pairs)

graphgen/models/partitioner/anchor_bfs_partitioner.py

@@ -1,6 +1,6 @@
 import random
 from collections import deque
-from typing import Any, List, Literal, Set
+from typing import Any, List, Literal, Set, Tuple
 
 from graphgen.bases import BaseGraphStorage
 from graphgen.bases.datatypes import Community
@@ -55,36 +55,9 @@ async def partition(
         for seed_node in seeds:
             if seed_node in used_n:
                 continue
-
-            comm_n: List[str] = []
-            comm_e: List[tuple[str, str]] = []
-            queue: deque[tuple[str, Any]] = deque([(NODE_UNIT, seed_node)])
-            cnt = 0
-
-            while queue and cnt < max_units_per_community:
-                k, it = queue.popleft()
-
-                if k == NODE_UNIT:
-                    if it in used_n:
-                        continue
-                    used_n.add(it)
-                    comm_n.append(it)
-                    cnt += 1
-                    for nei in adj[it]:
-                        e_key = frozenset((it, nei))
-                        if e_key not in used_e:
-                            queue.append((EDGE_UNIT, e_key))
-                else:  # EDGE_UNIT
-                    if it in used_e:
-                        continue
-                    used_e.add(it)
-                    u, v = it
-                    comm_e.append((u, v))
-                    cnt += 1
-                    for n in it:
-                        if n not in used_n:
-                            queue.append((NODE_UNIT, n))
-
+            comm_n, comm_e = await self._grow_community(
+                seed_node, adj, max_units_per_community, used_n, used_e
+            )
             if comm_n or comm_e:
                 communities.append(
                     Community(id=len(communities), nodes=comm_n, edges=comm_e)
@@ -105,3 +78,51 @@ async def _pick_anchor_ids(
             if self.anchor_type.lower() in node_type:
                 anchor_ids.add(node_id)
         return anchor_ids
+
+    @staticmethod
+    async def _grow_community(
+        seed: str,
+        adj: dict[str, List[str]],
+        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 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
+        :return: (list of node ids, list of edge tuples)
+        """
+        comm_n: List[str] = []
+        comm_e: List[Tuple[str, str]] = []
+        queue: deque[tuple[str, Any]] = deque([(NODE_UNIT, seed)])
+        cnt = 0
+
+        while queue and cnt < max_units:
+            k, it = queue.popleft()
+
+            if k == NODE_UNIT:
+                if it in used_n:
+                    continue
+                used_n.add(it)
+                comm_n.append(it)
+                cnt += 1
+                for nei in adj[it]:
+                    e_key = frozenset((it, nei))
+                    if e_key not in used_e:
+                        queue.append((EDGE_UNIT, e_key))
+            else:  # EDGE_UNIT
+                if it in used_e:
+                    continue
+                used_e.add(it)
+                u, v = it
+                comm_e.append((u, v))
+                cnt += 1
+                for n in it:
+                    if n not in used_n:
+                        queue.append((NODE_UNIT, n))
+
+        return comm_n, comm_e