diff --git a/recommend/utils/data_generator.py b/recommend/utils/graph_build.py
similarity index 92%
rename from recommend/utils/data_generator.py
rename to recommend/utils/graph_build.py
index d3ef9cf..becf5df 100644
--- a/recommend/utils/data_generator.py
+++ b/recommend/utils/graph_build.py
@@ -9,7 +9,7 @@
 SqlPassword = "123456"
 
 
-def fetch_data():
+def fetch_user_seed_data():
     conn = pymysql.connect(
         host=SqlURL,
         port=SqlPort,
@@ -74,8 +74,10 @@
             f.write(f"{uid}\t{items}\t{times}\n")
 
 
-def build_user_seed_graph():
-    download_rows, favorite_rows = fetch_data()
+def build_user_seed_graph(return_mapping=False):
+    download_rows, favorite_rows = fetch_user_seed_data()
     records, user_set, seed_set = process_records(download_rows, favorite_rows)
     user2idx, seed2idx = build_id_maps(user_set, seed_set)
     group_and_write(records, user2idx, seed2idx)
+    if return_mapping:
+        return user2idx, seed2idx
\ No newline at end of file